From 216a8f5190509e8819bc886e066ba44cb2817909 Mon Sep 17 00:00:00 2001 From: faresobeid <111092724+faresobeid@users.noreply.github.com> Date: Sun, 4 Oct 2026 11:52:48 +0000 Subject: [PATCH] refactor(trainer): one gradient manager for GPU and offload gradients --- src/prime_rl/trainer/optim/__init__.py | 12 +++--- src/prime_rl/trainer/optim/base.py | 52 +++++++++++++++++++++++ src/prime_rl/trainer/optim/offload.py | 9 +--- src/prime_rl/trainer/rl/train.py | 27 +++--------- src/prime_rl/trainer/sft/train.py | 57 +++++++------------------ src/prime_rl/trainer/utils.py | 59 +------------------------- 6 files changed, 82 insertions(+), 134 deletions(-) diff --git a/src/prime_rl/trainer/optim/__init__.py b/src/prime_rl/trainer/optim/__init__.py index 33755114091..e50a66a400b 100644 --- a/src/prime_rl/trainer/optim/__init__.py +++ b/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( diff --git a/src/prime_rl/trainer/optim/base.py b/src/prime_rl/trainer/optim/base.py index bbdf364f089..416700a4350 100644 --- a/src/prime_rl/trainer/optim/base.py +++ b/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 diff --git a/src/prime_rl/trainer/optim/offload.py b/src/prime_rl/trainer/optim/offload.py index cb5a692c020..29b38fa9d63 100644 --- a/src/prime_rl/trainer/optim/offload.py +++ b/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: diff --git a/src/prime_rl/trainer/rl/train.py b/src/prime_rl/trainer/rl/train.py index b363cc6543e..cfccfc2e8eb 100644 --- a/src/prime_rl/trainer/rl/train.py +++ b/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") diff --git a/src/prime_rl/trainer/sft/train.py b/src/prime_rl/trainer/sft/train.py index 5c42ffd5c23..045c6c7983c 100644 --- a/src/prime_rl/trainer/sft/train.py +++ b/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") diff --git a/src/prime_rl/trainer/utils.py b/src/prime_rl/trainer/utils.py index a17bb373a66..9e74a6d1623 100644 --- a/src/prime_rl/trainer/utils.py +++ b/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 (/checkpoints).