diff --git a/docs/guides/refit.md b/docs/guides/refit.md index e429b325b14..b3f87449693 100644 --- a/docs/guides/refit.md +++ b/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. diff --git a/examples/configs/recipes/llm/grpo-qwen3.5-35ba3b-4n4g-megatron-trtllm-mxfp8-noncolocated.yaml b/examples/configs/recipes/llm/grpo-qwen3.5-35ba3b-4n4g-megatron-trtllm-mxfp8-noncolocated.yaml new file mode 100644 index 00000000000..b0397bd2818 --- /dev/null +++ b/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 diff --git a/nemo_rl/models/generation/constants.py b/nemo_rl/models/generation/constants.py index 973ea9d0ae7..06564f055ef 100644 --- a/nemo_rl/models/generation/constants.py +++ b/nemo_rl/models/generation/constants.py @@ -22,3 +22,4 @@ SGLANG_BACKEND = "sglang" MEGATRON_BACKEND = "megatron" DYNAMO_BACKEND = "dynamo" +TRTLLM_BACKEND = "trtllm" diff --git a/nemo_rl/models/generation/dynamo/refit.py b/nemo_rl/models/generation/dynamo/refit.py index 72ea8865706..ad8e5988025 100644 --- a/nemo_rl/models/generation/dynamo/refit.py +++ b/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: diff --git a/nemo_rl/models/generation/interfaces.py b/nemo_rl/models/generation/interfaces.py index 8d489ee0da3..47204458008 100644 --- a/nemo_rl/models/generation/interfaces.py +++ b/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( diff --git a/nemo_rl/models/generation/trtllm/config.py b/nemo_rl/models/generation/trtllm/config.py index 68c5d3152f4..df2658d8478 100644 --- a/nemo_rl/models/generation/trtllm/config.py +++ b/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 diff --git a/nemo_rl/models/generation/trtllm/quantization/__init__.py b/nemo_rl/models/generation/trtllm/quantization/__init__.py new file mode 100644 index 00000000000..c1b4c32f805 --- /dev/null +++ b/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.""" diff --git a/nemo_rl/models/generation/trtllm/quantization/fp8.py b/nemo_rl/models/generation/trtllm/quantization/fp8.py new file mode 100644 index 00000000000..1588d654fe0 --- /dev/null +++ b/nemo_rl/models/generation/trtllm/quantization/fp8.py @@ -0,0 +1,532 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Routed-expert block-FP8 / MXFP8 refit support for TRT-LLM (MoE models).""" + +import math +import re +from collections.abc import Iterable, Sequence +from typing import Any + +import torch +import torch.nn.functional as F + +FP8_BLOCK_SIZE = (128, 128) +FP8_EXPERT_CHUNK_SIZE = 16 + +# MXFP8: E4M3 weights, one UE8M0 (power-of-two exponent byte) scale per 32 elements +# along K. Shares this module's plumbing with block-FP8 but not the caster. +MXFP8_BLOCK_SIZE = 32 +UE8M0_BIAS = 127 +E4M3_MAX = 448.0 + +# Scope: routed experts only. TRT-LLM only has a negative module filter, so +# everything else is listed here and stays BF16. +ROUTED_EXPERTS_ONLY_IGNORE_PATTERNS = ( + "model.layers.*.self_attn*", + "model.layers.*.linear_attn*", + "model.layers.*.mlp.gate", + "model.layers.*.mlp.shared_expert*", + "model.embed_tokens", + "model.norm", + "lm_head", + # MTP layer: same split as above (experts quantized, rest BF16). + "mtp.layers.0.self_attn*", + "mtp.layers.0.linear_attn*", + "mtp.layers.0.mlp.gate", + "mtp.layers.0.mlp.shared_expert*", + "mtp.fc*", + "mtp.norm*", + "mtp.pre_fc_norm_embedding*", + "mtp.pre_fc_norm_hidden*", +) + + +def build_quant_config(is_mx: bool) -> dict[str, Any]: + """Build the TRT-LLM ``quantization_config`` for block-FP8 or MXFP8.""" + return { + "activation_scheme": "dynamic", + "fmt": "e4m3", + "quant_method": "mxfp8" if is_mx else "fp8", + "weight_block_size": [1, MXFP8_BLOCK_SIZE] if is_mx else list(FP8_BLOCK_SIZE), + "modules_to_not_convert": list(ROUTED_EXPERTS_ONLY_IGNORE_PATTERNS), + } + + +FP8_BLOCK_QUANT_KWARGS: dict[str, Any] = build_quant_config(is_mx=False) +MXFP8_BLOCK_QUANT_KWARGS: dict[str, Any] = build_quant_config(is_mx=True) + + +_EXPERT_PREFIX: str = ( + r"(?:" + r"(?:(?:model\.)?(?:language_model\.)?)layers\.\d+" + r"|mtp\.layers\.\d+" + r")\.mlp\.experts" +) +_FUSED_EXPERT_RE = re.compile( + rf"^(?P{_EXPERT_PREFIX})\." + r"(?Pgate_up_proj|down_proj)$" +) +_SPLIT_EXPERT_RE = re.compile( + rf"^(?P{_EXPERT_PREFIX})\.\d+\." + r"(?:gate_proj|up_proj|down_proj)\.weight$" +) +_DENSE_MLP_RE = re.compile( + r"^(?:(?:model\.)?(?:language_model\.)?)layers\.\d+\.mlp\." + r"(?:gate_proj|up_proj|down_proj)\.weight$" +) +_SPLIT_LINEAR_ATTN_RE = re.compile( + r"^(?:.*\.)?layers\.\d+\.linear_attn\." + r"in_proj_(?:qkv|q|k|v|z|a|b)\..+$" +) + + +def validate_fused_expert_layout( + state_dict_info: dict[str, Any], +) -> None: + """Validate Qwen3.5's fused routed-expert layout before the first refit.""" + layouts: dict[str, dict[str, tuple[int, ...]]] = {} + for name, metadata in state_dict_info.items(): + match = _FUSED_EXPERT_RE.fullmatch(name) + if match is None: + continue + shape, _ = metadata + layouts.setdefault(match.group("prefix"), {})[match.group("projection")] = ( + tuple(shape) + ) + + if not layouts: + raise ValueError( + "No Qwen3.5 fused routed-expert weights were found; expected " + "gate_up_proj=[E,2I,H] and down_proj=[E,H,I]." + ) + + for prefix, projections in layouts.items(): + if set(projections) != {"gate_up_proj", "down_proj"}: + raise ValueError( + f"Qwen3.5 fused experts require both gate_up_proj and down_proj: " + f"{prefix} has {sorted(projections)}" + ) + gate_up_shape = projections["gate_up_proj"] + down_shape = projections["down_proj"] + valid = ( + len(gate_up_shape) == 3 + and len(down_shape) == 3 + and gate_up_shape[0] == down_shape[0] + and gate_up_shape[1] == 2 * down_shape[2] + and gate_up_shape[2] == down_shape[1] + ) + if not valid: + raise ValueError( + "Qwen3.5 fused expert weights must use gate_up_proj=[E,2I,H] " + f"and down_proj=[E,H,I]; {prefix} has " + f"gate_up_proj={gate_up_shape}, down_proj={down_shape}" + ) + + +# HF ``model_type`` values whose routed-expert layout this refit path supports. +SUPPORTED_MODEL_TYPES = ("qwen3_5_moe", "qwen3_moe") + + +def validate_model_type(model_type: str) -> None: + """Fail setup for models outside the supported list.""" + if model_type not in SUPPORTED_MODEL_TYPES: + raise ValueError( + f"precision='fp8' supports model types {list(SUPPORTED_MODEL_TYPES)}, " + f"got {model_type!r}; supporting a new model means extending " + "the routed-expert name patterns in quantization/fp8.py" + ) + + +def validate_routed_experts(state_dict_info: dict[str, Any]) -> None: + """Fail setup unless the FP8 filter will quantize something, correctly. + + Routed experts are either fused (``mlp.experts.{gate_up,down}_proj``, Qwen3.5) + or per-expert (``mlp.experts.{i}.{gate,up,down}_proj.weight``, Qwen3 MoE). + Only experts are quantized, so at least one must exist and dense MLP weights + are rejected. + """ + names = [str(name) for name in state_dict_info] + has_fused = any(_FUSED_EXPERT_RE.fullmatch(name) for name in names) + has_split = any(_SPLIT_EXPERT_RE.fullmatch(name) for name in names) + if not (has_fused or has_split): + raise ValueError( + "precision='fp8' found no routed-expert weights to quantize; " + "expected mlp.experts.{gate_up_proj,down_proj} or " + "mlp.experts.{i}.{gate,up,down}_proj.weight" + ) + dense = [name for name in names if _DENSE_MLP_RE.fullmatch(name)] + if dense: + raise ValueError( + "precision='fp8' quantizes only routed experts, but the model has " + f"dense MLP weights that would stay BF16 (e.g. {dense[0]})" + ) + if has_fused: + validate_fused_expert_layout(state_dict_info) + + +def configure_fp8_llm_kwargs( + llm_kwargs: dict[str, Any], + *, + is_mx: bool = False, +) -> None: + """Apply the block-FP8 / MXFP8 contract to TRT-LLM args. + + Conflicting quantization or load-format overrides raise at setup. + """ + base_kwargs = build_quant_config(is_mx) + label = "MXFP8" if is_mx else "block-FP8" + + model_kwargs = dict(llm_kwargs.get("model_kwargs") or {}) + existing_quant_config = model_kwargs.get("quantization_config") + if existing_quant_config is not None and dict(existing_quant_config) != base_kwargs: + raise ValueError( + f"precision='fp8' requires NeMo-RL's {label} quantization_config; " + "do not override it" + ) + + load_format = llm_kwargs.get("load_format") + if load_format is not None and load_format != "dummy": + raise ValueError( + "precision='fp8' requires load_format='dummy'; the initial BF16 " + "trainer refit populates the FP8 weights and scales" + ) + + quantization_config = dict(base_kwargs) + quantization_config["modules_to_not_convert"] = list( + base_kwargs["modules_to_not_convert"] + ) + model_kwargs["quantization_config"] = quantization_config + llm_kwargs["model_kwargs"] = model_kwargs + llm_kwargs["load_format"] = "dummy" + llm_kwargs["dtype"] = "bfloat16" + # Keep the block-FP8 Linear fallback on FP32 scales (not DeepGEMM's E8M0). + llm_kwargs["use_cute_dsl_blockscaling_mm"] = True + + +def configure_fp8_moe_backend( + llm_kwargs: dict[str, Any], moe_config_type: type[Any], *, is_mx: bool = False +) -> None: + """Force (or check) the MoE backend for the requested scale format. + + Block-FP8 needs TRTLLM (keeps FP32 scales). MXFP8 uses CUTLASS (default) or + CUTEDSL (Rubin). + """ + allowed = ("CUTLASS", "CUTEDSL") if is_mx else ("TRTLLM",) + default = allowed[0] + reason = ( + "precision='fp8' with is_mx=true (MXFP8 routed-expert scales) requires " + if is_mx + else "precision='fp8' with FP32 routed-expert scales requires " + ) + expected = " or ".join( + f"trtllm_kwargs.moe_config.backend={backend!r}" for backend in allowed + ) + + moe_config = llm_kwargs.get("moe_config") + if moe_config is None: + llm_kwargs["moe_config"] = moe_config_type(backend=default) + return + + if isinstance(moe_config, dict): + moe_config_kwargs = dict(moe_config) + configured_backend = str(moe_config_kwargs.get("backend", default)).upper() + if configured_backend not in allowed: + raise ValueError(f"{reason}{expected}, got {configured_backend!r}") + moe_config_kwargs["backend"] = configured_backend + llm_kwargs["moe_config"] = moe_config_type(**moe_config_kwargs) + return + + if isinstance(moe_config, moe_config_type): + configured_backend = str(moe_config.backend).upper() + if configured_backend not in allowed: + raise ValueError(f"{reason}{expected}, got {moe_config.backend!r}") + return + + raise TypeError( + "trtllm_kwargs.moe_config must be a dict or MoeConfig, got " + f"{type(moe_config).__name__}" + ) + + +def _has_fp8_block_scales(quant_config: Any) -> bool: + if quant_config is None: + return False + layer_mode = getattr(quant_config, "layer_quant_mode", None) + if layer_mode is None: + return False + return layer_mode.has_fp8_block_scales() is True + + +def is_fp8_model(quant_config: Any) -> bool: + """Return whether a TRT-LLM model uses 128x128 block FP8.""" + return _has_fp8_block_scales(quant_config) + + +def is_mxfp8_model(quant_config: Any) -> bool: + """Return whether a TRT-LLM model uses MXFP8 (E4M3 + UE8M0 1x32).""" + if quant_config is None: + return False + layer_mode = getattr(quant_config, "layer_quant_mode", None) + if layer_mode is None: + return False + has_mxfp8 = getattr(layer_mode, "has_mxfp8", None) + return has_mxfp8 is not None and has_mxfp8() is True + + +def is_quantized_expert_refit(quant_config: Any) -> bool: + """Whether refit must quantize routed experts before loading.""" + return is_fp8_model(quant_config) or is_mxfp8_model(quant_config) + + +def cast_tensor_to_fp8_blockwise( + data_hp: torch.Tensor, + weight_block_size: Sequence[int] = FP8_BLOCK_SIZE, +) -> tuple[torch.Tensor, torch.Tensor]: + """Quantize the last two dims into E4M3 with FP32 block scales. + + Leading dims are independent. ``weight_scale_inv`` is ``[..., out_block, + in_block]``; dequantize as ``fp8.float() * weight_scale_inv``. + """ + if data_hp.dim() < 2: + raise ValueError( + "cast_tensor_to_fp8_blockwise expects at least a 2-D tensor, got " + f"shape {tuple(data_hp.shape)}" + ) + if len(weight_block_size) != 2: + raise ValueError( + f"weight_block_size must contain two dimensions, got {weight_block_size}" + ) + block_m, block_n = (int(weight_block_size[0]), int(weight_block_size[1])) + if (block_m, block_n) != FP8_BLOCK_SIZE: + raise ValueError( + "TRT-LLM FP8_BLOCK_SCALES requires weight_block_size=[128, 128], " + f"got {[block_m, block_n]}" + ) + + batch_shape = tuple(data_hp.shape[:-2]) + rows, columns = data_hp.shape[-2:] + pad_rows = (-rows) % block_m + pad_columns = (-columns) % block_n + + data_fp32 = data_hp.to(torch.float32) + if pad_rows or pad_columns: + data_fp32 = F.pad(data_fp32, (0, pad_columns, 0, pad_rows), value=0.0) + + padded_rows, padded_columns = data_fp32.shape[-2:] + row_blocks = padded_rows // block_m + column_blocks = padded_columns // block_n + batch_size = math.prod(batch_shape) if batch_shape else 1 + + blocked = ( + data_fp32.reshape( + batch_size, + row_blocks, + block_m, + column_blocks, + block_n, + ) + .permute(0, 1, 3, 2, 4) + .contiguous() + .flatten(start_dim=3) + ) + + fp8_max = torch.finfo(torch.float8_e4m3fn).max + max_abs = torch.amax(torch.abs(blocked), dim=-1, keepdim=True) + valid_scale = torch.isfinite(max_abs) & (max_abs > 0) + scale_inv = torch.where( + valid_scale, + max_abs / fp8_max, + torch.ones_like(max_abs), + ) + quant_scale = torch.where( + valid_scale, + torch.reciprocal(scale_inv), + torch.ones_like(scale_inv), + ) + fp8_data = torch.clamp( + blocked * quant_scale, + min=-fp8_max, + max=fp8_max, + ).to(torch.float8_e4m3fn) + + fp8_data = ( + fp8_data.reshape( + batch_size, + row_blocks, + column_blocks, + block_m, + block_n, + ) + .permute(0, 1, 3, 2, 4) + .reshape(*batch_shape, padded_rows, padded_columns) + ) + fp8_data = fp8_data[..., :rows, :columns].contiguous() + scale_inv = scale_inv.squeeze(-1).reshape(*batch_shape, row_blocks, column_blocks) + return fp8_data, scale_inv.contiguous() + + +def cast_tensor_to_mxfp8_blockwise( + data_hp: torch.Tensor, + block_size: int = MXFP8_BLOCK_SIZE, +) -> tuple[torch.Tensor, torch.Tensor]: + """Quantize the last dim into E4M3 with UE8M0 1x{block} scales. + + Mirrors TRT-LLM's ``quant_bf16_to_mxfp8``: the power-of-two exponent maps + each block's amax just under the E4M3 max. Blocks span K only. + + Returns ``(e4m3 [..., K], ue8m0 uint8 [..., K // block])``. + """ + if data_hp.dim() < 2: + raise ValueError( + "cast_tensor_to_mxfp8_blockwise expects at least a 2-D tensor, got " + f"shape {tuple(data_hp.shape)}" + ) + columns = data_hp.shape[-1] + if columns % block_size != 0: + raise ValueError( + f"MXFP8 requires the last dim to be a multiple of {block_size}, got " + f"{columns}" + ) + + shape = tuple(data_hp.shape) + flat = data_hp.float().reshape(-1, columns) + blocked = flat.view(-1, columns // block_size, block_size) + amax = blocked.abs().amax(dim=-1).clamp_min(1e-12) + exponent = torch.ceil(torch.log2(amax / E4M3_MAX)) + scale_ue8m0 = (exponent + UE8M0_BIAS).clamp(0, 255).to(torch.uint8) + scale = torch.exp2(exponent).unsqueeze(-1) + quantized = (blocked / scale).to(torch.float8_e4m3fn).view(-1, columns) + return ( + quantized.reshape(shape).contiguous(), + scale_ue8m0.reshape(*shape[:-1], columns // block_size).contiguous(), + ) + + +def _insert_unique( + output: dict[str, torch.Tensor], name: str, tensor: torch.Tensor +) -> None: + if name in output: + raise ValueError(f"Duplicate refit weight after FP8 conversion: {name}") + output[name] = tensor + + +def _insert_quantized_projection( + output: dict[str, torch.Tensor], + name: str, + tensor: torch.Tensor, + *, + is_mx: bool = False, +) -> None: + if is_mx: + data, scale = cast_tensor_to_mxfp8_blockwise(tensor) + else: + data, scale = cast_tensor_to_fp8_blockwise(tensor) + _insert_unique(output, name, data) + # Both formats use `.weight_scale_inv` (TRT-LLM's MXFP8 loader probes it first). + scale_name = name.removesuffix(".weight") + ".weight_scale_inv" + _insert_unique(output, scale_name, scale) + + +def clone_mapper_staging_weights( + weights: dict[str, torch.Tensor], +) -> dict[str, torch.Tensor]: + """Clone tensors the Qwen3.5 mapper may retain across reload calls. + + IPC buffers are reused after each ACK, so retained views must own storage. + """ + return { + name: tensor.clone() + if _SPLIT_LINEAR_ATTN_RE.fullmatch(name) is not None + else tensor + for name, tensor in weights.items() + } + + +def _convert_fused_expert_weight( + output: dict[str, torch.Tensor], + *, + name: str, + tensor: torch.Tensor, + prefix: str, + projection: str, + is_mx: bool = False, +) -> None: + if tensor.dim() != 3: + raise ValueError( + f"Qwen3.5 fused expert weight {name} must be 3-D, got {tuple(tensor.shape)}" + ) + + num_experts = tensor.shape[0] + if projection == "gate_up_proj" and tensor.shape[1] % 2 != 0: + raise ValueError( + f"Qwen3.5 gate_up_proj dimension must be even, got {tuple(tensor.shape)}" + ) + + for start in range(0, num_experts, FP8_EXPERT_CHUNK_SIZE): + end = min(start + FP8_EXPERT_CHUNK_SIZE, num_experts) + if projection == "gate_up_proj": + intermediate_size = tensor.shape[1] // 2 + projections = ( + ("gate_proj", tensor[start:end, :intermediate_size, :]), + ("up_proj", tensor[start:end, intermediate_size:, :]), + ) + else: + projections = (("down_proj", tensor[start:end]),) + + for projection_name, projection_tensor in projections: + if is_mx: + fp8_data, scale_inv = cast_tensor_to_mxfp8_blockwise(projection_tensor) + else: + fp8_data, scale_inv = cast_tensor_to_fp8_blockwise(projection_tensor) + for chunk_index, expert_index in enumerate(range(start, end)): + weight_name = f"{prefix}.{expert_index}.{projection_name}.weight" + scale_name = weight_name.removesuffix(".weight") + ".weight_scale_inv" + _insert_unique(output, weight_name, fp8_data[chunk_index]) + _insert_unique(output, scale_name, scale_inv[chunk_index]) + + +def load_weights( + weight_list: Iterable[tuple[str, torch.Tensor]], + *, + is_mx: bool = False, +) -> dict[str, torch.Tensor]: + """Convert BF16 weights to block-FP8 or MXFP8 for the TRT-LLM engine. + + Only routed experts are quantized. Fused expert tensors (Qwen3.5) are expanded to per-expert HF names; other + weights pass through unchanged. + """ + output: dict[str, torch.Tensor] = {} + for name, tensor in weight_list: + if not isinstance(name, str): + raise TypeError( + f"TRT-LLM refit weight names must be strings, got {type(name).__name__}" + ) + weight_name = str(name) + fused_match = _FUSED_EXPERT_RE.fullmatch(weight_name) + is_split_expert = _SPLIT_EXPERT_RE.fullmatch(weight_name) is not None + if fused_match is not None: + _convert_fused_expert_weight( + output, + name=weight_name, + tensor=tensor, + prefix=fused_match.group("prefix"), + projection=fused_match.group("projection"), + is_mx=is_mx, + ) + elif is_split_expert: + _insert_quantized_projection(output, weight_name, tensor, is_mx=is_mx) + else: + _insert_unique(output, weight_name, tensor) + return output diff --git a/nemo_rl/models/generation/trtllm/trtllm_backend.py b/nemo_rl/models/generation/trtllm/trtllm_backend.py index 2f3e96e1647..d4b0c8f43ba 100644 --- a/nemo_rl/models/generation/trtllm/trtllm_backend.py +++ b/nemo_rl/models/generation/trtllm/trtllm_backend.py @@ -25,7 +25,6 @@ import gc import os -import traceback from typing import Any import torch @@ -33,6 +32,7 @@ from tensorrt_llm.executor.ray.utils import control_action_decorator from tensorrt_llm.llmapi.rlhf_utils import WorkerExtension +from nemo_rl.models.generation.trtllm.quantization import fp8 as fp8_quantization from nemo_rl.models.policy.utils import ( IPCProtocol, calculate_aligned_size, @@ -47,6 +47,25 @@ os.environ.setdefault("TRT_LLM_DISABLE_LOAD_WEIGHTS_IN_PARALLEL", "True") +def _require_fp8_refit_hooks(model_loader: Any) -> None: + """Require TRT-LLM hooks for transactional FP8 refits.""" + required_hooks = ( + "begin_update_weights", + "finalize_update_weights", + "abort_update_weights", + ) + missing_hooks = [ + hook_name + for hook_name in required_hooks + if not callable(getattr(model_loader, hook_name, None)) + ] + if missing_hooks: + raise RuntimeError( + "FP8 refit requires TRT-LLM weight-update hooks. " + f"Missing APIs: {missing_hooks}." + ) + + class NcclExtension(WorkerExtension): """NCCL-based weight update extension for TRT-LLM Ray workers. @@ -105,6 +124,10 @@ def stop_gpu_profiling(self) -> None: def prepare_refit_info(self, state_dict_info: dict[str, Any]) -> None: self.state_dict_info = state_dict_info + model = self.engine.model_engine.model + if fp8_quantization.is_quantized_expert_refit(model.model_config.quant_config): + fp8_quantization.validate_routed_experts(state_dict_info) + _require_fp8_refit_hooks(self.engine.model_engine.model_loader) # ------------------------------------------------------------------ # # NCCL weight receive + reload @@ -124,10 +147,12 @@ def update_weights_from_collective( If False, the swap happens at a scheduler step boundary with in-flight requests still in the engine (in-flight weight update). - recompute_kv: Only meaningful with ``drain=False``. If True, - preempt in-flight requests so they re-prefill under the new weights. - Otherwise, they keep decoding with their current KV cache. The - reusable prefix cache is cleared after every weight update. + recompute_kv: If True, call ``recompute_active_requests()`` after + the refit: in-flight requests are re-prefilled under the new + weights and the prefix cache is reset. If False, in-flight + requests keep decoding with their current KV cache. With + ``drain=True`` there are no in-flight requests, so True reduces + to a prefix-cache reset. """ assert hasattr(self, "state_dict_info") and self.state_dict_info is not None, ( "state_dict_info not set — call prepare_refit_info first" @@ -136,9 +161,20 @@ def update_weights_from_collective( model = model_engine.model def load_model_weight_func(weight_list): + if fp8_quantization.is_quantized_expert_refit( + model.model_config.quant_config + ): + weights = fp8_quantization.load_weights( + weight_list, + is_mx=fp8_quantization.is_mxfp8_model( + model.model_config.quant_config + ), + ) + else: + weights = dict(weight_list) model_engine.model_loader.reload( model, - dict(weight_list), + weights, allow_partial_loading=True, ) @@ -159,11 +195,14 @@ def load_model_weight_func(weight_list): post_unpack_func=load_model_weight_func, ) self.finalize_weight_update() + if recompute_kv: + # Re-prefill in-flight requests under the new weights; this + # also resets the prefix cache (NVIDIA/TensorRT-LLM#17937). + self.engine.recompute_active_requests() self.finish_weight_update() - except Exception as e: + except Exception: model_engine.model_loader.abort_update_weights() - print(f"Error in NcclExtension.update_weights_from_collective: {e}") - return False + raise return True @@ -234,6 +273,19 @@ def update_weights_via_ipc_zmq(self) -> bool: "Likely stale state_dict_info (wrong shape/dtype for some key)." ) + if fp8_quantization.is_quantized_expert_refit( + model.model_config.quant_config + ): + weights = fp8_quantization.load_weights( + weights.items(), + is_mx=fp8_quantization.is_mxfp8_model( + model.model_config.quant_config + ), + ) + # The mapper may retain split QKVZ/BA views across IPC chunks; + # detach them before ACK lets the trainer reuse the buffer. + weights = fp8_quantization.clone_mapper_staging_weights(weights) + model_engine.model_loader.reload( model, weights, @@ -253,13 +305,9 @@ def update_weights_via_ipc_zmq(self) -> bool: gc.collect() torch.cuda.empty_cache() return True - except Exception as e: + except Exception: model_engine.model_loader.abort_update_weights() - print( - f"Error in NcclExtension.update_weights_via_ipc_zmq: {e}\n" - f"{traceback.format_exc()}" - ) - return False + raise def cleanup_zmq(self) -> None: """Close ZMQ socket if open — called from worker shutdown.""" diff --git a/nemo_rl/models/generation/trtllm/trtllm_worker_async.py b/nemo_rl/models/generation/trtllm/trtllm_worker_async.py index 6126ee902f8..9aa47508ed5 100644 --- a/nemo_rl/models/generation/trtllm/trtllm_worker_async.py +++ b/nemo_rl/models/generation/trtllm/trtllm_worker_async.py @@ -334,6 +334,33 @@ def __init__( # they can override anything above for advanced tuning. llm_kwargs.update(extra_trtllm_kwargs) + if engine_cfg.get("is_mx") and engine_cfg["precision"] != "fp8": + raise ValueError("engine_cfg.is_mx=True requires precision='fp8'") + + if engine_cfg["precision"] == "fp8": + # Imported here: TRT-LLM is an optional dependency. + from tensorrt_llm.llmapi.llm_args import MoeConfig + from transformers import AutoConfig + + from nemo_rl.models.generation.trtllm.quantization.fp8 import ( + configure_fp8_llm_kwargs, + configure_fp8_moe_backend, + validate_model_type, + ) + + validate_model_type( + AutoConfig.from_pretrained( + self.model_name, trust_remote_code=True + ).model_type + ) + + is_mx = bool(engine_cfg.get("is_mx")) + configure_fp8_llm_kwargs(llm_kwargs, is_mx=is_mx) + + # Block-FP8 needs the TRTLLM MoE backend (DeepGEMM would resmooth to + # E8M0); MXFP8 uses CUTLASS or CuTe DSL. Other MoeConfig fields kept. + configure_fp8_moe_backend(llm_kwargs, MoeConfig, is_mx=is_mx) + # Propagate the nsight runtime_env down to TRT-LLM's internal Ray GPU # workers. The outer actor's @ray.remote nsight config does NOT inherit # into TRT-LLM's RayExecutor workers (ray_executor.py sets an explicit @@ -508,9 +535,10 @@ async def update_weights_from_collective_async( without draining in-flight requests (in-flight weight update). Default True preserves the original drain-first behavior. - recompute_kv: If True (and ``drain=False``), preempt all - in-flight requests after the refit so the scheduler - re-prefills them under the new weights. + recompute_kv: If True, re-prefill in-flight requests under the new + weights and reset the prefix cache after the refit + (``recompute_kv_cache_after_weight_updates``). If False, + in-flight requests keep their current KV cache. """ assert self.llm is not None try: @@ -518,10 +546,16 @@ async def update_weights_from_collective_async( "update_weights_from_collective", kwargs={"drain": drain, "recompute_kv": recompute_kv}, ) - worker_result = results[0] if results else True - if not worker_result: + if not results: + print("Error: TRT-LLM weight update returned no worker results.") + return False + failed_workers = [ + (rank, result) for rank, result in enumerate(results) if not result + ] + if failed_workers: print( - f"Error: TRT-LLM worker failed to update weights. Result: {worker_result}" + "Error: TRT-LLM workers failed to update weights. " + f"Results: {failed_workers}" ) return False return True @@ -530,16 +564,24 @@ async def update_weights_from_collective_async( import traceback traceback.print_exc() - return False + # Raise instead of returning False so the engine-side cause reaches + # the caller. + raise RuntimeError(f"TRT-LLM collective weight update failed: {e!r}") from e async def update_weights_via_ipc_zmq_async(self) -> bool: assert self.llm is not None try: results = await self.llm.collective_rpc("update_weights_via_ipc_zmq") - worker_result = results[0] if results else True - if not worker_result: + if not results: + print("Error: TRT-LLM IPC weight update returned no worker results.") + return False + failed_workers = [ + (rank, result) for rank, result in enumerate(results) if not result + ] + if failed_workers: print( - f"Error: TRT-LLM worker failed to update weights via IPC. Result: {worker_result}" + "Error: TRT-LLM workers failed to update weights via IPC. " + f"Results: {failed_workers}" ) return False return True @@ -548,7 +590,9 @@ async def update_weights_via_ipc_zmq_async(self) -> bool: import traceback traceback.print_exc() - return False + # Raise instead of returning False so the engine-side cause reaches + # the caller. + raise RuntimeError(f"TRT-LLM IPC weight update failed: {e!r}") from e async def report_device_id_async(self) -> list[str]: assert self.llm is not None diff --git a/nemo_rl/models/policy/interfaces.py b/nemo_rl/models/policy/interfaces.py index 075efbd8cb8..a9c7cfe6522 100644 --- a/nemo_rl/models/policy/interfaces.py +++ b/nemo_rl/models/policy/interfaces.py @@ -292,6 +292,7 @@ def broadcast_weights_for_collective( *, buffer_size_bytes: Optional[int] = None, num_buffers: Optional[int] = None, + use_legacy_packing: bool = False, ) -> list[ray.ObjectRef]: pass diff --git a/nemo_rl/models/policy/lm_policy.py b/nemo_rl/models/policy/lm_policy.py index f48b54bf1eb..6535a8e3d57 100644 --- a/nemo_rl/models/policy/lm_policy.py +++ b/nemo_rl/models/policy/lm_policy.py @@ -1420,6 +1420,7 @@ def broadcast_weights_for_collective( *, buffer_size_bytes: Optional[int] = None, num_buffers: Optional[int] = None, + use_legacy_packing: bool = False, ) -> list[ray.ObjectRef]: """Broadcast the weights for collective communication.""" futures = self.worker_group.run_all_workers_single_data( @@ -1428,6 +1429,7 @@ def broadcast_weights_for_collective( refit_timeout_s=refit_timeout_s, buffer_size_bytes=buffer_size_bytes, num_buffers=num_buffers, + use_legacy_packing=use_legacy_packing, ) # this function should co-work with vllm, so we should wait for all futures to complete outside return futures diff --git a/nemo_rl/models/policy/workers/automodel_policy_worker.py b/nemo_rl/models/policy/workers/automodel_policy_worker.py index cc582ddde3f..cecb2faa67d 100644 --- a/nemo_rl/models/policy/workers/automodel_policy_worker.py +++ b/nemo_rl/models/policy/workers/automodel_policy_worker.py @@ -1286,6 +1286,7 @@ def broadcast_weights_for_collective( *, buffer_size_bytes: Optional[int] = None, num_buffers: Optional[int] = None, + use_legacy_packing: bool = False, ) -> None: """Broadcast the weights for collective communication. @@ -1304,6 +1305,7 @@ def broadcast_weights_for_collective( kv_scales=kv_scales, buffer_size_bytes=buffer_size_bytes, num_buffers=num_buffers, + use_legacy_packing=use_legacy_packing, ) if guard.fired: # The aborted collective returned cleanly, so this is the only signal there is. @@ -1318,6 +1320,7 @@ def _broadcast_weights_for_collective( *, buffer_size_bytes: Optional[int] = None, num_buffers: Optional[int] = None, + use_legacy_packing: bool = False, ) -> None: if kv_scales is not None: raise NotImplementedError( @@ -1344,6 +1347,7 @@ def _broadcast_weights_for_collective( post_iter_func=automodel_post_iter_func, buffer_size_bytes=buffer_size_bytes, num_buffers=num_buffers, + use_legacy_packing=use_legacy_packing, ) # Manually move model to cpu for cpu offload case diff --git a/nemo_rl/models/policy/workers/megatron_policy_worker.py b/nemo_rl/models/policy/workers/megatron_policy_worker.py index b7c7efaaca8..93c25545ef2 100644 --- a/nemo_rl/models/policy/workers/megatron_policy_worker.py +++ b/nemo_rl/models/policy/workers/megatron_policy_worker.py @@ -3652,6 +3652,7 @@ def _iter_params_with_optional_kv_scales( # when pre-quantized refit is enabled for the parameter. for name, tensor in base_iter: yield from self._maybe_prequantize_param(name, tensor) + del tensor if include_draft and self.draft_model is not None: from nemo_rl.models.megatron.draft import export_eagle_weights_to_hf @@ -3918,6 +3919,7 @@ def broadcast_weights_for_collective( *, buffer_size_bytes: Optional[int] = None, num_buffers: Optional[int] = None, + use_legacy_packing: bool = False, ) -> None: """Broadcast the weights for collective communication. @@ -3938,6 +3940,7 @@ def broadcast_weights_for_collective( kv_scales=kv_scales, buffer_size_bytes=buffer_size_bytes, num_buffers=num_buffers, + use_legacy_packing=use_legacy_packing, ) if guard.fired: # The aborted collective returned cleanly, so this is the only signal there is. @@ -3952,6 +3955,7 @@ def _broadcast_weights_for_collective( *, buffer_size_bytes: Optional[int] = None, num_buffers: Optional[int] = None, + use_legacy_packing: bool = False, ) -> None: # param_iterator will return (name, tensor), we only need tensor. packed_broadcast_producer( @@ -3961,6 +3965,7 @@ def _broadcast_weights_for_collective( post_iter_func=lambda x: x[1], buffer_size_bytes=buffer_size_bytes, num_buffers=num_buffers, + use_legacy_packing=use_legacy_packing, ) def _build_layer_to_pp_stage( diff --git a/nemo_rl/utils/packed_tensor.py b/nemo_rl/utils/packed_tensor.py index ce725f42653..d420ebfbf00 100644 --- a/nemo_rl/utils/packed_tensor.py +++ b/nemo_rl/utils/packed_tensor.py @@ -45,6 +45,7 @@ def packed_broadcast_producer( *, buffer_size_bytes: int | None = None, num_buffers: int | None = None, + use_legacy_packing: bool = False, ): """Broadcast a list of tensors in a packed manner. @@ -55,6 +56,8 @@ def packed_broadcast_producer( post_iter_func: function to apply to each tensor before packing, should return a tensor buffer_size_bytes: packed-buffer target. Uses the NeMo-RL default when unset. num_buffers: number of alternating CUDA buffers. Uses the default when unset. + use_legacy_packing: pack an oversized tensor with the ones before it, as + vLLM's own consumer expects, instead of sending it in its own broadcast. Returns: None @@ -71,6 +74,7 @@ def packed_broadcast_producer( num_buffers = get_num_buffers() if num_buffers is None else num_buffers streams = [torch.cuda.Stream() for _ in range(num_buffers)] buffer_idx = 0 + pending_tensor = None packing_tensor_list = [[] for _ in range(num_buffers)] packing_tensor_sizes = [0 for _ in range(num_buffers)] @@ -103,34 +107,62 @@ def packed_broadcast_producer( packing_tensor_sizes[buffer_idx] = 0 # Pack the tensors while True: - # Apply backend specific post processing and then convert to linearized uint8 tensor. - # contiguous() is required because the upstream iterator may - # yield non-contiguous tensors that view(...) cannot handle. - tensor = post_iter_func(next(iterator)) - if tensor.device.type != "cuda": - # Everything here is concatenated into one buffer and - # broadcast over a CUDA collective, so a single host - # tensor anywhere in the stream fails the cat. The - # producer owns its buffer's device rather than - # trusting every upstream exporter to agree. - tensor = tensor.to(torch.cuda.current_device()) - tensor = tensor.contiguous().reshape(-1).view(torch.uint8) + if pending_tensor is None: + # Apply backend specific post processing and then convert to linearized uint8 tensor. + # contiguous() is required because the upstream iterator may + # yield non-contiguous tensors that view(...) cannot handle. + tensor = post_iter_func(next(iterator)) + if tensor.device.type != "cuda": + # Everything here is concatenated into one buffer and + # broadcast over a CUDA collective, so a single host + # tensor anywhere in the stream fails the cat. The + # producer owns its buffer's device rather than + # trusting every upstream exporter to agree. + tensor = tensor.to(torch.cuda.current_device()) + tensor = tensor.contiguous().reshape(-1).view(torch.uint8) + else: + tensor = pending_tensor + pending_tensor = None + if ( + not use_legacy_packing + and packing_tensor_list[buffer_idx] + and tensor.numel() > target_packed_tensor_size + ): + pending_tensor = tensor + break packing_tensor_list[buffer_idx].append(tensor) packing_tensor_sizes[buffer_idx] += tensor.numel() if packing_tensor_sizes[buffer_idx] > target_packed_tensor_size: break # Pack the tensors and call broadcast collective - packed_tensors[buffer_idx] = torch.cat( - packing_tensor_list[buffer_idx], dim=0 + packed_tensors[buffer_idx] = ( + packing_tensor_list[buffer_idx][0] + if len(packing_tensor_list[buffer_idx]) == 1 + else torch.cat(packing_tensor_list[buffer_idx], dim=0) ) export_done = torch.cuda.Event() export_done.record(streams[buffer_idx]) group.broadcast(packed_tensors[buffer_idx], src=src) + if ( + len(packing_tensor_list[buffer_idx]) == 1 + and packing_tensor_sizes[buffer_idx] > target_packed_tensor_size + ): + # Oversized tensor: drain its collective and free it before + # materializing the next one. + for idx, stream in enumerate(streams): + stream.synchronize() + packing_tensor_list[idx].clear() + packed_tensors[idx] = torch.empty( + 0, dtype=torch.uint8, device="cuda" + ) + del tensor except StopIteration: # do the last broadcast if there are remaining tensors if len(packing_tensor_list[buffer_idx]) > 0: - packed_tensors[buffer_idx] = torch.cat( - packing_tensor_list[buffer_idx], dim=0 + packed_tensors[buffer_idx] = ( + packing_tensor_list[buffer_idx][0] + if len(packing_tensor_list[buffer_idx]) == 1 + else torch.cat(packing_tensor_list[buffer_idx], dim=0) ) group.broadcast(packed_tensors[buffer_idx], src=src) break @@ -196,6 +228,7 @@ def restore_tensor( num_buffers = get_num_buffers() streams = [torch.cuda.Stream() for _ in range(num_buffers)] buffer_idx = 0 + pending_meta_data = None packing_tensor_meta_data = [[] for _ in range(num_buffers)] packing_tensor_sizes = [0 for _ in range(num_buffers)] @@ -218,8 +251,18 @@ def restore_tensor( try: # Form a packed tensor while True: - name, (shape, dtype) = next(iterator) - tensor_size = math.prod(shape) * dtype.itemsize + if pending_meta_data is None: + name, (shape, dtype) = next(iterator) + tensor_size = math.prod(shape) * dtype.itemsize + else: + name, shape, dtype, tensor_size = pending_meta_data + pending_meta_data = None + if ( + packing_tensor_meta_data[buffer_idx] + and tensor_size > target_packed_tensor_size + ): + pending_meta_data = (name, shape, dtype, tensor_size) + break packing_tensor_meta_data[buffer_idx].append( (name, shape, dtype, offsets[buffer_idx], tensor_size) ) @@ -238,6 +281,16 @@ def restore_tensor( packed_tensors[buffer_idx], packing_tensor_meta_data[buffer_idx], ) + if ( + len(packing_tensor_meta_data[buffer_idx]) == 1 + and packing_tensor_sizes[buffer_idx] > target_packed_tensor_size + ): + for idx, stream in enumerate(streams): + stream.synchronize() + packing_tensor_meta_data[idx].clear() + packed_tensors[idx] = torch.empty( + 0, dtype=torch.uint8, device="cuda" + ) except StopIteration: # Do the last broadcast if there are remaining tensors. if len(packing_tensor_meta_data[buffer_idx]) > 0: diff --git a/nemo_rl/weight_sync/collective_weight_synchronizer.py b/nemo_rl/weight_sync/collective_weight_synchronizer.py index 093800e01d6..6b8fa7d50ad 100644 --- a/nemo_rl/weight_sync/collective_weight_synchronizer.py +++ b/nemo_rl/weight_sync/collective_weight_synchronizer.py @@ -147,6 +147,7 @@ def sync_weights( refit_timeout_s=self._refit_timeout_s, buffer_size_bytes=sender_spec.buffer_size_bytes, num_buffers=sender_spec.num_buffers, + use_legacy_packing=sender_spec.use_legacy_packing, ) futures_inference = self._generation.update_weights_from_collective( refit_timeout_s=self._refit_timeout_s diff --git a/nemo_rl/weight_sync/factory.py b/nemo_rl/weight_sync/factory.py index d9e0795a3bd..3524b22c559 100644 --- a/nemo_rl/weight_sync/factory.py +++ b/nemo_rl/weight_sync/factory.py @@ -30,6 +30,7 @@ DYNAMO_BACKEND, MEGATRON_BACKEND, SGLANG_BACKEND, + TRTLLM_BACKEND, VLLM_BACKEND, ) from nemo_rl.weight_sync.checkpoint_engine_config import ( @@ -60,6 +61,8 @@ def validate_offload_policy_before_refit( supported_transports = ( (None, "nccl_reshard") if generation_backend == VLLM_BACKEND + else (None,) + if generation_backend == TRTLLM_BACKEND else (None, "mcore", "nccl_reshard") if generation_backend == MEGATRON_BACKEND else () @@ -67,7 +70,9 @@ def validate_offload_policy_before_refit( if refit_transport not in supported_transports: raise ValueError( "offload_policy_before_refit requires non-colocated vLLM or Megatron " - "generation with a collective, mcore, or nccl_reshard refit transport." + "generation with a collective, mcore, or nccl_reshard refit transport, " + "or non-colocated TensorRT-LLM generation with the default collective " + "transport." ) @@ -110,6 +115,7 @@ def create_weight_synchronizer( SGLANG_BACKEND, MEGATRON_BACKEND, DYNAMO_BACKEND, + TRTLLM_BACKEND, } if generation_backend not in _SUPPORTED_BACKENDS: raise ValueError( diff --git a/pyrefly.toml b/pyrefly.toml index f44eb587b2b..3ab7fcb45e3 100644 --- a/pyrefly.toml +++ b/pyrefly.toml @@ -254,6 +254,8 @@ project-includes = [ "nemo_rl/models/generation/trtllm/__init__.py", "nemo_rl/models/generation/trtllm/config.py", "nemo_rl/models/generation/trtllm/trtllm_disagg_server.py", + "nemo_rl/models/generation/trtllm/quantization/__init__.py", + "nemo_rl/models/generation/trtllm/quantization/fp8.py", "nemo_rl/models/generation/trtllm/trtllm_generation.py", "nemo_rl/models/generation/trtllm/trtllm_http_server.py", "nemo_rl/models/generation/vllm/__init__.py", diff --git a/tests/test_suites/llm/grpo-qwen3.5-35ba3b-4n4g-megatron-trtllm-mxfp8-noncolocated.sh b/tests/test_suites/llm/grpo-qwen3.5-35ba3b-4n4g-megatron-trtllm-mxfp8-noncolocated.sh new file mode 100755 index 00000000000..00860665be1 --- /dev/null +++ b/tests/test_suites/llm/grpo-qwen3.5-35ba3b-4n4g-megatron-trtllm-mxfp8-noncolocated.sh @@ -0,0 +1,57 @@ +#!/bin/bash +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd) +source "${SCRIPT_DIR}/common.env" + +# ===== BEGIN CONFIG ===== +NUM_NODES=4 +GPUS_PER_NODE=4 +SEGMENT_SIZE=2 +STEPS_PER_RUN=3 +MAX_STEPS=3 +NUM_RUNS=$(( (MAX_STEPS + STEPS_PER_RUN - 1) / STEPS_PER_RUN )) +NUM_MINUTES=60 +# ===== END CONFIG ===== + +exit_if_max_steps_reached + +cd ${PROJECT_ROOT} +uv run examples/run_grpo.py \ + --config ${CONFIG_PATH} \ + grpo.max_num_steps=${MAX_STEPS} \ + logger.log_dir=${LOG_DIR} \ + logger.wandb_enabled=true \ + logger.wandb.project=nemo-rl \ + logger.wandb.name=${EXP_NAME} \ + logger.monitor_gpus=true \ + logger.tensorboard_enabled=true \ + checkpointing.enabled=true \ + checkpointing.checkpoint_dir=${CKPT_DIR} \ + "$@" \ + 2>&1 | tee ${RUN_LOG} + +uv run tests/json_dump_tb_logs.py ${LOG_DIR} --output_path ${JSON_METRICS} + +if [[ $(jq 'to_entries | .[] | select(.key == "train/loss") | .value | keys | map(tonumber) | max' ${JSON_METRICS}) -ge ${MAX_STEPS} ]]; then + uv run tests/check_metrics.py ${JSON_METRICS} \ + 'median(data["train/token_mult_prob_error"]) < 1.05' \ + 'mean(data["train/gen_kl_error"]) < 0.004' \ + 'min(data["timing/train/prepare_for_generation/transfer_and_update_weights"]) > 0' \ + 'all_finite(data["train/loss"])' \ + 'max(data["train/reward"]) > 0.5' + + rm -rf "${CKPT_DIR}" +fi diff --git a/tests/test_suites/nightly_gb200.txt b/tests/test_suites/nightly_gb200.txt index 8a266aa7aac..0d7816e2e48 100644 --- a/tests/test_suites/nightly_gb200.txt +++ b/tests/test_suites/nightly_gb200.txt @@ -37,6 +37,8 @@ tests/test_suites/llm/grpo-qwen2.5-0.5b-1n4g-megatron-trtllm-noncolocated-async. # environment, so it needs no Apptainer or per-instance SIFs. tests/test_suites/llm/grpo-qwen3-1.7b-2n4g-megatron-trtllm-disagg-swe1.sh +tests/test_suites/llm/grpo-qwen3.5-35ba3b-4n4g-megatron-trtllm-mxfp8-noncolocated.sh + # Functional moonlight run tests/test_suites/llm/grpo-moonlight-16ba3b-4n4g-megatron.sh diff --git a/tests/unit/models/generation/test_dynamo_generation.py b/tests/unit/models/generation/test_dynamo_generation.py index 2417e589f25..7e1c5f39e3d 100644 --- a/tests/unit/models/generation/test_dynamo_generation.py +++ b/tests/unit/models/generation/test_dynamo_generation.py @@ -170,6 +170,7 @@ def test_runtime_start_world_size_sender_geometry_and_shutdown(monkeypatch) -> N assert sender.nccl_peer == "vllm" assert sender.buffer_size_bytes == VLLM_PACKED_BUFFER_SIZE_BYTES assert sender.num_buffers == VLLM_PACKED_NUM_BUFFERS + assert sender.use_legacy_packing assert generation.shutdown() assert generation.shutdown() assert calls.count("shutdown") == 1 diff --git a/tests/unit/models/generation/trtllm/test_trtllm_backend.py b/tests/unit/models/generation/trtllm/test_trtllm_backend.py index 9ab760be430..da2d4b5c599 100644 --- a/tests/unit/models/generation/trtllm/test_trtllm_backend.py +++ b/tests/unit/models/generation/trtllm/test_trtllm_backend.py @@ -67,24 +67,56 @@ def _ipc_extension(backend): return extension, model_loader, engine +def test_fp8_refit_hooks(monkeypatch): + from nemo_rl.models.generation.trtllm import trtllm_backend as backend + + extension, _, model, _, _ = _extension(backend) + model.model_config = SimpleNamespace(quant_config=object()) + extension.engine.model_engine.model_loader = SimpleNamespace() + monkeypatch.setattr(backend.fp8_quantization, "is_fp8_model", lambda _: True) + # A missing hook could leave only part of the model updated after a failure. + with pytest.raises(RuntimeError, match="weight-update hooks"): + extension.prepare_refit_info( + { + "model.layers.0.mlp.experts.gate_up_proj": ( + torch.Size([1, 2, 1]), + torch.bfloat16, + ), + "model.layers.0.mlp.experts.down_proj": ( + torch.Size([1, 1, 1]), + torch.bfloat16, + ), + } + ) + + @pytest.mark.parametrize( - ("drain", "recompute_kv"), - [(True, False), (False, False), (False, True)], + ("drain", "recompute_kv", "fp8"), + [ + (True, False, False), + (False, False, False), + (False, True, False), + (True, False, True), + ], ) def test_collective_refit_runs_at_async_engine_boundary( - monkeypatch, drain, recompute_kv + monkeypatch, drain, recompute_kv, fp8 ): from nemo_rl.models.generation.trtllm import trtllm_backend as backend extension, module, model, model_loader, engine = _extension(backend) call_order = [] + incoming = [("model.weight", torch.tensor([1.0]))] + converted = {"converted.weight": torch.tensor([2.0])} + convert = MagicMock(return_value=converted) + model.model_config = SimpleNamespace(quant_config=object()) def packed_consumer(*, iterator, group, src, post_unpack_func): call_order.append("broadcast") assert list(iterator) == list(extension.state_dict_info.items()) assert group is extension.model_update_group assert src == 0 - post_unpack_func([("model.weight", torch.tensor([1.0]))]) + post_unpack_func(incoming) module.pre_reload_weights.side_effect = lambda: call_order.append("pre") module.process_weights_after_loading.side_effect = lambda: call_order.append( @@ -103,8 +135,13 @@ def packed_consumer(*, iterator, group, src, post_unpack_func): model_engine.restore_compiled_model_after_refit.side_effect = ( lambda *_: call_order.append("restore") ) + engine.recompute_active_requests.side_effect = lambda: call_order.append( + "recompute" + ) engine.reset_prefix_cache.side_effect = lambda: call_order.append("reset_prefix") monkeypatch.setattr(backend, "packed_broadcast_consumer", packed_consumer) + monkeypatch.setattr(backend.fp8_quantization, "is_fp8_model", lambda _: fp8) + monkeypatch.setattr(backend.fp8_quantization, "load_weights", convert) monkeypatch.setattr( backend.torch.cuda, "synchronize", lambda: call_order.append("cuda_sync") ) @@ -120,12 +157,14 @@ def packed_consumer(*, iterator, group, src, post_unpack_func): engine.control_action.assert_called_once_with(drain=drain) model_loader.reload.assert_called_once_with( model, - {"model.weight": torch.tensor([1.0])}, + converted if fp8 else dict(incoming), allow_partial_loading=True, ) + assert convert.call_count == int(fp8) # WorkerExtension's lifecycle: begin (unwrap -> begin_update_weights -> # pre_reload_weights), the refit itself, finalize (finalize_update_weights - # -> process/post hooks), finish (reset_prefix_cache -> sync -> re-wrap). + # -> process/post hooks), recompute in-flight requests, finish + # (reset_prefix_cache -> sync -> re-wrap). assert call_order == [ "unwrap", "begin", @@ -135,11 +174,13 @@ def packed_consumer(*, iterator, group, src, post_unpack_func): "finalize", "process", "post", + *(["recompute"] if recompute_kv else []), "reset_prefix", "cuda_sync", "restore", ] model_loader.abort_update_weights.assert_not_called() + assert engine.recompute_active_requests.call_count == int(recompute_kv) engine.reset_prefix_cache.assert_called_once_with() @@ -153,7 +194,7 @@ def test_collective_refit_requires_metadata(): extension.update_weights_from_collective() -def test_collective_refit_returns_false_when_reload_fails(monkeypatch): +def test_collective_refit_aborts_and_raises_when_reload_fails(monkeypatch): from nemo_rl.models.generation.trtllm import trtllm_backend as backend extension, _, _, model_loader, engine = _extension(backend) @@ -165,18 +206,44 @@ def packed_consumer(*, post_unpack_func, **_): monkeypatch.setattr(backend, "packed_broadcast_consumer", packed_consumer) monkeypatch.setattr(backend.torch.cuda, "synchronize", lambda: None) - assert extension.update_weights_from_collective() is False + with pytest.raises(RuntimeError, match="reload failed"): + extension.update_weights_from_collective() model_loader.begin_update_weights.assert_called_once_with() model_loader.finalize_update_weights.assert_not_called() model_loader.abort_update_weights.assert_called_once_with() + engine.recompute_active_requests.assert_not_called() engine.reset_prefix_cache.assert_not_called() -def test_ipc_zmq_streams_chunk_and_reloads_with_aligned_offsets(monkeypatch): +def test_failed_refit_keeps_original_error_when_abort_fails(monkeypatch): + from nemo_rl.models.generation.trtllm import trtllm_backend as backend + + extension, _, _, model_loader, _ = _extension(backend) + + def packed_consumer(*, post_unpack_func, **_): + post_unpack_func([("model.weight", torch.tensor([1.0]))]) + + model_loader.reload.side_effect = RuntimeError("reload failed") + model_loader.abort_update_weights.side_effect = RuntimeError("abort failed") + monkeypatch.setattr(backend, "packed_broadcast_consumer", packed_consumer) + monkeypatch.setattr(backend.torch.cuda, "synchronize", lambda: None) + + with pytest.raises(RuntimeError, match="abort failed") as exc_info: + extension.update_weights_from_collective() + assert str(exc_info.value.__context__) == "reload failed" + + +@pytest.mark.parametrize("fp8", [False, True]) +def test_ipc_zmq_streams_chunk_and_reloads_with_aligned_offsets(monkeypatch, fp8): from nemo_rl.models.generation.trtllm import trtllm_backend as backend from nemo_rl.models.policy.utils import IPCProtocol extension, model_loader, engine = _ipc_extension(backend) + model = extension.engine.model_engine.model + model.model_config = SimpleNamespace(quant_config=object()) + convert = MagicMock(side_effect=lambda weights, **kwargs: dict(weights)) + monkeypatch.setattr(backend.fp8_quantization, "is_fp8_model", lambda _: fp8) + monkeypatch.setattr(backend.fp8_quantization, "load_weights", convert) buffer = torch.zeros(1024, dtype=torch.uint8) monkeypatch.setattr(backend, "rebuild_cuda_tensor_from_ipc", lambda h, d: buffer) @@ -201,12 +268,57 @@ def test_ipc_zmq_streams_chunk_and_reloads_with_aligned_offsets(monkeypatch): model_loader.begin_update_weights.assert_called_once_with() model_loader.finalize_update_weights.assert_called_once_with() model_loader.abort_update_weights.assert_not_called() + assert convert.call_count == int(fp8) + engine.recompute_active_requests.assert_not_called() engine.reset_prefix_cache.assert_called_once_with() # COMPLETE is ACKed after the final chunk. assert extension.zmq_socket.send.call_count == 2 -def test_ipc_zmq_offset_mismatch_returns_false_without_reload(monkeypatch): +def test_ipc_zmq_detaches_mapper_staging_tensor_before_buffer_ack(monkeypatch): + from nemo_rl.models.generation.trtllm import trtllm_backend as backend + from nemo_rl.models.policy.utils import IPCProtocol + + extension, model_loader, _ = _ipc_extension(backend) + model = extension.engine.model_engine.model + model.model_config = SimpleNamespace(quant_config=object()) + qkv_name = "model.language_model.layers.0.linear_attn.in_proj_qkv.weight" + z_name = "model.language_model.layers.0.linear_attn.in_proj_z.weight" + extension.state_dict_info = { + qkv_name: (torch.Size([2]), torch.float32), + z_name: (torch.Size([2]), torch.float32), + } + buffer = torch.zeros(512, dtype=torch.uint8) + staged_weights = [] + + def retain_reload_input(_, weights, **__): + staged_weights.extend(weights.values()) + + def overwrite_after_ack(*_): + if extension.zmq_socket.send.call_count == 1: + buffer.fill_(0xFF) + + model_loader.reload.side_effect = retain_reload_input + extension.zmq_socket.send.side_effect = overwrite_after_ack + extension.zmq_socket.recv_pyobj.side_effect = [ + ("ipc_handle", [qkv_name], 512), + ("ipc_handle", [z_name], 512), + IPCProtocol.COMPLETE, + ] + monkeypatch.setattr(backend, "rebuild_cuda_tensor_from_ipc", lambda h, d: buffer) + monkeypatch.setattr(backend.torch.cuda, "current_stream", lambda: MagicMock()) + monkeypatch.setattr(backend.torch.cuda, "empty_cache", lambda: None) + monkeypatch.setattr(backend.fp8_quantization, "is_fp8_model", lambda _: True) + + assert extension.update_weights_via_ipc_zmq() is True + + assert len(staged_weights) == 2 + assert torch.equal(staged_weights[0], torch.zeros(2, dtype=torch.float32)) + assert torch.isnan(staged_weights[1]).all() + assert extension.zmq_socket.send.call_count == 3 + + +def test_ipc_zmq_offset_mismatch_aborts_and_raises_without_reload(monkeypatch): from nemo_rl.models.generation.trtllm import trtllm_backend as backend from nemo_rl.models.policy.utils import IPCProtocol @@ -216,17 +328,19 @@ def test_ipc_zmq_offset_mismatch_returns_false_without_reload(monkeypatch): monkeypatch.setattr(backend, "rebuild_cuda_tensor_from_ipc", lambda h, d: buffer) monkeypatch.setattr(backend.torch.cuda, "current_stream", lambda: MagicMock()) - # used_bytes (999) != computed offset (1024) -> assertion -> caught -> False. + # used_bytes (999) != computed offset (1024). extension.zmq_socket.recv_pyobj.side_effect = [ ("ipc_handle", ["a", "b"], 999), IPCProtocol.COMPLETE, ] - assert extension.update_weights_via_ipc_zmq() is False + with pytest.raises(AssertionError, match="offset mismatch"): + extension.update_weights_via_ipc_zmq() model_loader.reload.assert_not_called() model_loader.begin_update_weights.assert_called_once_with() model_loader.finalize_update_weights.assert_not_called() model_loader.abort_update_weights.assert_called_once_with() + engine.recompute_active_requests.assert_not_called() engine.reset_prefix_cache.assert_not_called() @@ -246,3 +360,38 @@ def test_zmq_address_and_cleanup_are_per_gpu_and_idempotent(): socket.close.assert_called_once() context.destroy.assert_called_once() + + +def test_fp8_refit_quantizes_only_routed_experts(monkeypatch): + from nemo_rl.models.generation.trtllm import trtllm_backend as backend + + extension, _, model, model_loader, _ = _extension(backend) + model.model_config = SimpleNamespace(quant_config=SimpleNamespace()) + monkeypatch.setattr(backend.fp8_quantization, "is_fp8_model", lambda _: True) + quantized = { + "model.layers.0.mlp.experts.0.down_proj.weight": torch.randn(128, 128), + } + attention = {"model.layers.0.self_attn.q_proj.weight": torch.randn(128, 128)} + kept = { + "model.layers.0.mlp.gate.weight": torch.randn(4, 128), + "lm_head.weight": torch.randn(128, 128), + "model.embed_tokens.weight": torch.randn(128, 128), + } + incoming = {**quantized, **attention, **kept} + info = {name: (t.shape, t.dtype) for name, t in incoming.items()} + + extension.prepare_refit_info(info) + extension.state_dict_info = info + + def packed_consumer(*, post_unpack_func, **_): + post_unpack_func(list(incoming.items())) + + monkeypatch.setattr(backend, "packed_broadcast_consumer", packed_consumer) + monkeypatch.setattr(backend.torch.cuda, "synchronize", lambda: None) + assert extension.update_weights_from_collective() is True + + sent = model_loader.reload.call_args.args[1] + fp8 = torch.float8_e4m3fn + assert all(sent[name].dtype == fp8 for name in quantized) + assert all(sent[name] is incoming[name] for name in attention) + assert all(sent[name] is incoming[name] for name in kept) diff --git a/tests/unit/models/generation/trtllm/test_trtllm_fp8.py b/tests/unit/models/generation/trtllm/test_trtllm_fp8.py new file mode 100644 index 00000000000..6e5ff6c0acb --- /dev/null +++ b/tests/unit/models/generation/trtllm/test_trtllm_fp8.py @@ -0,0 +1,371 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import pytest +import torch + +from nemo_rl.models.generation.trtllm.quantization.fp8 import ( + E4M3_MAX, + FP8_BLOCK_QUANT_KWARGS, + MXFP8_BLOCK_QUANT_KWARGS, + MXFP8_BLOCK_SIZE, + UE8M0_BIAS, + cast_tensor_to_fp8_blockwise, + cast_tensor_to_mxfp8_blockwise, + configure_fp8_llm_kwargs, + configure_fp8_moe_backend, + load_weights, + validate_fused_expert_layout, + validate_model_type, + validate_routed_experts, +) + +pytestmark = pytest.mark.trtllm + + +def _block_matrix(values: list[list[float]]) -> torch.Tensor: + return torch.cat( + [ + torch.cat( + [torch.full((128, 128), value) for value in row], + dim=1, + ) + for row in values + ], + dim=0, + ) + + +def test_fused_expert_layout(): + prefix = "model.layers.0.mlp.experts" + valid = { + f"{prefix}.gate_up_proj": (torch.Size([256, 1024, 2048]), torch.bfloat16), + f"{prefix}.down_proj": (torch.Size([256, 2048, 512]), torch.bfloat16), + } + + validate_fused_expert_layout(valid) + + valid[f"{prefix}.gate_up_proj"] = ( + torch.Size([256, 2048, 1024]), + torch.bfloat16, + ) + with pytest.raises(ValueError, match=r"gate_up_proj=\[E,2I,H\]"): + validate_fused_expert_layout(valid) + + +def test_missing_expert_layout(): + with pytest.raises(ValueError, match="No Qwen3.5 fused routed-expert weights"): + validate_fused_expert_layout({}) + + +def test_fp8_config_preserves_overrides(): + llm_kwargs = { + "dtype": "fp8", + "model_kwargs": {"pretrained_config": {"num_hidden_layers": 4}}, + } + + configure_fp8_llm_kwargs(llm_kwargs) + + assert llm_kwargs["dtype"] == "bfloat16" + assert llm_kwargs["load_format"] == "dummy" + assert llm_kwargs["use_cute_dsl_blockscaling_mm"] is True + assert llm_kwargs["model_kwargs"]["pretrained_config"] == {"num_hidden_layers": 4} + assert llm_kwargs["model_kwargs"]["quantization_config"] == FP8_BLOCK_QUANT_KWARGS + + +@pytest.mark.parametrize( + "llm_kwargs", + [ + {"load_format": "auto"}, + {"model_kwargs": {"quantization_config": {"quant_method": "modelopt"}}}, + ], +) +def test_configure_fp8_llm_kwargs_rejects_unsupported_contract(llm_kwargs): + with pytest.raises(ValueError, match="precision='fp8'"): + configure_fp8_llm_kwargs(llm_kwargs) + + +class _MoeConfig: + def __init__(self, backend="AUTO", **kwargs) -> None: + self.backend = backend + self.kwargs = kwargs + + +def test_configure_fp8_moe_backend_preserves_other_fields(): + llm_kwargs = { + "moe_config": { + "backend": "trtllm", + "load_balancer_config": {"num_slots": 16}, + } + } + + configure_fp8_moe_backend(llm_kwargs, _MoeConfig) + + assert isinstance(llm_kwargs["moe_config"], _MoeConfig) + assert llm_kwargs["moe_config"].backend == "TRTLLM" + assert llm_kwargs["moe_config"].kwargs == { + "load_balancer_config": {"num_slots": 16} + } + + +@pytest.mark.parametrize( + "moe_config", + [ + {"backend": "DEEPGEMM"}, + _MoeConfig(backend="CUTLASS"), + ], +) +def test_configure_fp8_moe_backend_rejects_non_trtllm(moe_config): + with pytest.raises(ValueError, match="backend='TRTLLM'"): + configure_fp8_moe_backend({"moe_config": moe_config}, _MoeConfig) + + +def test_mxfp8_scale_is_ue8m0_exponent_and_round_trips(): + torch.manual_seed(0) + source = (torch.randn(4, 64, dtype=torch.float32) * 100.0).to(torch.bfloat16) + + quantized, scale_ue8m0 = cast_tensor_to_mxfp8_blockwise(source) + + assert quantized.dtype == torch.float8_e4m3fn + assert scale_ue8m0.dtype == torch.uint8 + assert scale_ue8m0.shape == (4, 64 // MXFP8_BLOCK_SIZE) + + blocked = source.float().reshape(4, -1, MXFP8_BLOCK_SIZE) + expected_exponent = torch.ceil(torch.log2(blocked.abs().amax(-1) / E4M3_MAX)) + assert torch.equal(scale_ue8m0, (expected_exponent + UE8M0_BIAS).to(torch.uint8)) + + scale = torch.exp2(scale_ue8m0.float() - UE8M0_BIAS).repeat_interleave( + MXFP8_BLOCK_SIZE, dim=-1 + ) + assert torch.allclose( + quantized.float() * scale, source.float(), rtol=1.0 / 16, atol=1e-3 + ) + assert quantized.float().abs().max() <= E4M3_MAX + + +@pytest.mark.parametrize("backend", ["CUTLASS", "cutedsl", "CuteDSL"]) +def test_configure_fp8_moe_backend_mxfp8_accepts_cutlass_and_cutedsl(backend): + """MXFP8 experts run on CUTLASS or on the fused FC1+FC2 CuTe DSL kernel.""" + llm_kwargs = {"moe_config": {"backend": backend, "load_balancer_config": {"n": 1}}} + configure_fp8_moe_backend(llm_kwargs, _MoeConfig, is_mx=True) + assert llm_kwargs["moe_config"].backend == backend.upper() + assert llm_kwargs["moe_config"].kwargs == {"load_balancer_config": {"n": 1}} + + llm_kwargs = {"moe_config": _MoeConfig(backend=backend)} + configure_fp8_moe_backend(llm_kwargs, _MoeConfig, is_mx=True) + assert llm_kwargs["moe_config"].backend == backend + + +def test_configure_fp8_moe_backend_mxfp8_defaults_to_cutlass(): + llm_kwargs = {} + configure_fp8_moe_backend(llm_kwargs, _MoeConfig, is_mx=True) + assert llm_kwargs["moe_config"].backend == "CUTLASS" + + +@pytest.mark.parametrize( + "moe_config", + [ + {"backend": "TRTLLM"}, + _MoeConfig(backend="DEEPGEMM"), + ], +) +def test_configure_fp8_moe_backend_mxfp8_rejects_other_backends(moe_config): + with pytest.raises(ValueError, match="backend='CUTLASS' or .*backend='CUTEDSL'"): + configure_fp8_moe_backend({"moe_config": moe_config}, _MoeConfig, is_mx=True) + + +def test_mxfp8_llm_kwargs_use_the_mx_quant_contract(): + llm_kwargs = {} + + configure_fp8_llm_kwargs(llm_kwargs, is_mx=True) + + quant = llm_kwargs["model_kwargs"]["quantization_config"] + assert quant["quant_method"] == "mxfp8" + assert quant["weight_block_size"] == [1, MXFP8_BLOCK_SIZE] + assert quant == MXFP8_BLOCK_QUANT_KWARGS + + +def test_block_fp8_scale_orientation_is_out_block_by_in_block(): + source = _block_matrix([[1.0, 2.0], [3.0, 4.0]]) + + fp8_data, scale_inv = cast_tensor_to_fp8_blockwise(source) + + assert fp8_data.dtype == torch.float8_e4m3fn + assert scale_inv.dtype == torch.float32 + assert scale_inv.shape == (2, 2) + assert torch.equal( + scale_inv, + torch.tensor([[1.0, 2.0], [3.0, 4.0]]) / 448.0, + ) + for row_block in range(2): + for column_block in range(2): + block = fp8_data[ + row_block * 128 : (row_block + 1) * 128, + column_block * 128 : (column_block + 1) * 128, + ] + dequantized = block.float() * scale_inv[row_block, column_block] + assert torch.equal( + dequantized, + source[ + row_block * 128 : (row_block + 1) * 128, + column_block * 128 : (column_block + 1) * 128, + ], + ) + + +def test_block_fp8_handles_batched_non_aligned_zero_weights(): + source = torch.zeros((2, 130, 259), dtype=torch.bfloat16) + + fp8_data, scale_inv = cast_tensor_to_fp8_blockwise(source) + + assert fp8_data.shape == source.shape + assert scale_inv.shape == (2, 2, 3) + assert torch.count_nonzero(fp8_data.float()) == 0 + assert torch.equal(scale_inv, torch.ones_like(scale_inv)) + + +def test_block_fp8_round_trip_preserves_intra_block_element_order(): + """Every element in the tests above is constant within its 128x128 block, + so a dropped/wrong inverse permute at the reshape back to [rows, cols] + would still pass them silently. Random (not periodic -- a ramp's + symmetry can coincidentally survive a scramble) values make intra-block + position matter.""" + torch.manual_seed(0) + source = (torch.randn(256, 256, dtype=torch.float32) * 10.0).to(torch.bfloat16) + + fp8_data, scale_inv = cast_tensor_to_fp8_blockwise(source) + + for row_block in range(2): + for column_block in range(2): + block = fp8_data[ + row_block * 128 : (row_block + 1) * 128, + column_block * 128 : (column_block + 1) * 128, + ] + dequantized = block.float() * scale_inv[row_block, column_block] + expected = source[ + row_block * 128 : (row_block + 1) * 128, + column_block * 128 : (column_block + 1) * 128, + ].float() + assert torch.allclose(dequantized, expected, rtol=1.0 / 8, atol=1.0) + + +def test_routed_expert_conversion(): + prefix = "model.language_model.layers.3.mlp.experts" + gate = torch.stack([_block_matrix([[1.0, 2.0]]), _block_matrix([[3.0, 4.0]])]) + up = torch.stack([_block_matrix([[5.0, 6.0]]), _block_matrix([[7.0, 8.0]])]) + down = torch.stack( + [_block_matrix([[9.0], [10.0]]), _block_matrix([[11.0], [12.0]])] + ) + gate_up = torch.cat((gate, up), dim=1).to(torch.bfloat16) + down = down.to(torch.bfloat16) + mtp_expert_name = "mtp.layers.0.mlp.experts.7.down_proj.weight" + mtp_expert = torch.randn(128, 128, dtype=torch.bfloat16) + passthrough = { + "model.layers.0.self_attn.q_proj.weight": torch.randn(128, 128), + "model.layers.0.linear_attn.in_proj_qkvz.weight": torch.randn(128, 128), + "model.layers.0.mlp.shared_expert.gate_proj.weight": torch.randn(128, 128), + "model.language_model.layers.3.mlp.gate.weight": torch.randn(2, 256), + } + + converted = load_weights( + [ + (f"{prefix}.gate_up_proj", gate_up), + (f"{prefix}.down_proj", down), + (mtp_expert_name, mtp_expert), + *passthrough.items(), + ] + ) + + assert f"{prefix}.gate_up_proj" not in converted + assert f"{prefix}.down_proj" not in converted + + source_projections = { + "gate_proj": gate, + "up_proj": up, + "down_proj": down, + } + for expert_index in range(2): + for projection_name, source in source_projections.items(): + weight_name = f"{prefix}.{expert_index}.{projection_name}.weight" + scale_name = weight_name.removesuffix(".weight") + ".weight_scale_inv" + expected_weight, expected_scale = cast_tensor_to_fp8_blockwise( + source[expert_index] + ) + assert torch.equal(converted[weight_name].float(), expected_weight.float()) + assert torch.equal(converted[scale_name], expected_scale) + assert converted[weight_name].dtype == torch.float8_e4m3fn + assert converted[scale_name].dtype == torch.float32 + assert converted[mtp_expert_name].dtype == torch.float8_e4m3fn + assert ( + converted["mtp.layers.0.mlp.experts.7.down_proj.weight_scale_inv"].dtype + == torch.float32 + ) + for name, tensor in passthrough.items(): + assert converted[name] is tensor + assert name.removesuffix(".weight") + ".weight_scale_inv" not in converted + + +def test_qwen3_moe_per_expert_weights_are_quantized(): + prefix = "model.layers.1.mlp.experts" + experts = { + f"{prefix}.{i}.{proj}.weight": torch.randn(128, 128, dtype=torch.bfloat16) + for i in range(2) + for proj in ("gate_proj", "up_proj", "down_proj") + } + passthrough = { + "model.layers.1.self_attn.q_proj.weight": torch.randn(128, 128), + "model.layers.1.mlp.gate.weight": torch.randn(2, 128), + } + validate_routed_experts( + {name: (tensor.shape, tensor.dtype) for name, tensor in experts.items()} + ) + + converted = load_weights([*experts.items(), *passthrough.items()]) + + for name in experts: + assert converted[name].dtype == torch.float8_e4m3fn + scale = converted[name.removesuffix(".weight") + ".weight_scale_inv"] + assert scale.dtype == torch.float32 + for name, tensor in passthrough.items(): + assert converted[name] is tensor + + +@pytest.mark.parametrize( + ("names", "match"), + [ + (["model.layers.0.self_attn.q_proj.weight"], "no routed-expert weights"), + ( + [ + "model.layers.0.mlp.experts.0.down_proj.weight", + "model.layers.1.mlp.down_proj.weight", + ], + "dense MLP weights", + ), + ], +) +def test_validate_routed_experts_rejects_unquantizable_models(names, match): + state_dict_info = {name: (torch.Size([128, 128]), torch.bfloat16) for name in names} + + with pytest.raises(ValueError, match=match): + validate_routed_experts(state_dict_info) + + +@pytest.mark.parametrize("model_type", ["qwen3_5_moe", "qwen3_moe"]) +def test_validate_model_type_accepts_supported_models(model_type): + validate_model_type(model_type) + + +def test_validate_model_type_rejects_other_models(): + with pytest.raises(ValueError, match="supports model types.*'llama'"): + validate_model_type("llama") diff --git a/tests/unit/models/generation/trtllm/test_trtllm_refit.py b/tests/unit/models/generation/trtllm/test_trtllm_refit.py index 99c87cc71e0..785fa521f39 100644 --- a/tests/unit/models/generation/trtllm/test_trtllm_refit.py +++ b/tests/unit/models/generation/trtllm/test_trtllm_refit.py @@ -9,7 +9,10 @@ pytestmark = pytest.mark.trtllm -def test_collective_refit_always_resets_prefix_cache(): +@pytest.mark.parametrize("recompute_kv", [False, True]) +def test_collective_refit_recomputes_active_requests_only_when_requested( + recompute_kv, +): from nemo_rl.models.generation.trtllm import trtllm_backend as backend extension = backend.NcclExtension.__new__(backend.NcclExtension) @@ -39,7 +42,7 @@ def test_collective_refit_always_resets_prefix_cache(): ): result = extension.update_weights_from_collective( drain=False, - recompute_kv=False, + recompute_kv=recompute_kv, ) assert result is True @@ -50,4 +53,12 @@ def test_collective_refit_always_resets_prefix_cache(): engine.model_engine.restore_compiled_model_after_refit.assert_called_once_with( engine.resource_manager ) - engine.reset_prefix_cache.assert_called_once_with() + # finish_weight_update always resets the prefix cache; recompute_kv adds a + # recompute of in-flight requests before it. + calls = [c[0] for c in engine.method_calls] + assert [ + c for c in calls if c in ("recompute_active_requests", "reset_prefix_cache") + ] == [ + *(["recompute_active_requests"] if recompute_kv else []), + "reset_prefix_cache", + ] diff --git a/tests/unit/models/generation/trtllm/test_trtllm_worker_async.py b/tests/unit/models/generation/trtllm/test_trtllm_worker_async.py index 4a7e7ebe1b6..b0bd79d7946 100644 --- a/tests/unit/models/generation/trtllm/test_trtllm_worker_async.py +++ b/tests/unit/models/generation/trtllm/test_trtllm_worker_async.py @@ -236,8 +236,11 @@ async def test_async_lifecycle_resets_cache_before_sleep_and_resumes_selected_ta @pytest.mark.asyncio -@pytest.mark.parametrize("result", [[True], [], [False]]) -async def test_async_collective_refit_propagates_worker_result(result): +@pytest.mark.parametrize( + "result", + [[True], [], [False], [True, False], [True, True]], +) +async def test_async_collective_refit_propagates_every_worker_result(result): worker = _worker() worker.llm.collective_rpc.return_value = result @@ -246,7 +249,7 @@ async def test_async_collective_refit_propagates_worker_result(result): recompute_kv=True, ) - assert succeeded is (result != [False]) + assert succeeded is (bool(result) and all(result)) worker.llm.collective_rpc.assert_awaited_once_with( "update_weights_from_collective", kwargs={"drain": False, "recompute_kv": True}, @@ -254,11 +257,16 @@ async def test_async_collective_refit_propagates_worker_result(result): @pytest.mark.asyncio -async def test_async_ipc_refit_returns_false_on_worker_exception(): +async def test_async_refit_surfaces_worker_exception(): worker = _worker() - worker.llm.collective_rpc.side_effect = RuntimeError("refit failed") + worker.llm.collective_rpc.side_effect = RuntimeError("worker is poisoned") - assert await worker.update_weights_via_ipc_zmq_async() is False + with pytest.raises(RuntimeError, match="worker is poisoned"): + await worker.update_weights_via_ipc_zmq_async() + with pytest.raises(RuntimeError, match="worker is poisoned"): + await worker.update_weights_from_collective_async( + drain=True, recompute_kv=False + ) class _RecordedConfig: @@ -423,3 +431,12 @@ def test_an_aggregated_engine_gets_no_cache_transceiver( kwargs = _engine_kwargs(recorded_async_llm, cfg) assert "cache_transceiver_config" not in kwargs + + +@pytest.mark.asyncio +@pytest.mark.parametrize("result", [[], [True, False]]) +async def test_async_ipc_refit_checks_every_worker_result(result): + worker = _worker() + worker.llm.collective_rpc.return_value = result + + assert await worker.update_weights_via_ipc_zmq_async() is False diff --git a/tests/unit/models/policy/test_lm_policy_collective.py b/tests/unit/models/policy/test_lm_policy_collective.py index 3f9212b0bdd..31763769e5e 100644 --- a/tests/unit/models/policy/test_lm_policy_collective.py +++ b/tests/unit/models/policy/test_lm_policy_collective.py @@ -301,6 +301,7 @@ def shutdown(self, **_kwargs): "refit_timeout_s": None, "buffer_size_bytes": 1024**3, "num_buffers": 2, + "use_legacy_packing": False, }, ) ] diff --git a/tests/unit/utils/test_packed_tensor.py b/tests/unit/utils/test_packed_tensor.py index 19c8204fa4e..1bc12e87f35 100644 --- a/tests/unit/utils/test_packed_tensor.py +++ b/tests/unit/utils/test_packed_tensor.py @@ -236,7 +236,9 @@ def post_unpack_func(tensor_list): for name, tensor in result: unpacked_tensors[name] = tensor - assert consumer_group.current_index == 3 + # Target is 2000 B; the two 2400 B tensors are isolated in their own + # batches: [800+40], [2400], [80], [2400], [2+80]. + assert consumer_group.current_index == 5 # Verify all parameters were unpacked assert len(unpacked_tensors) == len(params) @@ -284,6 +286,73 @@ def test_packed_broadcast_single_large_tensor(): assert mock_group.broadcasted_tensors[0].numel() == expected_size +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") +def test_packed_broadcast_isolates_large_tensor_after_small_tensor(): + """An oversized tensor must not force a copy of the preceding prefix.""" + params = [ + ("small_weight", torch.randn(16, dtype=torch.float32).cuda()), + ("large_weight", torch.randn(100, dtype=torch.float32).cuda()), + ] + producer_group = MockCommunicationGroup() + + with patch( + "nemo_rl.utils.packed_tensor.get_target_packed_tensor_size", return_value=100 + ): + packed_broadcast_producer( + iterator=iter(params), + group=producer_group, + src=0, + post_iter_func=lambda x: x[1], + ) + + consumer_group = MockConsumerCommunicationGroup( + producer_group.broadcasted_tensors + ) + unpacked_tensors = {} + + def post_unpack_func(tensor_list): + unpacked_tensors.update(tensor_list) + + packed_broadcast_consumer( + iterator=iter(create_mock_state_dict_info(params).items()), + group=consumer_group, + src=0, + post_unpack_func=post_unpack_func, + ) + + assert [tensor.numel() for tensor in producer_group.broadcasted_tensors] == [ + params[0][1].nbytes, + params[1][1].nbytes, + ] + for name, original_tensor in params: + assert torch.equal(unpacked_tensors[name], original_tensor) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") +def test_packed_broadcast_legacy_packing_keeps_oversized_tensor_in_prefix_buffer(): + """vLLM's own consumer (Dynamo) packs an oversized tensor with its prefix.""" + params = [ + ("small_weight", torch.randn(16, dtype=torch.float32).cuda()), + ("large_weight", torch.randn(100, dtype=torch.float32).cuda()), + ] + group = MockCommunicationGroup() + + with patch( + "nemo_rl.utils.packed_tensor.get_target_packed_tensor_size", return_value=100 + ): + packed_broadcast_producer( + iterator=iter(params), + group=group, + src=0, + post_iter_func=lambda x: x[1], + use_legacy_packing=True, + ) + + assert [t.numel() for t in group.broadcasted_tensors] == [ + sum(t.nbytes for _, t in params) + ] + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") def test_packed_broadcast_multiple_batches(): """Test that tensors are properly batched when exceeding target size.""" diff --git a/tests/unit/weight_sync/test_weight_synchronizer.py b/tests/unit/weight_sync/test_weight_synchronizer.py index 496564e9830..506a273fb68 100644 --- a/tests/unit/weight_sync/test_weight_synchronizer.py +++ b/tests/unit/weight_sync/test_weight_synchronizer.py @@ -803,6 +803,7 @@ def test_sync_weights_calls_broadcast_and_receive(self, mock_ray): refit_timeout_s=None, buffer_size_bytes=None, num_buffers=None, + use_legacy_packing=False, ) gen.update_weights_from_collective.assert_called_once() @@ -899,6 +900,7 @@ def test_backend_sender_contract_controls_geometry_and_world_size(self, mock_ray nccl_peer="vllm", buffer_size_bytes=1024**3, num_buffers=2, + use_legacy_packing=True, ) gen.get_inference_world_size.return_value = 8 sync = CollectiveWeightSynchronizer( @@ -922,6 +924,7 @@ def test_backend_sender_contract_controls_geometry_and_world_size(self, mock_ray refit_timeout_s=None, buffer_size_bytes=1024**3, num_buffers=2, + use_legacy_packing=True, ) @@ -1455,8 +1458,9 @@ def test_non_colocated_megatron_m2n_uses_effective_parallelism(self): @pytest.mark.parametrize(("configured", "expected_calls"), [(False, 0), (True, 1)]) @patch("nemo_rl.weight_sync.collective_weight_synchronizer.ray") + @pytest.mark.parametrize("backend", [VLLM_BACKEND, "trtllm"]) def test_non_colocated_vllm_propagates_refit_memory_release( - self, mock_ray, configured, expected_calls + self, mock_ray, configured, expected_calls, backend ): mock_ray.get.return_value = [True] policy = _mock_policy() @@ -1465,7 +1469,7 @@ def test_non_colocated_vllm_propagates_refit_memory_release( sync = create_weight_synchronizer( policy=policy, generation=_mock_generation(), - generation_backend=VLLM_BACKEND, + generation_backend=backend, colocated=False, train_cluster=_mock_cluster(), inference_cluster=_mock_cluster(), @@ -1505,6 +1509,8 @@ def test_refit_memory_release_rejects_dtensor_policy(self): None, "requires non-colocated vLLM or Megatron", ), + # TensorRT-LLM has no shard-to-shard refit transport. + ("trtllm", False, "nccl_reshard", "TensorRT-LLM generation"), ], ) def test_refit_memory_release_rejects_unsupported_transport(