perf: Reduce MXFP8 refit weight-transfer overheads

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:
 
 ```

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:

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

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

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

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

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

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:

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
 

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

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:

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):