perf(vllm): batch BF16 TRTLLM expert layout conversion

docs/design-docs/nccl-reshard-refit.md

@@ -39,6 +39,12 @@ single `ValueError` listing every violation. The current requirements are:
   `vllm_cfg.precision=fp8` and `vllm_cfg.is_mx=true`; the generation ranks
   quantize each received BF16 shard before installing it. Blockwise-FP8 train →
   MXFP8 gen is not supported.
+* BF16 FlashInfer TRTLLM MoE is supported through vLLM's native
+  layerwise-reload path. Its grouped expert weights must use expert-parallel
+  destination sharding with linear expert placement; tensor-sharded expert
+  destinations and round-robin placement are rejected. This path does not
+  support an FP8 KV cache or a co-trained MTP drafter; setup rejects both
+  combinations.
 * vLLM expert parallelism is supported with the NeMo RL convention
   `expert_parallel_size == tensor_parallel_size`. 
 * Generation-side, PP > 1 is not supported. 
@@ -67,18 +73,22 @@ nccl-reshard-refit implementation:
   Two FFN-named groups are explicitly excluded and ride the misc path instead:
   shared-expert weights (`*.shared_expert.*`, which fuse differently on the vLLM
   side) and co-trained MTP drafter weights (which vLLM keeps in a separate
-  drafter module updated through `load_weights`). MTP weights are recognized two
-  ways: bare-`mtp.`-prefix HF names (NemotronH, Qwen3.5) via
+  drafter module updated through `load_weights`). Co-trained MTP is not supported
+  with BF16 FlashInfer TRTLLM; this routing applies to other supported backend
+  combinations. MTP weights are recognized two ways: bare-`mtp.`-prefix HF names
+  (NemotronH, Qwen3.5) via
   `is_nccl_reshard_param()`, and DeepSeek-style MTP exported as trailing
   `model.layers.N` indices via provenance — the Megatron-side name carries an
   `mtp.` module segment (bare for LM bridges, `language_model.mtp.*` for the VL
   and EXAONE bridges), so the worker excludes those HF layers when building the
   metadata (`_collect_mtp_hf_layer_names()`).
 * **Misc path** — everything else (embeddings, attention projections, layernorms, the
-  MoE router, `lm_head`, FP8 `_scale_inv` siblings, FP8 KV-cache scales, …). These ride
-  a packed broadcast (conventional `packed_tensor.py` implementation) over the shared
-  `model_update_group` and are loaded on the generation side through the backend's
-  regular `load_weights` machinery.
+  MoE router, `lm_head`, FP8 `_scale_inv` siblings, FP8 KV-cache scales, …). FP8
+  KV-cache scales are supported only by backend combinations that allow an FP8 KV
+  cache; BF16 FlashInfer TRTLLM rejects that configuration at setup. These tensors
+  ride a packed broadcast (conventional `packed_tensor.py` implementation) over the
+  shared `model_update_group` and are loaded on the generation side through the
+  backend's regular `load_weights` machinery.
 
 The feature is integrated into the `nemo_rl/weight_sync/` framework:
 `create_weight_synchronizer(..., nccl_reshard_refit=True)` returns a
@@ -92,8 +102,8 @@ training starts:
 
 1. **`init_collective()`** — creates the `model_update_group`, a NCCL group spanning all
    training and generation ranks. The bulk path does not use it; it carries the misc
-   packed-broadcast (and FP8 KV-cache scales), identical to the conventional collective
-   transport.
+   packed-broadcast, including FP8 KV-cache scales for backend combinations that support
+   them, identical to the conventional collective transport.
 2. **`init_nccl_reshard_comm_group()`** — creates the bulk-path communicator(s): **one
    NCCL group per training PP stage**, each spanning that stage's training ranks plus
    *all* generation ranks (non-PP is simply `pp_size == 1`, a single group over
@@ -135,10 +145,11 @@ realized **locally**:
   (sent as-is); grouped MoE experts get a `pre` hook that stacks this rank's per-expert
   views into a `[num_local_experts, ...]` tensor fresh at each refit.
 * On the **generation side**, a direct parameter's `base` is the live vLLM parameter
-  (received into in place); a parameter that is a slice of a fused vLLM tensor (dense
-  `gate_up_proj`, grouped-expert `w13`/`w2`) gets a `pre` hook that allocates a receive
-  buffer for its region and a `post` hook that copies the received shard back into the
-  fused parameter.
+  (received into in place). Conventional fused parameters use `pre`/`post` hooks to
+  receive a component and copy it into the appropriate local region. BF16
+  FlashInfer TRTLLM grouped experts instead receive into canonical EP-local staging
+  tensors; `post` loads each logical expert with its global expert ID through vLLM's
+  native weight loader.
 
 ### Execution Flow: Refit Time
 
@@ -159,7 +170,9 @@ Every training step (with in-flight weight updates, concurrently with generation
   are distributed across `NRL_REFIT_NUM_STREAMS` CUDA streams so different stages'
   reshards overlap. For each parameter it runs `pre` (receive-buffer allocation), calls
   `xferdtensor(None, ..., dst, ..., group, stream)`, then `post` (copy back into the
-  fused parameter).
+  fused parameter or load staged TRTLLM experts). After every transfer completes, the
+  TRTLLM path finalizes vLLM's native layerwise reload once to restore the packed runtime
+  layout.
 
 ### The Misc Path
 
@@ -198,10 +211,11 @@ generation side maps those HF names onto whatever its own storage layout is.
   `nccl_reshard_refit()` send loop; the misc packed-broadcast producer.
 * **Generation side** (`vllm_backend.py`): building `hf_to_local_param_map` — mapping HF
   names onto vLLM's fused parameters (`qkv_proj`, `gate_up_proj`, grouped-expert
-  `w13_weight`/`w2_weight`) with `pre`/`post` hooks for the slice regions, which is
-  deliberately **shape-driven** so the same code handles generation TP and generation
-  EP; the comm bootstrap methods; the `nccl_reshard_refit()` receive loop; the misc
-  consumer feeding `load_weights`.
+  `w13_weight`/`w2_weight`) with `pre`/`post` hooks for slice regions or canonical
+  TRTLLM staging, which is deliberately **shape-driven** so the same code handles
+  supported generation parallelism; the comm bootstrap methods; the
+  `nccl_reshard_refit()` receive loop; the misc consumer feeding `load_weights`; and
+  backend-specific finalization after all weights arrive.
 
 **To extend to a new backend**, the only piece with genuinely new logic is
 `build_hf_to_local_param_map`. Everything else is boilerplate that follows a fixed
@@ -210,9 +224,10 @@ contract and can be copied from the existing backend almost verbatim.
 **The one backend-specific implementation — `build_hf_to_local_param_map`:** resolve
 each bulk HF name to your local storage as a `LocalParamSpec` — `base` for tensors
 sent/received as-is, and `pre`/`post` hooks wherever your layout requires staging
-(fused/merged tensors, layout conversions, grouped-expert stacking). This is the *only*
-place your backend's parameter layout is encoded; all cross-mesh byte movement is
-already handled by the shared metadata and `xferdtensor`.
+(fused/merged tensors, layout conversions, grouped-expert stacking). Backends that
+rebuild runtime storage may also need one transport-level finalizer after all specs have
+run. These are the only places the backend's parameter layout is encoded; all cross-mesh
+byte movement is already handled by the shared metadata and `xferdtensor`.
 
 (A new *training* backend additionally has to produce the HF-named metadata — names,
 global shapes, dtypes, and the parallelism description the agnostic builder consumes —

docs/fp8.md

@@ -73,6 +73,43 @@ attention, the router, and the language-model head in BF16:
                     - model.layers.*.mlp.gate
 ```
 
+`num_first_layers_in_bf16` and `num_last_layers_in_bf16` keep complete
+transformer layers in BF16. NeMo RL reads the model configuration and parameter
+names, so these options do not require a model-specific `model.layers` or
+`backbone.layers` prefix. The patterns above still apply to every middle layer.
+
+For example, this Nemotron 3.5 Lightning scope keeps the first two and last six
+layers in BF16. In the middle layers it quantizes only the non-shared routed
+experts. Attention, Mamba projections, routers, shared experts, latent
+projections, and MTP remain in BF16:
+
+```yaml
+policy:
+  generation:
+    vllm_cfg:
+      precision: fp8
+      is_mx: true
+      num_first_layers_in_bf16: 2
+      num_last_layers_in_bf16: 6
+      quantization_ignore_patterns:
+        - "*layers.*.mixer.qkv_proj"
+        - "*layers.*.mixer.o_proj"
+        - "*layers.*.mixer.in_proj"
+        - "*layers.*.mixer.out_proj"
+        - "*layers.*.mixer.up_proj"
+        - "*layers.*.mixer.down_proj"
+        - "*layers.*.mixer.gate"
+        - "*layers.*.mixer.shared_experts.*"
+        - "*layers.*.mixer.fc1_latent_proj"
+        - "*layers.*.mixer.fc2_latent_proj"
+        - "*mtp.*"
+```
+
+To quantize QKVO as well as the routed experts, remove the `qkv_proj` and
+`o_proj` entries. The first two and last six layers still remain entirely in
+BF16. Check the logged effective ignore list when adding a new model family;
+an ignore pattern that matches no module is usually a naming error.
+
 `lm_head` is always excluded from FP8 and MXFP8 quantization, even when it is
 not listed in `quantization_ignore_patterns` in the YAML configuration.
 Models with MTP layers must list their MTP module names explicitly, for example

docs/guides/models/qwen/qwen3-5.md

@@ -64,6 +64,7 @@ authoritative settings.
 | Qwen3.5-9B-Base | LLM | GRPO | Megatron | 1n8g | [`grpo-qwen3.5-9b-1n8g-megatron-fp8.yaml`](../../../../examples/configs/recipes/llm/grpo-qwen3.5-9b-1n8g-megatron-fp8.yaml) |
 | Qwen3.5-35B-A3B-Base | LLM | GRPO | Megatron | 2n8g | [`grpo-qwen3.5-35ba3b-2n8g-megatron-ep16tp2cp2.yaml`](../../../../examples/configs/recipes/llm/grpo-qwen3.5-35ba3b-2n8g-megatron-ep16tp2cp2.yaml) |
 | Qwen3.5-35B-A3B-Base | LLM | GRPO | Megatron | 2n8g | [`grpo-qwen3.5-35ba3b-2n8g-megatron-ep16tp2-fp8.yaml`](../../../../examples/configs/recipes/llm/grpo-qwen3.5-35ba3b-2n8g-megatron-ep16tp2-fp8.yaml) |
+| Qwen3.5-35B-A3B-Base | LLM | GRPO | Megatron | 6n4g | [`grpo-qwen3.5-35ba3b-6n4g-async-1off-bf16-trtllm.yaml`](../../../../examples/configs/recipes/llm/grpo-qwen3.5-35ba3b-6n4g-async-1off-bf16-trtllm.yaml) |
 | Qwen3.5-35B-A3B-Base | LLM | GRPO | AutoModel | 2n8g | [`grpo-qwen3.5-35ba3b-2n8g-automodel-ep16.yaml`](../../../../examples/configs/recipes/llm/grpo-qwen3.5-35ba3b-2n8g-automodel-ep16.yaml) |
 | Qwen3.5-35B-A3B-Base | LLM | GRPO | AutoModel | 4n8g | [`grpo-qwen3.5-35ba3b-dapo-4n8g-automodel.yaml`](../../../../examples/configs/recipes/llm/grpo-qwen3.5-35ba3b-dapo-4n8g-automodel.yaml) |
 | Qwen3.5-397B-A17B | LLM | GRPO | Megatron | 32n8g | [`grpo-qwen3.5-397ba17b-32n8g-megatron.v2.yaml`](../../../../examples/configs/recipes/llm/grpo-qwen3.5-397ba17b-32n8g-megatron.v2.yaml) |

examples/configs/recipes/llm/grpo-qwen3.5-35ba3b-6n4g-async-1off-bf16-trtllm.yaml

@@ -0,0 +1,35 @@
+defaults: grpo-qwen3.5-35ba3b-2n8g-megatron-ep16tp2cp2.yaml
+loss_fn:
+  use_importance_sampling_correction: true
+  truncated_importance_sampling_type: tis
+  truncated_importance_sampling_ratio: 2
+grpo:
+  async_grpo:
+    enabled: true
+    in_flight_weight_updates: true
+checkpointing:
+  checkpoint_dir: results/grpo-qwen3.5-35ba3b-6n4g-async-1off-bf16-trtllm
+policy:
+  generation:
+    refit_transport: nccl_reshard
+    colocated:
+      enabled: false
+      resources:
+        gpus_per_node: 4
+        num_nodes: 2
+    vllm_cfg:
+      async_engine: true
+      precision: bfloat16
+      expert_parallel_size: 4
+      gpu_memory_utilization: 0.8
+      enforce_eager: false
+    vllm_kwargs:
+      moe_backend: flashinfer_trtllm
+      expert_placement_strategy: linear
+logger:
+  wandb:
+    name: grpo-qwen3.5-35ba3b-6n4g-async-1off-bf16-trtllm
+cluster:
+  gpus_per_node: 4
+  num_nodes: 6
+  segment_size: 2

nemo_rl/models/generation/vllm/quantization/fp8.py

@@ -14,7 +14,7 @@
 
 import os
 import warnings
-from collections.abc import Iterable, Iterator, Sequence
+from collections.abc import Callable, Iterable, Iterator, Sequence
 from dataclasses import dataclass, field
 from typing import Any
 from unittest.mock import patch
@@ -264,9 +264,30 @@ def init_fp8(vllm_cfg, model_name, model_parallel_size):
         quantization_ignore_patterns = [
             pattern.strip() for pattern in quantization_ignore_patterns
         ]
+
+    num_first_layers_in_bf16 = vllm_cfg.get("num_first_layers_in_bf16", 0)
+    num_last_layers_in_bf16 = vllm_cfg.get("num_last_layers_in_bf16", 0)
+    get_text_config = getattr(config, "get_text_config", None)
+    text_config = (
+        get_text_config()
+        if callable(get_text_config)
+        else getattr(config, "text_config", config)
+    )
+    num_hidden_layers = text_config.num_hidden_layers
+    for field_name, value in (
+        ("num_first_layers_in_bf16", num_first_layers_in_bf16),
+        ("num_last_layers_in_bf16", num_last_layers_in_bf16),
+    ):
+        if not isinstance(value, int) or isinstance(value, bool):
+            raise ValueError(f"{field_name} must be an integer")
+        if not 0 <= value <= num_hidden_layers:
+            raise ValueError(
+                f"{field_name} must be between 0 and {num_hidden_layers}, got {value}"
+            )
+
     fp8_config_kwargs = {
-        "num_first_layers_in_bf16": vllm_cfg.get("num_first_layers_in_bf16", 0),
-        "num_last_layers_in_bf16": vllm_cfg.get("num_last_layers_in_bf16", 0),
+        "num_first_layers_in_bf16": num_first_layers_in_bf16,
+        "num_last_layers_in_bf16": num_last_layers_in_bf16,
         "model_parallel_size": model_parallel_size,
         "kv_cache_dtype": kv_cache_dtype,
         "use_fp8_weights": use_fp8_weights,
@@ -301,15 +322,13 @@ def init_fp8(vllm_cfg, model_name, model_parallel_size):
         monkey_patch_vllm_ray_executor(global_fp8_config)
 
     # create fp8 kwargs for vllm's LLM(...)
-    num_first_layers_in_bf16 = vllm_cfg.get("num_first_layers_in_bf16", 0)
-    num_last_layers_in_bf16 = vllm_cfg.get("num_last_layers_in_bf16", 0)
     if global_fp8_config.is_mx:
         fp8_block_quant_kwargs = dict(MXFP8_BLOCK_QUANT_KWARGS)
     else:
         fp8_block_quant_kwargs = dict(FP8_BLOCK_QUANT_KWARGS)
     if num_first_layers_in_bf16 > 0 or num_last_layers_in_bf16 > 0:
         with init_empty_weights():
-            model = AutoModel.from_config(config)
+            model = AutoModel.from_config(config, trust_remote_code=True)
         param_names = [name for name, _ in model.named_parameters()]
 
         bf16_params = []
@@ -321,8 +340,8 @@ def init_fp8(vllm_cfg, model_name, model_parallel_size):
             layers = [
                 l
                 for l in range(
-                    config.num_hidden_layers - num_last_layers_in_bf16,
-                    config.num_hidden_layers,
+                    num_hidden_layers - num_last_layers_in_bf16,
+                    num_hidden_layers,
                 )
             ]
             bf16_params.extend(_get_params_in_layers(param_names, layers))
@@ -338,7 +357,7 @@ def init_fp8(vllm_cfg, model_name, model_parallel_size):
         )
     if quantization_ignored_layer_kws:
         with init_empty_weights():
-            model = AutoModel.from_config(config)
+            model = AutoModel.from_config(config, trust_remote_code=True)
         param_names = [
             f"model.{name}".removesuffix(".weight").replace(
                 "model.backbone.", "backbone."
@@ -607,15 +626,17 @@ def get_quantized_weight_iterator(
 
 
 def load_weights(
-    weights: Iterable[tuple[str, torch.Tensor]], model_runner: Any
+    weights: Iterable[tuple[str, torch.Tensor]],
+    model_runner: Any,
+    *,
+    model_load_weights: Callable[..., object] | None = None,
 ) -> None:
     """Quantize weights for the legacy direct model-loading path."""
-    # Finally load the weights into vllm
-    model_runner.model.load_weights(
+    if model_load_weights is None:
+        model_load_weights = model_runner.model.load_weights
+    model_load_weights(
         get_quantized_weight_iterator(
-            weights,
-            model_runner,
-            refit_with_reload_api=False,
+            weights, model_runner, refit_with_reload_api=False
         )
     )