refactor(trainer): one gradient manager for GPU and offload gradients

src/prime_rl/trainer/optim/__init__.py

@@ -7,8 +7,8 @@
 
 from prime_rl.configs.trainer import OptimizerConfig, OptimizerInBackwardOffloadConfig
 from prime_rl.trainer.models.fusions import get_model_packed_parameters
+from prime_rl.trainer.optim.base import GPUGradientManager, OptimizerLike
 from prime_rl.trainer.optim.base import OffloadOptimizer as OffloadOptimizer
-from prime_rl.trainer.optim.base import OptimizerLike
 from prime_rl.trainer.optim.offload import (
     FullCPUOffloadOptimizer,
     GradientOffloadManager,
@@ -51,11 +51,11 @@ def setup_optimizer(
     config: OptimizerConfig,
     named_params: list[tuple[str, nn.Parameter]],
     parallel_dims: ParallelDims,
+    model: nn.Module,
     cpu_offload: bool = False,
     full_offload_config: OptimizerInBackwardOffloadConfig | None = None,
-    model: nn.Module | None = None,
     full_offload_dtype_policy: dict[int, tuple[torch.dtype, torch.dtype]] | None = None,
-) -> tuple[OptimizerLike, GradientOffloadManager | None]:
+) -> tuple[OptimizerLike, GPUGradientManager | GradientOffloadManager]:
     if cpu_offload and full_offload_config is not None:
         raise ValueError("State-only and full optimizer CPU offload cannot both be enabled")
     if full_offload_config is not None and config.type not in ("adamw", "sign_sgd"):
@@ -66,8 +66,6 @@ def setup_optimizer(
     optimizer_named_params = named_params
     master_weights = None
     if full_offload_config is not None:
-        if model is None:
-            raise ValueError("CPU optimizer offload requires the model")
         if full_offload_dtype_policy is None:
             raise ValueError("CPU optimizer offload requires an explicit per-parameter dtype policy")
         optimizer_named_params, master_weights = _create_cpu_master_weights(
@@ -95,9 +93,9 @@ def setup_optimizer(
 
     if cpu_offload:
         get_logger().info("Wrapping optimizer with CPUOffloadOptimizer for optimizer state CPU offloading")
-        return CPUOffloadOptimizer(optimizer), None
+        return CPUOffloadOptimizer(optimizer), GPUGradientManager(model)
 
-    return optimizer, None
+    return optimizer, GPUGradientManager(model)
 
 
 def _create_optimizer(

src/prime_rl/trainer/optim/base.py

@@ -1,6 +1,10 @@
 from abc import ABC, abstractmethod
+from collections import defaultdict
 from typing import TypeAlias
 
+import torch
+from torch import nn
+from torch.distributed.tensor import DTensor
 from torch.optim import Optimizer
 
 
@@ -26,3 +30,51 @@ def finish_model_only_checkpoint_load(self) -> None:
 
 
 OptimizerLike: TypeAlias = Optimizer | OffloadOptimizer
+
+
+class GPUGradientManager:
+    """Gradient scaling and clipping for gradients that stay on the GPU.
+
+    Has the same step/backward interface as `GradientOffloadManager`. The step's gradient
+    scale is applied once, after the final backward of the step.
+    """
+
+    def __init__(self, model: nn.Module):
+        self._model = model
+        self._gradient_scale = 1.0
+        self._final_backward = False
+
+    def begin_step(self, gradient_scale: float, *, overlap_optimizer: bool) -> None:
+        self._gradient_scale = gradient_scale
+
+    def begin_backward(self, *, final_backward: bool) -> None:
+        self._final_backward = final_backward
+
+    @torch.no_grad()
+    def finish_backward(self) -> None:
+        if not self._final_backward:
+            return
+        for param in self._model.parameters():
+            if param.grad is not None:
+                param.grad.mul_(self._gradient_scale)
+
+    @torch.no_grad()
+    def clip_grad_norm_(self, max_norm: float) -> torch.Tensor:
+        # Norm reductions must complete on each parameter mesh before combining
+        # dense, expert-parallel, and other model-parallel gradients.
+        mesh_parameters = defaultdict(list)
+        for param in self._model.parameters():
+            if param.grad is not None:
+                mesh = param.grad.device_mesh if isinstance(param.grad, DTensor) else None
+                mesh_parameters[mesh].append(param)
+        norms = []
+        for parameters in mesh_parameters.values():
+            norm = torch.nn.utils.get_total_norm([param.grad for param in parameters])
+            norms.append(norm.full_tensor() if isinstance(norm, DTensor) else norm)
+        grad_norm = torch.linalg.vector_norm(torch.stack(norms)) if norms else torch.tensor(0.0)
+        for parameters in mesh_parameters.values():
+            torch.nn.utils.clip_grads_with_norm_(parameters, max_norm, grad_norm)
+        return grad_norm.cuda() if grad_norm.device.type == "cpu" else grad_norm
+
+    def close(self) -> None:
+        pass

src/prime_rl/trainer/optim/offload.py

@@ -565,7 +565,7 @@ def _offload_hook(self, param: torch.Tensor) -> None:
             raise TypeError(f"Expected parameter in post-accumulate hook, got {type(param)}")
         self._offload_params([param])
 
-    def finish_backward(self, *, wait_for_copies: bool = True) -> None:
+    def finish_backward(self) -> None:
         remaining = [param for chunk in self._chunks for param in chunk if param.grad is not None]
         self._offload_params(remaining)
         with self._condition:
@@ -584,8 +584,6 @@ def finish_backward(self, *, wait_for_copies: bool = True) -> None:
             )
             self._backward_open = False
         self._transfer_requests.put(boundary)
-        if wait_for_copies:
-            self.wait()
 
     def wait(self) -> None:
         self._wait_for(
@@ -605,11 +603,6 @@ def optimizer_overlapped(self) -> bool:
     def gradient_scale(self) -> float:
         return self._gradient_scale
 
-    @torch.no_grad()
-    def scale_(self, factor: float) -> None:
-        self.wait()
-        self._gradient_scale *= factor
-
     @torch.no_grad()
     def clip_grad_norm_(self, max_norm: float) -> torch.Tensor:
         if self._overlap_optimizer:

src/prime_rl/trainer/rl/train.py

@@ -54,13 +54,8 @@
     GarbageCollection,
     MemoryProfiler,
     Tensors,
-    begin_backward,
-    clip_grad_norm_,
     filter_rl_trainer_tensor_stats_for_wandb,
-    finish_backward,
     get_ckpt_disk_metrics,
-    prepare_gradient_offload,
-    scale_gradients_,
     setup_full_cpu_optimizer_offload,
     setup_torch_distributed,
 )
@@ -335,11 +330,9 @@ def train(config: TrainerConfig):
         dp_cp_group = parallel_dims.get_mesh("dp_cp").get_group()
         dist.all_reduce(global_scales, op=dist.ReduceOp.SUM, group=dp_cp_group)
         rl_scale, ce_scale, ref_kl_scale = (scale if scale > 0 else 1 for scale in global_scales.tolist())
-        prepare_gradient_offload(
-            gradient_manager,
-            parallel_dims.fsdp_gradient_divide_factor,
-            overlap_optimizer=True,
-        )
+        # compute_loss already divides by the global token count. Undo FSDP's per-rank averaging
+        # across dp_cp so the final gradient is the true per-token mean over the global batch.
+        gradient_manager.begin_step(parallel_dims.fsdp_gradient_divide_factor, overlap_optimizer=True)
 
         logger.debug(f"Starting forward and backward pass ({batch_size=})")
         tensors = Tensors()  # Used to accumulate tensor statistics across micro-batches and ranks for logging
@@ -514,9 +507,9 @@ def train(config: TrainerConfig):
 
             # Backward pass
             with maybe_record_function("backward"):
-                begin_backward(gradient_manager, final_backward=micro_step == len(micro_batches) - 1)
+                gradient_manager.begin_backward(final_backward=micro_step == len(micro_batches) - 1)
                 loss.backward()
-                finish_backward(gradient_manager)
+                gradient_manager.finish_backward()
 
             mm_kwargs = None
 
@@ -581,15 +574,10 @@ def train(config: TrainerConfig):
 
         annotation_writer.flush()
 
-        # compute_loss already divided by the global token count. Undo FSDP's per-rank averaging
-        # across dp_cp so the final gradient is the true per-token mean over the global batch.
-        if gradient_manager is None:
-            scale_gradients_(None, model, parallel_dims.fsdp_gradient_divide_factor)
-
         # Optionally, clip the gradients
         grad_norm: torch.Tensor | None = None
         if config.optim.max_norm is not None:
-            grad_norm = clip_grad_norm_(gradient_manager, model, config.optim.max_norm)
+            grad_norm = gradient_manager.clip_grad_norm_(config.optim.max_norm)
 
         # Update the model parameters
         optimizer.step()
@@ -757,8 +745,7 @@ def train(config: TrainerConfig):
         ckpt_manager.save(progress.step, model, [optimizer], scheduler, progress)
         ckpt_manager.maybe_clean()
 
-    if gradient_manager is not None:
-        gradient_manager.close()
+    gradient_manager.close()
 
     logger.info(f"Peak memory: {max_peak_memory:.1f} GiB")
     logger.success("RL trainer finished")

src/prime_rl/trainer/sft/train.py

@@ -48,13 +48,8 @@
 from prime_rl.trainer.utils import (
     GarbageCollection,
     MemoryProfiler,
-    begin_backward,
-    clip_grad_norm_,
-    finish_backward,
     get_ckpt_disk_metrics,
-    prepare_gradient_offload,
     print_sample,
-    scale_gradients_,
     setup_full_cpu_optimizer_offload,
     setup_torch_distributed,
 )
@@ -431,25 +426,17 @@ def is_online_eval_step(step: int) -> bool:
             (is_first_step and config.val.eval_on_start)
             or (not is_first_step and progress.step % config.val.interval == 0)
         )
-        if gradient_manager is None:
-            micro_batches = (next(dataiter) for _ in range(grad_accum_steps))
-            step_local_token_count = torch.tensor(0, dtype=torch.int64, device="cuda")
-        else:
-            micro_batches = [next(dataiter) for _ in range(grad_accum_steps)]
-            local_token_count = sum(int(micro_batch["loss_mask"].sum()) for micro_batch in micro_batches)
-            global_step_token_count = torch.tensor(local_token_count, dtype=torch.int64, device="cuda")
-            dist.all_reduce(global_step_token_count, op=dist.ReduceOp.SUM, group=dp_cp_group)
-            global_token_count_val = global_step_token_count.item() // cp_size
-            grad_scale = (
-                parallel_dims.fsdp_gradient_divide_factor * grad_accum_steps / global_token_count_val
-                if global_token_count_val > 0
-                else 1.0
-            )
-            prepare_gradient_offload(
-                gradient_manager,
-                grad_scale,
-                overlap_optimizer=not run_validation_this_step,
-            )
+        micro_batches = [next(dataiter) for _ in range(grad_accum_steps)]
+        local_token_count = sum(int(micro_batch["loss_mask"].sum()) for micro_batch in micro_batches)
+        global_step_token_count = torch.tensor(local_token_count, dtype=torch.int64, device="cuda")
+        dist.all_reduce(global_step_token_count, op=dist.ReduceOp.SUM, group=dp_cp_group)
+        global_token_count_val = global_step_token_count.item() // cp_size
+        grad_scale = (
+            parallel_dims.fsdp_gradient_divide_factor * grad_accum_steps / global_token_count_val
+            if global_token_count_val > 0
+            else 1.0
+        )
+        gradient_manager.begin_step(grad_scale, overlap_optimizer=not run_validation_this_step)
 
         step_tokens_per_expert = 0
         for micro_step, micro_batch in enumerate(micro_batches):
@@ -459,10 +446,7 @@ def is_online_eval_step(step: int) -> bool:
                 )
 
             with maybe_record_function("forward"):
-                local_loss_sum, batch_token_count = compute_loss(micro_batch)
-
-            if gradient_manager is None:
-                step_local_token_count += batch_token_count
+                local_loss_sum, _ = compute_loss(micro_batch)
 
             if torch.isnan(local_loss_sum.detach()):
                 nan_loss_count += 1
@@ -473,9 +457,9 @@ def is_online_eval_step(step: int) -> bool:
                 scaled_loss = local_loss_sum / grad_accum_steps
 
             with maybe_record_function("backward"):
-                begin_backward(gradient_manager, final_backward=micro_step == grad_accum_steps - 1)
+                gradient_manager.begin_backward(final_backward=micro_step == grad_accum_steps - 1)
                 scaled_loss.backward()
-                finish_backward(gradient_manager)
+                gradient_manager.finish_backward()
 
             if is_moe_model:
                 micro_moe_stats, tokens_per_expert = get_global_moe_stats(model, ep_group, dp_cp_group)
@@ -488,14 +472,6 @@ def is_online_eval_step(step: int) -> bool:
         forward_backward_time = time.perf_counter() - forward_backward_start_time
         expert_load_stats = get_expert_load_stats(step_tokens_per_expert, dp_cp_group) if is_moe_model else {}
 
-        if gradient_manager is None:
-            global_step_token_count = step_local_token_count.clone()
-            dist.all_reduce(global_step_token_count, op=dist.ReduceOp.SUM, group=dp_cp_group)
-            global_token_count_val = global_step_token_count.item()
-            if global_token_count_val > 0:
-                grad_scale = parallel_dims.fsdp_gradient_divide_factor * grad_accum_steps / global_token_count_val
-                scale_gradients_(None, model, grad_scale)
-
         # Run validation after forward-backward (so torch.compile sees training graph first) but before
         # optimizer step (so eval_on_start evaluates untrained weights)
         if run_validation_this_step:
@@ -517,7 +493,7 @@ def is_online_eval_step(step: int) -> bool:
         grad_norm: torch.Tensor | None = None
         if config.optim.max_norm is not None:
             logger.debug(f"Clipping gradients with max norm {config.optim.max_norm}")
-            grad_norm = clip_grad_norm_(gradient_manager, model, config.optim.max_norm)
+            grad_norm = gradient_manager.clip_grad_norm_(config.optim.max_norm)
         logger.debug("Optimizer step")
         optimizer.step()
         optimizer.zero_grad()
@@ -684,8 +660,7 @@ def is_online_eval_step(step: int) -> bool:
         logger.info("Broadcasting final weights")
         weight_sender.broadcast(model, step=progress.step)
 
-    if gradient_manager is not None:
-        gradient_manager.close()
+    gradient_manager.close()
 
     logger.info(f"Peak memory: {max_peak_memory:.1f} GiB")
     logger.success("SFT trainer finished")

src/prime_rl/trainer/utils.py

@@ -12,8 +12,7 @@
 import torch.distributed as dist
 from rich import print as rich_print
 from rich.text import Text
-from torch import Tensor, nn
-from torch.distributed.tensor import DTensor
+from torch import Tensor
 from transformers.tokenization_utils import PreTrainedTokenizer
 
 from prime_rl.trainer.world import get_world
@@ -22,7 +21,6 @@
 
 if TYPE_CHECKING:
     from prime_rl.configs.trainer import OptimizerInBackwardOffloadConfig
-    from prime_rl.trainer.optim import GradientOffloadManager
 
 DEFAULT_TIMEOUT = timedelta(seconds=600)
 
@@ -54,61 +52,6 @@ def _collect(self, generation: int = 1):
         get_logger().debug(f"Collected garbage in {format_time(time.monotonic() - begin)}")
 
 
-def prepare_gradient_offload(
-    manager: "GradientOffloadManager | None",
-    gradient_scale: float,
-    *,
-    overlap_optimizer: bool,
-) -> None:
-    if manager is not None:
-        manager.begin_step(gradient_scale, overlap_optimizer=overlap_optimizer)
-
-
-def begin_backward(manager: "GradientOffloadManager | None", *, final_backward: bool) -> None:
-    if manager is not None:
-        manager.begin_backward(final_backward=final_backward)
-
-
-def finish_backward(manager: "GradientOffloadManager | None", *, wait_for_copies: bool = False) -> None:
-    if manager is not None:
-        manager.finish_backward(wait_for_copies=wait_for_copies)
-
-
-@torch.no_grad()
-def scale_gradients_(manager: "GradientOffloadManager | None", model: nn.Module, factor: float) -> None:
-    if manager is not None:
-        manager.scale_(factor)
-        return
-    for param in model.parameters():
-        if param.grad is not None:
-            param.grad.mul_(factor)
-
-
-def clip_grad_norm_(
-    manager: "GradientOffloadManager | None",
-    model: nn.Module,
-    max_norm: float,
-) -> Tensor:
-    if manager is not None:
-        grad_norm = manager.clip_grad_norm_(max_norm)
-    else:
-        # Norm reductions must complete on each parameter mesh before combining
-        # dense, expert-parallel, and other model-parallel gradients.
-        mesh_parameters = defaultdict(list)
-        for param in model.parameters():
-            if param.grad is not None:
-                mesh = param.grad.device_mesh if isinstance(param.grad, DTensor) else None
-                mesh_parameters[mesh].append(param)
-        norms = []
-        for parameters in mesh_parameters.values():
-            norm = torch.nn.utils.get_total_norm([param.grad for param in parameters])
-            norms.append(norm.full_tensor() if isinstance(norm, DTensor) else norm)
-        grad_norm = torch.linalg.vector_norm(torch.stack(norms)) if norms else torch.tensor(0.0)
-        for parameters in mesh_parameters.values():
-            torch.nn.utils.clip_grads_with_norm_(parameters, max_norm, grad_norm)
-    return grad_norm.cuda() if grad_norm.device.type == "cpu" else grad_norm
-
-
 def get_ckpt_disk_metrics(output_dir: Path) -> dict[str, float]:
     """
     Disk usage metrics for the checkpoint directory (<output_dir>/checkpoints).