diff --git a/docs/fp8.md b/docs/fp8.md index 51a57016f3e..25116647c23 100644 --- a/docs/fp8.md +++ b/docs/fp8.md @@ -85,6 +85,25 @@ MXFP8 configurations should use `quantization_ignore_patterns` instead. (`precision: "fp8"` without `is_mx`) has no pattern-based replacement yet and must continue to use `quantization_ignored_layer_kws`. +For MXFP8 rollout with Megatron training, trainer-side prequantization can +reduce the refit payload: + +```yaml +policy: + generation: + vllm_cfg: + precision: fp8 + is_mx: true + refit_prequantize: true +``` + +`refit_prequantize` is an MXFP8 refit optimization. It requires +`precision: fp8`, `is_mx: true`, and the Megatron policy backend. It moves +eligible weight quantization to the trainer and transfers E4M3 values plus E8M0 +scales instead of BF16 weights. It is rejected for blockwise FP8, BF16, NVFP4, +sparse-delta refit, and NCCL-Reshard refit. NVFP4 real-quant rollout uses its own +packed-weight refit protocol. + To train with FP8, you need to set the Megatron path and configure it using the following settings: ``` diff --git a/docs/guides/refit.md b/docs/guides/refit.md index 447a55ba892..b3e39a8fe11 100644 --- a/docs/guides/refit.md +++ b/docs/guides/refit.md @@ -38,6 +38,17 @@ workers can send weights. Sparse delta is currently limited to GRPO. NIXL is initialized by the GRPO and distillation setup paths; PPO currently requires colocated generation. +## Performance Options + +| Option | Scope | Effect | +|---|---|---| +| `policy.generation.vllm_cfg.refit_prequantize` | Megatron training with MXFP8 vLLM rollout | Quantizes eligible weights on the trainer and transfers E4M3 values plus E8M0 scales. Requires `precision: fp8` and `is_mx: true`; sparse delta and NCCL Reshard do not support it. | +| `policy.generation.vllm_cfg.refit_cache_loader_routes` | vLLM refit | Replays identity-validated weight-loader routes after the first refit. Disabled by default because loader behavior is model-dependent. | +| `policy.refit_buffer_size_gb` | Colocated CUDA IPC or non-colocated NCCL broadcast | Sets the packing threshold explicitly. For NCCL broadcast, the same byte value is sent to producer and consumer so their collective chunk boundaries match. | +| `policy.refit_persistent_ipc_buffers` | Colocated CUDA-IPC refit | Reuses the two trainer staging buffers across refits. A fixed `refit_buffer_size_gb` gives stable memory use. | +| `policy.megatron_cfg.refit_slim_offload_after` | Colocated Megatron refit | Avoids repeating grad-buffer offload and a second allocator cleanup after weights are transferred. | +| `policy.megatron_cfg.pinned_reference_swap` | Megatron reference-policy logprobs | Keeps the CPU reference copy in pinned memory for faster host-to-device swaps, at the cost of additional pinned host memory. | + ## Minimal Configuration Colocated refit needs no transport configuration: diff --git a/examples/configs/grpo_math_1B.yaml b/examples/configs/grpo_math_1B.yaml index 53c8c13abce..831094c8f35 100644 --- a/examples/configs/grpo_math_1B.yaml +++ b/examples/configs/grpo_math_1B.yaml @@ -175,6 +175,12 @@ policy: megatron_cfg: enabled: false + # When true, offload_after_refit skips rerunning the full offload_before_refit + # pass and only re-offloads the optimizer with a single allocator cleanup. + refit_slim_offload_after: false + # When true, stage the per-step reference-policy swap in persistent pinned + # CPU buffers instead of pageable copies (faster PCIe, more host RAM). + pinned_reference_swap: false # Arbitrary overrides applied recursively to the Megatron Bridge model # provider before model instantiation. Do not duplicate first-class fields. model_overrides: {} @@ -344,6 +350,10 @@ policy: # makes the training sequence length divisible by the tensor parallel size # this is useful for sequence parallel training make_sequence_length_divisible_by: ${policy.dtensor_cfg.tensor_parallel_size} + # Keep refit CUDA-IPC staging buffers allocated across refits (skips two large + # allocations plus a gc/empty_cache pair per refit). Best with a fixed + # refit_buffer_size_gb; buffers stay resident on the trainer GPU between refits. + refit_persistent_ipc_buffers: false max_grad_norm: 1.0 optimizer: @@ -420,6 +430,8 @@ policy: vllm_cfg: async_engine: false precision: ${policy.precision} + # MXFP8 + Megatron only: quantize on trainer and stream E4M3 values plus scales. + refit_prequantize: false kv_cache_dtype: "auto" logprobs_mode: processed_logprobs tensor_parallel_size: 1 diff --git a/examples/configs/recipes/llm/performance/grpo-qwen3-235b-16n4g-mxfp8-rollout.yaml b/examples/configs/recipes/llm/performance/grpo-qwen3-235b-16n4g-mxfp8-rollout.yaml index 9cd16f6d07a..c22646ac21e 100644 --- a/examples/configs/recipes/llm/performance/grpo-qwen3-235b-16n4g-mxfp8-rollout.yaml +++ b/examples/configs/recipes/llm/performance/grpo-qwen3-235b-16n4g-mxfp8-rollout.yaml @@ -7,6 +7,8 @@ policy: tensor_parallel_size: 4 precision: "fp8" is_mx: true + refit_prequantize: true + refit_cache_loader_routes: true quantization_ignore_patterns: - model.layers.*.self_attn.* - model.layers.*.mlp.gate diff --git a/examples/configs/recipes/llm/performance/grpo-qwen3-30ba3b-4n4g-async-1off-mxfp8-rollout.yaml b/examples/configs/recipes/llm/performance/grpo-qwen3-30ba3b-4n4g-async-1off-mxfp8-rollout.yaml index 468b0fda3a1..bf8376355d3 100644 --- a/examples/configs/recipes/llm/performance/grpo-qwen3-30ba3b-4n4g-async-1off-mxfp8-rollout.yaml +++ b/examples/configs/recipes/llm/performance/grpo-qwen3-30ba3b-4n4g-async-1off-mxfp8-rollout.yaml @@ -21,6 +21,8 @@ policy: gpu_memory_utilization: 0.8 precision: "fp8" is_mx: true + refit_prequantize: false + refit_cache_loader_routes: false quantization_ignore_patterns: - model.layers.*.self_attn.* - model.layers.*.mlp.gate diff --git a/examples/configs/recipes/llm/performance/grpo-qwen3-30ba3b-4n4g-mxfp8-rollout.yaml b/examples/configs/recipes/llm/performance/grpo-qwen3-30ba3b-4n4g-mxfp8-rollout.yaml index 6680144bba6..09a82f2b39d 100644 --- a/examples/configs/recipes/llm/performance/grpo-qwen3-30ba3b-4n4g-mxfp8-rollout.yaml +++ b/examples/configs/recipes/llm/performance/grpo-qwen3-30ba3b-4n4g-mxfp8-rollout.yaml @@ -41,6 +41,8 @@ policy: tensor_parallel_size: 1 precision: "fp8" is_mx: true + refit_prequantize: true + refit_cache_loader_routes: true quantization_ignore_patterns: - model.layers.*.self_attn.* - model.layers.*.mlp.gate diff --git a/examples/configs/recipes/llm/performance/grpo-qwen3-32b-4n4g-mxfp8-rollout.yaml b/examples/configs/recipes/llm/performance/grpo-qwen3-32b-4n4g-mxfp8-rollout.yaml index b0f0ebbc0af..8b7c59c0ca0 100644 --- a/examples/configs/recipes/llm/performance/grpo-qwen3-32b-4n4g-mxfp8-rollout.yaml +++ b/examples/configs/recipes/llm/performance/grpo-qwen3-32b-4n4g-mxfp8-rollout.yaml @@ -8,6 +8,8 @@ policy: vllm_cfg: precision: "fp8" is_mx: true + refit_prequantize: true + refit_cache_loader_routes: true quantization_ignore_patterns: - model.layers.*.self_attn.* - lm_head diff --git a/nemo_rl/algorithms/distillation.py b/nemo_rl/algorithms/distillation.py index 91e00346576..f3e1f2bffb0 100644 --- a/nemo_rl/algorithms/distillation.py +++ b/nemo_rl/algorithms/distillation.py @@ -87,6 +87,7 @@ checkpoint_engine_refit_config, ) from nemo_rl.weight_sync.factory import create_weight_synchronizer +from nemo_rl.weight_sync.interfaces import initialize_refit_metadata # =============================================================================== # Configuration @@ -628,8 +629,7 @@ def init_nemo_gym(): ) student_generation.weight_synchronizer.init_communicator() elif student_generation is not None: - state_dict_info = student_policy.prepare_refit_info() - student_generation.prepare_refit_info(state_dict_info) + initialize_refit_metadata(student_policy, student_generation) # if it is not colocated inference, initialize collective communication for update weights if not colocated_inference and checkpoint_engine_config is None: @@ -722,6 +722,7 @@ def distillation_train( val_at_start = master_config.distillation.val_at_start val_at_end = master_config.distillation.val_at_end colocated_inference = master_config.policy["generation"]["colocated"]["enabled"] + refit_buffer_size_gb = master_config.policy.get("refit_buffer_size_gb") max_epochs = ( master_config.distillation.max_num_epochs ) # max number of epochs to train for @@ -734,7 +735,10 @@ def distillation_train( print("\nšŸ” Running initial validation...", flush=True) if NEED_REFIT and POLICY_GENERATION_STALE: refit_policy_generation( - student_policy, student_generation, colocated_inference + student_policy, + student_generation, + colocated_inference, + _refit_buffer_size_gb=refit_buffer_size_gb, ) POLICY_GENERATION_STALE = False else: @@ -795,6 +799,7 @@ def distillation_train( student_policy, student_generation, colocated_inference, + _refit_buffer_size_gb=refit_buffer_size_gb, timer=timer, ) POLICY_GENERATION_STALE = False @@ -937,7 +942,10 @@ def distillation_train( ): if NEED_REFIT and POLICY_GENERATION_STALE: refit_policy_generation( - student_policy, student_generation, colocated_inference + student_policy, + student_generation, + colocated_inference, + _refit_buffer_size_gb=refit_buffer_size_gb, ) POLICY_GENERATION_STALE = False else: diff --git a/nemo_rl/algorithms/grpo.py b/nemo_rl/algorithms/grpo.py index ed97496ce5d..7e725175eda 100644 --- a/nemo_rl/algorithms/grpo.py +++ b/nemo_rl/algorithms/grpo.py @@ -151,6 +151,7 @@ checkpoint_engine_refit_config, ) from nemo_rl.weight_sync.factory import create_weight_synchronizer +from nemo_rl.weight_sync.interfaces import initialize_refit_metadata # =============================================================================== # Configuration @@ -1763,9 +1764,10 @@ def init_dynamo(): ) is None and _needs_hf_refit_handshake( backend, nccl_reshard_refit_enabled, colocated_inference ): - state_dict_info = policy.prepare_refit_info() if policy_generation is not None: - policy_generation.prepare_refit_info(state_dict_info) + initialize_refit_metadata(policy, policy_generation) + else: + policy.prepare_refit_info() # Spin up non-colocated OPD teacher worker groups AFTER policy / vLLM are # ready. Parallelizing with policy init races on Megatron-Bridge's HF->mcore @@ -2498,10 +2500,15 @@ def refit_policy_generation( with timer_context: # update weights update_success = False + configured_buffer_size_bytes = ( + None + if _refit_buffer_size_gb is None + else int(_refit_buffer_size_gb * 1024**3) + ) if colocated_inference: # get model param keys, which is grouped by size - if _refit_buffer_size_gb is not None: - buffer_size_bytes = int(_refit_buffer_size_gb * (1024**3)) + if configured_buffer_size_bytes is not None: + buffer_size_bytes = configured_buffer_size_bytes else: # Empirically sets ratio as 30% to maximize efficiency. # The remaining 70% is a necessary buffer reserved for the parameter all-gathering across the expert-parallelism dimension. @@ -2524,11 +2531,15 @@ def refit_policy_generation( results = ray.get(futures_inference) update_success = all(result for result in results if result is not None) else: - # update weights through nccl (vLLM) - futures_train = policy.broadcast_weights_for_collective( - kv_scales=kv_scales, - ) - futures_inference = policy_generation.update_weights_from_collective() + # update weights through nccl (vLLM) or megatron reshard + if isinstance(policy_generation, MegatronGeneration): + futures_train = policy.swap_weights_via_reshard(is_source=True) + futures_inference = policy_generation.update_weights_from_collective() + else: + futures_train = policy.broadcast_weights_for_collective( + kv_scales=kv_scales, + ) + futures_inference = policy_generation.update_weights_from_collective() # wait for all futures to complete ray.get(futures_train) results = ray.get(futures_inference) @@ -4292,6 +4303,7 @@ def async_grpo_train( val_at_start = master_config.grpo.val_at_start val_at_end = master_config.grpo.val_at_end colocated_inference = master_config.policy["generation"]["colocated"]["enabled"] + refit_buffer_size_gb = master_config.policy.get("refit_buffer_size_gb") stop_at_validation_threshold = master_config.grpo.stop_at_validation_threshold stop_at_validation_metric = master_config.grpo.stop_at_validation_metric @@ -4499,6 +4511,7 @@ def async_grpo_train( policy, policy_generation, colocated_inference, + _refit_buffer_size_gb=refit_buffer_size_gb, ) print("āœ… Policy generation refit completed successfully", flush=True) POLICY_GENERATION_STALE = False @@ -5157,6 +5170,7 @@ def async_grpo_train( policy, policy_generation, colocated_inference, + _refit_buffer_size_gb=refit_buffer_size_gb, ) POLICY_GENERATION_STALE = False diff --git a/nemo_rl/algorithms/grpo_sync.py b/nemo_rl/algorithms/grpo_sync.py index e3f1fe46a81..c562a394d82 100644 --- a/nemo_rl/algorithms/grpo_sync.py +++ b/nemo_rl/algorithms/grpo_sync.py @@ -437,6 +437,7 @@ def grpo_train_sync( val_period = master_config.grpo.val_period val_start_at = master_config.grpo.val_start_at colocated_inference = master_config.policy["generation"]["colocated"]["enabled"] + refit_buffer_size_gb = master_config.policy.get("refit_buffer_size_gb") stop_at_validation_threshold = master_config.grpo.stop_at_validation_threshold stop_at_validation_metric = master_config.grpo.stop_at_validation_metric @@ -503,7 +504,12 @@ def grpo_train_sync( memory_tracker.snapshot_start_of_stage("Initial validation", dir()) if POLICY_GENERATION_STALE: - refit_policy_generation(policy, policy_generation, colocated_inference) + refit_policy_generation( + policy, + policy_generation, + colocated_inference, + _refit_buffer_size_gb=refit_buffer_size_gb, + ) POLICY_GENERATION_STALE = False else: policy_generation.prepare_for_generation() @@ -629,6 +635,7 @@ def grpo_train_sync( policy, policy_generation, colocated_inference, + _refit_buffer_size_gb=refit_buffer_size_gb, timer=timer, kv_scales=kv_scales_cache if sync_kv_scales else None, ) @@ -1006,6 +1013,7 @@ def grpo_train_sync( policy, policy_generation, colocated_inference, + _refit_buffer_size_gb=refit_buffer_size_gb, kv_scales=kv_scales_cache if sync_kv_scales else None, ) POLICY_GENERATION_STALE = False diff --git a/nemo_rl/algorithms/ppo.py b/nemo_rl/algorithms/ppo.py index 4db2867a7ed..d1b86a83cf0 100644 --- a/nemo_rl/algorithms/ppo.py +++ b/nemo_rl/algorithms/ppo.py @@ -108,6 +108,7 @@ from nemo_rl.utils.timer import TimeoutChecker, Timer from nemo_rl.utils.venvs import make_actor_runtime_env from nemo_rl.weight_sync.factory import create_weight_synchronizer +from nemo_rl.weight_sync.interfaces import initialize_refit_metadata # =============================================================================== # Configuration @@ -926,10 +927,12 @@ def initialize_generation_with_policy( ray.get(futures_train + futures_inference) worker_init_timing_metrics["collective_init_time_s"] = time.perf_counter() - t0 + # prepare refit info (sglang initializes refit state via its weight synchronizer) if backend != "sglang": - state_dict_info = policy.prepare_refit_info() if policy_generation is not None: - policy_generation.prepare_refit_info(state_dict_info) + initialize_refit_metadata(policy, policy_generation) + else: + policy.prepare_refit_info() # Calculate total setup time total_setup_time = time.perf_counter() - setup_start_time @@ -1251,6 +1254,7 @@ def ppo_train( val_at_end = master_config.ppo.val_at_end val_period = master_config.ppo.val_period colocated_inference = master_config.policy["generation"]["colocated"]["enabled"] + refit_buffer_size_gb = master_config.policy.get("refit_buffer_size_gb") # Initialize advantage estimator adv_estimator = _create_advantage_estimator(master_config) @@ -1261,7 +1265,12 @@ def ppo_train( memory_tracker.snapshot_start_of_stage("Initial validation", dir()) if NEED_REFIT and POLICY_GENERATION_STALE: - refit_policy_generation(policy, policy_generation, colocated_inference) + refit_policy_generation( + policy, + policy_generation, + colocated_inference, + _refit_buffer_size_gb=refit_buffer_size_gb, + ) if not colocated_inference: # Colocated refit offloads policy inside # `refit_policy_generation`. Do it here so the value @@ -1357,6 +1366,7 @@ def ppo_train( policy, policy_generation, colocated_inference, + _refit_buffer_size_gb=refit_buffer_size_gb, timer=timer, kv_scales=kv_scales_cache if sync_kv_scales else None, ) @@ -1675,6 +1685,7 @@ def ppo_train( policy, policy_generation, colocated_inference, + _refit_buffer_size_gb=refit_buffer_size_gb, kv_scales=kv_scales_cache if sync_kv_scales else None, ) if not colocated_inference: diff --git a/nemo_rl/modelopt/models/generation/vllm_quant_backend.py b/nemo_rl/modelopt/models/generation/vllm_quant_backend.py index 1f9a5abe282..fa376b7dc61 100644 --- a/nemo_rl/modelopt/models/generation/vllm_quant_backend.py +++ b/nemo_rl/modelopt/models/generation/vllm_quant_backend.py @@ -16,7 +16,7 @@ import types from collections.abc import Iterator from contextlib import ExitStack, contextmanager -from typing import Any +from typing import Any, Optional import torch import vllm # noqa: F401 @@ -531,10 +531,17 @@ def _synchronize_before_ipc_data_ack(self) -> None: return super()._synchronize_before_ipc_data_ack() - def prepare_refit_info(self, state_dict_info: dict[str, Any]) -> None: - super().prepare_refit_info(state_dict_info) + def prepare_refit_info( + self, + state_dict_info: dict[str, Any], + serialized_fp8_config: Optional[dict[str, Any]] = None, + ) -> Optional[list[str]]: if not self._is_real_quant_model(): - return + return super().prepare_refit_info(state_dict_info, serialized_fp8_config) + + # Real quantization owns a separate refit handshake and must not import + # the legacy FP8 quantization path. + self.state_dict_info = state_dict_info self._get_modelopt_reload_roots() quant_config = ( self.model_runner.vllm_config.model_config.hf_config.quantization_config @@ -553,6 +560,7 @@ def prepare_refit_info(self, state_dict_info: dict[str, Any]) -> None: "Fused ModelOpt MoE refits require all experts local; " "vLLM expert parallelism is unsupported" ) + return None @contextmanager def _patch_named_parameters_to_include_buffers(self, model): diff --git a/nemo_rl/models/generation/__init__.py b/nemo_rl/models/generation/__init__.py index 52b80e6d317..56ffdc491c1 100644 --- a/nemo_rl/models/generation/__init__.py +++ b/nemo_rl/models/generation/__init__.py @@ -19,7 +19,10 @@ from nemo_rl.models.generation.interfaces import GenerationConfig from nemo_rl.models.generation.trtllm import TrtllmConfig from nemo_rl.models.generation.vllm import VllmConfig -from nemo_rl.models.generation.vllm.config import VLLM_SPARSE_REFIT_TRANSPORTS +from nemo_rl.models.generation.vllm.config import ( + VLLM_SPARSE_REFIT_TRANSPORTS, + validate_vllm_quantization_config, +) TokenizerType = PreTrainedTokenizerBase @@ -50,6 +53,7 @@ def configure_generation_config( if config["backend"] == "vllm": config = cast(VllmConfig, config) + validate_vllm_quantization_config(config) if config.get("real_quant"): export_cpu_offload = config.get("real_quant_export_cpu_offload") if not isinstance(export_cpu_offload, bool): diff --git a/nemo_rl/models/generation/dynamo/dynamo_generation.py b/nemo_rl/models/generation/dynamo/dynamo_generation.py index 1c488855c7d..4f338e7b048 100644 --- a/nemo_rl/models/generation/dynamo/dynamo_generation.py +++ b/nemo_rl/models/generation/dynamo/dynamo_generation.py @@ -621,12 +621,17 @@ def init_collective( train_world_size=train_world_size, ) - def prepare_refit_info(self, state_dict_info: dict[str, Any]) -> None: + def prepare_refit_info( + self, state_dict_info: Optional[dict[str, Any]] + ) -> Optional[list[str]]: """Serialize checkpoint-format tensor metadata for native vLLM refit.""" + if state_dict_info is None: + return None channel = self._refit_channel if channel is None: raise RuntimeError("Dynamo refit channel is unavailable") channel.prepare(state_dict_info) + return None def update_weights_via_ipc_zmq(self) -> list[ray.ObjectRef]: raise NotImplementedError( diff --git a/nemo_rl/models/generation/interfaces.py b/nemo_rl/models/generation/interfaces.py index 8691e735321..28fcb0b0bf4 100644 --- a/nemo_rl/models/generation/interfaces.py +++ b/nemo_rl/models/generation/interfaces.py @@ -445,8 +445,16 @@ def requires_kv_scale_sync(self) -> bool: """Whether the generation backend requires KV cache scales synchronization.""" return False - def prepare_refit_info(self, state_dict_info: dict[str, Any]) -> None: - """Prepare the info for refit.""" + def prepare_refit_info( + self, state_dict_info: Optional[dict[str, Any]] + ) -> Optional[list[str]]: + """Prepare the info for refit. + + Returns: + Optionally, the parameter names the backend wants pre-quantized on + the trainer before streaming (e.g. vllm_cfg.refit_prequantize); + None when no trainer-side pre-quantization is requested. + """ raise NotImplementedError def update_weights_via_ipc_zmq(self) -> list[ray.ObjectRef]: diff --git a/nemo_rl/models/generation/sglang/sglang_generation.py b/nemo_rl/models/generation/sglang/sglang_generation.py index 2dbdbb2990c..c5644eb601b 100644 --- a/nemo_rl/models/generation/sglang/sglang_generation.py +++ b/nemo_rl/models/generation/sglang/sglang_generation.py @@ -812,8 +812,10 @@ def init_collective( ) -> list[ray.ObjectRef]: return [] - def prepare_refit_info(self, state_dict_info: dict[str, Any]) -> None: - pass + def prepare_refit_info( + self, state_dict_info: dict[str, Any] | None + ) -> list[str] | None: + return None def update_weights_via_ipc_zmq(self) -> list[ray.ObjectRef]: return [] diff --git a/nemo_rl/models/generation/trtllm/trtllm_generation.py b/nemo_rl/models/generation/trtllm/trtllm_generation.py index 978303b5b06..89f153b48a8 100644 --- a/nemo_rl/models/generation/trtllm/trtllm_generation.py +++ b/nemo_rl/models/generation/trtllm/trtllm_generation.py @@ -435,13 +435,18 @@ def finish_generation(self, *args: Any, **kwargs: Any) -> bool: print(f"Error in finish_generation: {e}") return False - def prepare_refit_info(self, state_dict_info: dict[str, Any]) -> None: + def prepare_refit_info( + self, state_dict_info: Optional[dict[str, Any]] + ) -> Optional[list[str]]: + if state_dict_info is None: + return None futures = self.worker_group.run_all_workers_single_data( "prepare_refit_info_async", state_dict_info=state_dict_info, run_rank_0_only_axes=["tensor_parallel"], ) ray.get(futures) + return None def start_gpu_profiling(self) -> None: """Grpo profiling protocol: start nsys capture on the GPU workers.""" diff --git a/nemo_rl/models/generation/vllm/checkpoint_engine.py b/nemo_rl/models/generation/vllm/checkpoint_engine.py index 7168a9d3f3d..eefa08f2553 100644 --- a/nemo_rl/models/generation/vllm/checkpoint_engine.py +++ b/nemo_rl/models/generation/vllm/checkpoint_engine.py @@ -147,18 +147,18 @@ async def _update_weights_from_checkpoint_engine_async(self) -> bool: load_time = 0.0 start_time = time.time() - async for weight_batch in self.checkpoint_engine.receive_weight_batches(): - loaded_batches += 1 - loaded_tensors += len(weight_batch) - loaded_bytes += sum(weight.nbytes for _name, weight in weight_batch) - - load_start = time.time() - self._load_weights(weight_batch) - torch.cuda.current_stream().synchronize() - load_time += time.time() - load_start - del weight_batch - - self._maybe_process_fp8_kv_cache() + with self._weight_update_lifecycle("checkpoint_engine") as finalize: + async for weight_batch in self.checkpoint_engine.receive_weight_batches(): + loaded_batches += 1 + loaded_tensors += len(weight_batch) + loaded_bytes += sum(weight.nbytes for _name, weight in weight_batch) + + load_start = time.time() + self._load_weights(weight_batch) + torch.cuda.current_stream().synchronize() + load_time += time.time() - load_start + del weight_batch + finalize() total_time = time.time() - start_time loaded_gib = loaded_bytes / (1024 * 1024 * 1024) diff --git a/nemo_rl/models/generation/vllm/config.py b/nemo_rl/models/generation/vllm/config.py index dabf01d0b1e..5ac90257261 100644 --- a/nemo_rl/models/generation/vllm/config.py +++ b/nemo_rl/models/generation/vllm/config.py @@ -66,6 +66,13 @@ class VllmSpecificArgs(TypedDict): cap_max_tokens_to_context: NotRequired[bool] # Use ModelOpt MXFP8 quantization when precision is fp8. is_mx: NotRequired[bool] + # With is_mx, quantize weights to MXFP8 on the trainer during refit and + # stream E4M3 data plus scales (~47% smaller payload) instead of BF16; + # the vLLM worker then skips its per-refit re-quantization. Requires the + # Megatron policy backend. + refit_prequantize: NotRequired[bool] + # Cache and replay stable vLLM weight-loader routes across refits. + refit_cache_loader_routes: NotRequired[bool] # Deprecated in 0.8. Use quantization_ignore_patterns instead. quantization_ignored_layer_kws: NotRequired[list[str]] # MXFP8 exclusion patterns forwarded through vLLM's quantization config. @@ -191,6 +198,34 @@ class VllmConfig(GenerationConfig): real_quant_ignore: NotRequired[list[str]] +def validate_vllm_quantization_config(config: VllmConfig) -> None: + """Reject quantization options that would otherwise be silently ignored.""" + vllm_cfg = config.get("vllm_cfg") + if vllm_cfg is None: + return + refit_prequantize = vllm_cfg.get("refit_prequantize") + if refit_prequantize is not None and not isinstance(refit_prequantize, bool): + raise ValueError( + "policy.generation.vllm_cfg.refit_prequantize must be a boolean." + ) + if refit_prequantize and not ( + vllm_cfg.get("precision") == "fp8" and vllm_cfg.get("is_mx") is True + ): + raise ValueError( + "policy.generation.vllm_cfg.refit_prequantize requires " + "precision='fp8' and is_mx=true." + ) + if refit_prequantize and config.get("refit_transport") == "nccl_reshard": + raise ValueError( + "policy.generation.vllm_cfg.refit_prequantize is not supported with " + "nccl_reshard; that transport owns its weight-format conversion." + ) + for field in ("refit_cache_loader_routes",): + value = vllm_cfg.get(field) + if value is not None and not isinstance(value, bool): + raise ValueError(f"policy.generation.vllm_cfg.{field} must be a boolean.") + + def resolve_vllm_video_config(config: VllmConfig) -> VllmVideoConfig | None: """Validate and return the optional vLLM video sampling contract.""" raw_video_config = config["vllm_cfg"].get("video") @@ -261,6 +296,7 @@ def materialize_vllm_video_config( def normalize_vllm_refit_config(config: VllmConfig) -> VllmRefitConfig | None: """Validate the selected refit transport and resolve its scoped defaults.""" + validate_vllm_quantization_config(config) if cast(dict[str, Any], config).get("checkpoint_engine") is not None: raise ValueError( "policy.generation.checkpoint_engine was replaced by " diff --git a/nemo_rl/models/generation/vllm/quantization/fp8.py b/nemo_rl/models/generation/vllm/quantization/fp8.py index fd2bf721924..b8e5d8124c2 100644 --- a/nemo_rl/models/generation/vllm/quantization/fp8.py +++ b/nemo_rl/models/generation/vllm/quantization/fp8.py @@ -15,7 +15,8 @@ import os import warnings from collections.abc import Sequence -from dataclasses import dataclass, field +from dataclasses import asdict, dataclass, field +from typing import Any from unittest.mock import patch import ray @@ -33,6 +34,9 @@ from nemo_rl.models.generation.vllm.quantization.mxfp8_utils import ( pad_flashinfer_scale_k, ) +from nemo_rl.models.generation.vllm.worker_utils import ( + refit_cache_loader_routes_enabled, +) logger = init_logger(__name__) @@ -61,6 +65,9 @@ class FP8Config: kv_cache_dtype: str = "auto" use_fp8_weights: bool = True # Whether model weights are quantized to FP8 is_mx: bool = False + # Weights arrive from the trainer already MXFP8-quantized (E4M3 data plus + # *_scale_from_checkpoint entries), so load_weights skips re-quantization. + refit_prequantize: bool = False @dataclass() @@ -70,11 +77,16 @@ class FP8State: seen_params: set = field(default_factory=lambda: set()) fp8_param_names: set = field(default_factory=lambda: set()) vllm_patches: list = field(default_factory=lambda: []) + # Full refit manifest names from prepare_refit_info. load_weights receives + # transfer batches split by buffer size, so a weight and its + # *_scale_from_checkpoint entry may arrive in different batches; presence + # must be validated against the full manifest, not the current batch. + refit_manifest_names: set | None = None # Global FP8 config that can be accessed by patched vLLM functions # initialized by 'init_fp8_cfg()' -global_fp8_config: FP8Config = None +global_fp8_config: FP8Config | None = None # Global FP8 state that holds runtime fp8 objects fp8_state: FP8State = FP8State() @@ -96,7 +108,62 @@ def my_run_engine_core(*args, **kwargs): return original_run_engine_core(*args, **kwargs) +def serialize_fp8_config() -> dict[str, Any] | None: + if global_fp8_config is None: + return None + return asdict(global_fp8_config) + + +def install_fp8_config(config: dict[str, Any] | None) -> None: + if config is None: + return + global global_fp8_config + global_fp8_config = FP8Config(**config) + + +def set_refit_manifest_names(names: set[str] | None) -> None: + fp8_state.refit_manifest_names = names + + +def _patch_ray_executor_v2_worker(ray_executor_v2: Any, fp8_config: FP8Config) -> None: + """Install FP8 patches inside RayExecutorV2 workers before model loading.""" + original_ray_worker_proc = ray_executor_v2.RayWorkerProc + if getattr(original_ray_worker_proc, "_nrl_fp8_patched", False): + original_ray_worker_proc._nrl_fp8_config = fp8_config + return + + class NRLFP8RayWorkerProc(original_ray_worker_proc): + _nrl_fp8_patched = True + _nrl_fp8_config = fp8_config + + def initialize_worker( + self, + local_rank: int, + env_vars: dict[str, str], + driver_env_vars: dict[str, str] | None = None, + assigned_physical_gpu_ids: list[int] | None = None, + ) -> Any: + global fp8_patches_applied + if not fp8_patches_applied: + apply_fp8_patches(None, type(self)._nrl_fp8_config) + return super().initialize_worker( + local_rank, + env_vars, + driver_env_vars, + assigned_physical_gpu_ids, + ) + + ray_executor_v2.RayWorkerProc = NRLFP8RayWorkerProc + + def monkey_patch_vllm_ray_executor(fp8_config): + try: + from vllm.v1.executor import ray_executor_v2 + except ImportError: + pass + else: + _patch_ray_executor_v2_worker(ray_executor_v2, fp8_config) + if fp8_config.model_parallel_size > 1: # we patch vllm's collective_rpc so that before vllm initalizes the model on each rank, we execute # a ray remote that patches each worker with the required fp8 vllm patches @@ -167,6 +234,12 @@ def apply_fp8_patches(self, fp8_config): process_weights_after_loading_mxfp8_moe, ) ) + fp8_state.vllm_patches.append( + patch( + "vllm.model_executor.layers.quantization.modelopt.ModelOptMxFp8FusedMoE.apply_monolithic", + apply_monolithic_mxfp8_moe, + ) + ) # These patches add support for pow2, e8 dynamic activation scalings factors which are believed to have higher # SNR compared to plain fp32 scaling factors. This feature is still under active research. @@ -243,6 +316,7 @@ def init_fp8(vllm_cfg, model_name, model_parallel_size): } if is_mx: fp8_config_kwargs["is_mx"] = True + fp8_config_kwargs["refit_prequantize"] = bool(vllm_cfg.get("refit_prequantize")) if vllm_cfg.get("pow2_weight_scaling_factors") is False: raise ValueError("only pow2 weight scaling factors are supported for MXFP8") if vllm_cfg.get("pow2_activation_scaling_factors") is False: @@ -518,6 +592,8 @@ def quantize_mxfp8_weight(weight: torch.Tensor) -> tuple[torch.Tensor, torch.Ten def load_weights(weights, model_runner): global global_fp8_config + weights = list(weights) + weight_names = {name for name, _tensor in weights} weights_quantized = [] model = model_runner.model @@ -543,9 +619,15 @@ def load_weights(weights, model_runner): and experts_module.w13_weight.dtype == torch.float8_e4m3fn and experts_module.w2_weight.dtype == torch.float8_e4m3fn ): + if v.dtype == torch.float8_e4m3fn: + # Trainer-side prequantized slab: already E4M3 with its + # scale streamed separately; pass through untouched. + weights_quantized.append((k, v)) + continue if global_fp8_config.is_mx: raise NotImplementedError( - "MXFP8 refit does not support grouped MoE expert weights." + "MXFP8 refit does not support quantizing grouped MoE " + "expert weights on the fly; enable refit_prequantize." ) weights_quantized.extend(_expand_grouped_moe_expert_to_fp8(k, v)) else: @@ -554,6 +636,29 @@ def load_weights(weights, model_runner): if not _is_fp8_weight(k, model): weights_quantized.append((k, v)) continue + if v.dtype == torch.float8_e4m3fn: + if global_fp8_config.is_mx and not global_fp8_config.refit_prequantize: + raise ValueError( + "MXFP8 E4M3 refit weights require refit_prequantize=true; " + "other FP8 trainer scale layouts are not compatible." + ) + # Transfer batches split by buffer size, so the matching scale may + # arrive in an earlier or later batch; validate against the full + # refit manifest when available, not just the current batch. + scale_name = k + "_scale_from_checkpoint" + manifest = fp8_state.refit_manifest_names + if ( + global_fp8_config.is_mx + and scale_name not in weight_names + and (manifest is None or scale_name not in manifest) + ): + raise ValueError( + f"Prequantized MXFP8 weight {k!r} is missing {scale_name!r}." + ) + # Prequantized MXFP8 sends the matching *_scale_from_checkpoint + # entry separately. Non-MXFP8 blockwise FP8 sends *_scale_inv. + weights_quantized.append([k, v]) + continue # Cast the weight into fp8 and its scale factor if global_fp8_config.is_mx: param_lp, param_scale = quantize_mxfp8_weight(v) @@ -562,15 +667,22 @@ def load_weights(weights, model_runner): v.to(torch.float), weight_block_size=FP8_BLOCK_QUANT_KWARGS["weight_block_size"], ) - param_scale = torch.squeeze(param_scale, dim=-1) if global_fp8_config.is_mx: weights_quantized.append([k, param_lp]) weights_quantized.append([k + "_scale_from_checkpoint", param_scale]) else: + param_scale = torch.squeeze(param_scale, dim=-1) weights_quantized.append([k, param_lp]) weights_quantized.append([k + "_scale_inv", param_scale]) - # Finally load the weights into vllm - model.load_weights(weights_quantized) + # Finally load the weights into vllm. Deferred: importing vllm_backend at + # module top would cycle through the nemo_rl generation package init. + from nemo_rl.models.generation.vllm.vllm_backend import load_weights_maybe_cached + + load_weights_maybe_cached( + model, + weights_quantized, + cache_loader_routes=refit_cache_loader_routes_enabled(model_runner.vllm_config), + ) def cast_tensor_to_fp8_blockwise( @@ -932,7 +1044,7 @@ def process_weights_after_loading_mxfp8_linear(self, layer) -> None: def create_weights_mxfp8_moe( self, - layer: torch.nn.Module, + layer: RoutedExperts, num_experts: int, hidden_size: int, intermediate_size_per_partition: int, @@ -941,9 +1053,8 @@ def create_weights_mxfp8_moe( ): """Create ModelOpt MXFP8 MoE weights without assigning read-only vLLM attrs. - vLLM 0.20.0's ModelOpt MXFP8 path writes hidden/intermediate sizes onto - FusedMoE, but those are read-only properties backed by moe_config. Keep the - upstream allocation behavior while relying on those existing properties. + Keep vLLM's allocation behavior while preserving the checkpoint layout + required by repeated refits. """ # vLLM 0.25 moved FusedMoeWeightScaleSupported out of fused_moe.layer; # it is re-exported from the fused_moe package. @@ -958,6 +1069,8 @@ def create_weights_mxfp8_moe( from vllm.model_executor.parameter import ModelWeightParameter from vllm.model_executor.utils import set_weight_attrs + assert layer.intermediate_size_per_partition == intermediate_size_per_partition + assert layer.hidden_size == hidden_size layer.orig_dtype = params_dtype if hidden_size % MXFP8_BLOCK_SIZE != 0: raise ValueError( @@ -1095,6 +1208,78 @@ def process_weights_after_loading_moe(self, layer) -> None: ) +def _round_up(value: int, multiple: int) -> int: + return ((value + multiple - 1) // multiple) * multiple + + +def _pad_tensor_dim( + tensor: torch.Tensor, dim: int, padded_size: int, pad_value: int = 0 +) -> torch.Tensor: + current_size = tensor.shape[dim] + if current_size == padded_size: + return tensor + padded_shape = list(tensor.shape) + padded_shape[dim] = padded_size + padded = torch.zeros(padded_shape, dtype=tensor.dtype, device=tensor.device) + if pad_value != 0: + padded.fill_(pad_value) + padded.narrow(dim, 0, current_size).copy_(tensor) + return padded + + +def _pad_w13_shards( + tensor: torch.Tensor, + intermediate_size_factor: int, + padded_intermediate_size: int, + pad_value: int = 0, +) -> torch.Tensor: + """Pad the intermediate dim of a [E, factor*I, K] tensor to [E, factor*I_pad, K]. + + Pads each of the factor shards separately so gated W13 halves stay aligned + for swap_w13_to_w31 and reorder_rows_for_gated_act_gemm. + """ + num_experts, total_rows, cols = tensor.shape + rows_per_shard = total_rows // intermediate_size_factor + if rows_per_shard == padded_intermediate_size: + return tensor + shards = tensor.reshape(num_experts, intermediate_size_factor, rows_per_shard, cols) + shards = _pad_tensor_dim(shards, 2, padded_intermediate_size, pad_value) + return shards.reshape( + num_experts, intermediate_size_factor * padded_intermediate_size, cols + ) + + +def _clamp_mxfp8_scale(scale: torch.Tensor) -> torch.Tensor: + return torch.where(scale == 0, torch.ones_like(scale), scale) + + +def _set_mxfp8_apply_tensor(layer, name: str, value: torch.Tensor) -> None: + existing = getattr(layer, name, None) + if existing is not None and existing.shape == value.shape: + # Keep storage stable across refits (CUDA graphs capture pointers). + existing.copy_(value) + else: + # Clone: value may view a scratch buffer shared across layers. + setattr(layer, name, torch.nn.Parameter(value.clone(), requires_grad=False)) + + +# Shared gather destinations keyed by (tag, shape, device). They persist across +# refits so the batched shuffle allocates nothing after the first pass; their +# contents are rewritten on every call, so a sleep-mode discard is harmless. +mxfp8_shuffle_scratch_buffers: dict[ + tuple[str, tuple[int, ...], torch.device], torch.Tensor +] = {} + + +def _mxfp8_scratch(tag: str, shape: torch.Size, device: torch.device) -> torch.Tensor: + key = (tag, tuple(shape), device) + buf = mxfp8_shuffle_scratch_buffers.get(key) + if buf is None: + buf = torch.empty(shape, dtype=torch.uint8, device=device) + mxfp8_shuffle_scratch_buffers[key] = buf + return buf + + def _mxfp8_moe_row_permutations( layer, w13_weight: torch.Tensor, @@ -1102,7 +1287,16 @@ def _mxfp8_moe_row_permutations( is_gated: bool, epilogue_tile_m: int, ) -> tuple[torch.Tensor, torch.Tensor]: - """Return row permutations equivalent to the FlashInfer shuffle calls.""" + """Composed row-index permutations for the batched TRTLLM MoE shuffle. + + shuffle_matrix_a / shuffle_matrix_sf_a and reorder_rows_for_gated_act_gemm + are input-independent row permutations (and the sf row indices equal the + weight row indices for the same row count), so composing the indices once + reproduces the per-expert call sequence as a single gather per tensor. + Cached on the layer as CPU tensors; device copies are made per call because + tensors allocated while vLLM loads the model land in the sleep-mode weights + pool, whose contents are discarded at sleep_level=2. + """ perm_w13 = getattr(layer, "_mxfp8_shuffle_perm_w13", None) perm_w2 = getattr(layer, "_mxfp8_shuffle_perm_w2", None) if perm_w13 is None or perm_w2 is None: @@ -1131,7 +1325,15 @@ def _shuffle_mxfp8_moe_batched( is_gated: bool, epilogue_tile_m: int, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - """Shuffle stacked expert values and scales with four batched gathers.""" + """Apply the TRTLLM row shuffles as one batched gather per stacked tensor. + + Bit-identical to the per-expert loop (`_shuffle_mxfp8_moe_per_expert`): + the row gathers run as a single dim-1 index_select over the whole + [E, M, K] tensor and the scale swizzle as one 3D block_scale_interleave, + whose flat output equals the stacked per-expert 2D outputs. Gathers write + into shared scratch buffers; callers copy_ into the persistent + destinations (w13/w2 weights may alias their own gather source). + """ from flashinfer import block_scale_interleave from vllm.model_executor.layers.quantization.utils.mxfp8_utils import ( MXFP8_SCALE_DTYPE, @@ -1142,24 +1344,39 @@ def _shuffle_mxfp8_moe_batched( layer, w13_weight, w2_weight, is_gated, epilogue_tile_m ) num_experts = w13_weight.shape[0] + w13_u8 = w13_weight.view(torch.uint8) w2_u8 = w2_weight.view(torch.uint8) - w13_shuffled = torch.index_select(w13_u8, 1, perm_w13) - w2_shuffled = torch.index_select(w2_u8, 1, perm_w2) - - w13_scale_u8 = pad_flashinfer_scale_k(w13_scale.view(torch.uint8)) - w2_scale_u8 = pad_flashinfer_scale_k(w2_scale.view(torch.uint8)) - assert w13_scale_u8.shape[1] % 128 == 0 - assert w2_scale_u8.shape[1] % 128 == 0 - w13_scale_gathered = torch.index_select(w13_scale_u8, 1, perm_w13) - w2_scale_gathered = torch.index_select(w2_scale_u8, 1, perm_w2) + w13_shuffled = torch.index_select( + w13_u8, 1, perm_w13, out=_mxfp8_scratch("w13", w13_u8.shape, w13_u8.device) + ) + w2_shuffled = torch.index_select( + w2_u8, 1, perm_w2, out=_mxfp8_scratch("w2", w2_u8.shape, w2_u8.device) + ) + + w13_sf_u8 = pad_flashinfer_scale_k(w13_scale.view(torch.uint8)) + w2_sf_u8 = pad_flashinfer_scale_k(w2_scale.view(torch.uint8)) + # Same constraint shuffle_matrix_sf_a asserts on the per-expert path. + assert w13_sf_u8.shape[1] % 128 == 0 and w2_sf_u8.shape[1] % 128 == 0 + w13_sf_gathered = torch.index_select( + w13_sf_u8, + 1, + perm_w13, + out=_mxfp8_scratch("w13_sf", w13_sf_u8.shape, w13_sf_u8.device), + ) + w2_sf_gathered = torch.index_select( + w2_sf_u8, + 1, + perm_w2, + out=_mxfp8_scratch("w2_sf", w2_sf_u8.shape, w2_sf_u8.device), + ) w13_scale_shuffled = ( - block_scale_interleave(w13_scale_gathered) + block_scale_interleave(w13_sf_gathered) .view(MXFP8_SCALE_DTYPE) .view(num_experts, -1) ) w2_scale_shuffled = ( - block_scale_interleave(w2_scale_gathered) + block_scale_interleave(w2_sf_gathered) .view(MXFP8_SCALE_DTYPE) .view(num_experts, -1) ) @@ -1179,7 +1396,7 @@ def _shuffle_mxfp8_moe_per_expert( is_gated: bool, epilogue_tile_m: int, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - """Run the original per-expert FlashInfer shuffle as a reference path.""" + """Per-expert reference shuffle, kept for NRL_MXFP8_SHUFFLE_VERIFY.""" from flashinfer import ( reorder_rows_for_gated_act_gemm, shuffle_matrix_a, @@ -1232,57 +1449,103 @@ def _shuffle_mxfp8_moe_per_expert( ) -def process_weights_after_loading_mxfp8_moe(self, layer) -> None: - """Shuffle weights and scales into FlashInfer TRTLLM MXFP8 layout.""" +def process_weights_after_loading_mxfp8_moe(self, layer: RoutedExperts) -> None: + """Shuffle weights and scales into FlashInfer TRTLLM MXFP8 layout. + + FlashInfer's shuffle_matrix_sf_a requires M divisible by 128, so the + intermediate dim is padded to 128 (Nemotron-3-Nano: 1856 -> 1920, + 928 -> 1024 at TP2). The TRTLLM kernel also produced non-finite output + for Nano's hidden 2816 unless padded to the next 512 boundary (3072). + When padding is needed, the checkpoint-layout parameters keep their + original shapes so refit keeps loading into them unchanged, and the + padded+shuffled kernel inputs are maintained as separate *_for_apply + tensors consumed by apply_monolithic_mxfp8_moe. Already-aligned models + keep the original in-place shuffle behavior. + """ from vllm.model_executor.layers.fused_moe import FusedMoeWeightScaleSupported from vllm.model_executor.layers.fused_moe.oracle.fp8 import Fp8MoeBackend from vllm.model_executor.layers.quantization.utils.flashinfer_utils import ( swap_w13_to_w31, ) + from vllm.model_executor.layers.quantization.utils.mxfp8_utils import ( + MXFP8_BLOCK_SIZE, + ) from vllm.model_executor.parameter import ModelWeightParameter from vllm.model_executor.utils import set_weight_attrs - if self.mxfp8_backend != Fp8MoeBackend.FLASHINFER_TRTLLM: + if ( + self.mxfp8_backend != Fp8MoeBackend.FLASHINFER_TRTLLM + or not self.experts_cls.is_monolithic() + ): raise NotImplementedError( - "MXFP8 MoE refit layout conversion only supports FLASHINFER_TRTLLM; " - f"got {self.mxfp8_backend}." + "NeMo-RL's padded and batched MXFP8 refit path requires the " + "monolithic FlashInfer TRTLLM backend." ) + layer.weight_block_size = self.weight_block_size epilogue_tile_m = 128 is_gated = self.moe.is_act_and_mul + intermediate_size_factor = 2 if is_gated else 1 + + # Sizes come from the checkpoint-layout params, not layer attributes: + # layer.intermediate_size_per_partition is switched to the padded value + # below, while these params keep their original shapes across refits. + original_intermediate_size = layer.w2_weight.shape[2] + original_hidden_size = layer.w13_weight.shape[2] + padded_intermediate_size = _round_up(original_intermediate_size, 128) + padded_hidden_size = _round_up(original_hidden_size, 512) + needs_padding = ( + padded_intermediate_size != original_intermediate_size + or padded_hidden_size != original_hidden_size + ) + + first_load = not hasattr(layer, "w13_weight_scale_from_checkpoint") w13_weight = layer.w13_weight.data - if not hasattr(layer, "w13_weight_scale_from_checkpoint"): + if first_load: w13_scale = layer.w13_weight_scale.data + w2_scale = layer.w2_weight_scale.data else: w13_scale = layer.w13_weight_scale_from_checkpoint.data + w2_scale = layer.w2_weight_scale_from_checkpoint.data + w2_weight = layer.w2_weight.data + + if needs_padding: + w13_weight = _pad_w13_shards( + w13_weight, intermediate_size_factor, padded_intermediate_size + ) + w13_weight = _pad_tensor_dim(w13_weight, 2, padded_hidden_size) + w13_scale = _pad_w13_shards( + w13_scale, intermediate_size_factor, padded_intermediate_size, pad_value=1 + ) + w13_scale = _pad_tensor_dim( + w13_scale, 2, padded_hidden_size // MXFP8_BLOCK_SIZE, pad_value=1 + ) + w2_weight = _pad_tensor_dim(w2_weight, 1, padded_hidden_size) + w2_weight = _pad_tensor_dim(w2_weight, 2, padded_intermediate_size) + w2_scale = _pad_tensor_dim(w2_scale, 1, padded_hidden_size, pad_value=1) + w2_scale = _pad_tensor_dim( + w2_scale, 2, padded_intermediate_size // MXFP8_BLOCK_SIZE, pad_value=1 + ) + # Zero E8M0 scale bytes destabilize the TRTLLM kernel; clamp to byte 1. + w13_scale = _clamp_mxfp8_scale(w13_scale) + w2_scale = _clamp_mxfp8_scale(w2_scale) + if is_gated: # FI TRTLLM gated kernels use W31 ordering. Model checkpoints store # gated projection as W13, so convert once before shuffling. w13_weight = swap_w13_to_w31(w13_weight) w13_scale = swap_w13_to_w31(w13_scale) - w2_weight = layer.w2_weight.data - if not hasattr(layer, "w2_weight_scale_from_checkpoint"): - w2_scale = layer.w2_weight_scale.data - else: - w2_scale = layer.w2_weight_scale_from_checkpoint.data - shuffled = _shuffle_mxfp8_moe_batched( - layer, - w13_weight, - w2_weight, - w13_scale, - w2_scale, - is_gated, - epilogue_tile_m, - ) ( w13_weight_shuffled, w2_weight_shuffled, w13_scale_shuffled, w2_scale_shuffled, - ) = shuffled + ) = _shuffle_mxfp8_moe_batched( + layer, w13_weight, w2_weight, w13_scale, w2_scale, is_gated, epilogue_tile_m + ) - if not hasattr(layer, "w13_weight_scale_from_checkpoint"): + if first_load: layer.w13_weight_scale_from_checkpoint = ModelWeightParameter( data=layer.w13_weight_scale.data, input_dim=2, @@ -1315,24 +1578,64 @@ def process_weights_after_loading_mxfp8_moe(self, layer) -> None: layer.w2_weight_scale_from_checkpoint, {"quant_method": FusedMoeWeightScaleSupported.BLOCK.value}, ) - layer.w13_weight_scale = torch.nn.Parameter( - w13_scale_shuffled, requires_grad=False - ) - layer.w2_weight_scale = torch.nn.Parameter( - w2_scale_shuffled, requires_grad=False + + if needs_padding: + # Checkpoint-layout params (w13_weight, w2_weight, w13_weight_scale, + # w2_weight_scale and the *_from_checkpoint aliases) stay untouched so + # refit keeps loading unpadded tensors into them. The kernel consumes + # the *_for_apply tensors instead. + layer.mxfp8_unpadded_hidden_size = original_hidden_size + layer.mxfp8_padded_hidden_size = padded_hidden_size + layer.mxfp8_unpadded_intermediate_size_per_partition = ( + original_intermediate_size ) + layer.mxfp8_padded_intermediate_size_per_partition = padded_intermediate_size + # vLLM 0.25 stores this size on both RoutedExperts and its FusedMoEConfig. + # Keep them aligned so routing/kernel setup sees the padded value. + # Hidden size stays original: apply pads x and narrows the output back. + layer.intermediate_size_per_partition = padded_intermediate_size + layer.moe_config.intermediate_size_per_partition = padded_intermediate_size + self.moe.intermediate_size_per_partition = padded_intermediate_size + _set_mxfp8_apply_tensor(layer, "w13_weight_for_apply", w13_weight_shuffled) + _set_mxfp8_apply_tensor(layer, "w2_weight_for_apply", w2_weight_shuffled) + _set_mxfp8_apply_tensor(layer, "w13_scale_for_apply", w13_scale_shuffled) + _set_mxfp8_apply_tensor(layer, "w2_scale_for_apply", w2_scale_shuffled) else: - layer.w13_weight_scale.copy_(w13_scale_shuffled) - layer.w2_weight_scale.copy_(w2_scale_shuffled) - layer.w13_weight.copy_(w13_weight_shuffled) - layer.w2_weight.copy_(w2_weight_shuffled) + if first_load: + layer.w13_weight_scale = torch.nn.Parameter( + w13_scale_shuffled, requires_grad=False + ) + layer.w2_weight_scale = torch.nn.Parameter( + w2_scale_shuffled, requires_grad=False + ) + else: + layer.w13_weight_scale.copy_(w13_scale_shuffled) + layer.w2_weight_scale.copy_(w2_scale_shuffled) + layer.w13_weight.copy_(w13_weight_shuffled) + layer.w2_weight.copy_(w2_weight_shuffled) + + runtime_w13_scale = getattr(layer, "w13_scale_for_apply", layer.w13_weight_scale) + runtime_w2_scale = getattr(layer, "w2_scale_for_apply", layer.w2_weight_scale) + assert self.moe is layer.moe_config if self.moe_kernel is None: - from vllm.model_executor.layers.quantization.fp8 import make_fp8_moe_kernel + from vllm.model_executor.layers.fused_moe.oracle.fp8 import ( + make_fp8_moe_kernel, + make_fp8_moe_quant_config, + ) - self.moe_quant_config = self.get_fused_moe_quant_config(layer) - assert self.moe_quant_config is not None - assert self.experts_cls is not None + self.moe_quant_config = make_fp8_moe_quant_config( + fp8_backend=self.mxfp8_backend, + w1_scale=runtime_w13_scale, + w2_scale=runtime_w2_scale, + a1_scale=None, + a2_scale=None, + block_shape=self.weight_block_size, + swiglu_limit=getattr(layer, "swiglu_limit", None), + gemm1_alpha=getattr(layer, "swiglu_alpha", None), + gemm1_beta=getattr(layer, "swiglu_beta", None), + layer=layer, + ) self.moe_kernel = make_fp8_moe_kernel( moe_quant_config=self.moe_quant_config, moe_config=self.moe, @@ -1341,6 +1644,54 @@ def process_weights_after_loading_mxfp8_moe(self, layer) -> None: routing_tables=layer._expert_routing_tables(), layer=layer, ) + else: + assert self.moe_quant_config is not None + assert self.moe_quant_config.w1_scale is runtime_w13_scale + assert self.moe_quant_config.w2_scale is runtime_w2_scale + + +def apply_monolithic_mxfp8_moe( + self, + layer: RoutedExperts, + x: torch.Tensor, + router_logits: torch.Tensor, + input_ids: torch.Tensor | None = None, +) -> torch.Tensor: + """Forward for the FlashInfer TRTLLM MXFP8 MoE with hidden-dim padding. + + Uses vLLM 0.25.1's modular MoE kernel with the *_for_apply tensors built by + process_weights_after_loading_mxfp8_moe when padding is active, pads x's + hidden dim to mxfp8_padded_hidden_size before the kernel, and narrows the + output back. + """ + assert self.is_monolithic + assert self.moe_kernel is not None + unpadded_hidden_size = x.shape[-1] + padded_hidden_size = getattr( + layer, "mxfp8_padded_hidden_size", unpadded_hidden_size + ) + if unpadded_hidden_size < padded_hidden_size: + x = torch.nn.functional.pad( + x, (0, padded_hidden_size - unpadded_hidden_size), value=0.0 + ) + + output = self.moe_kernel.apply_monolithic( + x, + getattr(layer, "w13_weight_for_apply", layer.w13_weight), + getattr(layer, "w2_weight_for_apply", layer.w2_weight), + router_logits, + activation=layer.activation, + global_num_experts=layer.global_num_experts, + expert_map=layer.expert_map, + apply_router_weight_on_input=layer.apply_router_weight_on_input, + num_expert_group=layer.num_expert_group, + topk_group=layer.topk_group, + e_score_correction_bias=layer.e_score_correction_bias, + routed_scaling_factor=layer.routed_scaling_factor, + ) + if output.shape[-1] != unpadded_hidden_size: + output = output[..., :unpadded_hidden_size].contiguous() + return output def process_weights_after_loading_kv(self, layer) -> None: diff --git a/nemo_rl/models/generation/vllm/quantization/fp8_train_utils.py b/nemo_rl/models/generation/vllm/quantization/fp8_train_utils.py index ac4db666cfe..c4d2df7e290 100644 --- a/nemo_rl/models/generation/vllm/quantization/fp8_train_utils.py +++ b/nemo_rl/models/generation/vllm/quantization/fp8_train_utils.py @@ -13,6 +13,88 @@ # limitations under the License. +import torch + +MXFP8_BLOCK_SIZE = 32 +MXFP8_VALUE_DTYPE = torch.float8_e4m3fn + + +def _mxfp8_e4m3_quantize_torch( + x: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + """Reference MXFP8 quantization with row-major scales. + + Replicates vLLM's _mxfp8_e4m3_quantize_torch (Apache-2.0, + vllm/model_executor/layers/quantization/utils/mxfp8_utils.py) for trainer + processes without a vLLM install: for each block of 32 elements along the + last dimension, a shared e8m0 scale (biased exponent of the block amax) + and float8_e4m3fn values. + """ + assert x.shape[-1] % MXFP8_BLOCK_SIZE == 0, ( + f"MXFP8 requires the last dim to be divisible by {MXFP8_BLOCK_SIZE}, got {x.shape}" + ) + orig_shape = x.shape + num_blocks = x.shape[-1] // MXFP8_BLOCK_SIZE + + x_fp32 = x.to(torch.float32) + x_blocked = x_fp32.view(*orig_shape[:-1], num_blocks, MXFP8_BLOCK_SIZE) + + amax = x_blocked.abs().amax(dim=-1) + amax = amax.clamp(min=torch.finfo(torch.float32).tiny) + scale_biased = torch.floor(torch.log2(amax)) + 127.0 + scale_biased = scale_biased.clamp(0, 254) + scales_uint8 = scale_biased.to(torch.uint8) + + descale = torch.exp2(scale_biased - 127.0) + x_scaled = x_blocked / descale.unsqueeze(-1) + + x_fp8 = x_scaled.view(orig_shape).to(MXFP8_VALUE_DTYPE) + + if x.ndim == 2: + scales_uint8 = scales_uint8.view(x.shape[0], -1) + elif x.ndim == 3: + scales_uint8 = scales_uint8.view(x.shape[0], x.shape[1], -1) + + return x_fp8, scales_uint8 + + +def mxfp8_e4m3_quantize_for_refit( + x: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + """Quantize a weight to MXFP8 on the trainer for pre-quantized refit. + + Mirrors the receiver path in quantization/fp8.py load_weights + (mxfp8_e4m3_quantize + scale reshape) so the streamed E4M3 data and + *_scale_from_checkpoint scales load bit-identically without receiver-side + re-quantization. Uses the same flashinfer kernel as vLLM on Blackwell and + the torch reference elsewhere. + """ + x_q = x_scales = None + if x.is_cuda and torch.cuda.get_device_capability(x.device) >= (10, 0): + try: + from flashinfer import mxfp8_quantize as flashinfer_mxfp8_quantize + except ImportError as exc: + raise RuntimeError( + "Trainer-side MXFP8 refit prequantization on sm100+ requires " + "FlashInfer so it matches the vLLM receiver quantization path." + ) from exc + else: + x_q, x_scales = flashinfer_mxfp8_quantize( + x, is_sf_swizzled_layout=False, alignment=32 + ) + if x_scales.ndim == 1 and x.ndim == 2: + x_scales = x_scales.view(x.size(0), -1) + if x_q is None or x_scales is None: + x_q, x_scales = _mxfp8_e4m3_quantize_torch(x) + x_scales = x_scales.reshape(*x.shape[:-1], x.shape[-1] // MXFP8_BLOCK_SIZE) + # Match the receiver path's zero-scale clamp: an E8M0 byte of 0 (2^-127) + # destabilizes the TRTLLM kernels, and pre-quantized tensors skip the + # receiver-side quantize branch where the clamp normally runs. + # pyrefly: ignore # no-matching-overload + x_scales = x_scales.masked_fill(x_scales == 0, 1) + return x_q, x_scales + + def get_vllm_qkv_scale_names(layer_idx: int) -> dict[str, str]: """Get vLLM-compatible parameter names for Q/K/V FP8 scales. diff --git a/nemo_rl/models/generation/vllm/vllm_backend.py b/nemo_rl/models/generation/vllm/vllm_backend.py index 5ccc315a185..dd4574127d2 100644 --- a/nemo_rl/models/generation/vllm/vllm_backend.py +++ b/nemo_rl/models/generation/vllm/vllm_backend.py @@ -13,11 +13,12 @@ # limitations under the License. import gc import logging +import os import re import socket from collections.abc import Callable, Iterable, Iterator, Sequence from contextlib import contextmanager -from typing import Any, Literal +from typing import Any, Literal, Optional import torch import zmq @@ -27,6 +28,9 @@ preinit_nixl_from_vllm_config, resolve_rollout_rank, ) +from nemo_rl.models.generation.vllm.worker_utils import ( + refit_cache_loader_routes_enabled, +) from nemo_rl.models.policy.utils import ( IPCProtocol, calculate_aligned_size, @@ -201,6 +205,139 @@ def fix_gemma3_vision_weight_name(key: str) -> str: ) +class _RefitLoaderCache: + """Recorded weight_loader calls for refit weight names. + + vLLM's model.load_weights re-resolves every weight name through the + model's stacked/expert parameter mappings on each call; for large MoE + refits that is millions of substring checks per refit over a key set that + is static after the prepare_refit_info handshake. This cache records, per + name, the (loader, param, args, kwargs) of every weight_loader call the + first time a name is loaded and replays them directly afterwards. + """ + + def __init__(self) -> None: + self.calls: dict[str, list[tuple[Any, torch.nn.Parameter, tuple, dict]]] = {} + # Names whose loads never reached a wrapped weight_loader (skipped, + # transformed before dispatch, or default-loaded); these keep going + # through model.load_weights. + self.uncached: set[str] = set() + self.snapshot: dict[str, torch.nn.Parameter] = {} + + def reset(self) -> None: + self.calls.clear() + self.uncached.clear() + self.snapshot.clear() + + +def _cached_params_still_valid(model: Any, cache: _RefitLoaderCache) -> bool: + current = dict(model.named_parameters()) + return all(current.get(name) is param for name, param in cache.snapshot.items()) + + +def _record_loader_calls( + model: Any, cache: _RefitLoaderCache, weights: list[tuple[str, torch.Tensor]] +) -> set[str]: + """Run model.load_weights once while recording every weight_loader call. + + Incoming weights are matched to loader calls by tensor object identity, + so only loads that pass the original tensor through a parameter's + weight_loader attribute are captured; everything else lands in + cache.uncached. Returns the loaded names from model.load_weights. + """ + from vllm.model_executor.model_loader.weight_utils import default_weight_loader + + weight_names = {id(weight): name for name, weight in weights} + recorded: dict[str, list] = {} + originals: list[tuple[torch.nn.Parameter, Any]] = [] + + def make_recorder(loader): + def recorder(param, loaded_weight, *args, **kwargs): + name = weight_names.get(id(loaded_weight)) + if name is not None: + recorded.setdefault(name, []).append((loader, param, args, kwargs)) + return loader(param, loaded_weight, *args, **kwargs) + + return recorder + + try: + for param_name, param in model.named_parameters(): + loader = getattr(param, "weight_loader", None) + # Leave default_weight_loader params unwrapped: model + # load_weights implementations dispatch on + # `weight_loader == default_weight_loader` with a different + # argument list. + if loader is None or loader is default_weight_loader: + continue + cache.snapshot[param_name] = param + originals.append((param, loader)) + param.weight_loader = make_recorder(loader) + loaded = model.load_weights(weights=weights) + finally: + for param, loader in originals: + param.weight_loader = loader + + for name, _ in weights: + calls = recorded.get(name) + if calls is None: + cache.uncached.add(name) + else: + cache.calls[name] = calls + return loaded if loaded is not None else set() + + +def load_weights_maybe_cached( + model: Any, + weights: list[tuple[str, torch.Tensor]], + *, + cache_loader_routes: bool, +) -> set[str]: + """Load weights, optionally replaying cached loader routes. + + Cached parameter identities are re-validated against named_parameters() + on every call, so a process_weights_after_loading pass that replaces + parameter objects drops the cache instead of loading into orphans. + Returns the set of loaded weight names, mirroring model.load_weights. + """ + if not cache_loader_routes: + return model.load_weights(weights=weights) + + cache = getattr(model, "_nrl_refit_loader_cache", None) + if cache is None: + cache = _RefitLoaderCache() + model._nrl_refit_loader_cache = cache + + replay = [] + fallback = [] + record = [] + for name, weight in weights: + if name in cache.calls: + replay.append((name, weight)) + elif name in cache.uncached: + fallback.append((name, weight)) + else: + record.append((name, weight)) + + if replay and not _cached_params_still_valid(model, cache): + cache.reset() + return model.load_weights(weights=weights) + + loaded: set[str] = set() + for name, weight in replay: + for loader, param, args, kwargs in cache.calls[name]: + # Expert loaders return False for non-local shards; a name only + # counts as loaded when some call does not report failure. + if loader(param, weight, *args, **kwargs) is not False: + loaded.add(name) + if record: + loaded |= _record_loader_calls(model, cache, record) + if fallback: + fallback_loaded = model.load_weights(weights=fallback) + if fallback_loaded is not None: + loaded |= fallback_loaded + return loaded + + def _read_mtp_layer_weights_from_checkpoint( model_path: str, mtp_layer_indices: set[int] ) -> list[tuple[str, torch.Tensor]]: @@ -219,7 +356,6 @@ def _read_mtp_layer_weights_from_checkpoint( tensors on CPU. """ import json - import os from safetensors import safe_open @@ -267,7 +403,13 @@ def _load_full_hf_weights( ) -> None: """Load HF weights and detach any deferred reload tensors from transport storage.""" if not getattr(self, "_nrl_layerwise_reload_active", False): - self.model_runner.model.load_weights(weights=policy_weights) + load_weights_maybe_cached( + self.model_runner.model, + policy_weights, + cache_loader_routes=refit_cache_loader_routes_enabled( + self.model_runner.vllm_config + ), + ) return source_storage_ptrs = { @@ -412,13 +554,24 @@ def maybe_init_zmq(self): self.zmq_socket.setsockopt(zmq.LINGER, 0) self.zmq_socket.connect(self.get_zmq_address()) - def prepare_refit_info(self, state_dict_info: dict[str, Any]) -> None: + def prepare_refit_info( + self, + state_dict_info: dict[str, Any], + serialized_fp8_config: Optional[dict[str, Any]] = None, + ) -> Optional[list[str]]: """Prepare state dict metadata for weight refitting and IPC streaming. Args: state_dict_info (dict): A dictionary containing the info for refit. e.g. {tensor_name: (shape, dtype)} + Returns: + When MXFP8 trainer-side pre-quantization is enabled + (vllm_cfg.refit_prequantize), the list of parameter names this + worker will quantize at load time; the trainer quantizes exactly + these and streams E4M3 data plus *_scale_from_checkpoint scales. + None otherwise. + Raises: RuntimeError: If the model realizes the unquantized FlashInfer TRTLLM MoE backend while a co-trained MTP drafter is enabled (unsupported @@ -427,6 +580,29 @@ def prepare_refit_info(self, state_dict_info: dict[str, Any]) -> None: self._validate_native_layerwise_refit() self.state_dict_info = state_dict_info # pyrefly: ignore[implicitly-defined-attribute] This class does not define __init__ so assignments like this should be ignored + # Non-FP8 runs serialize no FP8 config; skip the fp8 module (and its + # heavyweight vLLM imports) entirely so quant backends stubbed without + # the full vLLM surface can still prepare refit info. + if serialized_fp8_config is None: + return None + + from nemo_rl.models.generation.vllm.quantization import fp8 + + fp8.install_fp8_config(serialized_fp8_config) + fp8.set_refit_manifest_names(set(state_dict_info)) + if not ( + fp8.global_fp8_config is not None + and fp8.global_fp8_config.is_mx + and fp8.global_fp8_config.refit_prequantize + and fp8.is_fp8_model(self.model_runner.vllm_config) + ): + return None + return [ + name + for name in state_dict_info + if fp8._is_fp8_weight(name, self.model_runner.model) + ] + def prepare_sparse_delta_refit_info( self, state_dict_info: dict[str, tuple[tuple[int, ...], torch.dtype]] ) -> list[str]: @@ -442,28 +618,6 @@ def _uses_fp8_kv_cache(self) -> bool: kv_cache_dtype = getattr(cache_config, "cache_dtype", None) return kv_cache_dtype is not None and "fp8" in str(kv_cache_dtype).lower() - def _maybe_process_fp8_kv_cache(self) -> None: - """Process weights after loading for FP8 KV cache (static scales).""" - if not self._uses_fp8_kv_cache(): - return - - # FP8 KV cache: process KV scales after weight loading - from vllm.config import set_current_vllm_config - from vllm.model_executor.model_loader.utils import ( - process_weights_after_loading, - ) - - # Get target device for processing - target_device = next(self.model_runner.model.parameters()).device - - # Call process_weights_after_loading to handle KV scales - with set_current_vllm_config(self.model_runner.vllm_config): - process_weights_after_loading( - self.model_runner.model, - self.model_runner.model_config, - target_device, - ) - @staticmethod def _split_policy_and_draft_weights( weights: list[tuple[str, torch.Tensor]], @@ -837,9 +991,8 @@ def finalize() -> None: self._maybe_process_mtp_drafter_after_loading() yield finalize - # Preserve the IPC lifetime boundary: the COMPLETE ACK is sent before - # this optional second pass, just as it was before lifecycle hooks. - self._maybe_process_fp8_kv_cache() + # KV-cache scales are covered by the full process_weights_after_loading + # pass in finalize(); no second pass is needed. def _weight_update_errors_are_fatal(self) -> bool: """Whether transport errors should propagate instead of returning False.""" diff --git a/nemo_rl/models/generation/vllm/vllm_generation.py b/nemo_rl/models/generation/vllm/vllm_generation.py index a8cb78bd13c..27838a0e0ae 100644 --- a/nemo_rl/models/generation/vllm/vllm_generation.py +++ b/nemo_rl/models/generation/vllm/vllm_generation.py @@ -41,7 +41,10 @@ GenerationInterface, GenerationOutputSpec, ) -from nemo_rl.models.generation.vllm.config import VllmConfig +from nemo_rl.models.generation.vllm.config import ( + VllmConfig, + validate_vllm_quantization_config, +) from nemo_rl.models.generation.vllm.utils import ( aggregate_spec_decode_counters, compute_spec_decode_metrics, @@ -109,6 +112,8 @@ def __init__( workers_per_node: Workers per node override defer_model_load: If True, defer model loading for overlapped init """ + validate_vllm_quantization_config(config) + # Store config self.cfg = config self._defer_model_load = defer_model_load @@ -1042,8 +1047,19 @@ def shutdown(self) -> bool: print(f"Error during policy shutdown: {e}") return False - def prepare_refit_info(self, state_dict_info: dict[str, Any]) -> None: - """Prepare the info for refit.""" + def prepare_refit_info( + self, state_dict_info: Optional[dict[str, Any]] + ) -> Optional[list[str]]: + """Prepare the info for refit. + + Returns: + When MXFP8 trainer-side pre-quantization is enabled + (vllm_cfg.refit_prequantize), the parameter names the engine wants + quantized on the trainer before streaming. None otherwise. + """ + if state_dict_info is None: + return None + # Choose the appropriate method based on async_engine setting method_name = ( "prepare_refit_info_async" @@ -1058,8 +1074,13 @@ def prepare_refit_info(self, state_dict_info: dict[str, Any]) -> None: run_rank_0_only_axes=["tensor_parallel", "pipeline_parallel"], ) - # Wait for all futures to complete - ray.get(futures) + # Union the fp8-eligible parameter names across workers: replicas are + # equivalent, but with pipeline parallelism each worker only reports + # the parameters of its local shard. + names = sorted( + {name for result in ray.get(futures) if result for name in result} + ) + return names or None def update_weights_via_ipc_zmq(self) -> list[ray.ObjectRef]: """Update weights of the policy using IPC handles via ZMQ socket.""" diff --git a/nemo_rl/models/generation/vllm/vllm_worker.py b/nemo_rl/models/generation/vllm/vllm_worker.py index 15448d97225..a401b08adfe 100644 --- a/nemo_rl/models/generation/vllm/vllm_worker.py +++ b/nemo_rl/models/generation/vllm/vllm_worker.py @@ -54,6 +54,7 @@ register_torchcodec_vllm_video_loader, ) from nemo_rl.models.generation.vllm.worker_utils import ( + configure_refit_runtime, resolve_data_parallel_local_rank, resolve_distributed_executor_backend, ) @@ -413,6 +414,7 @@ def _load_model(self, bundle_indices, seed): "please run at least once with the environment variable NRL_FORCE_REBUILD_VENVS=true set to force the rebuild of the environment." ) vllm_kwargs: dict[str, Any] = copy.deepcopy(self.cfg.get("vllm_kwargs", {})) + configure_refit_runtime(self.cfg["vllm_cfg"], vllm_kwargs) checkpoint_engine_config = checkpoint_engine_refit_config(self.cfg) if checkpoint_engine_config is not None: from nemo_rl.models.generation.vllm.checkpoint_engine import ( @@ -1195,9 +1197,24 @@ def report_device_id(self) -> list[str]: ) return cast(list[str], list_of_worker_results) - def prepare_refit_info(self, state_dict_info: dict[str, Any]) -> None: - """Prepare the info for refit.""" - self.llm.collective_rpc("prepare_refit_info", args=(state_dict_info,)) + def prepare_refit_info( + self, state_dict_info: dict[str, Any] + ) -> Optional[list[str]]: + """Prepare the info for refit. + + Returns the parameter names the engine wants pre-quantized on the + trainer (vllm_cfg.refit_prequantize), or None. + """ + from nemo_rl.models.generation.vllm.quantization import fp8 + + results = self.llm.collective_rpc( + "prepare_refit_info", + args=(state_dict_info, fp8.serialize_fp8_config()), + ) + # Union across the engine's TP/PP workers: with pipeline parallelism + # each shard only classifies its local parameters as fp8-eligible. + names = sorted({name for result in results if result for name in result}) + return names or None @wrap_with_nvtx_name("vllm_genertion_worker/update_weights_via_ipc_zmq") def update_weights_via_ipc_zmq(self) -> bool: diff --git a/nemo_rl/models/generation/vllm/vllm_worker_async.py b/nemo_rl/models/generation/vllm/vllm_worker_async.py index e791fa4a023..ee14bc47d6d 100644 --- a/nemo_rl/models/generation/vllm/vllm_worker_async.py +++ b/nemo_rl/models/generation/vllm/vllm_worker_async.py @@ -1403,9 +1403,20 @@ async def report_device_id_async(self) -> list[str]: return cast(list[str], list_of_worker_results) - async def prepare_refit_info_async(self, state_dict_info: dict[str, Any]) -> None: + async def prepare_refit_info_async( + self, state_dict_info: dict[str, Any] + ) -> Optional[list[str]]: """Async version of prepare_refit_info.""" - await self.llm.collective_rpc("prepare_refit_info", args=(state_dict_info,)) + from nemo_rl.models.generation.vllm.quantization import fp8 + + results = await self.llm.collective_rpc( + "prepare_refit_info", + args=(state_dict_info, fp8.serialize_fp8_config()), + ) + # Union across the engine's TP/PP workers: with pipeline parallelism + # each shard only classifies its local parameters as fp8-eligible. + names = sorted({name for result in results if result for name in result}) + return names or None async def _reset_encoder_cache_after_weight_update(self) -> None: """Invalidate weight-dependent multimodal encoder outputs when enabled.""" diff --git a/nemo_rl/models/generation/vllm/worker_utils.py b/nemo_rl/models/generation/vllm/worker_utils.py index 831fbde0dd5..2c22170e8fd 100644 --- a/nemo_rl/models/generation/vllm/worker_utils.py +++ b/nemo_rl/models/generation/vllm/worker_utils.py @@ -12,6 +12,33 @@ # See the License for the specific language governing permissions and # limitations under the License. +from collections.abc import Mapping +from typing import Any + +_REFIT_CACHE_LOADER_ROUTES_KEY = "nemo_rl_refit_cache_loader_routes" + + +def configure_refit_runtime( + vllm_cfg: Mapping[str, Any], vllm_kwargs: dict[str, Any] +) -> None: + """Forward NeMo-RL refit options through vLLM's worker config.""" + additional_config = dict(vllm_kwargs.get("additional_config") or {}) + additional_config[_REFIT_CACHE_LOADER_ROUTES_KEY] = vllm_cfg.get( + "refit_cache_loader_routes", False + ) + vllm_kwargs["additional_config"] = additional_config + + +def refit_cache_loader_routes_enabled(vllm_config: Any) -> bool: + """Return the configured loader-route cache setting in a vLLM worker.""" + additional_config = getattr(vllm_config, "additional_config", None) or {} + if _REFIT_CACHE_LOADER_ROUTES_KEY not in additional_config: + return False + value = additional_config[_REFIT_CACHE_LOADER_ROUTES_KEY] + if not isinstance(value, bool): + raise TypeError(f"{_REFIT_CACHE_LOADER_ROUTES_KEY} must be a boolean") + return value + def resolve_distributed_executor_backend( tensor_parallel_size: int, diff --git a/nemo_rl/models/megatron/setup.py b/nemo_rl/models/megatron/setup.py index 399633cea52..a161fdc3aa2 100644 --- a/nemo_rl/models/megatron/setup.py +++ b/nemo_rl/models/megatron/setup.py @@ -2176,6 +2176,10 @@ def composed_peft_hook(model: list[MegatronModule]) -> list[MegatronModule]: ) reference_state_dict = {} + # NotRequired key: absent means disabled, default lives in the exemplar YAML. + pinned_reference_swap = bool( + config["megatron_cfg"].get("pinned_reference_swap") + ) if should_load_checkpoint or use_peft: reference_model = reference_model[0] @@ -2183,13 +2187,23 @@ def composed_peft_hook(model: list[MegatronModule]) -> list[MegatronModule]: # Store reference state dict on CPU for name, item in reference_model.state_dict().items(): if isinstance(item, torch.Tensor): - cpu_item = item.detach().to( - device="cpu", non_blocking=True, copy=True - ) + if pinned_reference_swap: + # Pinned so use_reference_model can upload the reference + # weights with non-blocking H2D copies each step. + cpu_item = torch.empty( + item.shape, dtype=item.dtype, device="cpu", pin_memory=True + ) + cpu_item.copy_(item.detach(), non_blocking=True) + else: + cpu_item = item.detach().to( + device="cpu", non_blocking=True, copy=True + ) del item else: cpu_item = item reference_state_dict[name] = cpu_item + if pinned_reference_swap: + torch.cuda.synchronize() print("Reference model loaded") else: print("Reference model not loaded") diff --git a/nemo_rl/models/policy/__init__.py b/nemo_rl/models/policy/__init__.py index 7d14674a392..245bed08991 100644 --- a/nemo_rl/models/policy/__init__.py +++ b/nemo_rl/models/policy/__init__.py @@ -336,6 +336,17 @@ class MegatronConfig(TypedDict): # 1 is the minimum recommendation for RL since we almost always need to offload before beginning generation. # Setting to 0 is faster, but you are more likely to run out of GPU memory. In SFT/DPO, the default is 0. empty_unused_memory_level: int + # When True, offload_after_refit skips rerunning the full offload_before_refit + # pass (grad-buffer moves, cache clears, and a second gc.collect/empty_cache) + # and only re-offloads the optimizer with a single allocator cleanup. + refit_slim_offload_after: NotRequired[bool] + # When True, the reference-policy swap in use_reference_model stages the + # active weights in persistent pinned CPU buffers and keeps the reference + # state dict pinned from setup, so the three per-step full-model PCIe + # transfers run as non-blocking pinned copies instead of pageable ones. + # Costs 2x model-size pinned host RAM per worker. Default False keeps the + # current pageable-copy behavior. + pinned_reference_swap: NotRequired[bool] activation_checkpointing: bool # Recompute granularity: "full" recomputes all activations, "selective" recomputes # only specific modules (see recompute_modules). "selective" typically saves ~10-18GB @@ -590,6 +601,11 @@ class PolicyConfig(TypedDict): # This sets the clipping norm for the DTensorPolicyWorkers (Megatron's is called clip_grad) max_grad_norm: NotRequired[float | int | None] refit_buffer_size_gb: NotRequired[float | int] + # Keep the CUDA-IPC ping-pong staging buffers allocated across refits instead of + # reallocating and freeing them (plus a gc.collect/empty_cache pair) every refit. + # Works best with a fixed refit_buffer_size_gb; the buffers stay resident on the + # trainer GPU between refits (2x half-buffer bytes). + refit_persistent_ipc_buffers: NotRequired[bool] optimizer: NotRequired[PytorchOptimizerConfig | None] scheduler: NotRequired[ list[SinglePytorchSchedulerConfig | SinglePytorchMilestonesConfig] diff --git a/nemo_rl/models/policy/interfaces.py b/nemo_rl/models/policy/interfaces.py index ea1d64dc1ac..d6d5032f7c1 100644 --- a/nemo_rl/models/policy/interfaces.py +++ b/nemo_rl/models/policy/interfaces.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. from abc import ABC, abstractmethod -from typing import Any, Optional, TypedDict +from typing import TYPE_CHECKING, Any, Optional, TypedDict import ray import torch @@ -22,6 +22,9 @@ from nemo_rl.models.generation.interfaces import GenerationDatumSpec from nemo_rl.utils.timer import Timer +if TYPE_CHECKING: + from nemo_rl.models.policy import PolicyConfig + class LogprobOutputSpec(TypedDict): """logprobs: Tensor of log probabilities.""" @@ -164,6 +167,8 @@ def shutdown(self) -> bool: class ColocatablePolicyInterface(PolicyInterface): + cfg: "PolicyConfig" + @abstractmethod def init_collective( self, @@ -191,6 +196,18 @@ def offload_to_cpu(self) -> None: def prepare_refit_info(self) -> Optional[dict[str, Any]]: pass + def enable_refit_prequantize( + self, param_names: list[str] + ) -> Optional[dict[str, Any]]: + """Quantize the listed params on the trainer during refit streaming. + + Returns: + Refit info updated with the quantized dtypes and scale entries. + """ + raise NotImplementedError( + "enable_refit_prequantize is not implemented for this policy worker" + ) + @abstractmethod def stream_weights_via_ipc_zmq( self, diff --git a/nemo_rl/models/policy/lm_policy.py b/nemo_rl/models/policy/lm_policy.py index b07b30a7f62..2791370dc15 100644 --- a/nemo_rl/models/policy/lm_policy.py +++ b/nemo_rl/models/policy/lm_policy.py @@ -1001,6 +1001,20 @@ def prepare_refit_info(self) -> Optional[dict[str, Any]]: # Only get the first worker's info since all workers will have the same result return results[0] + def enable_refit_prequantize( + self, param_names: list[str] + ) -> Optional[dict[str, Any]]: + """Enable trainer-side MXFP8 quantization of the listed params for refit. + + Returns: + dict: Refit info updated with quantized dtypes and scale entries. + """ + futures = self.worker_group.run_all_workers_single_data( + "enable_refit_prequantize", param_names=param_names + ) + results = ray.get(futures) + return results[0] + def finish_inference(self) -> None: """Offload policy model to CPU after inference.""" futures = self.worker_group.run_all_workers_single_data("finish_inference") diff --git a/nemo_rl/models/policy/utils.py b/nemo_rl/models/policy/utils.py index c2f61bf6b03..0507ffac1a8 100644 --- a/nemo_rl/models/policy/utils.py +++ b/nemo_rl/models/policy/utils.py @@ -360,7 +360,12 @@ def calculate_aligned_size(size_bytes: int, alignment: int = 512) -> int: def stream_weights_via_ipc_zmq_impl( - params_generator, buffer_size_bytes: int, zmq_socket, rank: int, worker_name: str + params_generator, + buffer_size_bytes: int, + zmq_socket, + rank: int, + worker_name: str, + buffer_cache: dict[str, Any] | None = None, ) -> None: """Shared implementation for streaming weights via IPC ZMQ with improved memory management. @@ -373,6 +378,11 @@ def stream_weights_via_ipc_zmq_impl( zmq_socket: ZMQ socket for communication rank: Worker rank for logging worker_name: Name of the worker for logging + buffer_cache: Optional caller-owned dict that keeps the ping-pong staging + buffers alive across calls. When provided, buffers are reused between + refits (skipping two large allocations plus a gc.collect/empty_cache + pair per refit) as long as the buffer size and device are unchanged. + The buffers then stay resident on the GPU between refits. """ # Divide total buffer size by 2 because we use two individual buffers (ping-pong) for overlapping communication. buffer_size_bytes = buffer_size_bytes // 2 @@ -420,6 +430,11 @@ def release_staging_buffers() -> None: """Release acyclic IPC buffers without scanning the worker object graph.""" nonlocal buffer_a, buffer_b, current_buffer + if buffer_cache is not None: + # Persistent-buffer mode: the buffers intentionally outlive this + # call (and the refit), so both the mid-stream reclaim and the + # final cleanup are skipped. + return had_buffers = buffer_a is not None or buffer_b is not None current_buffer = None buffer_a = None @@ -439,8 +454,30 @@ def release_staging_buffers() -> None: buffer_device = tensor.device if buffer_device.type == "cpu" and torch.cuda.is_available(): buffer_device = torch.device("cuda", torch.cuda.current_device()) - buffer_a = allocate_buffer(buffer_device) - buffer_b = allocate_buffer(buffer_device) + if ( + buffer_cache is not None + and "a" in buffer_cache + and buffer_cache.get("device") == buffer_device + ): + # Reuse the first-refit size even when dynamic sizing asks + # for more: parameters are identical every refit, and with + # vLLM sleep level 2 the free-memory-based request inflates + # while rollout weights are discarded - allocating (and + # keeping) larger buffers then OOMs vLLM's wake_up remap. + buffer_a = buffer_cache["a"] + buffer_b = buffer_cache["b"] + buffer_size_bytes = buffer_cache["size"] + else: + buffer_a = allocate_buffer(buffer_device) + buffer_b = allocate_buffer(buffer_device) + if buffer_cache is not None: + buffer_cache.clear() + buffer_cache.update( + a=buffer_a, + b=buffer_b, + size=buffer_size_bytes, + device=buffer_device, + ) current_buffer = buffer_a aligned_size = calculate_aligned_size(tensor.nbytes) diff --git a/nemo_rl/models/policy/workers/dtensor_policy_worker.py b/nemo_rl/models/policy/workers/dtensor_policy_worker.py index 89e5c95ee28..e44402fef94 100644 --- a/nemo_rl/models/policy/workers/dtensor_policy_worker.py +++ b/nemo_rl/models/policy/workers/dtensor_policy_worker.py @@ -258,6 +258,9 @@ def __init__( configure_dynamo_cache() self.cfg = config + # Staging-buffer cache for refit weight streaming; only populated when + # cfg["refit_persistent_ipc_buffers"] is enabled. + self._refit_ipc_buffer_cache: dict[str, Any] = {} # torch distributed init. Envars for rank, world_size, and master_addr and master_port are set from the ray remote call torch.distributed.init_process_group(backend="nccl") self.rank = torch.distributed.get_rank() @@ -1891,6 +1894,11 @@ def stream_weights_via_ipc_zmq( zmq_socket=self.zmq_socket, rank=self.rank, worker_name=str(self), + buffer_cache=( + self._refit_ipc_buffer_cache + if self.cfg.get("refit_persistent_ipc_buffers") + else None + ), ) def _checkpoint_engine_params( diff --git a/nemo_rl/models/policy/workers/dtensor_policy_worker_v2.py b/nemo_rl/models/policy/workers/dtensor_policy_worker_v2.py index 9f38da334ae..41c16e54466 100644 --- a/nemo_rl/models/policy/workers/dtensor_policy_worker_v2.py +++ b/nemo_rl/models/policy/workers/dtensor_policy_worker_v2.py @@ -265,6 +265,9 @@ def __init__( # Store configuration self.cfg = config + # Staging-buffer cache for refit weight streaming; only populated when + # cfg["refit_persistent_ipc_buffers"] is enabled. + self._refit_ipc_buffer_cache: dict[str, Any] = {} # Reconstruct tokenizer/processor locally to avoid pickling across # incompatible transformers versions (v4 head node ā†’ v5 worker). @@ -1144,6 +1147,11 @@ def stream_weights_via_ipc_zmq( zmq_socket=self.zmq_socket, rank=self.rank, worker_name=str(self), + buffer_cache=( + self._refit_ipc_buffer_cache + if self.cfg.get("refit_persistent_ipc_buffers") + else None + ), ) @torch.no_grad() diff --git a/nemo_rl/models/policy/workers/megatron_policy_worker.py b/nemo_rl/models/policy/workers/megatron_policy_worker.py index 6f0bef2e237..3bfe4929914 100644 --- a/nemo_rl/models/policy/workers/megatron_policy_worker.py +++ b/nemo_rl/models/policy/workers/megatron_policy_worker.py @@ -462,6 +462,25 @@ def __init__( self.cfg = config self._router_replay_enabled = router_replay_enabled(config) + # Staging-buffer cache for refit weight streaming; only populated when + # cfg["refit_persistent_ipc_buffers"] is enabled. + self._refit_ipc_buffer_cache: dict[str, Any] = {} + # HF param names to MXFP8-quantize on the trainer during refit; set via + # enable_refit_prequantize() when vllm_cfg.refit_prequantize is on. + self._refit_prequant_names: set[str] = set() + # HF param metadata cached by prepare_refit_info so that + # enable_refit_prequantize can derive updated metadata without + # re-running the export + quantize pass. + self._refit_param_info_hf: Optional[ + dict[str, tuple[torch.Size, torch.dtype]] + ] = None + # Pinned host staging for the reference-policy swap; only populated when + # megatron_cfg["pinned_reference_swap"] is enabled. Buffer contents are + # only live within a single use_reference_model call (every copy + # synchronizes before control leaves it) and each call re-reads + # model.state_dict(), so offload_before/after_refit moving or replacing + # param storages between calls cannot collide with these buffers. + self._pinned_swap_save_buffers: dict[str, torch.Tensor] = {} self._nixl_preinit_agent = maybe_preinit_nixl_checkpoint_engine(config) # Set rank for non-collocated to check which ranks to broadcast from @@ -1948,6 +1967,7 @@ def _apply_state_dict_to_model( source_state_dict: dict, *, raise_if_key_missing: bool = False, + non_blocking: bool = False, ) -> None: """Apply a state dict to self.model in-place. @@ -1959,6 +1979,8 @@ def _apply_state_dict_to_model( source_state_dict: State dict to apply (e.g. reference_state_dict or saved model_state_dict). raise_if_key_missing: If True, raise when a key in self.model.state_dict() is missing from source_state_dict; if False, skip such keys. + non_blocking: Passed to the in-place copies; callers staging from + pinned CPU memory must synchronize afterwards. """ for state_dict_key, param_or_buf in self.model.state_dict().items(): if ( @@ -1979,7 +2001,7 @@ def _apply_state_dict_to_model( isinstance(source_value, torch.Tensor) and param_or_buf.shape == source_value.shape ): - param_or_buf.copy_(source_value) + param_or_buf.copy_(source_value, non_blocking=non_blocking) continue # Case 2: _extra_state (shape mismatch or non-Tensor) ā†’ set_extra_state() @@ -2010,12 +2032,39 @@ def use_reference_model(self): self.disable_forward_pre_hook() with torch.no_grad(): + # NotRequired key: absent means disabled, default lives in the exemplar YAML. + use_pinned_swap = bool( + self.cfg["megatron_cfg"].get("pinned_reference_swap") + ) + # Save original references model_state_dict = {} for name, item in self.model.state_dict().items(): if isinstance(item, torch.Tensor): - item = item.detach().to(device="cpu", non_blocking=True, copy=True) + # extra_state tensors stay on fresh pageable copies: + # set_extra_state() may retain the tensor it is given, and + # a reused pinned buffer would mutate it on the next swap. + if use_pinned_swap and "extra_state" not in name: + buf = self._pinned_swap_save_buffers.get(name) + if buf is None: + buf = torch.empty( + item.shape, + dtype=item.dtype, + device="cpu", + pin_memory=True, + ) + self._pinned_swap_save_buffers[name] = buf + buf.copy_(item.detach(), non_blocking=True) + item = buf + else: + item = item.detach().to( + device="cpu", non_blocking=True, copy=True + ) model_state_dict[name] = item + if use_pinned_swap: + # D2H saves must land before the reference apply overwrites + # the params they read from. + torch.cuda.synchronize() # Swap reference state into self.model. Use _apply_state_dict_to_model # (rather than load_state_dict) so FP8 _extra_state with mismatched shape @@ -2023,7 +2072,10 @@ def use_reference_model(self): self._apply_state_dict_to_model( self.reference_state_dict, raise_if_key_missing=True, + non_blocking=use_pinned_swap, ) + if use_pinned_swap: + torch.cuda.synchronize() if self.cfg["megatron_cfg"]["empty_unused_memory_level"] >= 1: gc.collect() @@ -2055,7 +2107,10 @@ def use_reference_model(self): self._apply_state_dict_to_model( model_state_dict, raise_if_key_missing=True, + non_blocking=use_pinned_swap, ) + if use_pinned_swap: + torch.cuda.synchronize() if self.cfg["megatron_cfg"]["empty_unused_memory_level"] >= 1: gc.collect() @@ -2162,7 +2217,7 @@ def get_topk_logits( @torch.no_grad() @wrap_with_nvtx_name("megatron_policy_worker/prepare_refit_info") - def prepare_refit_info(self) -> None: + def prepare_refit_info(self) -> Optional[dict[str, Any]]: """Prepare state dict metadata for weight refitting and IPC streaming.""" self.refit_param_info_mcore = self._calculate_refit_param_info() @@ -2171,8 +2226,87 @@ def prepare_refit_info(self) -> None: for name, tensor in self._iter_params_with_optional_kv_scales(): refit_param_info_hf[name] = (tensor.shape, tensor.dtype) + self._refit_param_info_hf = refit_param_info_hf + return refit_param_info_hf + + def enable_refit_prequantize(self, param_names: list[str]) -> dict[str, Any]: + """Quantize the listed HF params to MXFP8 on the trainer during refit. + + Derives the updated metadata from the shapes cached by + prepare_refit_info instead of re-running the export + quantize pass: + MXFP8 keeps the value shape and adds one uint8 E8M0 scale per 32-wide + block along the last dim, so no tensor needs to be materialized here. + + Args: + param_names: fp8-eligible parameter names reported by the vLLM + workers (see VllmInternalWorkerExtension.prepare_refit_info). + + Returns: + Updated refit metadata: the listed params become float8_e4m3fn and + each gains a *_scale_from_checkpoint uint8 entry. + """ + if self._is_fp8_export(): + raise ValueError( + "vllm_cfg.refit_prequantize requires BF16 trainer-exported weights; " + "Megatron blockwise FP8 parameter storage uses a different scale layout." + ) + if self._refit_param_info_hf is None: + raise RuntimeError( + "enable_refit_prequantize requires prepare_refit_info to have " + "run first so the HF parameter metadata is available." + ) + + from nemo_rl.models.generation.vllm.quantization.fp8_train_utils import ( + MXFP8_BLOCK_SIZE, + ) + + self._refit_prequant_names = set(param_names) + + refit_param_info_hf = {} + for name, (shape, dtype) in self._refit_param_info_hf.items(): + if name not in self._refit_prequant_names: + refit_param_info_hf[name] = (shape, dtype) + continue + if dtype == torch.float8_e4m3fn: + raise ValueError( + "vllm_cfg.refit_prequantize requires BF16 trainer-exported weights; " + f"{name} is already stored as E4M3 with a non-MXFP8 scale layout." + ) + if shape[-1] % MXFP8_BLOCK_SIZE != 0: + raise ValueError( + f"MXFP8 requires the last dim to be divisible by " + f"{MXFP8_BLOCK_SIZE}; {name} has shape {tuple(shape)}." + ) + scale_shape = torch.Size((*shape[:-1], shape[-1] // MXFP8_BLOCK_SIZE)) + refit_param_info_hf[name] = (shape, torch.float8_e4m3fn) + refit_param_info_hf[name + "_scale_from_checkpoint"] = ( + scale_shape, + torch.uint8, + ) return refit_param_info_hf + def _maybe_prequantize_param( + self, name: str, tensor: torch.Tensor + ) -> Iterator[tuple[str, torch.Tensor]]: + if name not in self._refit_prequant_names: + yield name, tensor + return + if tensor.dtype == torch.float8_e4m3fn: + raise ValueError( + "vllm_cfg.refit_prequantize requires BF16 trainer-exported weights; " + f"{name} is already stored as E4M3 with a non-MXFP8 scale layout." + ) + + # Deferred: pulls in the heavy nemo_rl...generation.vllm package init, + # which trainer workers only need when prequantized refit is enabled. + from nemo_rl.models.generation.vllm.quantization.fp8_train_utils import ( + mxfp8_e4m3_quantize_for_refit, + ) + + param_lp, param_scale = mxfp8_e4m3_quantize_for_refit(tensor) + yield name, param_lp + yield name + "_scale_from_checkpoint", param_scale + def _collect_mtp_metrics( self, metrics: dict[str, Any], @@ -2387,9 +2521,10 @@ def _iter_params_with_optional_kv_scales( conversion_tasks=conversion_tasks, # used for metadata caching ) - # Yield the original parameters first. + # Yield the original parameters first, MXFP8-quantizing on the trainer + # when pre-quantized refit is enabled for the parameter. for name, tensor in base_iter: - yield name, tensor + yield from self._maybe_prequantize_param(name, tensor) if include_draft and self.draft_model is not None: from nemo_rl.models.megatron.draft import export_eagle_weights_to_hf @@ -2687,6 +2822,11 @@ def stream_weights_via_ipc_zmq( zmq_socket=self.zmq_socket, rank=self.rank, worker_name=str(self), + buffer_cache=( + self._refit_ipc_buffer_cache + if self.cfg.get("refit_persistent_ipc_buffers") + else None + ), ) @torch.no_grad() @@ -3039,7 +3179,8 @@ def nccl_reshard_refit(self, kv_scales=None): ``kv_scales`` (FP8 KV cache): the per-layer k/v(/q) scales ride the misc packed-broadcast as plain scale tensors (the is_nccl_reshard_param whitelist excludes ``.k_scale``/``.v_scale``/``.q_scale`` -> misc); the gen side finalizes - them via ``_maybe_process_fp8_kv_cache``. No out-of-band channel needed. + them via the ``process_weights_after_loading`` pass in the weight-update + finalizer. No out-of-band channel needed. """ # hf_to_local_param_map is built once in prepare_nccl_reshard_refit_info; # weight values change but the name ā†’ spec mapping is stable across @@ -3294,56 +3435,7 @@ def offload_before_refit(self): self._clear_fp8_caches() if self.cfg["megatron_cfg"].get("clear_memory_caches_before_refit", False): - # Clear RotaryEmbedding's @lru_cache(maxsize=32). The cache accumulates one - # entry per unique (max_seq_len, offset, packed_seq) seen, and each entry is - # a GPU tensor (the concatenated sin/cos embedding). With training + logprob - # runs at different sequence lengths, the cache fills quickly and the tensors - # anchor large CUDA segments. - try: - from megatron.core.models.common.embeddings.rotary_pos_embedding import ( - RotaryEmbedding, - ) - - RotaryEmbedding.forward.cache_clear() - except Exception: - pass - - # Clear MoE token dispatcher persistent routing tensors. - # - # MoETokenDispatcher is a plain Python class (NOT an nn.Module), so iterating - # self.model.modules() never yields it. We must access it via the token_dispatcher - # attribute on MoELayer nn.Module objects. - # - # When recompute_mlp=True and fp8=True, - # transformer_layer._forward_mlp wraps self.mlp (the MoE layer) with te_checkpoint. - # te_checkpoint._CheckpointFunction.backward recomputes the forward with - # torch.enable_grad(), which causes dispatch_preprocess to store - # dispatcher.probs = routing_probs (with grad_fn, under enable_grad) - # This creates a reference cycle: - # _CheckpointFunctionBackward ā†’ ctx ā†’ ctx.run_function=mlp - # ā†’ mlp.token_dispatcher.probs ā†’ probs.grad_fn ā†’ ... ā†’ _CheckpointFunctionBackward - # - # Breaking this cycle by nulling dispatcher.probs frees BOTH: - # - the routing tensors - # - the te_checkpoint ctx saved tensors - try: - for module in self.model.modules(): - if not hasattr(module, "token_dispatcher"): - continue - dispatcher = module.token_dispatcher - if dispatcher is None: - continue - for attr in ( - "probs", # AllToAll + AllGather - "routing_map", # AllToAll - "reversed_local_input_permutation_mapping", # AllToAll - "local_probs", # AllGather - "local_map", # AllGather - ): - if isinstance(getattr(dispatcher, attr, None), torch.Tensor): - setattr(dispatcher, attr, None) - except Exception: - pass + self._clear_rope_and_moe_dispatcher_caches() torch.randn(1).cuda() # wake up torch allocator if ( @@ -3365,6 +3457,59 @@ def offload_before_refit(self): ) no_grad.__exit__(None, None, None) + def _clear_rope_and_moe_dispatcher_caches(self) -> None: + """Clear rotary-embedding and MoE dispatcher caches repopulated by forwards.""" + # Clear RotaryEmbedding's @lru_cache(maxsize=32). The cache accumulates one + # entry per unique (max_seq_len, offset, packed_seq) seen, and each entry is + # a GPU tensor (the concatenated sin/cos embedding). With training + logprob + # runs at different sequence lengths, the cache fills quickly and the tensors + # anchor large CUDA segments. + try: + from megatron.core.models.common.embeddings.rotary_pos_embedding import ( + RotaryEmbedding, + ) + + RotaryEmbedding.forward.cache_clear() + except Exception: + pass + + # Clear MoE token dispatcher persistent routing tensors. + # + # MoETokenDispatcher is a plain Python class (NOT an nn.Module), so iterating + # self.model.modules() never yields it. We must access it via the token_dispatcher + # attribute on MoELayer nn.Module objects. + # + # When recompute_mlp=True and fp8=True, + # transformer_layer._forward_mlp wraps self.mlp (the MoE layer) with te_checkpoint. + # te_checkpoint._CheckpointFunction.backward recomputes the forward with + # torch.enable_grad(), which causes dispatch_preprocess to store + # dispatcher.probs = routing_probs (with grad_fn, under enable_grad) + # This creates a reference cycle: + # _CheckpointFunctionBackward ā†’ ctx ā†’ ctx.run_function=mlp + # ā†’ mlp.token_dispatcher.probs ā†’ probs.grad_fn ā†’ ... ā†’ _CheckpointFunctionBackward + # + # Breaking this cycle by nulling dispatcher.probs frees BOTH: + # - the routing tensors + # - the te_checkpoint ctx saved tensors + try: + for module in self.model.modules(): + if not hasattr(module, "token_dispatcher"): + continue + dispatcher = module.token_dispatcher + if dispatcher is None: + continue + for attr in ( + "probs", # AllToAll + AllGather + "routing_map", # AllToAll + "reversed_local_input_permutation_mapping", # AllToAll + "local_probs", # AllGather + "local_map", # AllGather + ): + if isinstance(getattr(dispatcher, attr, None), torch.Tensor): + setattr(dispatcher, attr, None) + except Exception: + pass + @wrap_with_nvtx_name("megatron_policy_worker/offload_after_refit") def offload_after_refit(self): """Offload as much as possible on the CPU.""" @@ -3389,7 +3534,28 @@ def offload_after_refit(self): ) self.model.eval() torch.randn(1).cuda() # wake up torch allocator - self.offload_before_refit() # rerun the old offload function + if self.cfg.get("megatron_cfg", {}).get("refit_slim_offload_after"): + # Grad buffers were already offloaded by offload_before_refit at + # the start of the refit, so skip the full rerun (grad-buffer moves + # and a second gc/empty_cache pair). Cache clears must still honor + # their knobs: callers may have run a forward pass since (e.g. + # teacher logits in distillation), repopulating TE fp8 workspaces, + # the rotary-embedding lru_cache, and MoE dispatcher tensors. + if self.fp8_cfg and self.fp8_cfg.get("force_clear_fp8_caches", False): + self._clear_fp8_caches() + if self.cfg["megatron_cfg"].get("clear_memory_caches_before_refit", False): + self._clear_rope_and_moe_dispatcher_caches() + if ( + hasattr(self, "optimizer") + and self.optimizer is not None + and not self.optimizer_cpu_offload + and self.offload_optimizer_for_refit + ): + self.move_optimizer("cpu") + gc.collect() + torch.cuda.empty_cache() + else: + self.offload_before_refit() # rerun the old offload function allocated = torch.cuda.memory_allocated() / (1024**3) # Convert to GB reserved = torch.cuda.memory_reserved() / (1024**3) # Convert to GB diff --git a/nemo_rl/utils/packed_tensor.py b/nemo_rl/utils/packed_tensor.py index baa9b65af92..d742d045c4d 100644 --- a/nemo_rl/utils/packed_tensor.py +++ b/nemo_rl/utils/packed_tensor.py @@ -59,6 +59,8 @@ def packed_broadcast_producer( None """ + if buffer_size_bytes is not None and buffer_size_bytes <= 0: + raise ValueError("buffer_size_bytes must be > 0") target_packed_tensor_size = ( get_target_packed_tensor_size() if buffer_size_bytes is None @@ -127,7 +129,12 @@ def packed_broadcast_producer( def packed_broadcast_consumer( - iterator, group, src, post_unpack_func, *, num_buffers: int | None = None + iterator, + group, + src, + post_unpack_func, + *, + num_buffers: int | None = None, ): """Consume a packed tensor and unpack it into a list of tensors. diff --git a/nemo_rl/weight_sync/checkpoint_engine_weight_synchronizer.py b/nemo_rl/weight_sync/checkpoint_engine_weight_synchronizer.py index 946547adbd6..8b97c008a3a 100644 --- a/nemo_rl/weight_sync/checkpoint_engine_weight_synchronizer.py +++ b/nemo_rl/weight_sync/checkpoint_engine_weight_synchronizer.py @@ -20,7 +20,10 @@ from nemo_rl.models.generation.interfaces import CheckpointEngineConfig from nemo_rl.utils.timer import Timer -from nemo_rl.weight_sync.interfaces import WeightSynchronizer +from nemo_rl.weight_sync.interfaces import ( + WeightSynchronizer, + initialize_refit_metadata, +) _MEBIBYTE = 1024 * 1024 @@ -67,7 +70,7 @@ class CheckpointEngineWeightSynchronizer(WeightSynchronizer): _bucket_size_bytes: int | None = None def init_communicator(self) -> None: - self._generation.prepare_refit_info(self._policy.prepare_refit_info()) + initialize_refit_metadata(self._policy, self._generation) self._ensure_checkpoint_engine_ready() @property diff --git a/nemo_rl/weight_sync/collective_weight_synchronizer.py b/nemo_rl/weight_sync/collective_weight_synchronizer.py index ff730d076e9..a77d7eed8aa 100644 --- a/nemo_rl/weight_sync/collective_weight_synchronizer.py +++ b/nemo_rl/weight_sync/collective_weight_synchronizer.py @@ -34,7 +34,10 @@ import ray from nemo_rl.utils.timer import Timer -from nemo_rl.weight_sync.interfaces import WeightSynchronizer +from nemo_rl.weight_sync.interfaces import ( + WeightSynchronizer, + initialize_refit_metadata, +) class CollectiveWeightSynchronizer(WeightSynchronizer): @@ -105,8 +108,7 @@ def init_communicator(self) -> None: # prepare_refit_info is called before init_collective. This matches # distillation.py ordering. Neither call depends on the other today, # but we document this as the canonical ordering for future reference. - state_dict_info = self._policy.prepare_refit_info() - self._generation.prepare_refit_info(state_dict_info) + initialize_refit_metadata(self._policy, self._generation) ip, port = self._train_cluster.get_master_address_and_port() train_world_size = self._train_cluster.world_size() diff --git a/nemo_rl/weight_sync/interfaces.py b/nemo_rl/weight_sync/interfaces.py index e2314d9170c..eb308ecdf2a 100644 --- a/nemo_rl/weight_sync/interfaces.py +++ b/nemo_rl/weight_sync/interfaces.py @@ -39,10 +39,39 @@ """ from abc import ABC, abstractmethod -from typing import Optional +from typing import TYPE_CHECKING, Optional from nemo_rl.utils.timer import Timer +if TYPE_CHECKING: + from nemo_rl.models.generation.interfaces import GenerationInterface + from nemo_rl.models.policy.interfaces import ColocatablePolicyInterface + + +def initialize_refit_metadata( + policy: "ColocatablePolicyInterface", generation: "GenerationInterface" +) -> None: + """Negotiate the wire-format metadata used by policy-to-generation refit.""" + state_dict_info = policy.prepare_refit_info() + prequant_names = generation.prepare_refit_info(state_dict_info) + if not prequant_names: + return + + megatron_cfg = policy.cfg.get("megatron_cfg") + if megatron_cfg is None or not megatron_cfg["enabled"]: + raise ValueError( + "vllm_cfg.refit_prequantize requires the Megatron policy backend " + "(policy.megatron_cfg.enabled=true); the DTensor workers do not " + "implement trainer-side pre-quantized refit." + ) + + updated_info = policy.enable_refit_prequantize(prequant_names) + if updated_info is None: + raise RuntimeError( + "Trainer-side refit prequantization did not return updated metadata." + ) + generation.prepare_refit_info(updated_info) + class WeightSynchronizer(ABC): """Abstract base class for weight synchronization between policy and generation. diff --git a/nemo_rl/weight_sync/ipc_weight_synchronizer.py b/nemo_rl/weight_sync/ipc_weight_synchronizer.py index b23f4906268..728a7613442 100644 --- a/nemo_rl/weight_sync/ipc_weight_synchronizer.py +++ b/nemo_rl/weight_sync/ipc_weight_synchronizer.py @@ -34,7 +34,10 @@ import ray from nemo_rl.utils.timer import Timer -from nemo_rl.weight_sync.interfaces import WeightSynchronizer +from nemo_rl.weight_sync.interfaces import ( + WeightSynchronizer, + initialize_refit_metadata, +) class IPCWeightSynchronizer(WeightSynchronizer): @@ -109,8 +112,7 @@ def is_stale(self) -> bool: return self._stale def init_communicator(self) -> None: - state_dict_info = self._policy.prepare_refit_info() - self._generation.prepare_refit_info(state_dict_info) + initialize_refit_metadata(self._policy, self._generation) def shutdown(self) -> None: pass diff --git a/pyrefly.toml b/pyrefly.toml index 7e4b0e6a0e7..a871ae53bdc 100644 --- a/pyrefly.toml +++ b/pyrefly.toml @@ -8,6 +8,7 @@ replace-imports-with-any = [ "transformers.*", "tensorrt_llm.*", "vllm.*", + "flashinfer.*", "dynamo.*", "math_verify.*", "sympy.*", diff --git a/tests/functional/L1_Functional_Tests_GB200_MXFP8.sh b/tests/functional/L1_Functional_Tests_GB200_MXFP8.sh index 725bd279df5..a868da7d2fd 100644 --- a/tests/functional/L1_Functional_Tests_GB200_MXFP8.sh +++ b/tests/functional/L1_Functional_Tests_GB200_MXFP8.sh @@ -34,6 +34,8 @@ run_test() { fi } +run_test fast uv run --no-sync pytest -q \ + tests/unit/models/generation/test_mxfp8_prequant.py run_test uv run --no-sync bash ./tests/functional/grpo_vllm_mxfp8_rollout_gb200.sh run_test uv run --no-sync bash ./tests/functional/grpo_megatron_mxfp8_refit_gb200.sh diff --git a/tests/functional/grpo_vllm_mxfp8_rollout_gb200.sh b/tests/functional/grpo_vllm_mxfp8_rollout_gb200.sh index 677853dc6cd..76217aceff6 100644 --- a/tests/functional/grpo_vllm_mxfp8_rollout_gb200.sh +++ b/tests/functional/grpo_vllm_mxfp8_rollout_gb200.sh @@ -13,6 +13,7 @@ LOG_DIR=$EXP_DIR/logs JSON_METRICS=$EXP_DIR/metrics.json RUN_LOG=$EXP_DIR/run.log export PYTHONPATH=${PROJECT_ROOT}:${PYTHONPATH:-} +export NRL_MXFP8_SHUFFLE_VERIFY=1 rm -rf "$EXP_DIR" "$LOG_DIR" mkdir -p "$EXP_DIR" "$LOG_DIR" @@ -37,9 +38,12 @@ uv run coverage run -a --data-file="$PROJECT_ROOT/tests/.coverage" --source="$PR policy.train_micro_batch_size=1 \ policy.logprob_batch_size=1 \ policy.max_total_sequence_length=256 \ + policy.dtensor_cfg.enabled=false \ + policy.megatron_cfg.enabled=true \ policy.generation.max_new_tokens=128 \ policy.generation.vllm_cfg.precision=fp8 \ ++policy.generation.vllm_cfg.is_mx=true \ + policy.generation.vllm_cfg.refit_prequantize=true \ policy.generation.vllm_cfg.kv_cache_dtype=auto \ policy.generation.vllm_cfg.max_model_len=256 \ policy.generation.vllm_cfg.gpu_memory_utilization=0.5 \ diff --git a/tests/test_mxfp8_rollout_recipes.py b/tests/test_mxfp8_rollout_recipes.py index 1c12462e9bc..7d1d2d8b14e 100644 --- a/tests/test_mxfp8_rollout_recipes.py +++ b/tests/test_mxfp8_rollout_recipes.py @@ -114,6 +114,7 @@ "segment_size": 4, "async_engine": None, "moe_backend": "flashinfer_trtllm", + "refit_optimizations": True, "train_global_batch_size": 2048, "ignore_patterns": [ "model.layers.*.self_attn.*", @@ -140,6 +141,7 @@ "segment_size": 4, "async_engine": None, "moe_backend": "flashinfer_trtllm", + "refit_optimizations": True, "ignore_patterns": [ "model.layers.*.self_attn.*", "lm_head", @@ -163,6 +165,7 @@ "async_engine": True, "tensor_parallel_size": 4, "moe_backend": "flashinfer_trtllm", + "refit_optimizations": True, "ignore_patterns": [ "model.layers.*.self_attn.*", "model.layers.*.mlp.gate", @@ -246,6 +249,10 @@ def test_mxfp8_rollout_recipe_matrix(case_name: str, expected: dict) -> None: assert config["logger"]["wandb"]["name"] == case_name assert vllm_cfg["precision"] == "fp8" assert vllm_cfg["is_mx"] is True + if expected.get("refit_optimizations"): + assert vllm_cfg["refit_prequantize"] is True + assert vllm_cfg["refit_cache_loader_routes"] is True + assert config["policy"]["megatron_cfg"]["enabled"] is True assert "quantization_ignored_layer_kws" not in vllm_cfg assert vllm_cfg["quantization_ignore_patterns"] == expected["ignore_patterns"] assert cluster["num_nodes"] == expected["nodes"] @@ -276,6 +283,42 @@ def test_mxfp8_rollout_recipe_matrix(case_name: str, expected: dict) -> None: assert config["policy"]["generation"]["colocated"]["enabled"] is False +@pytest.mark.parametrize( + "case_name", + ( + "grpo-qwen3-30ba3b-4n4g-mxfp8-rollout", + "grpo-qwen3-32b-4n4g-mxfp8-rollout", + "grpo-qwen3-235b-16n4g-mxfp8-rollout", + ), +) +def test_sync_qwen3_mxfp8_rollout_recipes_enable_refit_optimizations( + case_name: str, +) -> None: + config = _load_resolved_yaml(PERF_CONFIG_DIR / f"{case_name}.yaml") + vllm_cfg = config["policy"]["generation"]["vllm_cfg"] + + assert vllm_cfg["refit_prequantize"] is True + assert vllm_cfg["refit_cache_loader_routes"] is True + + +@pytest.mark.parametrize( + "case_name", + ( + "grpo-qwen3-30ba3b-4n4g-async-1off-mxfp8-rollout", + "grpo-qwen3-32b-8n4g-async-1off-mxfp8-rollout", + "grpo-qwen3-235b-32n4g-async-1off-mxfp8-rollout", + ), +) +def test_async_qwen3_mxfp8_rollout_recipes_skip_sync_refit_optimizations( + case_name: str, +) -> None: + config = _load_resolved_yaml(PERF_CONFIG_DIR / f"{case_name}.yaml") + vllm_cfg = config["policy"]["generation"]["vllm_cfg"] + + assert not vllm_cfg.get("refit_prequantize", False) + assert not vllm_cfg.get("refit_cache_loader_routes", False) + + def test_mxfp8_rollout_recipes_are_in_gb200_performance_suite() -> None: recipe_names = {path.stem for path in PERF_CONFIG_DIR.glob("*-mxfp8-rollout.yaml")} assert recipe_names == set(MXFP8_CASES) diff --git a/tests/unit/algorithms/test_ppo.py b/tests/unit/algorithms/test_ppo.py index 8f51fa4f9fa..96965236c29 100644 --- a/tests/unit/algorithms/test_ppo.py +++ b/tests/unit/algorithms/test_ppo.py @@ -1462,6 +1462,7 @@ def get_placement_groups(self): policy.init_collective.return_value = ["policy-future"] value_model = MagicMock() generation = MagicMock() + generation.prepare_refit_info.return_value = None generation.init_collective.return_value = ["generation-future"] policy_factory = MagicMock(return_value=policy) value_factory = MagicMock(return_value=value_model) diff --git a/tests/unit/environments/test_code_environment.py b/tests/unit/environments/test_code_environment.py index cfdb96cc089..7e2a7440069 100644 --- a/tests/unit/environments/test_code_environment.py +++ b/tests/unit/environments/test_code_environment.py @@ -54,6 +54,7 @@ "vllm_cfg": { "async_engine": False, "precision": "bfloat16", + "refit_cache_loader_routes": False, "tensor_parallel_size": 1, "pipeline_parallel_size": 1, "expert_parallel_size": 1, diff --git a/tests/unit/environments/test_retriever.py b/tests/unit/environments/test_retriever.py index 9e0c27e62f1..9d53aa15870 100644 --- a/tests/unit/environments/test_retriever.py +++ b/tests/unit/environments/test_retriever.py @@ -53,6 +53,7 @@ "vllm_cfg": { "async_engine": False, "precision": "bfloat16", + "refit_cache_loader_routes": False, "tensor_parallel_size": 1, "pipeline_parallel_size": 1, "expert_parallel_size": 1, diff --git a/tests/unit/experience/test_rollouts.py b/tests/unit/experience/test_rollouts.py index 4a810da668b..905e516f682 100644 --- a/tests/unit/experience/test_rollouts.py +++ b/tests/unit/experience/test_rollouts.py @@ -1082,6 +1082,7 @@ def initial_multi_step_calculator_batch(rollout_tokenizer): "vllm_cfg": { "async_engine": False, "precision": "bfloat16", + "refit_cache_loader_routes": False, "tensor_parallel_size": 1, "pipeline_parallel_size": 1, "expert_parallel_size": 1, @@ -1946,10 +1947,9 @@ def _postprocess_group(**kwargs): monkeypatch.setattr( rollouts_mod, "collect_multimodal_payload_metrics", - lambda payload, boundary, enabled: payload_calls.append( - (payload, boundary, enabled) - ) - or {}, + lambda payload, boundary, enabled: ( + payload_calls.append((payload, boundary, enabled)) or {} + ), ) monkeypatch.setattr( rollouts_mod, "print_multimodal_payload_metrics", lambda metrics: None diff --git a/tests/unit/models/generation/sglang/test_sglang_generation.py b/tests/unit/models/generation/sglang/test_sglang_generation.py index c6e9db783c1..4e91a26be20 100644 --- a/tests/unit/models/generation/sglang/test_sglang_generation.py +++ b/tests/unit/models/generation/sglang/test_sglang_generation.py @@ -46,6 +46,12 @@ pytestmark = pytest.mark.sglang +def test_prepare_refit_info_accepts_missing_metadata(): + generation = SGLangGeneration.__new__(SGLangGeneration) + + assert generation.prepare_refit_info(None) is None + + @pytest.fixture(scope="module") def ray_cluster(): """Initialise Ray once for this module's tests.""" diff --git a/tests/unit/models/generation/test_mxfp8_prequant.py b/tests/unit/models/generation/test_mxfp8_prequant.py new file mode 100644 index 00000000000..5c762b8ca46 --- /dev/null +++ b/tests/unit/models/generation/test_mxfp8_prequant.py @@ -0,0 +1,207 @@ +# 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 sys + +import pytest +import torch + +from nemo_rl.models.generation.vllm.quantization.fp8_train_utils import ( + MXFP8_BLOCK_SIZE, + _mxfp8_e4m3_quantize_torch, + mxfp8_e4m3_quantize_for_refit, +) + +pytestmark = pytest.mark.vllm + + +def _dequantize(x_fp8: torch.Tensor, scales: torch.Tensor) -> torch.Tensor: + num_blocks = x_fp8.shape[-1] // MXFP8_BLOCK_SIZE + x_blocked = x_fp8.to(torch.float32).view( + *x_fp8.shape[:-1], num_blocks, MXFP8_BLOCK_SIZE + ) + descale = torch.exp2(scales.to(torch.float32) - 127.0) + return (x_blocked * descale.unsqueeze(-1)).view(*x_fp8.shape) + + +@pytest.mark.parametrize("shape", [(64, 128), (7, 96), (4, 16, 64)]) +def test_torch_reference_shapes_and_roundtrip(shape): + torch.manual_seed(0) + x = torch.randn(*shape, dtype=torch.bfloat16) + + x_fp8, scales = _mxfp8_e4m3_quantize_torch(x) + + assert x_fp8.shape == x.shape + assert x_fp8.dtype == torch.float8_e4m3fn + assert scales.dtype == torch.uint8 + expected_scale_shape = (*shape[:-1], shape[-1] // MXFP8_BLOCK_SIZE) + assert tuple(scales.shape) == expected_scale_shape + + x_dq = _dequantize(x_fp8, scales) + x32 = x.to(torch.float32) + abs_err = (x_dq - x32).abs() + block_amax = ( + x32.abs() + .reshape(*shape[:-1], shape[-1] // MXFP8_BLOCK_SIZE, MXFP8_BLOCK_SIZE) + .amax(dim=-1, keepdim=True) + .expand(*shape[:-1], shape[-1] // MXFP8_BLOCK_SIZE, MXFP8_BLOCK_SIZE) + .reshape(shape) + ) + # e4m3 has 3 mantissa bits: elements within the block's representable range + # must round-trip to ~12.5% relative error; elements far below the block + # amax may legitimately quantize to zero, so bound them by absolute error + # instead of a ratio (keeps the test deterministic across torch versions). + representable = x32.abs() >= block_amax / 64 + rel_err = (abs_err / x32.abs().clamp(min=1e-6))[representable] + assert rel_err.median() < 0.05 + assert rel_err.max() < 0.25 + assert (abs_err[~representable] <= block_amax[~representable] / 32).all() + + +def test_last_dim_not_divisible_raises(): + x = torch.randn(8, MXFP8_BLOCK_SIZE + 1, dtype=torch.bfloat16) + with pytest.raises(AssertionError): + _mxfp8_e4m3_quantize_torch(x) + + +def test_refit_quantize_preserves_single_scale_block_dimension(): + x = torch.randn(8, MXFP8_BLOCK_SIZE, dtype=torch.bfloat16) + + _, scales = mxfp8_e4m3_quantize_for_refit(x) + + assert scales.shape == (8, 1) + + +def test_blackwell_refit_prequantization_requires_flashinfer(monkeypatch): + class FakeBlackwellTensor: + is_cuda = True + device = "cuda" + + monkeypatch.setattr(torch.cuda, "get_device_capability", lambda _device: (10, 0)) + monkeypatch.setitem(sys.modules, "flashinfer", None) + + with pytest.raises(RuntimeError, match=r"sm100\+ requires FlashInfer"): + mxfp8_e4m3_quantize_for_refit(FakeBlackwellTensor()) + + +@pytest.mark.skipif( + not torch.cuda.is_available() or torch.cuda.get_device_capability() < (10, 0), + reason=( + "requires sm100+; below it both sides fall back to the shared torch " + "reference and the comparison is vacuous" + ), +) +def test_refit_quantize_matches_receiver_path(): + """Bitwise parity with the vLLM receiver path (mxfp8_e4m3_quantize + squeeze).""" + vllm_mxfp8 = pytest.importorskip( + "vllm.model_executor.layers.quantization.utils.mxfp8_utils" + ) + + torch.manual_seed(0) + x = torch.randn(256, 512, dtype=torch.bfloat16, device="cuda") + x[0].zero_() + + ref_lp, ref_scale = vllm_mxfp8.mxfp8_e4m3_quantize(x) + ref_scale = torch.squeeze(ref_scale, dim=-1) + assert torch.any(ref_scale == 0) + ref_scale = torch.where(ref_scale == 0, torch.ones_like(ref_scale), ref_scale) + + got_lp, got_scale = mxfp8_e4m3_quantize_for_refit(x) + + assert got_lp.dtype == ref_lp.dtype + assert torch.equal(got_lp.view(torch.uint8), ref_lp.view(torch.uint8)) + assert got_scale.dtype == ref_scale.dtype + assert got_scale.shape == ref_scale.shape + assert torch.equal(got_scale.reshape(-1), ref_scale.reshape(-1)) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") +def test_refit_quantize_matches_receiver_quantize_mxfp8_weight(): + """Sender prequantization and the receiver helper must agree bit-for-bit. + + The trainer streams E4M3 data + *_scale_from_checkpoint produced by + mxfp8_e4m3_quantize_for_refit; weights the receiver quantizes itself go + through quantize_mxfp8_weight. Refit correctness relies on the two + implementations producing identical bits for the same input. + """ + from nemo_rl.models.generation.vllm.quantization.fp8 import quantize_mxfp8_weight + + torch.manual_seed(0) + x = torch.randn(256, 512, dtype=torch.bfloat16, device="cuda") + x[0].zero_() + + recv_lp, recv_scale = quantize_mxfp8_weight(x) + sent_lp, sent_scale = mxfp8_e4m3_quantize_for_refit(x) + + assert sent_lp.dtype == recv_lp.dtype + assert torch.equal(sent_lp.view(torch.uint8), recv_lp.view(torch.uint8)) + assert sent_scale.dtype == recv_scale.dtype + assert sent_scale.shape == recv_scale.shape + assert torch.equal(sent_scale, recv_scale) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") +@pytest.mark.parametrize( + "is_gated,intermediate_size,hidden_size", + [ + # Aligned: both scale K dims (hidden/32=8, intermediate/32=4) are %4. + (True, 128, 256), + # w2 scale K = 192/32 = 6, so pad_flashinfer_scale_k pads it to 8. + (True, 192, 128), + # Non-gated (single w13 shard), aligned. + (False, 128, 256), + ], +) +def test_batched_moe_shuffle_matches_per_expert( + is_gated, intermediate_size, hidden_size +): + """Bitwise parity of the batched TRTLLM MoE shuffle with the per-expert loop.""" + pytest.importorskip("flashinfer") + fp8 = pytest.importorskip("nemo_rl.models.generation.vllm.quantization.fp8") + + from types import SimpleNamespace + + torch.manual_seed(0) + num_experts = 4 + w13_rows = (2 if is_gated else 1) * intermediate_size + + def rand_bytes(*shape): + return torch.randint(0, 256, shape, dtype=torch.uint8, device="cuda") + + w13_weight = rand_bytes(num_experts, w13_rows, hidden_size).view( + torch.float8_e4m3fn + ) + w2_weight = rand_bytes(num_experts, hidden_size, intermediate_size).view( + torch.float8_e4m3fn + ) + w13_scale = rand_bytes(num_experts, w13_rows, hidden_size // MXFP8_BLOCK_SIZE) + w2_scale = rand_bytes( + num_experts, hidden_size, intermediate_size // MXFP8_BLOCK_SIZE + ) + + layer = SimpleNamespace() # holds the cached row permutations + epilogue_tile_m = 128 + batched = fp8._shuffle_mxfp8_moe_batched( + layer, w13_weight, w2_weight, w13_scale, w2_scale, is_gated, epilogue_tile_m + ) + reference = fp8._shuffle_mxfp8_moe_per_expert( + w13_weight, w2_weight, w13_scale, w2_scale, is_gated, epilogue_tile_m + ) + + for got, want, name in zip( + batched, reference, ("w13_weight", "w2_weight", "w13_scale", "w2_scale") + ): + assert got.shape == want.shape, name + assert got.dtype == want.dtype, name + assert torch.equal(got.view(torch.uint8), want.view(torch.uint8)), name diff --git a/tests/unit/models/generation/test_vllm_backend.py b/tests/unit/models/generation/test_vllm_backend.py index ee5e6699ab7..658dec26e41 100644 --- a/tests/unit/models/generation/test_vllm_backend.py +++ b/tests/unit/models/generation/test_vllm_backend.py @@ -666,6 +666,114 @@ def test_sparse_delta_refit_rejected_for_native_trtllm_backend(): ext.update_weights_from_decoded_sparse_payload(b"") +@pytest.mark.vllm +@pytest.mark.parametrize("enabled", [False, True]) +def test_prepare_refit_info_reports_only_fp8_weights(monkeypatch, enabled): + from nemo_rl.models.generation.vllm import vllm_backend + from nemo_rl.models.generation.vllm.quantization import fp8 + + ext = vllm_backend.VllmInternalWorkerExtension.__new__( + vllm_backend.VllmInternalWorkerExtension + ) + model = object() + config = object() + ext.model_runner = SimpleNamespace(model=model, vllm_config=config) + state_dict_info = { + "model.linear.weight": ((2, 2), torch.bfloat16), + "model.norm.weight": ((2,), torch.bfloat16), + } + is_fp8_model = MagicMock(return_value=True) + checked_names = [] + + def is_fp8_weight(name, candidate_model): + assert candidate_model is model + checked_names.append(name) + return name == "model.linear.weight" + + source_config = fp8.FP8Config( + is_mx=True, + refit_prequantize=enabled, + use_fp8_weights=True, + ) + monkeypatch.setattr(fp8, "global_fp8_config", source_config) + serialized_config = fp8.serialize_fp8_config() + monkeypatch.setattr(fp8, "global_fp8_config", None) + monkeypatch.setattr(fp8, "is_fp8_model", is_fp8_model) + monkeypatch.setattr(fp8, "_is_fp8_weight", is_fp8_weight) + + result = ext.prepare_refit_info(state_dict_info, serialized_config) + + assert ext.state_dict_info is state_dict_info + assert fp8.global_fp8_config == source_config + if enabled: + assert result == ["model.linear.weight"] + is_fp8_model.assert_called_once_with(config) + assert checked_names == list(state_dict_info) + else: + assert result is None + is_fp8_model.assert_not_called() + assert checked_names == [] + + +@pytest.mark.vllm +def test_sync_prepare_refit_info_unions_worker_names(monkeypatch): + from nemo_rl.models.generation.vllm.quantization import fp8 + from nemo_rl.models.generation.vllm.vllm_worker import ( + VllmGenerationWorkerImpl, + ) + + worker = VllmGenerationWorkerImpl.__new__(VllmGenerationWorkerImpl) + state_dict_info = {"model.weight": ((2, 2), torch.bfloat16)} + serialized_config = {"is_mx": True, "refit_prequantize": True} + monkeypatch.setattr(fp8, "serialize_fp8_config", lambda: serialized_config) + worker.llm = SimpleNamespace( + collective_rpc=MagicMock( + return_value=[ + None, + ["model.b.weight", "model.a.weight"], + ["model.a.weight"], + ] + ) + ) + + assert worker.prepare_refit_info(state_dict_info) == [ + "model.a.weight", + "model.b.weight", + ] + worker.llm.collective_rpc.assert_called_once_with( + "prepare_refit_info", + args=(state_dict_info, serialized_config), + ) + + +@pytest.mark.vllm +@pytest.mark.asyncio +async def test_async_prepare_refit_info_unions_worker_names(monkeypatch): + from nemo_rl.models.generation.vllm.quantization import fp8 + from nemo_rl.models.generation.vllm.vllm_worker_async import ( + VllmAsyncGenerationWorkerImpl, + ) + + worker = VllmAsyncGenerationWorkerImpl.__new__(VllmAsyncGenerationWorkerImpl) + state_dict_info = {"model.weight": ((2, 2), torch.bfloat16)} + serialized_config = {"is_mx": True, "refit_prequantize": True} + monkeypatch.setattr(fp8, "serialize_fp8_config", lambda: serialized_config) + worker.llm = SimpleNamespace( + collective_rpc=AsyncMock( + return_value=[None, ["model.b.weight"], ["model.a.weight"]] + ) + ) + + assert await worker.prepare_refit_info_async(state_dict_info) == [ + "model.a.weight", + "model.b.weight", + ] + worker.llm.collective_rpc.assert_awaited_once_with( + "prepare_refit_info", + args=(state_dict_info, serialized_config), + ) + + @pytest.mark.vllm @pytest.mark.parametrize("with_mtp", [False, True]) def test_update_weights_from_collective_processes_weights_after_loading( @@ -721,7 +829,6 @@ def packed_broadcast_consumer( post_unpack_func([("model.weight", "weight-value")]) ext._load_weights = load_weights - ext._maybe_process_fp8_kv_cache = lambda: call_order.append("kv") monkeypatch.setattr( vllm_backend, "packed_broadcast_consumer", packed_broadcast_consumer ) @@ -745,7 +852,7 @@ def packed_broadcast_consumer( if with_mtp: expected_process_calls.append((draft_model, draft_model_config, ext.device)) expected_call_order.extend(["config_enter", "process_mtp", "config_exit"]) - expected_call_order.extend(["kv", "gc", "empty_cache"]) + expected_call_order.extend(["gc", "empty_cache"]) assert process_calls == expected_process_calls assert call_order == expected_call_order diff --git a/tests/unit/models/generation/test_vllm_checkpoint_engine.py b/tests/unit/models/generation/test_vllm_checkpoint_engine.py index 868b00efaea..66df07ffd72 100644 --- a/tests/unit/models/generation/test_vllm_checkpoint_engine.py +++ b/tests/unit/models/generation/test_vllm_checkpoint_engine.py @@ -15,6 +15,8 @@ """Tests for vLLM checkpoint-engine worker lifecycle helpers.""" import asyncio +from collections.abc import Callable, Iterator +from contextlib import contextmanager from types import SimpleNamespace from unittest.mock import AsyncMock, MagicMock @@ -95,7 +97,20 @@ async def receive_weight_batches(self): worker._load_weights = lambda batch: events.append( ("load", [name for name, _weight in batch]) ) - worker._maybe_process_fp8_kv_cache = lambda: events.append(("fp8",)) + + @contextmanager + def weight_update_lifecycle( + transport: str, + ) -> Iterator[Callable[[], None]]: + events.append(("setup", transport)) + + def finalize() -> None: + events.append(("finalize",)) + + yield finalize + events.append(("teardown", transport)) + + worker._weight_update_lifecycle = weight_update_lifecycle monkeypatch.setattr( torch.cuda, "current_stream", @@ -104,11 +119,13 @@ async def receive_weight_batches(self): assert asyncio.run(worker._update_weights_from_checkpoint_engine_async()) is True assert events == [ + ("setup", "checkpoint_engine"), ("load", ["a"]), ("sync",), ("load", ["b", "c"]), ("sync",), - ("fp8",), + ("finalize",), + ("teardown", "checkpoint_engine"), ] diff --git a/tests/unit/models/generation/test_vllm_config.py b/tests/unit/models/generation/test_vllm_config.py new file mode 100644 index 00000000000..e4043878d93 --- /dev/null +++ b/tests/unit/models/generation/test_vllm_config.py @@ -0,0 +1,152 @@ +# 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. + +from types import SimpleNamespace +from typing import cast + +import pytest + +from nemo_rl.models.generation import configure_generation_config +from nemo_rl.models.generation.interfaces import GenerationConfig +from nemo_rl.models.generation.vllm.config import ( + VllmConfig, + validate_vllm_quantization_config, +) + + +@pytest.mark.parametrize( + "generation_config", + [ + { + "vllm_cfg": { + "precision": "fp8", + "is_mx": False, + "refit_prequantize": True, + } + }, + { + "vllm_cfg": { + "precision": "fp8", + "refit_prequantize": True, + } + }, + { + "vllm_cfg": { + "precision": "bfloat16", + "refit_prequantize": True, + }, + "quant_cfg": "examples/modelopt/quant_configs/nvfp4_a16.yaml", + "real_quant": True, + }, + ], +) +def test_refit_prequantize_requires_mxfp8(generation_config: dict) -> None: + with pytest.raises( + ValueError, + match="refit_prequantize requires precision='fp8' and is_mx=true", + ): + validate_vllm_quantization_config(cast(VllmConfig, generation_config)) + + +def test_refit_prequantize_must_be_boolean() -> None: + generation_config = cast( + VllmConfig, + { + "vllm_cfg": { + "precision": "fp8", + "is_mx": True, + "refit_prequantize": "false", + } + }, + ) + + with pytest.raises(ValueError, match="refit_prequantize must be a boolean"): + validate_vllm_quantization_config(generation_config) + + +def test_refit_prequantize_accepts_mxfp8() -> None: + generation_config = cast( + VllmConfig, + { + "vllm_cfg": { + "precision": "fp8", + "is_mx": True, + "refit_prequantize": True, + } + }, + ) + + validate_vllm_quantization_config(generation_config) + + +def test_refit_prequantize_rejects_nccl_reshard() -> None: + generation_config = cast( + VllmConfig, + { + "refit_transport": "nccl_reshard", + "vllm_cfg": { + "precision": "fp8", + "is_mx": True, + "refit_prequantize": True, + }, + }, + ) + + with pytest.raises(ValueError, match="not supported with nccl_reshard"): + validate_vllm_quantization_config(generation_config) + + +@pytest.mark.parametrize( + "field", + ["refit_cache_loader_routes"], +) +def test_refit_optimization_flags_must_be_boolean(field: str) -> None: + generation_config = cast( + VllmConfig, + { + "vllm_cfg": { + "precision": "fp8", + "is_mx": True, + field: "true", + } + }, + ) + + with pytest.raises(ValueError, match=rf"{field} must be a boolean"): + validate_vllm_quantization_config(generation_config) + + +def test_refit_prequantize_validation_allows_omitted_vllm_cfg() -> None: + generation_config = cast(VllmConfig, {"quant_cfg": None}) + + validate_vllm_quantization_config(generation_config) + + +def test_configure_generation_config_validates_refit_prequantize() -> None: + generation_config = cast( + GenerationConfig, + { + "backend": "vllm", + "stop_token_ids": None, + "stop_strings": None, + "vllm_cfg": { + "precision": "bfloat16", + "refit_prequantize": True, + }, + }, + ) + tokenizer = SimpleNamespace(pad_token_id=0, eos_token_id=1) + + with pytest.raises(ValueError, match="requires precision='fp8' and is_mx=true"): + configure_generation_config(generation_config, tokenizer) diff --git a/tests/unit/models/generation/test_vllm_fp8_quantization.py b/tests/unit/models/generation/test_vllm_fp8_quantization.py index 34240d34f83..0e10f28eb96 100644 --- a/tests/unit/models/generation/test_vllm_fp8_quantization.py +++ b/tests/unit/models/generation/test_vllm_fp8_quantization.py @@ -14,7 +14,9 @@ import types from pathlib import Path +from typing import Any +import cloudpickle import pytest import torch import yaml @@ -68,6 +70,7 @@ def test_init_fp8_uses_mxfp8_quantization_config(fp8_module, monkeypatch): "kv_cache_dtype": "auto", "async_engine": False, "is_mx": True, + "refit_cache_loader_routes": True, "use_deep_gemm": True, }, "dummy-model", @@ -91,6 +94,60 @@ def test_init_fp8_uses_mxfp8_quantization_config(fp8_module, monkeypatch): assert "VLLM_USE_DEEP_GEMM_E8M0" not in fp8.os.environ +def test_ray_executor_v2_worker_applies_fp8_patches_before_model_load( + fp8_module, monkeypatch +): + fp8 = fp8_module + monkeypatch.setattr(fp8, "_test_applied_configs", [], raising=False) + config = fp8.FP8Config( + use_fp8_weights=True, + model_parallel_size=2, + is_mx=True, + ) + + class FakeRayWorkerProc: + def initialize_worker( + self, + local_rank, + env_vars, + driver_env_vars=None, + assigned_physical_gpu_ids=None, + ): + assert fp8.fp8_patches_applied + return ( + local_rank, + env_vars, + driver_env_vars, + assigned_physical_gpu_ids, + ) + + def fake_apply_fp8_patches(_self, fp8_config): + fp8._test_applied_configs.append(fp8_config) + fp8.fp8_patches_applied = True + + monkeypatch.setattr(fp8, "apply_fp8_patches", fake_apply_fp8_patches) + ray_executor_v2 = types.SimpleNamespace(RayWorkerProc=FakeRayWorkerProc) + fp8._patch_ray_executor_v2_worker(ray_executor_v2, config) + patched_worker_cls = cloudpickle.loads( + cloudpickle.dumps(ray_executor_v2.RayWorkerProc) + ) + + result = patched_worker_cls().initialize_worker( + 1, + {"WORKER_ENV": "1"}, + {"DRIVER_ENV": "1"}, + assigned_physical_gpu_ids=[2, 3], + ) + + assert fp8._test_applied_configs == [config] + assert result == ( + 1, + {"WORKER_ENV": "1"}, + {"DRIVER_ENV": "1"}, + [2, 3], + ) + + def test_init_fp8_passes_modelopt_ignore_patterns_without_hf_expansion( fp8_module, monkeypatch ): @@ -579,7 +636,7 @@ def track_index_select(*args, **kwargs): monkeypatch.setattr(torch, "index_select", original_index_select) assert len(index_select_out_tensors) == 4 - assert all(tensor is None for tensor in index_select_out_tensors) + assert all(tensor is not None for tensor in index_select_out_tensors) reference = fp8._shuffle_mxfp8_moe_per_expert( w13_weight, @@ -609,12 +666,35 @@ def test_process_mxfp8_moe_refit_uses_batched_flashinfer_shuffle( is_mx=True, ) - w13_weight = torch.nn.Parameter(torch.zeros(2, 4, 3), requires_grad=False) - w2_weight = torch.nn.Parameter(torch.zeros(2, 3, 2), requires_grad=False) - w13_scale = torch.nn.Parameter(torch.zeros(2, 4, 1), requires_grad=False) - w2_scale = torch.nn.Parameter(torch.zeros(2, 3, 1), requires_grad=False) - w13_scale_from_checkpoint = torch.ones_like(w13_scale) - w2_scale_from_checkpoint = torch.ones_like(w2_scale) + hidden_size = 512 + intermediate_size = 128 + w13_rows = intermediate_size * (2 if is_gated else 1) + w13_weight = torch.nn.Parameter( + torch.zeros(2, w13_rows, hidden_size, dtype=torch.float8_e4m3fn), + requires_grad=False, + ) + w2_weight = torch.nn.Parameter( + torch.zeros(2, hidden_size, intermediate_size, dtype=torch.float8_e4m3fn), + requires_grad=False, + ) + w13_scale = torch.nn.Parameter( + torch.zeros(2, w13_rows * (hidden_size // 32), dtype=torch.uint8), + requires_grad=False, + ) + w2_scale = torch.nn.Parameter( + torch.zeros(2, hidden_size * (intermediate_size // 32), dtype=torch.uint8), + requires_grad=False, + ) + w13_scale_from_checkpoint = torch.ones( + 2, w13_rows, hidden_size // 32, dtype=torch.uint8 + ) + w2_scale_from_checkpoint = torch.ones( + 2, hidden_size, intermediate_size // 32, dtype=torch.uint8 + ) + moe_config = types.SimpleNamespace( + is_act_and_mul=is_gated, + intermediate_size_per_partition=intermediate_size, + ) layer = types.SimpleNamespace( w13_weight=w13_weight, w2_weight=w2_weight, @@ -626,14 +706,20 @@ def test_process_mxfp8_moe_refit_uses_batched_flashinfer_shuffle( w2_weight_scale_from_checkpoint=types.SimpleNamespace( data=w2_scale_from_checkpoint ), + moe_config=moe_config, ) moe_kernel = object() - moe_quant_config = object() + moe_quant_config = types.SimpleNamespace( + w1_scale=w13_scale, + w2_scale=w2_scale, + ) quant_method = types.SimpleNamespace( - moe=types.SimpleNamespace(is_act_and_mul=is_gated), + moe=moe_config, moe_kernel=moe_kernel, moe_quant_config=moe_quant_config, mxfp8_backend=Fp8MoeBackend.FLASHINFER_TRTLLM, + experts_cls=types.SimpleNamespace(is_monolithic=lambda: True), + weight_block_size=[32, 32], ) shuffled = ( torch.full_like(w13_weight, 1), @@ -851,6 +937,7 @@ def test_init_fp8_rejects_non_pow2_mxfp8_scales(fp8_module, monkeypatch, field, "kv_cache_dtype": "auto", "async_engine": False, "is_mx": True, + "refit_cache_loader_routes": False, field: False, }, "dummy-model", @@ -915,9 +1002,467 @@ def fake_patch(path, _replacement): "ModelOptMxFp8FusedMoE.process_weights_after_loading" in path for path in patched_paths ) + assert any( + "ModelOptMxFp8FusedMoE.apply_monolithic" in path for path in patched_paths + ) assert all(patcher.started for patcher in fp8.fp8_state.vllm_patches) +def test_get_module_from_param_name_resolves_vllm_025_routed_experts( + fp8_module: types.ModuleType, +) -> None: + from vllm.model_executor.layers.fused_moe.routed_experts import RoutedExperts + from vllm.model_executor.layers.fused_moe.runner.moe_runner import MoERunner + + fp8 = fp8_module + routed_experts = RoutedExperts.__new__(RoutedExperts) + torch.nn.Module.__init__(routed_experts) + runner = MoERunner.__new__(MoERunner) + torch.nn.Module.__init__(runner) + runner.routed_experts = routed_experts + model = types.SimpleNamespace(packed_modules_mapping={}, experts=runner) + + assert ( + fp8._get_module_from_param_name(model, "experts.w13_weight") is routed_experts + ) + + +def test_load_weights_preserves_prequantized_mxfp8_and_clamps_scales( + fp8_module, monkeypatch +): + from vllm.model_executor.layers.quantization.utils import mxfp8_utils + + from nemo_rl.models.generation.vllm import vllm_backend + + fp8 = fp8_module + fp8.global_fp8_config = types.SimpleNamespace( + is_mx=True, + refit_prequantize=True, + ) + native = torch.ones(2, 2, dtype=torch.bfloat16) + prequantized = torch.ones(2, 2, dtype=torch.float8_e4m3fn) + prequantized_scales = torch.ones(2, 1, dtype=torch.uint8) + receiver_quantized = torch.full((2, 64), 2.0, dtype=torch.bfloat16) + receiver_fp8 = torch.ones(2, 64, dtype=torch.float8_e4m3fn) + receiver_scales = torch.tensor([[0, 7], [3, 0]], dtype=torch.uint8) + loaded = [] + + monkeypatch.setattr( + fp8, + "_is_fp8_weight", + lambda name, _model: name.endswith(".weight"), + ) + monkeypatch.setattr( + mxfp8_utils, + "mxfp8_e4m3_quantize", + lambda tensor, **_kwargs: ( + ( + receiver_fp8, + receiver_scales, + ) + if tensor is receiver_quantized + else pytest.fail("unexpected receiver quantization input") + ), + ) + monkeypatch.setattr( + vllm_backend, + "load_weights_maybe_cached", + lambda model, weights, *, cache_loader_routes: loaded.extend(weights), + ) + model = object() + + fp8.load_weights( + [ + ("model.native", native), + ("model.prequantized.weight", prequantized), + ( + "model.prequantized.weight_scale_from_checkpoint", + prequantized_scales, + ), + ("model.receiver.weight", receiver_quantized), + ], + types.SimpleNamespace( + model=model, + vllm_config=types.SimpleNamespace(additional_config={}), + ), + ) + + assert loaded[0][0] == "model.native" + assert loaded[0][1] is native + assert loaded[1][0] == "model.prequantized.weight" + assert loaded[1][1] is prequantized + assert loaded[2][0] == "model.prequantized.weight_scale_from_checkpoint" + assert loaded[2][1] is prequantized_scales + assert loaded[3][0] == "model.receiver.weight" + # quantize_mxfp8_weight reshapes to the checkpoint layout, so compare + # contents rather than object identity. + assert loaded[3][1].dtype == torch.float8_e4m3fn + assert torch.equal(loaded[3][1].view(torch.uint8), receiver_fp8.view(torch.uint8)) + assert loaded[4][0] == "model.receiver.weight_scale_from_checkpoint" + torch.testing.assert_close( + loaded[4][1], + torch.tensor([[1, 7], [3, 1]], dtype=torch.uint8), + ) + + +def test_load_weights_rejects_unnegotiated_mxfp8_payload(fp8_module, monkeypatch): + fp8 = fp8_module + fp8.global_fp8_config = types.SimpleNamespace( + is_mx=True, + refit_prequantize=False, + ) + monkeypatch.setattr(fp8, "_is_fp8_weight", lambda _name, _model: True) + + with pytest.raises(ValueError, match="refit_prequantize=true"): + fp8.load_weights( + [("model.weight", torch.ones(2, 2, dtype=torch.float8_e4m3fn))], + types.SimpleNamespace( + model=object(), + vllm_config=types.SimpleNamespace(additional_config={}), + ), + ) + + +def test_load_weights_rejects_prequantized_mxfp8_without_scale(fp8_module, monkeypatch): + fp8 = fp8_module + fp8.global_fp8_config = types.SimpleNamespace( + is_mx=True, + refit_prequantize=True, + ) + monkeypatch.setattr(fp8, "_is_fp8_weight", lambda _name, _model: True) + + with pytest.raises(ValueError, match="missing.*scale_from_checkpoint"): + fp8.load_weights( + [("model.weight", torch.ones(2, 2, dtype=torch.float8_e4m3fn))], + types.SimpleNamespace( + model=object(), + vllm_config=types.SimpleNamespace(additional_config={}), + ), + ) + + +def test_load_weights_preserves_non_mx_blockwise_fp8_payload(fp8_module, monkeypatch): + from nemo_rl.models.generation.vllm import vllm_backend + + fp8 = fp8_module + fp8.global_fp8_config = types.SimpleNamespace( + is_mx=False, + refit_prequantize=False, + ) + weight = torch.ones(2, 2, dtype=torch.float8_e4m3fn) + scale = torch.ones(1, dtype=torch.float32) + loaded = [] + monkeypatch.setattr( + fp8, + "_is_fp8_weight", + lambda name, _model: name == "model.weight", + ) + monkeypatch.setattr( + vllm_backend, + "load_weights_maybe_cached", + lambda model, weights, *, cache_loader_routes: loaded.extend(weights), + ) + + fp8.load_weights( + [("model.weight", weight), ("model.weight_scale_inv", scale)], + types.SimpleNamespace( + model=object(), + vllm_config=types.SimpleNamespace(additional_config={}), + ), + ) + + assert loaded[0][0] == "model.weight" + assert loaded[0][1] is weight + assert loaded[1][0] == "model.weight_scale_inv" + assert loaded[1][1] is scale + + +def test_mxfp8_padding_helpers_preserve_values_and_fill_padding( + fp8_module: types.ModuleType, +) -> None: + fp8 = fp8_module + tensor = torch.arange(12).reshape(2, 2, 3) + + assert fp8._round_up(1856, 128) == 1920 + assert fp8._pad_tensor_dim(tensor, 1, 2) is tensor + + padded = fp8._pad_tensor_dim(tensor, 1, 4, pad_value=7) + torch.testing.assert_close(padded[:, :2], tensor) + torch.testing.assert_close(padded[:, 2:], torch.full((2, 2, 3), 7)) + + w13 = torch.arange(24).reshape(2, 4, 3) + assert fp8._pad_w13_shards(w13, 2, 2) is w13 + + padded_w13 = fp8._pad_w13_shards(w13, 2, 3, pad_value=9) + expected_w13 = torch.tensor( + [ + [[0, 1, 2], [3, 4, 5], [9, 9, 9], [6, 7, 8], [9, 10, 11], [9, 9, 9]], + [ + [12, 13, 14], + [15, 16, 17], + [9, 9, 9], + [18, 19, 20], + [21, 22, 23], + [9, 9, 9], + ], + ] + ) + torch.testing.assert_close(padded_w13, expected_w13) + torch.testing.assert_close( + fp8._clamp_mxfp8_scale(torch.tensor([0, 2, 0], dtype=torch.uint8)), + torch.tensor([1, 2, 1], dtype=torch.uint8), + ) + + +def test_set_mxfp8_apply_tensor_reuses_matching_storage( + fp8_module: types.ModuleType, +) -> None: + fp8 = fp8_module + layer = torch.nn.Module() + + fp8._set_mxfp8_apply_tensor(layer, "weight_for_apply", torch.ones(2, 3)) + first = layer.weight_for_apply + first_data_ptr = first.data_ptr() + + fp8._set_mxfp8_apply_tensor(layer, "weight_for_apply", torch.full((2, 3), 4.0)) + + assert layer.weight_for_apply is first + assert layer.weight_for_apply.data_ptr() == first_data_ptr + torch.testing.assert_close(layer.weight_for_apply, torch.full((2, 3), 4.0)) + + +def test_process_mxfp8_moe_pads_kernel_tensors_without_changing_checkpoint_layout( + fp8_module: types.ModuleType, + monkeypatch: pytest.MonkeyPatch, +) -> None: + from vllm.model_executor.layers.fused_moe.oracle import fp8 as fp8_oracle + from vllm.model_executor.layers.fused_moe.oracle.fp8 import Fp8MoeBackend + + fp8 = fp8_module + captured: dict[str, Any] = {} + kernel_builds = 0 + + def fake_make_quant_config(**kwargs: Any) -> Any: + captured["quant_config_kwargs"] = kwargs + return types.SimpleNamespace( + w1_scale=kwargs["w1_scale"], + w2_scale=kwargs["w2_scale"], + ) + + def fake_make_kernel(**kwargs: Any) -> Any: + nonlocal kernel_builds + kernel_builds += 1 + captured["kernel_kwargs"] = kwargs + return types.SimpleNamespace() + + def fake_batched_shuffle( + layer: torch.nn.Module, + w13_weight: torch.Tensor, + w2_weight: torch.Tensor, + w13_scale: torch.Tensor, + w2_scale: torch.Tensor, + is_gated: bool, + epilogue_tile_m: int, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + captured.update( + { + "layer": layer, + "w13_weight": w13_weight, + "w2_weight": w2_weight, + "w13_scale": w13_scale, + "w2_scale": w2_scale, + "is_gated": is_gated, + "epilogue_tile_m": epilogue_tile_m, + } + ) + return tuple( + tensor.clone() for tensor in (w13_weight, w2_weight, w13_scale, w2_scale) + ) + + monkeypatch.setattr(fp8, "_shuffle_mxfp8_moe_batched", fake_batched_shuffle) + monkeypatch.setattr(fp8_oracle, "make_fp8_moe_quant_config", fake_make_quant_config) + monkeypatch.setattr(fp8_oracle, "make_fp8_moe_kernel", fake_make_kernel) + fp8.global_fp8_config = fp8.FP8Config() + + layer = torch.nn.Module() + layer.w13_weight = torch.nn.Parameter( + torch.arange(30, dtype=torch.float32).reshape(2, 3, 5), + requires_grad=False, + ) + layer.w2_weight = torch.nn.Parameter( + torch.arange(30, dtype=torch.float32).reshape(2, 5, 3), + requires_grad=False, + ) + layer.w13_weight_scale = torch.nn.Parameter( + torch.zeros(2, 3, 1, dtype=torch.uint8), + requires_grad=False, + ) + layer.w2_weight_scale = torch.nn.Parameter( + torch.zeros(2, 5, 1, dtype=torch.uint8), + requires_grad=False, + ) + layer.w13_weight_scale_from_checkpoint = torch.nn.Parameter( + torch.zeros(2, 3, 1, dtype=torch.uint8), + requires_grad=False, + ) + layer.w2_weight_scale_from_checkpoint = torch.nn.Parameter( + torch.zeros(2, 5, 1, dtype=torch.uint8), + requires_grad=False, + ) + moe_config = types.SimpleNamespace( + intermediate_size_per_partition=3, + is_act_and_mul=False, + ) + layer.moe_config = moe_config + layer._expert_routing_tables = lambda: None + quant_method = types.SimpleNamespace( + mxfp8_backend=Fp8MoeBackend.FLASHINFER_TRTLLM, + experts_cls=types.SimpleNamespace(is_monolithic=lambda: True), + weight_block_size=[1, 32], + moe=moe_config, + moe_kernel=None, + moe_quant_config=None, + ) + original_w13 = layer.w13_weight.detach().clone() + original_w2 = layer.w2_weight.detach().clone() + + fp8.process_weights_after_loading_mxfp8_moe(quant_method, layer) + + assert captured["layer"] is layer + assert captured["is_gated"] is False + assert captured["epilogue_tile_m"] == 128 + assert captured["w13_weight"].shape == (2, 128, 512) + assert captured["w2_weight"].shape == (2, 512, 128) + assert captured["w13_scale"].shape == (2, 128, 16) + assert captured["w2_scale"].shape == (2, 512, 4) + assert torch.count_nonzero(captured["w13_scale"] == 0) == 0 + assert torch.count_nonzero(captured["w2_scale"] == 0) == 0 + + torch.testing.assert_close(layer.w13_weight, original_w13) + torch.testing.assert_close(layer.w2_weight, original_w2) + assert layer.mxfp8_unpadded_hidden_size == 5 + assert layer.mxfp8_padded_hidden_size == 512 + assert layer.mxfp8_unpadded_intermediate_size_per_partition == 3 + assert layer.mxfp8_padded_intermediate_size_per_partition == 128 + assert layer.intermediate_size_per_partition == 128 + assert layer.moe_config.intermediate_size_per_partition == 128 + assert quant_method.moe.intermediate_size_per_partition == 128 + assert layer.w13_weight_for_apply.shape == (2, 128, 512) + assert layer.w2_weight_for_apply.shape == (2, 512, 128) + assert layer.w13_scale_for_apply.shape == (2, 128, 16) + assert layer.w2_scale_for_apply.shape == (2, 512, 4) + assert layer.weight_block_size == [1, 32] + assert captured["quant_config_kwargs"]["w1_scale"] is layer.w13_scale_for_apply + assert captured["quant_config_kwargs"]["w2_scale"] is layer.w2_scale_for_apply + assert captured["kernel_kwargs"]["moe_config"] is moe_config + assert captured["kernel_kwargs"]["routing_tables"] is None + assert kernel_builds == 1 + + w13_scale_for_apply = layer.w13_scale_for_apply + w2_scale_for_apply = layer.w2_scale_for_apply + fp8.process_weights_after_loading_mxfp8_moe(quant_method, layer) + + assert layer.w13_scale_for_apply is w13_scale_for_apply + assert layer.w2_scale_for_apply is w2_scale_for_apply + assert quant_method.moe_quant_config.w1_scale is w13_scale_for_apply + assert quant_method.moe_quant_config.w2_scale is w2_scale_for_apply + assert kernel_builds == 1 + + +def test_process_mxfp8_moe_rejects_non_trtllm_backend_before_mutation( + fp8_module: types.ModuleType, +) -> None: + from vllm.model_executor.layers.fused_moe.oracle.fp8 import Fp8MoeBackend + + fp8 = fp8_module + quant_method = types.SimpleNamespace( + mxfp8_backend=Fp8MoeBackend.DEEPGEMM, + experts_cls=types.SimpleNamespace(is_monolithic=lambda: True), + ) + layer = types.SimpleNamespace(marker=object()) + + with pytest.raises( + NotImplementedError, + match="requires the monolithic FlashInfer TRTLLM backend", + ): + fp8.process_weights_after_loading_mxfp8_moe(quant_method, layer) + + assert not hasattr(layer, "weight_block_size") + + +def test_apply_monolithic_mxfp8_moe_uses_vllm_025_moe_config( + fp8_module: types.ModuleType, +) -> None: + from vllm.model_executor.layers.fused_moe.activation import MoEActivation + + fp8 = fp8_module + captured: dict[str, Any] = {} + + def fake_apply( + x: torch.Tensor, + w13_weight: torch.Tensor, + w2_weight: torch.Tensor, + router_logits: torch.Tensor, + **kwargs: Any, + ) -> torch.Tensor: + captured.update( + { + "x": x, + "w13_weight": w13_weight, + "w2_weight": w2_weight, + "router_logits": router_logits, + } + ) + captured.update(kwargs) + return torch.zeros_like(x, dtype=torch.bfloat16) + + kernel = types.SimpleNamespace(apply_monolithic=fake_apply) + quant_method = types.SimpleNamespace( + is_monolithic=True, + moe_kernel=kernel, + ) + runtime_w13 = torch.empty(4, 128, 512, dtype=torch.float8_e4m3fn) + runtime_w2 = torch.empty(4, 512, 128, dtype=torch.float8_e4m3fn) + layer = types.SimpleNamespace( + activation=MoEActivation.RELU2_NO_MUL, + global_num_experts=32, + expert_map=None, + apply_router_weight_on_input=False, + num_expert_group=0, + topk_group=0, + routed_scaling_factor=1.0, + e_score_correction_bias=None, + w13_weight=torch.empty(4, 128, 512, dtype=torch.float8_e4m3fn), + w2_weight=torch.empty(4, 512, 128, dtype=torch.float8_e4m3fn), + w13_weight_for_apply=runtime_w13, + w2_weight_for_apply=runtime_w2, + mxfp8_padded_hidden_size=512, + ) + x = torch.ones(2, 64, dtype=torch.bfloat16) + router_logits = torch.zeros(2, 32, dtype=torch.bfloat16) + + output = fp8.apply_monolithic_mxfp8_moe( + quant_method, + layer, + x, + router_logits, + ) + + assert captured["x"].shape == (2, 512) + assert captured["w13_weight"] is runtime_w13 + assert captured["w2_weight"] is runtime_w2 + assert captured["router_logits"] is router_logits + assert captured["activation"] == MoEActivation.RELU2_NO_MUL + assert captured["global_num_experts"] == 32 + assert captured["expert_map"] is None + assert captured["apply_router_weight_on_input"] is False + assert captured["num_expert_group"] == 0 + assert captured["topk_group"] == 0 + assert captured["e_score_correction_bias"] is None + assert captured["routed_scaling_factor"] == 1.0 + assert output.shape == x.shape + + def test_process_weights_after_loading_copies_in_place_on_refit(monkeypatch): """Refit runs this every step; rebinding .data each time fragments memory. diff --git a/tests/unit/models/generation/test_vllm_generation.py b/tests/unit/models/generation/test_vllm_generation.py index 579d927d049..1a1bc61b0d8 100644 --- a/tests/unit/models/generation/test_vllm_generation.py +++ b/tests/unit/models/generation/test_vllm_generation.py @@ -74,6 +74,7 @@ "stop_strings": None, "vllm_cfg": { "precision": "bfloat16", + "refit_cache_loader_routes": False, "tensor_parallel_size": 1, "pipeline_parallel_size": 1, "expert_parallel_size": 1, @@ -143,6 +144,15 @@ } +@pytest.mark.vllm +def test_prepare_refit_info_skips_missing_metadata(): + generation = VllmGeneration.__new__(VllmGeneration) + generation.worker_group = MagicMock() + + assert generation.prepare_refit_info(None) is None + generation.worker_group.run_all_workers_single_data.assert_not_called() + + def test_context_capped_max_new_tokens(): assert ( _context_capped_max_new_tokens( diff --git a/tests/unit/models/generation/test_vllm_large_model.py b/tests/unit/models/generation/test_vllm_large_model.py index 9b18a446ebf..bf603642b27 100644 --- a/tests/unit/models/generation/test_vllm_large_model.py +++ b/tests/unit/models/generation/test_vllm_large_model.py @@ -45,6 +45,7 @@ "stop_strings": None, "vllm_cfg": { "precision": "bfloat16", + "refit_cache_loader_routes": False, "tensor_parallel_size": 8, "pipeline_parallel_size": 2, "expert_parallel_size": 1, diff --git a/tests/unit/models/generation/test_vllm_modelopt_real_quant_config.py b/tests/unit/models/generation/test_vllm_modelopt_real_quant_config.py index ae9f60d57a7..11443db281f 100644 --- a/tests/unit/models/generation/test_vllm_modelopt_real_quant_config.py +++ b/tests/unit/models/generation/test_vllm_modelopt_real_quant_config.py @@ -1165,9 +1165,9 @@ def make_model(expert_map): batched_forwarded = [] extension = _make_real_quant_extension(backend, make_model(None), []) + _patch_real_quant_load(monkeypatch, backend, batched_forwarded) extension.prepare_refit_info(state_dict_info) extension._nrl_w13_num_shards_by_prefix = {prefix: 1} - _patch_real_quant_load(monkeypatch, backend, batched_forwarded) assert ( extension._load_weights( [ diff --git a/tests/unit/models/generation/test_vllm_quant_backend.py b/tests/unit/models/generation/test_vllm_quant_backend.py index 7e88f4c993a..b375c7cf8d3 100644 --- a/tests/unit/models/generation/test_vllm_quant_backend.py +++ b/tests/unit/models/generation/test_vllm_quant_backend.py @@ -70,6 +70,7 @@ def _make_vllm_config( "quant_cfg": quant_cfg, "vllm_cfg": { "precision": "bfloat16", + "refit_cache_loader_routes": False, "tensor_parallel_size": 1, "pipeline_parallel_size": 1, "expert_parallel_size": 1, diff --git a/tests/unit/models/generation/test_vllm_refit_loader.py b/tests/unit/models/generation/test_vllm_refit_loader.py index baa769e91a2..7878eac51a8 100644 --- a/tests/unit/models/generation/test_vllm_refit_loader.py +++ b/tests/unit/models/generation/test_vllm_refit_loader.py @@ -145,6 +145,161 @@ def test_refit_load_weights_uses_full_weight_path_by_default(): assert loaded == [("model.weight", weight)] +@pytest.mark.vllm +def test_refit_loader_cache_records_replays_and_falls_back(monkeypatch): + from vllm.model_executor.model_loader.weight_utils import default_weight_loader + + from nemo_rl.models.generation.vllm.vllm_backend import ( + load_weights_maybe_cached, + ) + + events = [] + + def remote_loader(param, loaded_weight, *args, **kwargs): + events.append(("remote", loaded_weight, args, kwargs)) + return False + + def local_loader(param, loaded_weight, *args, **kwargs): + events.append(("local", loaded_weight, args, kwargs)) + with torch.no_grad(): + param.copy_(loaded_weight) + return None + + class Model: + def __init__(self): + self.remote = torch.nn.Parameter(torch.zeros(1), requires_grad=False) + self.local = torch.nn.Parameter(torch.zeros(1), requires_grad=False) + self.default = torch.nn.Parameter(torch.zeros(1), requires_grad=False) + self.remote.weight_loader = remote_loader + self.local.weight_loader = local_loader + self.default.weight_loader = default_weight_loader + self.load_calls = [] + + def named_parameters(self): + return [ + ("remote", self.remote), + ("local", self.local), + ("default", self.default), + ] + + def load_weights(self, *, weights): + self.load_calls.append([name for name, _weight in weights]) + loaded = set() + for name, weight in weights: + if name == "expert": + results = [ + self.remote.weight_loader( + self.remote, weight, "w1", expert_id=0 + ), + self.local.weight_loader(self.local, weight, "w1", expert_id=1), + ] + if any(result is not False for result in results): + loaded.add(name) + else: + self.default.weight_loader(self.default, weight) + loaded.add(name) + return loaded + + model = Model() + first_expert = torch.tensor([1.0]) + first_default = torch.tensor([2.0]) + second_expert = torch.tensor([3.0]) + second_default = torch.tensor([4.0]) + + assert load_weights_maybe_cached( + model, + [("expert", first_expert), ("default", first_default)], + cache_loader_routes=True, + ) == {"expert", "default"} + assert load_weights_maybe_cached( + model, + [("expert", second_expert), ("default", second_default)], + cache_loader_routes=True, + ) == {"expert", "default"} + + cache = model._nrl_refit_loader_cache + assert model.load_calls == [["expert", "default"], ["default"]] + assert cache.uncached == {"default"} + assert set(cache.calls) == {"expert"} + assert len(cache.calls["expert"]) == 2 + first_loader, first_param, first_args, first_kwargs = cache.calls["expert"][0] + second_loader, second_param, second_args, second_kwargs = cache.calls["expert"][1] + assert first_loader is remote_loader + assert first_param is model.remote + assert first_args == ("w1",) + assert first_kwargs == {"expert_id": 0} + assert second_loader is local_loader + assert second_param is model.local + assert second_args == ("w1",) + assert second_kwargs == {"expert_id": 1} + assert [event[0] for event in events] == ["remote", "local", "remote", "local"] + assert events[2][1] is second_expert + assert events[3][1] is second_expert + assert model.remote.weight_loader is remote_loader + assert model.local.weight_loader is local_loader + torch.testing.assert_close(model.local, second_expert) + torch.testing.assert_close(model.default, second_default) + + +@pytest.mark.vllm +def test_refit_loader_cache_invalidates_replaced_parameter(monkeypatch): + from nemo_rl.models.generation.vllm.vllm_backend import ( + load_weights_maybe_cached, + ) + + events = [] + + def make_loader(label): + def loader(param, loaded_weight): + events.append(label) + with torch.no_grad(): + param.copy_(loaded_weight) + + return loader + + class Model: + def __init__(self): + self.remote = torch.nn.Parameter(torch.zeros(1), requires_grad=False) + self.local = torch.nn.Parameter(torch.zeros(1), requires_grad=False) + self.remote.weight_loader = make_loader("remote") + self.local.weight_loader = make_loader("local") + self.load_calls = [] + + def named_parameters(self): + return [("remote", self.remote), ("local", self.local)] + + def load_weights(self, *, weights): + self.load_calls.append([name for name, _weight in weights]) + for _name, weight in weights: + self.remote.weight_loader(self.remote, weight) + self.local.weight_loader(self.local, weight) + return {name for name, _weight in weights} + + model = Model() + first = torch.tensor([1.0]) + second = torch.tensor([2.0]) + + assert load_weights_maybe_cached( + model, [("expert", first)], cache_loader_routes=True + ) == {"expert"} + cache = model._nrl_refit_loader_cache + old_local = model.local + model.local = torch.nn.Parameter(torch.zeros(1), requires_grad=False) + model.local.weight_loader = make_loader("replacement") + + assert load_weights_maybe_cached( + model, [("expert", second)], cache_loader_routes=True + ) == {"expert"} + + assert model.load_calls == [["expert"], ["expert"]] + assert events == ["remote", "local", "remote", "replacement"] + torch.testing.assert_close(old_local, first) + torch.testing.assert_close(model.local, second) + assert cache.calls == {} + assert cache.uncached == set() + assert cache.snapshot == {} + + @pytest.mark.vllm def test_refit_load_weights_dispatches_to_sharded_path_when_enabled(): from nemo_rl.models.generation.vllm.vllm_backend import ( diff --git a/tests/unit/models/generation/test_vllm_worker_helpers.py b/tests/unit/models/generation/test_vllm_worker_helpers.py index 29d61934447..9ef0e6de842 100644 --- a/tests/unit/models/generation/test_vllm_worker_helpers.py +++ b/tests/unit/models/generation/test_vllm_worker_helpers.py @@ -14,14 +14,41 @@ """Tests for vLLM worker helper functions.""" +from types import SimpleNamespace + import pytest from nemo_rl.models.generation.vllm.worker_utils import ( + configure_refit_runtime, + refit_cache_loader_routes_enabled, resolve_data_parallel_local_rank, resolve_distributed_executor_backend, ) +@pytest.mark.parametrize("enabled", [False, True]) +def test_refit_loader_cache_round_trips_through_additional_config(enabled): + vllm_kwargs = {"additional_config": {"existing": "value"}} + + configure_refit_runtime( + {"refit_cache_loader_routes": enabled}, + vllm_kwargs, + ) + + assert vllm_kwargs["additional_config"]["existing"] == "value" + vllm_config = SimpleNamespace(additional_config=vllm_kwargs["additional_config"]) + assert refit_cache_loader_routes_enabled(vllm_config) is enabled + + +def test_refit_loader_cache_defaults_to_disabled(): + vllm_kwargs = {} + + configure_refit_runtime({}, vllm_kwargs) + + vllm_config = SimpleNamespace(additional_config=vllm_kwargs["additional_config"]) + assert refit_cache_loader_routes_enabled(vllm_config) is False + + @pytest.mark.parametrize( ("tp", "pp", "ep", "expected"), [ diff --git a/tests/unit/models/generation/trtllm/test_trtllm_generation.py b/tests/unit/models/generation/trtllm/test_trtllm_generation.py index 1123ebae58a..3424e65082c 100644 --- a/tests/unit/models/generation/trtllm/test_trtllm_generation.py +++ b/tests/unit/models/generation/trtllm/test_trtllm_generation.py @@ -185,8 +185,8 @@ async def worker_result(): } ) - generation.worker_group.run_single_worker_single_data.side_effect = ( - lambda **_: worker_result() + generation.worker_group.run_single_worker_single_data.side_effect = lambda **_: ( + worker_result() ) data = BatchedDataDict( { @@ -297,6 +297,30 @@ def test_generation_lifecycle_routes_by_colocation( ) +def test_prepare_refit_info_skips_missing_metadata_and_dispatches_present_metadata( + monkeypatch, +): + generation = _bare_generation() + ray_get = MagicMock() + monkeypatch.setattr(trtllm_generation.ray, "get", ray_get) + + assert generation.prepare_refit_info(None) is None + generation.worker_group.run_all_workers_single_data.assert_not_called() + ray_get.assert_not_called() + + state_dict_info = {"weight": {"shape": [2, 2]}} + futures = [SimpleNamespace()] + generation.worker_group.run_all_workers_single_data.return_value = futures + + assert generation.prepare_refit_info(state_dict_info) is None + generation.worker_group.run_all_workers_single_data.assert_called_once_with( + "prepare_refit_info_async", + state_dict_info=state_dict_info, + run_rank_0_only_axes=["tensor_parallel"], + ) + ray_get.assert_called_once_with(futures) + + @pytest.mark.parametrize( ("in_flight", "recompute_kv", "expected_drain"), [(False, False, True), (True, False, False), (True, True, False)], diff --git a/tests/unit/models/policy/test_megatron_worker.py b/tests/unit/models/policy/test_megatron_worker.py index 1755cbac3f3..c868f52a2b2 100644 --- a/tests/unit/models/policy/test_megatron_worker.py +++ b/tests/unit/models/policy/test_megatron_worker.py @@ -17,7 +17,7 @@ import time from pathlib import Path from types import SimpleNamespace -from typing import Optional +from typing import Any, Optional from unittest.mock import MagicMock import numpy as np @@ -495,6 +495,379 @@ def test_megatron_move_model_does_not_serialize_extra_state(): assert model.scale.device.type == "cpu" +def test_checkpoint_engine_prequant_handshake_exports_mxfp8_weights(): + from nemo_rl.models.policy.workers.checkpoint_engine import ( + MegatronCheckpointEngineSendMixin, + ) + from nemo_rl.models.policy.workers.megatron_policy_worker import ( + MegatronPolicyWorkerImpl, + ) + from nemo_rl.weight_sync.checkpoint_engine_weight_synchronizer import ( + CheckpointEngineWeightSynchronizer, + ) + + class _PrequantCheckpointWorker(MegatronCheckpointEngineSendMixin): + enable_refit_prequantize = MegatronPolicyWorkerImpl.enable_refit_prequantize + _is_fp8_export = MegatronPolicyWorkerImpl._is_fp8_export + _iter_params_with_optional_kv_scales = ( + MegatronPolicyWorkerImpl._iter_params_with_optional_kv_scales + ) + _maybe_prequantize_param = MegatronPolicyWorkerImpl._maybe_prequantize_param + + name = "model.layers.0.mlp.down_proj.weight" + weight = torch.randn(64, 64, dtype=torch.bfloat16) + worker = _PrequantCheckpointWorker() + worker._refit_prequant_names = set() + worker._refit_param_info_hf = None + worker.fp8_cfg = None + worker.model = object() + worker.draft_model = None + worker.refit_conversion_tasks = [] + worker.cfg = {"megatron_cfg": {"enabled": True}} + worker.megatron_bridge = SimpleNamespace( + export_hf_weights=lambda *_args, **_kwargs: iter([(name, weight)]) + ) + + def _prepare_refit_info() -> dict[str, Any]: + worker._refit_param_info_hf = {name: (weight.shape, weight.dtype)} + return worker._refit_param_info_hf + + worker.prepare_refit_info = _prepare_refit_info + worker.checkpoint_engine = SimpleNamespace(get_target_weight_layout=lambda: None) + + class _Generation: + def __init__(self) -> None: + self.refit_info: list[dict[str, Any]] = [] + + def prepare_refit_info( + self, state_dict_info: dict[str, Any] | None + ) -> list[str] | None: + assert state_dict_info is not None + self.refit_info.append(state_dict_info) + return [name] if len(self.refit_info) == 1 else None + + generation = _Generation() + synchronizer = CheckpointEngineWeightSynchronizer(worker, generation, {}) + synchronizer._ensure_checkpoint_engine_ready = lambda: None + + synchronizer.init_communicator() + exported = dict(worker._checkpoint_engine_weight_iterator()) + + scale_name = name + "_scale_from_checkpoint" + assert generation.refit_info[1][name][1] == torch.float8_e4m3fn + assert generation.refit_info[1][scale_name][1] == torch.uint8 + assert exported[name].dtype == torch.float8_e4m3fn + assert exported[scale_name].dtype == torch.uint8 + assert exported[scale_name].shape == (64, 2) + + +def test_reference_model_pinned_swap_restores_state_and_reuses_buffer(monkeypatch): + from nemo_rl.models.policy.workers.megatron_policy_worker import ( + MegatronPolicyWorkerImpl, + ) + + class Model(torch.nn.Module): + def __init__(self): + super().__init__() + self.weight = torch.nn.Parameter(torch.tensor([1.0])) + self.register_buffer("extra_state_cache", torch.tensor([2.0])) + + worker = object.__new__(MegatronPolicyWorkerImpl) + worker.model = Model() + worker.reference_state_dict = { + "weight": torch.tensor([11.0]), + "extra_state_cache": torch.tensor([12.0]), + } + worker.cfg = { + "megatron_cfg": { + "pinned_reference_swap": True, + "empty_unused_memory_level": 0, + } + } + worker._pinned_swap_save_buffers = {} + worker.should_disable_forward_pre_hook = False + worker.sampling_params = None + allocations = [] + original_empty = torch.empty + + def empty(*args, **kwargs): + allocations.append(kwargs) + kwargs = {**kwargs, "pin_memory": False} + return original_empty(*args, **kwargs) + + synchronize = MagicMock() + monkeypatch.setattr(torch, "empty", empty) + monkeypatch.setattr(torch.cuda, "synchronize", synchronize) + + cached_buffer = None + for _ in range(2): + with worker.use_reference_model(): + torch.testing.assert_close(worker.model.weight, torch.tensor([11.0])) + torch.testing.assert_close( + worker.model.extra_state_cache, torch.tensor([12.0]) + ) + torch.testing.assert_close(worker.model.weight, torch.tensor([1.0])) + torch.testing.assert_close(worker.model.extra_state_cache, torch.tensor([2.0])) + if cached_buffer is None: + cached_buffer = worker._pinned_swap_save_buffers["weight"] + else: + assert worker._pinned_swap_save_buffers["weight"] is cached_buffer + + assert list(worker._pinned_swap_save_buffers) == ["weight"] + assert allocations == [ + { + "dtype": torch.float32, + "device": "cpu", + "pin_memory": True, + } + ] + assert synchronize.call_count == 6 + + +def test_clear_rope_and_moe_dispatcher_caches_clears_tensor_state(monkeypatch): + from megatron.core.models.common.embeddings import rotary_pos_embedding + + from nemo_rl.models.policy.workers.megatron_policy_worker import ( + MegatronPolicyWorkerImpl, + ) + + cache_clear = MagicMock() + + def forward(): + return None + + forward.cache_clear = cache_clear + monkeypatch.setattr( + rotary_pos_embedding, + "RotaryEmbedding", + SimpleNamespace(forward=forward), + ) + dispatcher = SimpleNamespace( + probs=torch.ones(1), + routing_map=torch.ones(1), + reversed_local_input_permutation_mapping=torch.ones(1), + local_probs=torch.ones(1), + local_map=torch.ones(1), + non_tensor="keep", + ) + worker = object.__new__(MegatronPolicyWorkerImpl) + worker.model = SimpleNamespace( + modules=lambda: [ + SimpleNamespace(), + SimpleNamespace(token_dispatcher=None), + SimpleNamespace(token_dispatcher=dispatcher), + ] + ) + + worker._clear_rope_and_moe_dispatcher_caches() + + cache_clear.assert_called_once_with() + assert dispatcher.probs is None + assert dispatcher.routing_map is None + assert dispatcher.reversed_local_input_permutation_mapping is None + assert dispatcher.local_probs is None + assert dispatcher.local_map is None + assert dispatcher.non_tensor == "keep" + + +def test_clear_rope_and_moe_dispatcher_caches_is_best_effort(monkeypatch): + from megatron.core.models.common.embeddings import rotary_pos_embedding + + from nemo_rl.models.policy.workers.megatron_policy_worker import ( + MegatronPolicyWorkerImpl, + ) + + def forward(): + return None + + forward.cache_clear = MagicMock(side_effect=RuntimeError("rotary cache")) + monkeypatch.setattr( + rotary_pos_embedding, + "RotaryEmbedding", + SimpleNamespace(forward=forward), + ) + worker = object.__new__(MegatronPolicyWorkerImpl) + worker.model = SimpleNamespace( + modules=MagicMock(side_effect=RuntimeError("module traversal")) + ) + + worker._clear_rope_and_moe_dispatcher_caches() + + +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float8_e4m3fn]) +def test_maybe_prequantize_param_passthrough_when_not_selected(dtype): + from nemo_rl.models.policy.workers.megatron_policy_worker import ( + MegatronPolicyWorkerImpl, + ) + + worker = object.__new__(MegatronPolicyWorkerImpl) + name = "model.weight" + worker._refit_prequant_names = set() + tensor = torch.ones(2, 2, dtype=dtype) + + result = list(worker._maybe_prequantize_param(name, tensor)) + + assert len(result) == 1 + assert result[0][0] == name + assert result[0][1] is tensor + + +def test_maybe_prequantize_param_rejects_fp8_trainer_storage(): + from nemo_rl.models.policy.workers.megatron_policy_worker import ( + MegatronPolicyWorkerImpl, + ) + + worker = object.__new__(MegatronPolicyWorkerImpl) + name = "model.weight" + worker._refit_prequant_names = {name} + tensor = torch.ones(2, 2, dtype=torch.float8_e4m3fn) + + with pytest.raises(ValueError, match="BF16 trainer-exported weights"): + list(worker._maybe_prequantize_param(name, tensor)) + + +def test_enable_refit_prequantize_rejects_blockwise_fp8_storage(): + from nemo_rl.models.policy.workers.megatron_policy_worker import ( + MegatronPolicyWorkerImpl, + ) + + worker = object.__new__(MegatronPolicyWorkerImpl) + worker.fp8_cfg = { + "fp8_param": True, + "fp8_recipe": "blockwise", + } + + with pytest.raises(ValueError, match="BF16 trainer-exported weights"): + worker.enable_refit_prequantize(["model.weight"]) + + +def test_enable_refit_prequantize_requires_prepare_refit_info(): + from nemo_rl.models.policy.workers.megatron_policy_worker import ( + MegatronPolicyWorkerImpl, + ) + + worker = object.__new__(MegatronPolicyWorkerImpl) + worker.fp8_cfg = None + worker._refit_param_info_hf = None + + with pytest.raises(RuntimeError, match="prepare_refit_info"): + worker.enable_refit_prequantize(["model.weight"]) + + +def test_enable_refit_prequantize_derives_metadata_without_export(): + from nemo_rl.models.policy.workers.megatron_policy_worker import ( + MegatronPolicyWorkerImpl, + ) + + worker = object.__new__(MegatronPolicyWorkerImpl) + worker.fp8_cfg = None + worker._refit_param_info_hf = { + "model.a.weight": (torch.Size([4, 64]), torch.bfloat16), + "model.b.weight": (torch.Size([4, 64]), torch.bfloat16), + } + + def _fail_iter(*_args, **_kwargs): + raise AssertionError("metadata derivation must not re-export weights") + + worker._iter_params_with_optional_kv_scales = _fail_iter + + info = worker.enable_refit_prequantize(["model.a.weight"]) + + assert info["model.a.weight"] == (torch.Size([4, 64]), torch.float8_e4m3fn) + assert info["model.a.weight_scale_from_checkpoint"] == ( + torch.Size([4, 2]), + torch.uint8, + ) + assert info["model.b.weight"] == (torch.Size([4, 64]), torch.bfloat16) + assert worker._refit_prequant_names == {"model.a.weight"} + + +def test_enable_refit_prequantize_rejects_indivisible_last_dim(): + from nemo_rl.models.policy.workers.megatron_policy_worker import ( + MegatronPolicyWorkerImpl, + ) + + worker = object.__new__(MegatronPolicyWorkerImpl) + worker.fp8_cfg = None + worker._refit_param_info_hf = { + "model.weight": (torch.Size([4, 48]), torch.bfloat16), + } + + with pytest.raises(ValueError, match="divisible"): + worker.enable_refit_prequantize(["model.weight"]) + + +@pytest.mark.parametrize( + "slim,offload_optimizer", + [(False, True), (True, True), (True, False)], +) +def test_offload_after_refit_routes_cleanup_by_mode( + monkeypatch, slim, offload_optimizer +): + from nemo_rl.models.policy.workers.megatron_policy_worker import ( + MegatronPolicyWorkerImpl, + ) + + worker = object.__new__(MegatronPolicyWorkerImpl) + model = SimpleNamespace(eval=MagicMock()) + worker.model = model + worker.move_model = MagicMock(return_value=model) + worker.cfg = { + "megatron_cfg": { + "refit_slim_offload_after": slim, + "clear_memory_caches_before_refit": True, + } + } + worker.fp8_cfg = {"force_clear_fp8_caches": True} + worker._clear_fp8_caches = MagicMock() + worker._clear_rope_and_moe_dispatcher_caches = MagicMock() + worker.optimizer = object() + worker.optimizer_cpu_offload = False + worker.offload_optimizer_for_refit = offload_optimizer + worker.move_optimizer = MagicMock() + worker.offload_before_refit = MagicMock() + worker.finalize_async_save = MagicMock() + collect = MagicMock() + empty_cache = MagicMock() + monkeypatch.setattr( + torch, + "randn", + lambda *_args, **_kwargs: SimpleNamespace(cuda=lambda: None), + ) + monkeypatch.setattr(torch.cuda, "memory_allocated", lambda: 0) + monkeypatch.setattr(torch.cuda, "memory_reserved", lambda: 0) + monkeypatch.setattr(torch.cuda, "empty_cache", empty_cache) + monkeypatch.setattr(torch.cuda.nvtx, "range_push", lambda *_args: None) + monkeypatch.setattr(torch.cuda.nvtx, "range_pop", lambda: None) + monkeypatch.setattr( + "nemo_rl.models.policy.workers.megatron_policy_worker.gc.collect", + collect, + ) + + worker.offload_after_refit() + + worker.finalize_async_save.assert_called_once_with() + worker.move_model.assert_called_once_with(model, "cpu") + model.eval.assert_called_once_with() + if slim: + worker._clear_fp8_caches.assert_called_once_with() + worker._clear_rope_and_moe_dispatcher_caches.assert_called_once_with() + if offload_optimizer: + worker.move_optimizer.assert_called_once_with("cpu") + else: + worker.move_optimizer.assert_not_called() + collect.assert_called_once_with() + empty_cache.assert_called_once_with() + worker.offload_before_refit.assert_not_called() + else: + worker.offload_before_refit.assert_called_once_with() + worker._clear_fp8_caches.assert_not_called() + worker._clear_rope_and_moe_dispatcher_caches.assert_not_called() + worker.move_optimizer.assert_not_called() + collect.assert_not_called() + empty_cache.assert_not_called() + + def test_megatron_prepare_for_training_restores_optimizer(): from nemo_rl.models.policy.workers.megatron_policy_worker import ( MegatronPolicyWorkerImpl, diff --git a/tests/unit/models/policy/test_policy_validation.py b/tests/unit/models/policy/test_policy_validation.py index 6a4c6a8d3ab..67e4c858086 100644 --- a/tests/unit/models/policy/test_policy_validation.py +++ b/tests/unit/models/policy/test_policy_validation.py @@ -35,6 +35,30 @@ def test_shutdown_succeeds_before_worker_group_is_initialized(capsys) -> None: assert capsys.readouterr().out == "" +def test_enable_refit_prequantize_forwards_names_and_returns_metadata( + monkeypatch, +) -> None: + policy = Policy.__new__(Policy) + futures = [object()] + updated_info = { + "model.weight": ((2, 2), "float8_e4m3fn"), + "model.weight_scale_from_checkpoint": ((2, 1), "uint8"), + } + policy.worker_group = MagicMock() + policy.worker_group.run_all_workers_single_data.return_value = futures + ray_get = MagicMock(return_value=[updated_info]) + monkeypatch.setattr("nemo_rl.models.policy.lm_policy.ray.get", ray_get) + + result = policy.enable_refit_prequantize(["model.weight"]) + + assert result is updated_info + policy.worker_group.run_all_workers_single_data.assert_called_once_with( + "enable_refit_prequantize", + param_names=["model.weight"], + ) + ray_get.assert_called_once_with(futures) + + def create_mock_cluster(world_size: int): """Create a mock cluster with the specified world size.""" cluster = MagicMock() diff --git a/tests/unit/reference_configs/distillation_math.yaml b/tests/unit/reference_configs/distillation_math.yaml index 3eec82229b2..45af0a0539f 100644 --- a/tests/unit/reference_configs/distillation_math.yaml +++ b/tests/unit/reference_configs/distillation_math.yaml @@ -205,6 +205,7 @@ policy: &POLICY_BASE vllm_cfg: async_engine: false precision: ${...precision} + refit_cache_loader_routes: false # Replay stable vLLM loader routes across refits. kv_cache_dtype: "auto" logprobs_mode: processed_logprobs tensor_parallel_size: 1 diff --git a/tests/unit/reference_configs/eval.yaml b/tests/unit/reference_configs/eval.yaml index 95813a0a433..53d1a5dc138 100644 --- a/tests/unit/reference_configs/eval.yaml +++ b/tests/unit/reference_configs/eval.yaml @@ -22,6 +22,7 @@ generation: vllm_cfg: async_engine: false precision: "bfloat16" + refit_cache_loader_routes: false tensor_parallel_size: 1 pipeline_parallel_size: 1 expert_parallel_size: 1 diff --git a/tests/unit/reference_configs/grpo_math_1B.yaml b/tests/unit/reference_configs/grpo_math_1B.yaml index c51186f9aea..12a818efb6e 100644 --- a/tests/unit/reference_configs/grpo_math_1B.yaml +++ b/tests/unit/reference_configs/grpo_math_1B.yaml @@ -181,6 +181,8 @@ policy: megatron_cfg: enabled: false + refit_slim_offload_after: false + pinned_reference_swap: false model_overrides: {} checkpoint: async_save: true @@ -332,6 +334,7 @@ policy: # makes the training sequence length divisible by the tensor parallel size # this is useful for sequence parallel training make_sequence_length_divisible_by: ${policy.dtensor_cfg.tensor_parallel_size} + refit_persistent_ipc_buffers: false max_grad_norm: 1.0 optimizer: @@ -405,6 +408,9 @@ policy: vllm_cfg: async_engine: false precision: ${policy.precision} + # MXFP8 + Megatron only: quantize on trainer and stream E4M3 values plus scales. + refit_prequantize: false + refit_cache_loader_routes: false # Replay stable vLLM loader routes across refits. kv_cache_dtype: "auto" logprobs_mode: processed_logprobs tensor_parallel_size: 1 diff --git a/tests/unit/reference_configs/ppo_math_1B_megatron.yaml b/tests/unit/reference_configs/ppo_math_1B_megatron.yaml index 71a9251a559..04c1b446782 100644 --- a/tests/unit/reference_configs/ppo_math_1B_megatron.yaml +++ b/tests/unit/reference_configs/ppo_math_1B_megatron.yaml @@ -252,6 +252,7 @@ policy: vllm_cfg: async_engine: false precision: ${policy.precision} + refit_cache_loader_routes: false # Replay stable vLLM loader routes across refits. kv_cache_dtype: "auto" logprobs_mode: processed_logprobs tensor_parallel_size: 1 diff --git a/tests/unit/utils/test_packed_tensor.py b/tests/unit/utils/test_packed_tensor.py index ca4cf6e9548..3eabbd47df1 100644 --- a/tests/unit/utils/test_packed_tensor.py +++ b/tests/unit/utils/test_packed_tensor.py @@ -23,6 +23,18 @@ ) +@pytest.mark.parametrize("buffer_size_bytes", [0, -1]) +def test_producer_rejects_nonpositive_buffer_size(buffer_size_bytes): + with pytest.raises(ValueError, match="buffer_size_bytes must be > 0"): + packed_broadcast_producer( + iterator=iter([]), + group=None, + src=0, + post_iter_func=lambda x: x, + buffer_size_bytes=buffer_size_bytes, + ) + + class MockCommunicationGroup: """Mock communication group for testing broadcast operations.""" diff --git a/tests/unit/weight_sync/test_checkpoint_engine_weight_synchronizer.py b/tests/unit/weight_sync/test_checkpoint_engine_weight_synchronizer.py index f7e8a7c43e1..67408fe05b9 100644 --- a/tests/unit/weight_sync/test_checkpoint_engine_weight_synchronizer.py +++ b/tests/unit/weight_sync/test_checkpoint_engine_weight_synchronizer.py @@ -14,7 +14,7 @@ """Tests for checkpoint-engine weight synchronization and factory routing.""" -from unittest.mock import MagicMock, patch +from unittest.mock import MagicMock, call, patch import pytest @@ -33,6 +33,7 @@ def _mock_policy(**overrides): policy = MagicMock() + policy.cfg = {"megatron_cfg": {"enabled": True}} policy.offload_before_refit.return_value = None policy.offload_after_refit.return_value = None policy.prepare_refit_info.return_value = {"layer_0": {"shape": [4096, 4096]}} @@ -151,6 +152,54 @@ def _checkpoint_sync( class TestCheckpointEngineWeightSynchronizer: + def test_init_communicator_completes_prequant_handshake(self): + sync = _checkpoint_sync(MagicMock()) + sync._ensure_checkpoint_engine_ready = MagicMock() + state_dict_info = {"layer_0": {"shape": [4096, 4096]}} + prequant_names = ["layer_0"] + updated_info = {"layer_0": {"shape": [4096, 4096], "dtype": "float8_e4m3fn"}} + sync._policy.prepare_refit_info.return_value = state_dict_info + sync._generation.prepare_refit_info.side_effect = [prequant_names, None] + sync._policy.enable_refit_prequantize.return_value = updated_info + + sync.init_communicator() + + sync._policy.enable_refit_prequantize.assert_called_once_with(prequant_names) + assert sync._generation.prepare_refit_info.call_args_list == [ + call(state_dict_info), + call(updated_info), + ] + sync._ensure_checkpoint_engine_ready.assert_called_once_with() + + def test_init_communicator_rejects_missing_prequant_metadata(self): + sync = _checkpoint_sync(MagicMock()) + sync._ensure_checkpoint_engine_ready = MagicMock() + sync._policy.prepare_refit_info.return_value = { + "layer_0": {"shape": [4096, 4096]} + } + sync._generation.prepare_refit_info.return_value = ["layer_0"] + sync._policy.enable_refit_prequantize.return_value = None + + with pytest.raises( + RuntimeError, + match="did not return updated metadata", + ): + sync.init_communicator() + + sync._ensure_checkpoint_engine_ready.assert_not_called() + + def test_init_communicator_rejects_prequantization_without_megatron(self): + sync = _checkpoint_sync(MagicMock()) + sync._ensure_checkpoint_engine_ready = MagicMock() + sync._policy.cfg = {"megatron_cfg": {"enabled": False}} + sync._generation.prepare_refit_info.return_value = ["layer_0"] + + with pytest.raises(ValueError, match="requires the Megatron policy backend"): + sync.init_communicator() + + sync._policy.enable_refit_prequantize.assert_not_called() + sync._ensure_checkpoint_engine_ready.assert_not_called() + @patch("nemo_rl.weight_sync.checkpoint_engine_weight_synchronizer.ray") def test_bucket_uses_minimum_total_memory_and_is_cached(self, mock_ray, capsys): config = _checkpoint_engine_cfg(bucket_memory_ratio=0.125) diff --git a/tests/unit/weight_sync/test_vllm_remote_sparse_weight_synchronizer.py b/tests/unit/weight_sync/test_vllm_remote_sparse_weight_synchronizer.py index 3724f93835f..673938232fe 100644 --- a/tests/unit/weight_sync/test_vllm_remote_sparse_weight_synchronizer.py +++ b/tests/unit/weight_sync/test_vllm_remote_sparse_weight_synchronizer.py @@ -105,6 +105,16 @@ def test_validate_remote_sparse_refit_accepts_supported_scope(): ({"refit_cfg": {"sparse": {"storage": {"s3_bucket": None}}}}, {}), ({"quant_cfg": "fp8"}, {}), ({"vllm_cfg": {"precision": "fp8", "kv_cache_dtype": "auto"}}, {}), + ( + { + "vllm_cfg": { + "precision": "bfloat16", + "kv_cache_dtype": "auto", + "refit_prequantize": True, + } + }, + {}, + ), ( {"vllm_cfg": {"precision": "bfloat16", "kv_cache_dtype": "fp8_e4m3"}}, {}, diff --git a/tests/unit/weight_sync/test_weight_synchronizer.py b/tests/unit/weight_sync/test_weight_synchronizer.py index 77219047268..a3523d2fc75 100644 --- a/tests/unit/weight_sync/test_weight_synchronizer.py +++ b/tests/unit/weight_sync/test_weight_synchronizer.py @@ -14,7 +14,7 @@ """Unit tests for the WeightSynchronizer abstraction and its implementations.""" -from unittest.mock import MagicMock, patch +from unittest.mock import MagicMock, call, patch import pytest @@ -29,7 +29,10 @@ CollectiveWeightSynchronizer, ) from nemo_rl.weight_sync.factory import create_weight_synchronizer -from nemo_rl.weight_sync.interfaces import WeightSynchronizer +from nemo_rl.weight_sync.interfaces import ( + WeightSynchronizer, + initialize_refit_metadata, +) from nemo_rl.weight_sync.ipc_weight_synchronizer import ( IPCWeightSynchronizer, ) @@ -106,6 +109,73 @@ class IncompleteSync(WeightSynchronizer): IncompleteSync() # type: ignore[abstract] +# --------------------------------------------------------------------------- +# Shared refit metadata handshake +# --------------------------------------------------------------------------- + + +class TestInitializeRefitMetadata: + def test_returns_when_generation_does_not_request_prequantization(self): + policy = _mock_policy() + generation = _mock_generation() + state_dict_info = policy.prepare_refit_info.return_value + + initialize_refit_metadata(policy, generation) + + generation.prepare_refit_info.assert_called_once_with(state_dict_info) + policy.enable_refit_prequantize.assert_not_called() + + def test_rejects_prequantization_without_megatron(self): + policy = _mock_policy() + policy.cfg = {"megatron_cfg": {"enabled": False}} + generation = _mock_generation() + generation.prepare_refit_info.return_value = ["layer_0"] + + with pytest.raises(ValueError, match="requires the Megatron policy backend"): + initialize_refit_metadata(policy, generation) + + policy.enable_refit_prequantize.assert_not_called() + + def test_refreshes_generation_metadata_after_prequantization(self): + policy = _mock_policy() + policy.cfg = {"megatron_cfg": {"enabled": True}} + generation = _mock_generation() + state_dict_info = policy.prepare_refit_info.return_value + updated_info = { + "layer_0": { + "shape": [4096, 4096], + "dtype": "float8_e4m3fn", + }, + "layer_0_scale_from_checkpoint": { + "shape": [4096, 128], + "dtype": "uint8", + }, + } + generation.prepare_refit_info.side_effect = [["layer_0"], None] + policy.enable_refit_prequantize.return_value = updated_info + + initialize_refit_metadata(policy, generation) + + policy.enable_refit_prequantize.assert_called_once_with(["layer_0"]) + assert generation.prepare_refit_info.call_args_list == [ + call(state_dict_info), + call(updated_info), + ] + + def test_rejects_missing_prequantized_metadata(self): + policy = _mock_policy() + policy.cfg = {"megatron_cfg": {"enabled": True}} + policy.enable_refit_prequantize.return_value = None + generation = _mock_generation() + generation.prepare_refit_info.return_value = ["layer_0"] + + with pytest.raises( + RuntimeError, + match="did not return updated metadata", + ): + initialize_refit_metadata(policy, generation) + + # --------------------------------------------------------------------------- # IPCWeightSynchronizer # --------------------------------------------------------------------------- @@ -190,6 +260,29 @@ def test_init_communicator(self): policy.prepare_refit_info.assert_called_once() gen.prepare_refit_info.assert_called_once() + def test_init_communicator_completes_prequantization_handshake(self): + policy = _mock_policy() + policy.cfg = {"megatron_cfg": {"enabled": True}} + updated_info = { + "layer_0": { + "shape": [4096, 4096], + "dtype": "float8_e4m3fn", + } + } + policy.enable_refit_prequantize.return_value = updated_info + gen = _mock_generation() + state_dict_info = policy.prepare_refit_info.return_value + gen.prepare_refit_info.side_effect = [["layer_0"], None] + sync = IPCWeightSynchronizer(policy, gen) + + sync.init_communicator() + + policy.enable_refit_prequantize.assert_called_once_with(["layer_0"]) + assert gen.prepare_refit_info.call_args_list == [ + call(state_dict_info), + call(updated_info), + ] + @patch("nemo_rl.weight_sync.ipc_weight_synchronizer.ray") def test_phase_restoration_on_transfer_failure(self, mock_ray): """offload_after_refit and kv_cache prep run even when transfer raises.""" @@ -543,6 +636,30 @@ def test_sync_weights_passes_kv_scales(self, mock_ray): call_kwargs = policy.broadcast_weights_for_collective.call_args assert call_kwargs.kwargs["kv_scales"] == kv_scales + @patch("nemo_rl.weight_sync.collective_weight_synchronizer.ray") + def test_sync_weights_forwards_fixed_buffer_size(self, mock_ray): + mock_ray.get.return_value = [True] + policy = _mock_policy() + gen = _mock_generation() + sync = CollectiveWeightSynchronizer( + policy, + gen, + _mock_cluster(), + _mock_cluster(), + refit_buffer_size_gb=1.5, + ) + + sync.sync_weights() + + expected_bytes = int(1.5 * 1024**3) + policy.broadcast_weights_for_collective.assert_called_once_with( + kv_scales=None, + buffer_size_bytes=expected_bytes, + ) + gen.update_weights_from_collective.assert_called_once_with( + buffer_size_bytes=expected_bytes + ) + @patch("nemo_rl.weight_sync.collective_weight_synchronizer.ray") def test_sync_weights_raises_on_failure(self, mock_ray): mock_ray.get.side_effect = [ @@ -580,6 +697,37 @@ def test_init_communicator_sets_up_collective(self, mock_ray): "10.0.0.1", 29500, 6, train_world_size=4 ) + @patch("nemo_rl.weight_sync.collective_weight_synchronizer.ray") + def test_init_communicator_prequantizes_before_collective_setup(self, mock_ray): + mock_ray.get.return_value = [True] + policy = _mock_policy() + policy.cfg = {"megatron_cfg": {"enabled": True}} + updated_info = { + "layer_0": { + "shape": [4096, 4096], + "dtype": "float8_e4m3fn", + } + } + policy.enable_refit_prequantize.return_value = updated_info + gen = _mock_generation() + state_dict_info = policy.prepare_refit_info.return_value + gen.prepare_refit_info.side_effect = [["layer_0"], None] + sync = CollectiveWeightSynchronizer( + policy, + gen, + _mock_cluster(world_size=4), + _mock_cluster(world_size=2), + ) + + sync.init_communicator() + + policy.enable_refit_prequantize.assert_called_once_with(["layer_0"]) + assert gen.prepare_refit_info.call_args_list == [ + call(state_dict_info), + call(updated_info), + ] + policy.init_collective.assert_called_once() + @patch("nemo_rl.weight_sync.collective_weight_synchronizer.ray") def test_backend_sender_contract_controls_geometry_and_world_size(self, mock_ray): mock_ray.get.return_value = [True] @@ -855,8 +1003,10 @@ def test_non_colocated_vllm_returns_collective(self): colocated=False, train_cluster=_mock_cluster(), inference_cluster=_mock_cluster(), + refit_buffer_size_gb=1.5, ) assert isinstance(sync, CollectiveWeightSynchronizer) + assert sync._buffer_size_bytes == int(1.5 * 1024**3) def test_non_colocated_dynamo_returns_collective(self): sync = create_weight_synchronizer(