feat(sglang): mxfp8 quantized rollout and refit

examples/configs/grpo_math_1B_sglang.yaml

@@ -22,8 +22,17 @@ policy:
       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
+      cuda_graph_backend_prefill: breakable
       disable_cuda_graph: false
+      # Weight precision for rollout/refit. scheme=bf16 (default) sends BF16
+      # HF tensors; scheme=mxfp8 boots SGLang from a matching quantized
+      # HF checkpoint and quantizes refit tensors online.
+      quantization:
+        scheme: bf16
+        extra_high_precision_layers_hf: []
+        num_layers_at_start_in_bf16: 0
+        num_layers_at_end_in_bf16: 0
+        modules_to_not_convert: []
       sglang_server_config:
         needs_offload: true
         cpu_weight_backup: true
@@ -32,8 +41,16 @@ policy:
         pause_generation_mode: retract
         num_gpus: 2
         num_gpus_per_engine: ${policy.generation.sglang_cfg.tp_size}
+        # "ipc" for colocated inference, "broadcast" for disaggregate GPUs.
+        weight_transfer_mode: ipc
       sglang_router_config:
         use_external_router: false
+      # Fault tolerance (RolloutHealthMonitor). Off by default; when enabled,
+      # a daemon thread health-checks each engine and restarts hung/dead actors.
+      use_fault_tolerance: false
+      rollout_health_check_interval: 60
+      rollout_health_check_timeout: 60
+      rollout_health_check_first_wait: 60
     colocated:
       enabled: true
 

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

@@ -0,0 +1,63 @@
+defaults: ../../grpo_math_1B.yaml
+
+grpo:
+  max_num_steps: 450
+  val_batch_size: 128
+
+checkpointing:
+  checkpoint_dir: results/grpo-qwen2.5-math-1.5b-instruct-1n4g-fsdp2tp1-sglang
+
+policy:
+  model_name: Qwen/Qwen2.5-Math-1.5B-Instruct
+  tokenizer:
+    name: Qwen/Qwen2.5-Math-1.5B-Instruct
+  dynamic_batching:
+    enabled: true
+  sequence_packing:
+    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: 1
+      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
+      cuda_graph_backend_prefill: breakable
+      disable_cuda_graph: false
+      sglang_server_config:
+        needs_offload: true
+        cpu_weight_backup: true
+        sglang_server_concurrency: 1024
+        pause_generation_mode: retract
+        num_gpus: 4
+        num_gpus_per_engine: ${policy.generation.sglang_cfg.tp_size}
+        weight_transfer_mode: ipc
+      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-1n4g-fsdp2tp1-sglang
+  wandb_enabled: true
+  tensorboard_enabled: true
+  wandb:
+    project: nemo-rl-sglang-1n4g
+    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,72 @@
+defaults: ../../grpo_math_1B.yaml
+
+grpo:
+  max_num_steps: 450
+  val_batch_size: 128
+
+checkpointing:
+  checkpoint_dir: results/grpo-qwen2.5-math-1.5b-instruct-1n4g-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
+    tensor_model_parallel_size: 1
+    pipeline_model_parallel_size: 1
+    scheduler:
+      lr_warmup_iters: 50
+  dtensor_cfg:
+    enabled: false
+  dynamic_batching:
+    enabled: true
+  sequence_packing:
+    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: 1
+      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
+      cuda_graph_backend_prefill: breakable
+      disable_cuda_graph: false
+      sglang_server_config:
+        needs_offload: true
+        cpu_weight_backup: true
+        sglang_server_concurrency: 1024
+        pause_generation_mode: retract
+        num_gpus: 4
+        num_gpus_per_engine: ${policy.generation.sglang_cfg.tp_size}
+        weight_transfer_mode: ipc
+      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-1n4g-megatrontp1-sglang
+  wandb_enabled: true
+  tensorboard_enabled: true
+  wandb:
+    project: nemo-rl-sglang-1n4g
+    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-1n8g-fsdp2tp1-sglang.yaml

@@ -29,7 +29,7 @@ policy:
       max_running_requests: null
       mem_fraction_static: 0.6
       skip_server_warmup: true
-      disable_piecewise_cuda_graph: true
+      cuda_graph_backend_prefill: breakable
       disable_cuda_graph: false
       sglang_server_config:
         needs_offload: true
@@ -38,9 +38,11 @@ policy:
         pause_generation_mode: retract
         num_gpus: 8
         num_gpus_per_engine: ${policy.generation.sglang_cfg.tp_size}
+        weight_transfer_mode: ipc
       sglang_router_config:
         use_external_router: false
         use_distributed_post: true
+      use_fault_tolerance: false
 data:
   max_input_seq_length: 512
 logger:

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

@@ -33,7 +33,7 @@ policy:
       max_running_requests: null
       mem_fraction_static: 0.6
       skip_server_warmup: true
-      disable_piecewise_cuda_graph: true
+      cuda_graph_backend_prefill: breakable
       disable_cuda_graph: false
       sglang_server_config:
         needs_offload: true
@@ -42,9 +42,11 @@ policy:
         pause_generation_mode: retract
         num_gpus: 16
         num_gpus_per_engine: ${policy.generation.sglang_cfg.tp_size}
+        weight_transfer_mode: ipc
       sglang_router_config:
         use_external_router: false
         use_distributed_post: true
+      use_fault_tolerance: false
 
 data:
   max_input_seq_length: 512

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

@@ -0,0 +1,70 @@
+defaults: ../../grpo_math_1B.yaml
+grpo:
+  max_num_steps: 450
+  val_batch_size: 128
+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
+    tensor_model_parallel_size: 1
+    pipeline_model_parallel_size: 1
+    scheduler:
+      lr_warmup_iters: 50
+  dtensor_cfg:
+    enabled: false
+  dynamic_batching:
+    enabled: true
+  sequence_packing:
+    enabled: false
+  make_sequence_length_divisible_by: 1
+  generation:
+    backend: sglang
+    max_new_tokens: 512
+    colocated:
+      enabled: false
+      resources:
+        gpus_per_node: 8
+        num_nodes: 1
+    sglang_cfg:
+      model_path: ${policy.model_name}
+      dtype: ${policy.precision}
+      context_length: 512
+      allow_auto_truncate: true
+      tp_size: 2
+      dp_size: 1
+      pp_size: 1
+      ep_size: 1
+      random_seed: 42
+      max_running_requests: null
+      mem_fraction_static: 0.6
+      skip_server_warmup: true
+      cuda_graph_backend_prefill: breakable
+      disable_cuda_graph: false
+      sglang_server_config:
+        needs_offload: true
+        cpu_weight_backup: true
+        sglang_server_concurrency: 1024
+        pause_generation_mode: retract
+        num_gpus: 8
+        num_gpus_per_engine: ${policy.generation.sglang_cfg.tp_size}
+        weight_transfer_mode: broadcast
+      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-megatrontp1-sglang-noncolocated
+  wandb_enabled: true
+  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