feat(weights): sparse deltas from the full-offload optimizer step, without a snapshot

src/prime_rl/trainer/optim/offload.py

@@ -719,6 +719,9 @@ def __init__(
         self.optimizer = optimizer
         self._initialized = False
         self._master_weights = master_weights
+        # Called on the H2D stream with each parameter's local shard and its new values (pinned host) instead of
+        # the copy, which it must do; the shard still holds the previous values
+        self.weight_hook: Callable[[torch.Tensor, torch.Tensor], None] | None = None
         self._chunks = self._build_chunks()
         # Reuse the transfer streams across steps: fresh streams each step land
         # their H2D/D2H staging in new per-stream allocator pools, growing
@@ -797,7 +800,10 @@ def _update_compute_weights(
                     raise TypeError(
                         f"Full-offload H2D source dtype mismatch: expected {master.compute_dtype}, got {source.dtype}"
                     )
-                local_param.copy_(source, non_blocking=True)
+                if self.weight_hook is None:
+                    local_param.copy_(source, non_blocking=True)
+                else:
+                    self.weight_hook(local_param, source)
 
     def _step_native_cpu_adamw_chunk(
         self,

src/prime_rl/trainer/rl/train.py

@@ -9,13 +9,14 @@
 # Import environment before any other imports
 # ruff: noqa: I001
 
-from prime_rl.transports.weights import prune_broadcasts_beyond, setup_weight_sender
+from prime_rl.transports.weights import SparseFileSystemWeightSender, prune_broadcasts_beyond, setup_weight_sender
 from prime_rl.utils.act_offloading import maybe_activation_offloading
 import torch
 import torch.distributed as dist
 from torch.profiler import profile, ProfilerActivity, record_function
 from prime_rl.trainer.ckpt import Progress, setup_ckpt_manager
 from prime_rl.trainer.optim import setup_optimizer
+from prime_rl.trainer.optim.offload import FullCPUOffloadOptimizer
 from prime_rl.trainer.scheduler import setup_scheduler
 from prime_rl.configs.trainer import TrainerConfig
 from prime_rl.trainer.rl.data import DataLoader, FakeDataLoader, prepare_micro_batch
@@ -192,6 +193,8 @@ def train(config: TrainerConfig):
             config.model.lora,
         )
         logger.debug(f"Initialized weight broadcast in {format_time(time.perf_counter() - t0)}")
+        if isinstance(weight_sender, SparseFileSystemWeightSender) and isinstance(optimizer, FullCPUOffloadOptimizer):
+            weight_sender.diff_in_optimizer(optimizer)
 
     if parallel_dims.cp_enabled:
         setup_context_parallel(model, config.model, parallel_dims)