perf: Reduce MXFP8 refit weight-transfer overheads

examples/configs/grpo_math_1B.yaml

@@ -158,6 +158,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
     checkpoint:
       async_save: true
       ckpt_assume_constant_structure: true
@@ -306,6 +312,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:

nemo_rl/algorithms/distillation.py

@@ -36,7 +36,7 @@
     DistillationLossDataDict,
     DistillationLossFn,
 )
-from nemo_rl.algorithms.utils import set_seed
+from nemo_rl.algorithms.utils import maybe_enable_refit_prequantize, set_seed
 from nemo_rl.data import DataConfig
 from nemo_rl.data.collate_fn import rl_collate_fn
 from nemo_rl.data.datasets import AllTaskProcessedDataset
@@ -620,7 +620,9 @@ 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)
+        maybe_enable_refit_prequantize(
+            student_policy, student_generation, state_dict_info, master_config.policy
+        )
 
     # if it is not colocated inference, initialize collective communication for update weights
     if not colocated_inference and checkpoint_engine_config is None:

nemo_rl/algorithms/grpo.py

@@ -54,6 +54,7 @@
     calculate_baseline_and_std_per_prompt,
     get_gdpo_reward_component_keys,
     log_generation_metrics_to_wandb,
+    maybe_enable_refit_prequantize,
     print_efficiency_summary,
     print_performance_metrics,
     set_seed,
@@ -1384,7 +1385,9 @@ def init_trtllm():
     else:
         state_dict_info = policy.prepare_refit_info()
         if policy_generation is not None:
-            policy_generation.prepare_refit_info(state_dict_info)
+            maybe_enable_refit_prequantize(
+                policy, policy_generation, state_dict_info, master_config.policy
+            )
 
     # Spin up non-colocated OPD teacher worker groups AFTER policy / vLLM are
     # ready. Parallelizing with policy init races on Megatron-Bridge's HF->mcore

nemo_rl/algorithms/ppo.py

@@ -47,7 +47,11 @@
     RewardShapingConfig,
     apply_reward_shaping,
 )
-from nemo_rl.algorithms.utils import print_performance_metrics, set_seed
+from nemo_rl.algorithms.utils import (
+    maybe_enable_refit_prequantize,
+    print_performance_metrics,
+    set_seed,
+)
 from nemo_rl.data import DataConfig
 from nemo_rl.data.collate_fn import rl_collate_fn
 from nemo_rl.data.datasets import AllTaskProcessedDataset
@@ -658,7 +662,9 @@ def initialize_generation_with_policy(
     # prepare refit info
     state_dict_info = policy.prepare_refit_info()
     if policy_generation is not None:
-        policy_generation.prepare_refit_info(state_dict_info)
+        maybe_enable_refit_prequantize(
+            policy, policy_generation, state_dict_info, master_config.policy
+        )
 
     # Calculate total setup time
     total_setup_time = time.perf_counter() - setup_start_time

nemo_rl/algorithms/utils.py

@@ -16,7 +16,7 @@
 import random
 import warnings
 from functools import partial, wraps
-from typing import Any, Optional
+from typing import TYPE_CHECKING, Any, Optional
 
 import numpy as np
 import torch
@@ -27,10 +27,16 @@
 )
 
 from nemo_rl.data.chat_templates import COMMON_CHAT_TEMPLATES
-from nemo_rl.models.policy import TokenizerConfig
+from nemo_rl.models.policy import PolicyConfig, TokenizerConfig
 from nemo_rl.utils.fastokens import maybe_patch_fastokens
 from nemo_rl.utils.logger import Logger
 
+if TYPE_CHECKING:
+    # Runtime import would cycle: policy.interfaces pulls in algorithms.loss,
+    # which imports back into algorithms.utils.
+    from nemo_rl.models.generation.interfaces import GenerationInterface
+    from nemo_rl.models.policy.interfaces import ColocatablePolicyInterface
+
 
 def get_gdpo_reward_component_keys(batch) -> list[str]:
     """Return batch keys that are named reward components (e.g. reward/correctness) in sorted order."""
@@ -1032,3 +1038,30 @@ def print_efficiency_summary(
     loggable["efficiency/total_wall_time_s"] = total_wall_time_s
 
     return loggable
+
+
+def maybe_enable_refit_prequantize(
+    policy: "ColocatablePolicyInterface",
+    policy_generation: "GenerationInterface",
+    state_dict_info: Optional[dict[str, Any]],
+    policy_config: PolicyConfig,
+) -> None:
+    """Complete the trainer-side pre-quantized refit handshake if requested.
+
+    The generation backend's prepare_refit_info returns the fp8-eligible
+    parameter names when vllm_cfg.refit_prequantize is enabled; the trainer
+    then quantizes exactly those during refit and the receiver's unpack
+    metadata is refreshed to the quantized dtypes plus scale entries.
+    """
+    prequant_names = policy_generation.prepare_refit_info(state_dict_info)
+    if not prequant_names:
+        return
+    megatron_cfg = policy_config.get("megatron_cfg")
+    if not (megatron_cfg and 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)
+    policy_generation.prepare_refit_info(updated_info)

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
@@ -510,10 +510,15 @@ 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]
+    ) -> Optional[list[str]]:
         if not self._is_real_quant_model():
-            return
+            return super().prepare_refit_info(state_dict_info)
+
+        # 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
@@ -532,6 +537,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):

nemo_rl/models/generation/interfaces.py

@@ -330,8 +330,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]:

nemo_rl/models/generation/sglang/sglang_generation.py

@@ -731,8 +731,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 []

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."""

nemo_rl/models/generation/vllm/checkpoint_engine.py

@@ -143,18 +143,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)

nemo_rl/models/generation/vllm/config.py

@@ -37,6 +37,11 @@ class VllmSpecificArgs(TypedDict):
     precision: NotRequired[str]
     # 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]
     kv_cache_dtype: Literal["auto", "fp8", "fp8_e4m3"]
     enforce_eager: NotRequired[bool]
     enable_return_routed_experts: NotRequired[bool]