feat(trtllm): Qwen3.5 routed-expert block-FP8/MXFP8 refit

docs/guides/refit.md

@@ -77,6 +77,69 @@ refit workload.
 | `policy.generation.vllm_cfg.refit_prequantize` | Megatron policy with MXFP8 vLLM rollout | Quantizes supported weights before transfer. See [Optional trainer-side MXFP8 refit](../fp8.md#optional-trainer-side-mxfp8-refit) for requirements and exclusions. |
 | `policy.refit_buffer_size_gb` | Refit paths that stage weights in buffers | Sets a fixed buffer size in GiB. NCCL broadcast uses the same size on both sides to keep chunk boundaries equal. |
 
+### TensorRT-LLM Routed-Expert FP8
+
+TensorRT-LLM supports online block-FP8 refit for MoE rollouts. Policy
+training remains in BF16, and this setting does not change optimizer-state
+precision. During each refit, NeMo RL converts weights to E4M3 with 128x128
+blocks and FP32 scales. Only the routed-expert weights are converted;
+attention, routers, shared experts, embeddings, and the output head remain in
+BF16.
+
+This path starts from a BF16 policy checkpoint. Direct loading of a
+pre-quantized ModelOpt FP8 checkpoint is outside its scope.
+
+Enable the path with:
+
+```yaml
+policy:
+  generation:
+    trtllm_cfg:
+      precision: fp8
+```
+
+Supported models are Qwen3.5 MoE (`qwen3_5_moe`) and Qwen3 MoE (`qwen3_moe`);
+setup fails for any other model type. The path quantizes routed-expert weights,
+either fused (`mlp.experts.gate_up_proj` / `down_proj`, the Qwen3.5 layout) or
+per-expert (`mlp.experts.{i}.{gate,up,down}_proj.weight`, the Qwen3 MoE layout).
+Setup also fails if no routed-expert weight is found, or if the model has dense
+MLP weights that would stay BF16. Supporting another model means extending the
+routed-expert name patterns in `quantization/fp8.py`.
+
+Verified end to end: Qwen3.5 MoE with MXFP8 (`Qwen/Qwen3.5-35B-A3B-Base` and
+`Qwen/Qwen3.5-397B-A17B`).
+
+NeMo RL initializes the TensorRT-LLM model with `load_format` set
+to `dummy`, then populates it from the first BF16 policy refit. By default
+(`is_mx: false`), the TensorRT-LLM MoE backend (`moe_config.backend: TRTLLM`)
+is required to preserve FP32 block scales; MXFP8/E8M0 scales are not used.
+
+#### MXFP8 variant
+
+Set `is_mx: true` to quantize routed experts to MXFP8 (E4M3 weights with
+UE8M0 1x32 block scales) instead of 128x128 block-FP8 with FP32 scales:
+
+```yaml
+policy:
+  generation:
+    trtllm_cfg:
+      precision: fp8
+      is_mx: true
+```
+
+This selects the CUTLASS MoE backend (the default) or the CuTe DSL backend
+(`moe_config.backend: CUTEDSL`, for Rubin) instead of the `TRTLLM` backend that
+block-FP8 uses; any other backend fails at setup. MXFP8 also constrains which
+GPUs it can run on: TensorRT-LLM's
+CUTLASS MoE gates `QuantAlgo.MXFP8` on `sm_constraint in {100, 103}`, so it does
+not run everywhere the default block-FP8 path does.
+
+The installed TensorRT-LLM must provide the incremental-refit lifecycle APIs
+`begin_update_weights`, `finalize_update_weights`, `abort_update_weights`, and
+`WorkerExtension.finalize_weight_update`. NeMo RL fails during setup if any of
+the three model-loader hooks is missing; `finalize_weight_update` is called
+directly and is not checked up front.
+
 ## Minimal Configuration
 
 Colocated vLLM and SGLang refit need no transport configuration:
@@ -120,8 +183,9 @@ policy:
 ```
 
 This option requires the Megatron policy backend. It applies to non-colocated
-vLLM collective and `nccl_reshard` transports, plus non-colocated Megatron
-generation. Unsupported combinations fail during synchronizer setup. It is
+vLLM collective and `nccl_reshard` transports, non-colocated TensorRT-LLM
+collective refit, plus non-colocated Megatron generation. Unsupported
+combinations fail during synchronizer setup. It is
 disabled by default because CPU offload adds transfer overhead when the export
 already fits in trainer GPU memory.
 

examples/configs/recipes/llm/grpo-qwen3.5-35ba3b-4n4g-megatron-trtllm-mxfp8-noncolocated.yaml

@@ -0,0 +1,61 @@
+defaults: ../../grpo_math_1B_trtllm.yaml
+grpo:
+  num_prompts_per_step: 16
+  num_generations_per_prompt: 8
+loss_fn:
+  use_importance_sampling_correction: true
+  truncated_importance_sampling_type: tis
+  truncated_importance_sampling_ratio: 2
+checkpointing:
+  checkpoint_dir: results/grpo-qwen3.5-35ba3b-4n4g-megatron-trtllm-mxfp8-noncolocated
+policy:
+  model_name: Qwen/Qwen3.5-35B-A3B-Base
+  train_global_batch_size: 128
+  train_micro_batch_size: 1
+  logprob_batch_size: 1
+  max_total_sequence_length: 4096
+  make_sequence_length_divisible_by: 16
+  optimizer: null
+  scheduler: null
+  automodel_cfg:
+    enabled: false
+  megatron_cfg:
+    enabled: true
+    tensor_model_parallel_size: 2
+    expert_model_parallel_size: 4
+    sequence_parallel: true
+    apply_rope_fusion: false
+    activation_checkpointing: true
+    defer_fp32_logits: true
+    moe_router_dtype: fp32
+    freeze_config:
+      freeze_vision_model: true
+      freeze_vision_projection: true
+      freeze_language_model: false
+  generation:
+    colocated:
+      enabled: false
+      resources:
+        gpus_per_node: 4
+        num_nodes: 2
+    trtllm_cfg:
+      tensor_parallel_size: 4
+      moe_expert_parallel_size: 4
+      gpu_memory_utilization: 0.7
+      precision: fp8
+      is_mx: true
+    trtllm_kwargs:
+      kv_cache_config:
+        mamba_ssm_cache_dtype: bfloat16
+        # TODO: migrate to KVCM V2
+        use_kv_cache_manager_v2: false
+      disable_mm_encoder: true
+logger:
+  log_dir: logs/grpo-qwen3.5-35ba3b-4n4g-megatron-trtllm-mxfp8-noncolocated
+  wandb:
+    project: nemo-rl
+    name: grpo-qwen3.5-35ba3b-4n4g-megatron-trtllm-mxfp8-noncolocated
+cluster:
+  gpus_per_node: 4
+  num_nodes: 4
+  segment_size: 2

nemo_rl/models/generation/constants.py

@@ -22,3 +22,4 @@
 SGLANG_BACKEND = "sglang"
 MEGATRON_BACKEND = "megatron"
 DYNAMO_BACKEND = "dynamo"
+TRTLLM_BACKEND = "trtllm"

nemo_rl/models/generation/dynamo/refit.py

@@ -152,6 +152,7 @@ def sender_spec(self) -> CollectiveSenderSpec:
             nccl_peer="vllm",
             buffer_size_bytes=VLLM_PACKED_BUFFER_SIZE_BYTES,
             num_buffers=VLLM_PACKED_NUM_BUFFERS,
+            use_legacy_packing=True,
         )
 
     def prepare(self, state_dict_info: dict[str, Any] | None) -> None:

nemo_rl/models/generation/interfaces.py

@@ -442,6 +442,7 @@ class CollectiveSenderSpec:
     nccl_peer: str = "nemo"
     buffer_size_bytes: int | None = None
     num_buffers: int | None = None
+    use_legacy_packing: bool = False
 
 
 def agreed_media_columns(

nemo_rl/models/generation/trtllm/config.py

@@ -84,6 +84,8 @@ class TrtllmSpecificArgs(TrtllmEngineKnobs):
     model_name: NotRequired[str]
     max_model_len: int
     precision: str
+    # With precision="fp8", use MXFP8 (UE8M0 1x32 scales) instead of 128x128 block-FP8.
+    is_mx: NotRequired[bool]
     expose_http_server: NotRequired[bool]
     async_engine: NotRequired[bool]
     # Duplicated from grpo.async_grpo because the generation backend never sees

nemo_rl/models/generation/trtllm/quantization/__init__.py

@@ -0,0 +1,3 @@
+# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved.
+
+"""Quantization helpers for TRT-LLM generation."""