feat: add delta-compressed collective refit

docs/index.md

@@ -322,6 +322,7 @@ design-docs/uv.md
 design-docs/dependency-management.md
 design-docs/chat-datasets.md
 design-docs/generation.md
+design-docs/sparse-delta-refit.md
 design-docs/checkpointing.md
 design-docs/loss-functions.md
 design-docs/fsdp2-parallel-plan.md

examples/configs/grpo_math_1B.yaml

@@ -339,6 +339,8 @@ policy:
     top_k: null
     stop_token_ids: null
     stop_strings: null
+    refit_transport: null  # Set to "vllm_s3_sparse" or "vllm_zmq_sparse" for remote sparse-delta refit.
+    refit_cfg: null  # Optional tuning and storage settings for remote sparse-delta refit.
     mcore_generation_config:
       async_engine: false
       max_model_len: ${policy.max_total_sequence_length}  # Engine-side max sequence length.
@@ -376,6 +378,9 @@ policy:
       num_first_layers_in_bf16: 0
       enable_vllm_metrics_logger: true # Set to true to enable vLLM internal metrics logger, turn off for better performance
       vllm_metrics_logger_interval: 0.5 # Interval in seconds to collect vLLM logger metrics
+      http_refit_api_key_env_var: null # Optional env var containing the internal refit API key.
+      http_refit_server_port: null # Optional fixed port for Kubernetes targetPorts.
+      zmq_refit_server_port: null # Optional fixed ZeroMQ relay port for Kubernetes targetPorts.
     vllm_kwargs: {}
     colocated:
       # true: generation shares training GPUs

examples/configs/recipes/llm/grpo-qwen3-30ba3b-4n8g-megatron-zmq-deltaweight-noncolocated.yaml

@@ -0,0 +1,41 @@
+defaults: ./performance/grpo-qwen3-30ba3b-4n8g.yaml
+
+grpo:
+  num_prompts_per_step: 16
+  num_generations_per_prompt: 8
+  max_num_steps: 50
+  val_period: 1000
+
+checkpointing:
+  checkpoint_dir: results/grpo-qwen3-30ba3b-4n8g-megatron-zmq-deltaweight-noncolocated
+
+policy:
+  train_global_batch_size: 128
+  generation_batch_size: 16
+  max_total_sequence_length: 2048
+  sequence_packing:
+    train_mb_tokens: 2048
+    logprob_mb_tokens: 4096
+  megatron_cfg:
+    activation_checkpointing: true
+    env_vars:
+      PYTORCH_CUDA_ALLOC_CONF: expandable_segments:True
+  generation:
+    refit_transport: vllm_zmq_sparse
+    refit_cfg:
+      delta_compression:
+        encoding: xor
+      verify_samples_per_payload: 0
+      baseline:
+        in_memory: false
+    colocated:
+      enabled: false
+      resources:
+        gpus_per_node: 8
+        num_nodes: 2
+
+logger:
+  log_dir: logs/grpo-qwen3-30ba3b-4n8g-megatron-zmq-deltaweight-noncolocated
+  wandb:
+    project: nemo-rl-refit
+    name: grpo-qwen3-30ba3b-4n8g-megatron-zmq-deltaweight-noncolocated

nemo_rl/algorithms/distillation.py

@@ -217,6 +217,14 @@ def setup(
     assert generation_config is not None, (
         "A generation config in the PolicyConfig is required for distillation"
     )
+    if (
+        generation_config["backend"] == "vllm"
+        and cast(VllmConfig, generation_config).get("refit_transport") is not None
+    ):
+        raise ValueError(
+            "Remote sparse refit is currently supported only by GRPO; distillation "
+            "support is tracked in https://github.com/NVIDIA-NeMo/RL/issues/3275."
+        )
 
     # Disallow SP + packing for dtensor path
     for cfg, who in ((policy_config, "student"), (teacher_config, "teacher")):

nemo_rl/algorithms/grpo.py

@@ -101,6 +101,7 @@
 from nemo_rl.models.generation.sglang.config import SGLangConfig
 from nemo_rl.models.generation.sglang.sglang_generation import SGLangGeneration
 from nemo_rl.models.generation.vllm import VllmConfig, VllmGeneration
+from nemo_rl.models.generation.vllm.config import normalize_vllm_refit_config
 from nemo_rl.models.megatron.router_replay import (
     configure_vllm_for_router_replay,
     router_replay_enabled,
@@ -354,6 +355,8 @@ def setup(
     assert generation_config is not None, (
         "A generation config in the PolicyConfig is required for GRPO"
     )
+    if generation_config["backend"] == "vllm":
+        normalize_vllm_refit_config(cast(VllmConfig, generation_config))
 
     # Set seed for all random number generators
     set_seed(grpo_config["seed"])
@@ -907,6 +910,9 @@ def _spinup_nemo_gym(base_urls, model_name):
     # vllm model loading prefers clean environment, initialize policy_generation before policy in colocated mode
     backend = generation_config["backend"]
     generation_config["model_name"] = policy_config["model_name"]  # Needed for vLLM
+    remote_transport = None
+    remote_synchronizer_cls = None
+    remote_baseline_init_refs: list[Any] = []
 
     # Dictionary to store worker initialization timing stats for logging
     worker_init_timing_metrics = {}
@@ -967,6 +973,11 @@ def init_policy():
             init_optimizer=True,
             init_reference_model=init_reference_model,
         )
+        if remote_transport is not None:
+            assert remote_synchronizer_cls is not None
+            remote_baseline_init_refs.extend(
+                remote_synchronizer_cls.start_baseline(p, remote_transport)
+            )
         return p, time.perf_counter() - t0
 
     def init_vllm():
@@ -1095,6 +1106,21 @@ def initialize_generation_with_policy(
     elif backend == "vllm":
         # vLLM generation: setup config, then initialize with policy
         generation_config = cast(VllmConfig, generation_config)
+        if generation_config.get("refit_transport") is not None:
+            # Keep optional remote transport dependencies off the default path.
+            from nemo_rl.weight_sync.vllm_remote_sparse_weight_synchronizer import (
+                VllmRemoteSparseWeightSynchronizer,
+                validate_vllm_remote_sparse_refit,
+            )
+
+            remote_transport = validate_vllm_remote_sparse_refit(
+                generation_config,
+                colocated=colocated_inference,
+                megatron_enabled=policy_config["megatron_cfg"]["enabled"],
+            )
+            assert remote_transport is not None
+            remote_synchronizer_cls = VllmRemoteSparseWeightSynchronizer
+
         if generation_config["vllm_cfg"]["precision"] == "fp8":
             assert loss_config.use_importance_sampling_correction, (
                 "Importance sampling must be enabled for vLLM FP8 generation for good convergence!"
@@ -1232,7 +1258,7 @@ def init_vllm_then_policy():
     policy.print_node_ip_and_gpu_id()
 
     # if it is not colocated inference, initialize collective communication for update weights
-    if not colocated_inference:
+    if not colocated_inference and remote_transport is None:
         t0 = time.perf_counter()
         ip, port = train_cluster.get_master_address_and_port()
         print(f"Using ip: {ip}, port: {port} for collective communication", flush=True)
@@ -1271,9 +1297,30 @@ def init_vllm_then_policy():
             ray.get(futures_train + futures_inference)
         worker_init_timing_metrics["collective_init_time_s"] = time.perf_counter() - t0
 
-    state_dict_info = policy.prepare_refit_info()
-    if policy_generation is not None:
-        policy_generation.prepare_refit_info(state_dict_info)
+    if remote_transport is not None:
+        t0 = time.perf_counter()
+        assert isinstance(policy_generation, VllmGeneration)
+        assert remote_synchronizer_cls is not None
+        refit_config = generation_config["refit_cfg"]
+        assert refit_config is not None
+        policy_generation.weight_synchronizer = remote_synchronizer_cls(
+            policy,
+            policy_generation,
+            transport=remote_transport,
+            api_key_env_var=generation_config["vllm_cfg"].get(
+                "http_refit_api_key_env_var"
+            ),
+            request_timeout_s=refit_config.request_timeout_s,
+            baseline_init_refs=remote_baseline_init_refs,
+        )
+        policy_generation.weight_synchronizer.init_communicator()
+        worker_init_timing_metrics[f"vllm_{remote_transport}_sparse_init_time_s"] = (
+            time.perf_counter() - t0
+        )
+    else:
+        state_dict_info = policy.prepare_refit_info()
+        if policy_generation is not None:
+            policy_generation.prepare_refit_info(state_dict_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
@@ -2033,7 +2080,7 @@ def refit_policy_generation(
     _refit_buffer_size_gb: Optional[float] = None,
     timer: Optional[Timer] = None,
     kv_scales: Optional[dict[str, float]] = None,
-) -> None:
+) -> dict[str, float]:
     """Refit the policy generation interface with the latest policy weights.
 
     Args:
@@ -2043,7 +2090,14 @@ def refit_policy_generation(
             the buffer size is computed from remaining memory.
         timer: Optional Timer used to time the prepare/transfer/update phase
         kv_scales: Optional dictionary of KV cache scales for FP8 quantization.
+
+    Returns:
+        Scalar metrics reported by the selected weight synchronizer.
     """
+    synchronizer = getattr(policy_generation, "weight_synchronizer", None)
+    if synchronizer is not None:
+        return synchronizer.sync_weights(timer=timer, kv_scales=kv_scales) or {}
+
     # Megatron generation backend needs explicit suspend/resume around refits.
     if isinstance(policy_generation, MegatronGeneration):
         policy_generation.suspend_for_refit()
@@ -2144,6 +2198,16 @@ def refit_policy_generation(
     if isinstance(policy_generation, MegatronGeneration):
         policy_generation.resume_after_refit()
 
+    return {}
+
+
+def _initial_policy_generation_stale(
+    policy_generation: GenerationInterface, completed_steps: int
+) -> bool:
+    """Skip a fresh run's redundant sync when the synchronizer is already current."""
+    synchronizer = getattr(policy_generation, "weight_synchronizer", None)
+    return completed_steps > 0 or synchronizer is None or synchronizer.is_stale
+
 
 def _log_mixed_rewards_and_advantages_information(
     logger: Logger,
@@ -2333,7 +2397,6 @@ def grpo_train(
         isinstance(policy_generation, MegatronGeneration)
         and master_config.policy["generation"]["colocated"]["enabled"]
     )
-    POLICY_GENERATION_STALE = True  # tracks if generation needs a refit before running
     assert policy_generation is not None
 
     # Check if we need to sync KV cache scales
@@ -2343,6 +2406,9 @@ def grpo_train(
     # common config/state times
     current_step = grpo_save_state["current_step"]  # current step within an epoch
     total_steps = grpo_save_state["total_steps"]  # total steps across all epochs
+    POLICY_GENERATION_STALE = _initial_policy_generation_stale(
+        policy_generation, total_steps
+    )
     max_num_steps = master_config.grpo[
         "max_num_steps"
     ]  # max number of steps to train for
@@ -2413,6 +2479,7 @@ def grpo_train(
 
         # Run grpo/dapo training loop (single-turn)
         for batch in wrapped_dataloader:
+            refit_metrics: dict[str, float] = {}
             # A central place to store logging data that won't be deleted until the loop ends
             metrics_logging_data = dict()
             metrics = dict()
@@ -2489,7 +2556,7 @@ def grpo_train(
                                 calibration_data, include_q=True
                             )["layers"]
 
-                        refit_policy_generation(
+                        refit_metrics = refit_policy_generation(
                             policy,
                             policy_generation,
                             colocated_inference,
@@ -2944,7 +3011,7 @@ def grpo_train(
                 ):
                     memory_tracker.snapshot_start_of_stage("Validation", dir())
                     if NEED_REFIT and POLICY_GENERATION_STALE:
-                        refit_policy_generation(
+                        refit_metrics = refit_policy_generation(
                             policy,
                             policy_generation,
                             colocated_inference,
@@ -3296,6 +3363,8 @@ def grpo_train(
                 train_results, metrics, timing_metrics, master_config
             )
 
+            if refit_metrics:
+                logger.log_metrics(refit_metrics, total_steps + 1, prefix="refit")
             logger.log_metrics(metrics, total_steps + 1, prefix="train")
             logger.log_metrics(
                 performance_metrics, total_steps + 1, prefix="performance"
@@ -3627,11 +3696,11 @@ def async_grpo_train(
         isinstance(policy_generation, MegatronGeneration)
         and master_config.policy["generation"]["colocated"]["enabled"]
     )
-    POLICY_GENERATION_STALE = True
     assert policy_generation is not None
 
     # Training state
     step = grpo_save_state["current_step"]
+    POLICY_GENERATION_STALE = _initial_policy_generation_stale(policy_generation, step)
     weight_version = step  # Tracks refitted weight versions
     consumed_samples = grpo_save_state["consumed_samples"]
     total_valid_tokens = grpo_save_state.get(
@@ -3895,6 +3964,7 @@ def async_grpo_train(
     # Main training loop
     try:
         while step < master_config.grpo["max_num_steps"]:
+            refit_metrics: dict[str, float] = {}
             print(
                 f"\n{'=' * 25} Step {step + 1}/{master_config.grpo['max_num_steps']} {'=' * 25}"
             )
@@ -4287,7 +4357,7 @@ def async_grpo_train(
                     # Only the actual refit/weight transfer should be counted as weight_sync
                     print("🔄 Performing policy generation refit...")
                     with timer.time("weight_sync"):
-                        refit_policy_generation(
+                        refit_metrics = refit_policy_generation(
                             policy,
                             policy_generation,
                             colocated_inference,
@@ -4318,7 +4388,7 @@ def async_grpo_train(
                         trajectory_collector.pause.remote()
 
                         if NEED_REFIT and POLICY_GENERATION_STALE:
-                            refit_policy_generation(
+                            refit_metrics = refit_policy_generation(
                                 policy, policy_generation, colocated_inference
                             )
                             POLICY_GENERATION_STALE = False
@@ -4645,6 +4715,8 @@ def async_grpo_train(
                 merged_efficiency, total_wall_time, step + 1
             )
 
+            if refit_metrics:
+                logger.log_metrics(refit_metrics, step + 1, prefix="refit")
             logger.log_metrics(performance_metrics, step + 1, prefix="performance")
             logger.log_metrics(metrics, step + 1, prefix="train")
             logger.log_metrics(efficiency_loggable, step + 1, prefix="")

nemo_rl/algorithms/ppo.py

@@ -228,6 +228,14 @@ def setup(
     assert generation_config is not None, (
         "A generation config in the PolicyConfig is required for PPO"
     )
+    if (
+        generation_config["backend"] == "vllm"
+        and cast(VllmConfig, generation_config).get("refit_transport") is not None
+    ):
+        raise ValueError(
+            "Remote sparse refit is currently supported only by GRPO; PPO support "
+            "is tracked in https://github.com/NVIDIA-NeMo/RL/issues/3275."
+        )
 
     if "megatron_cfg" in policy_config and policy_config["megatron_cfg"]["enabled"]:
         policy_megatron_config = cast(MegatronConfig, policy_config["megatron_cfg"])

nemo_rl/models/generation/__init__.py

@@ -45,7 +45,9 @@ def configure_generation_config(
     if config["backend"] == "vllm":
         config = cast(VllmConfig, config)
         # set load_format
-        config["vllm_cfg"]["load_format"] = "auto" if is_eval else "dummy"
+        config["vllm_cfg"]["load_format"] = (
+            "auto" if is_eval or config.get("refit_transport") else "dummy"
+        )
         speculative_config = config.get("vllm_kwargs", {}).get("speculative_config")
         if speculative_config and not is_eval and not has_refit_draft_weights:
             # Speculative decoding needs real draft weights at startup, since the