feat(sglang): megatron backend weight refit for sglang rollouts

docker/Dockerfile

@@ -73,6 +73,7 @@ apt-get install -y --no-install-recommends \
     less \
     vim \
     ccache \
+    protobuf-compiler \
 
 # Nsight
 apt install -y --no-install-recommends gnupg
@@ -107,6 +108,13 @@ RUN GITHUB_ARTIFACTORY=github.com \
     && cp -r "${CMAKE_INSTALLER}/share/"* /usr/local/share/ \
     && rm -rf "${CMAKE_INSTALLER}" "${CMAKE_INSTALLER}.tar.gz"
 
+# Rust toolchain for setuptools-rust extensions in SGLang.
+ENV PATH="/root/.cargo/bin:$PATH"
+RUN curl --proto '=https' --tlsv1.2 --retry 3 --retry-delay 2 -sSf https://sh.rustup.rs \
+        | sh -s -- -y --no-modify-path --profile minimal \
+    && rustc --version \
+    && cargo --version
+
 # Install uv and python
 ARG UV_VERSION=0.11.28
 ARG PYTHON_VERSION=3.13.14

docs/design-docs/checkpoint-engines.md

@@ -5,9 +5,10 @@ They let GRPO move policy weights directly from policy workers to generation
 workers without using the driver as a model-sized staging point.
 
 The first built-in backend is NIXL. The current implementation targets policy
-workers refitting non-colocated vLLM generation workers. Colocated generation
-still uses the existing IPC/HTTP refit paths, and non-colocated generation
-without checkpoint engines still uses the existing NCCL collective path.
+workers refitting non-colocated vLLM generation workers. Without checkpoint
+engines, vLLM uses ZMQ/CUDA IPC when colocated and NCCL when non-colocated.
+SGLang uses Ray CUDA IPC when colocated and its own NCCL weight-update group
+with a Megatron policy when non-colocated.
 
 The user-facing guide is [Checkpoint-Engine Refit](../guides/checkpoint-engine-refit.md).
 

docs/guides/checkpoint-engine-refit.md

@@ -10,8 +10,9 @@ Use it only for non-colocated vLLM generation:
 - `policy.generation.colocated.enabled=false`
 - `policy.generation.refit_transport=nixl`
 
-Colocated generation still uses IPC/HTTP refit. Non-colocated generation without
-checkpoint-engine refit still uses the NCCL collective update path.
+Without checkpoint-engine refit, colocated vLLM uses ZMQ/CUDA IPC and
+non-colocated vLLM uses NCCL. SGLang uses Ray CUDA IPC when colocated and its
+own NCCL weight-update group with a Megatron policy when non-colocated.
 
 `examples/configs/grpo_math_8B_megatron_nixl.yaml` is a complete two-node
 example built as an overlay on the standard 8B Megatron recipe.

docs/guides/refit.md

@@ -8,8 +8,8 @@ topology first, then select one non-colocated transport with
 
 | Topology | `refit_transport` | Transport | Use when |
 |---|---|---|---|
-| `colocated.enabled: true` | `null` | CUDA IPC or HTTP | Policy and rollout workers share GPUs; vLLM uses IPC and SGLang uses HTTP. |
-| `colocated.enabled: false` | `null` | NCCL broadcast | You want the default full-weight path without extra dependencies. |
+| `colocated.enabled: true` | `null` | CUDA IPC | Policy and rollout workers share GPUs; vLLM uses ZMQ plus CUDA IPC handles and SGLang uses Ray CUDA IPC. |
+| `colocated.enabled: false` | `null` | NCCL broadcast | vLLM uses the default full-weight collective; SGLang uses its own weight-update group. |
 | `colocated.enabled: false` | `nccl_reshard` | NCCL reshard | Provides the best performance for large models (>100s B models). |
 | `colocated.enabled: false` | `vllm_zmq_sparse` | Sparse delta over ZeroMQ | The link is bandwidth-limited and workers can reach a relay over TCP. |
 | `colocated.enabled: false` | `vllm_s3_sparse` | Sparse delta through S3 | Workers communicate through shared object storage. |
@@ -23,17 +23,20 @@ delta and NIXL cannot both be active.
 
 | Transport | Generation backend | Policy backend | Quantization and MoE |
 |---|---|---|---|
-| Colocated IPC/HTTP | vLLM or SGLang | DTensor or Megatron | Uses the generation backend's standard loader. |
+| Colocated CUDA IPC | vLLM or SGLang | DTensor or Megatron | Uses the generation backend's standard loader. |
 | NCCL | vLLM or Megatron | DTensor or Megatron | Uses the standard full-weight loader. |
+| SGLang NCCL weight-update group | SGLang | Megatron | Trainer rank 0 broadcasts finalized HF tensors to the engine leaders. |
 | NCCL reshard | vLLM | Megatron | Requires matching BF16 or blockwise FP8 precision; Megatron ETP must be 1. Currently supporting Megatron+vLLM backends. |
 | Sparse delta | vLLM | Megatron | BF16/FP16, unquantized rollout only. |
 | NIXL, full weights | vLLM | DTensor or Megatron | Supports the standard full-weight FP8 loader. DTensor FP8 KV-cache scale transfer is not yet supported. |
 | NIXL, sharded experts | vLLM | DTensor or Megatron | Unquantized BF16/FP16 Triton MoE only; FP8/MXFP8 and dynamic expert placement are rejected. |
 
-Non-colocated SGLang generation is not supported. The NIXL restrictions are on
-the generation backend; both Megatron and DTensor policy workers can send
-weights. Sparse delta is currently limited to GRPO. NIXL is initialized by the
-GRPO and distillation setup paths; PPO currently requires colocated generation.
+Non-colocated SGLang generation is supported with a Megatron policy and
+`refit_transport: null`; it creates its own NCCL weight-update group. The NIXL
+restrictions are on the generation backend; both Megatron and DTensor policy
+workers can send weights. Sparse delta is currently limited to GRPO. NIXL is
+initialized by the GRPO and distillation setup paths; PPO currently requires
+colocated generation.
 
 ## Minimal Configuration
 

examples/configs/grpo_math_1B_sglang.yaml

@@ -8,6 +8,8 @@ policy:
     backend: "sglang"
     use_async_rollouts: false
     sglang_cfg:
+      quantization:
+        scheme: bf16
       model_path: ${policy.model_name}
       dtype: ${policy.precision}
       context_length: 512
@@ -20,15 +22,16 @@ policy:
       max_running_requests: null
       mem_fraction_static: 0.6
       skip_server_warmup: true
-      # Piecewise CUDA graph currently crashes with "illegal memory access"
-      # (likely torch 2.10 + sglang incompatibility). Keep disabled until upstream fix.
-      disable_piecewise_cuda_graph: true
+      # Prefill graphs off: "tc_piecewise" crashes with "illegal memory access" on
+      # torch 2.10, and "breakable" (upstream's CUDA default) is rejected under the
+      # memory saver that needs_offload: true enables below.
+      cuda_graph_backend_prefill: disabled
       disable_cuda_graph: false
       sglang_server_config:
         needs_offload: true
         cpu_weight_backup: true
         sglang_server_concurrency: 1024
-        # pause generation has no effect here
+        # Weight refit supports abort or retract; in_place preserves stale KV.
         pause_generation_mode: retract
         num_gpus: 2
         num_gpus_per_engine: ${policy.generation.sglang_cfg.tp_size}

examples/configs/recipes/llm/grpo-qwen2.5-math-1.5b-instruct-1n4g-fsdp2tp1-sglang.yaml

@@ -0,0 +1,20 @@
+defaults: ./grpo-qwen2.5-math-1.5b-instruct-2n8g-fsdp2tp1-sglang.yaml
+
+checkpointing:
+  checkpoint_dir: results/grpo-qwen2.5-math-1.5b-instruct-1n4g-fsdp2tp1-sglang
+
+policy:
+  generation:
+    sglang_cfg:
+      tp_size: 1
+      sglang_server_config:
+        num_gpus: 4
+
+logger:
+  log_dir: logs/grpo-qwen2.5-math-1.5b-instruct-1n4g-fsdp2tp1-sglang
+  wandb:
+    name: grpo-qwen2.5-math-1.5b-instruct-1n4g-fsdp2tp1-sglang
+
+cluster:
+  gpus_per_node: 4
+  num_nodes: 1

examples/configs/recipes/llm/grpo-qwen2.5-math-1.5b-instruct-1n4g-megatrontp1-sglang.yaml

@@ -0,0 +1,20 @@
+defaults: ./grpo-qwen2.5-math-1.5b-instruct-2n8g-megatrontp1-sglang.yaml
+
+checkpointing:
+  checkpoint_dir: results/grpo-qwen2.5-math-1.5b-instruct-1n4g-megatrontp1-sglang
+
+policy:
+  generation:
+    sglang_cfg:
+      tp_size: 1
+      sglang_server_config:
+        num_gpus: 4
+
+logger:
+  log_dir: logs/grpo-qwen2.5-math-1.5b-instruct-1n4g-megatrontp1-sglang
+  wandb:
+    name: grpo-qwen2.5-math-1.5b-instruct-1n4g-megatrontp1-sglang
+
+cluster:
+  gpus_per_node: 4
+  num_nodes: 1

examples/configs/recipes/llm/grpo-qwen2.5-math-1.5b-instruct-2n4g-megatrontp1-sglang-noncolocated.yaml

@@ -0,0 +1,22 @@
+defaults: ./grpo-qwen2.5-math-1.5b-instruct-2n8g-megatrontp1-sglang-noncolocated.yaml
+
+checkpointing:
+  checkpoint_dir: results/grpo-qwen2.5-math-1.5b-instruct-2n4g-megatrontp1-sglang-noncolocated
+
+policy:
+  generation:
+    colocated:
+      resources:
+        gpus_per_node: 4
+    sglang_cfg:
+      tp_size: 1
+      sglang_server_config:
+        num_gpus: 4
+
+logger:
+  log_dir: logs/grpo-qwen2.5-math-1.5b-instruct-2n4g-megatrontp1-sglang-noncolocated
+  wandb:
+    name: grpo-qwen2.5-math-1.5b-instruct-2n4g-megatrontp1-sglang-noncolocated
+
+cluster:
+  gpus_per_node: 4

examples/configs/recipes/llm/grpo-qwen2.5-math-1.5b-instruct-2n8g-fsdp2tp1-sglang.yaml

@@ -1,12 +1,8 @@
-defaults: ../../grpo_math_1B.yaml
-
+defaults: ../../grpo_math_1B_sglang.yaml
 grpo:
   max_num_steps: 450
-  val_batch_size: 128
-
 checkpointing:
   checkpoint_dir: results/grpo-qwen2.5-math-1.5b-instruct-2n8g-fsdp2tp1-sglang
-
 policy:
   model_name: Qwen/Qwen2.5-Math-1.5B-Instruct
   tokenizer:
@@ -17,46 +13,21 @@ policy:
     enabled: false
   make_sequence_length_divisible_by: 1
   generation:
-    backend: "sglang"
     max_new_tokens: 512
     sglang_cfg:
-      model_path: ${policy.model_name}
-      dtype: ${policy.precision}
-      context_length: 512
-      allow_auto_truncate: true
       tp_size: 4
-      dp_size: 1
-      # pp_size must be 1 (sglang pipeline parallelism not yet supported here).
-      pp_size: 1
-      ep_size: 1
-      random_seed: 42
-      max_running_requests: null
-      mem_fraction_static: 0.6
-      skip_server_warmup: true
-      disable_piecewise_cuda_graph: true
-      disable_cuda_graph: false
       sglang_server_config:
-        needs_offload: true
-        cpu_weight_backup: true
-        sglang_server_concurrency: 1024
-        pause_generation_mode: retract
         num_gpus: 16
-        num_gpus_per_engine: ${policy.generation.sglang_cfg.tp_size}
       sglang_router_config:
-        use_external_router: false
         use_distributed_post: true
-
 data:
   max_input_seq_length: 512
-
 logger:
   log_dir: logs/grpo-qwen2.5-math-1.5b-instruct-2n8g-fsdp2tp1-sglang
-  wandb_enabled: true
   tensorboard_enabled: true
   wandb:
     project: nemo-rl
     name: grpo-qwen2.5-math-1.5b-instruct-2n8g-fsdp2tp1-sglang
-
 cluster:
   gpus_per_node: 8
   num_nodes: 2

examples/configs/recipes/llm/grpo-qwen2.5-math-1.5b-instruct-2n8g-megatrontp1-sglang-noncolocated.yaml

@@ -0,0 +1,45 @@
+defaults: ../../grpo_math_1B_sglang.yaml
+grpo:
+  max_num_steps: 450
+checkpointing:
+  checkpoint_dir: results/grpo-qwen2.5-math-1.5b-instruct-2n8g-megatrontp1-sglang-noncolocated
+policy:
+  model_name: Qwen/Qwen2.5-Math-1.5B-Instruct
+  tokenizer:
+    name: Qwen/Qwen2.5-Math-1.5B-Instruct
+  optimizer: null
+  megatron_cfg:
+    enabled: true
+    scheduler:
+      lr_warmup_iters: 50
+  dtensor_cfg:
+    enabled: false
+  dynamic_batching:
+    enabled: true
+  sequence_packing:
+    enabled: false
+  make_sequence_length_divisible_by: 1
+  generation:
+    max_new_tokens: 512
+    colocated:
+      enabled: false
+      resources:
+        gpus_per_node: 8
+        num_nodes: 1
+    sglang_cfg:
+      tp_size: 2
+      sglang_server_config:
+        num_gpus: 8
+      sglang_router_config:
+        use_distributed_post: true
+data:
+  max_input_seq_length: 512
+logger:
+  log_dir: logs/grpo-qwen2.5-math-1.5b-instruct-2n8g-megatrontp1-sglang-noncolocated
+  tensorboard_enabled: true
+  wandb:
+    project: nemo-rl
+    name: grpo-qwen2.5-math-1.5b-instruct-2n8g-megatrontp1-sglang-noncolocated
+cluster:
+  gpus_per_node: 8
+  num_nodes: 2

examples/configs/recipes/llm/grpo-qwen2.5-math-1.5b-instruct-2n8g-megatrontp1-sglang.yaml

@@ -0,0 +1,40 @@
+defaults: ../../grpo_math_1B_sglang.yaml
+grpo:
+  max_num_steps: 450
+checkpointing:
+  checkpoint_dir: results/grpo-qwen2.5-math-1.5b-instruct-2n8g-megatrontp1-sglang
+policy:
+  model_name: Qwen/Qwen2.5-Math-1.5B-Instruct
+  tokenizer:
+    name: Qwen/Qwen2.5-Math-1.5B-Instruct
+  optimizer: null
+  megatron_cfg:
+    enabled: true
+    scheduler:
+      lr_warmup_iters: 50
+  dtensor_cfg:
+    enabled: false
+  dynamic_batching:
+    enabled: true
+  sequence_packing:
+    enabled: false
+  make_sequence_length_divisible_by: 1
+  generation:
+    max_new_tokens: 512
+    sglang_cfg:
+      tp_size: 4
+      sglang_server_config:
+        num_gpus: 16
+      sglang_router_config:
+        use_distributed_post: true
+data:
+  max_input_seq_length: 512
+logger:
+  log_dir: logs/grpo-qwen2.5-math-1.5b-instruct-2n8g-megatrontp1-sglang
+  tensorboard_enabled: true
+  wandb:
+    project: nemo-rl
+    name: grpo-qwen2.5-math-1.5b-instruct-2n8g-megatrontp1-sglang
+cluster:
+  gpus_per_node: 8
+  num_nodes: 2