diff --git a/docs/scaling.md b/docs/scaling.md index d10ef6f444..2c7620f780 100644 --- a/docs/scaling.md +++ b/docs/scaling.md @@ -146,6 +146,16 @@ optim_cpu_offload = true # already the default Mutually exclusive with `fsdp_cpu_offload`. Also incompatible with `trainer.max_concurrent_runs > 1` (multi-tenant training) — set `optim_cpu_offload = false` for multi-run. Muon doesn't support `fsdp_cpu_offload` but does support `optim_cpu_offload`. +The optimizer step is performed per-transformer-layer ("chunked"): each layer's optimizer states are moved to GPU, the step runs for that layer only, and the states move back to CPU before the next layer. This bounds peak GPU optimizer-state memory to a single layer's worth, preventing OOM when `weight + grad + all_opt_states` exceeds available VRAM even without activations. H2D and D2H transfers are pipelined on dedicated CUDA streams so the next layer's states are prefetched while the current layer computes: + +```toml +[trainer.model] +optim_cpu_offload_chunked = true +optim_cpu_offload_stream = true +``` + +Both default to `false`. Set `optim_cpu_offload_chunked = true` to enable per-layer chunking. Set `optim_cpu_offload_stream = true` to additionally overlap H2D/D2H transfers with optimizer compute on dedicated CUDA streams. When chunked is enabled but stream is not, the step uses a simple sequential move→step→move loop. + ### LM Head Chunking The vanilla LM head materializes a `[batch * seq, vocab]` logits tensor on every step — a major memory tax when the vocabulary is large (often >100K). `fused_lm_head_token_chunk_size` swaps in a custom fused linear + logprob/entropy kernel that streams through `chunk_size` tokens at a time, avoiding the materialization. It defaults to `1024` for RL training: diff --git a/packages/prime-rl-configs/src/prime_rl/configs/trainer.py b/packages/prime-rl-configs/src/prime_rl/configs/trainer.py index 11ecb92064..4ea29cbfb6 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/trainer.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/trainer.py @@ -171,6 +171,12 @@ class ModelConfig(BaseModelConfig): optim_cpu_offload: bool = True """Offload only optimizer states (momentum, variance) to CPU, keeping weights on GPU. Avoids the H2D all-gather overhead of FSDP CPU offload while still saving GPU memory.""" + optim_cpu_offload_chunked: bool = False + """When optimizer CPU offload is enabled, perform the optimizer step per-transformer-layer instead of all-at-once. Each layer's states are moved to GPU, the step runs for that layer only, and states move back to CPU before the next layer. Bounds peak GPU optimizer-state memory to one layer's worth, preventing OOM when weight + grad + all_opt_states exceeds available VRAM.""" + + optim_cpu_offload_stream: bool = False + """When chunked optimizer CPU offload is enabled, overlap H2D (state→GPU) and D2H (state→CPU) transfers with optimizer compute using dedicated CUDA streams. Disabling simplifies the transfer logic to a sequential move→step→move loop at the cost of serializing transfers with compute.""" + reshard_after_forward: bool = True """Reshard the model after each forward pass.""" diff --git a/src/prime_rl/trainer/ckpt.py b/src/prime_rl/trainer/ckpt.py index 83728c8d9a..c4c92da7ed 100644 --- a/src/prime_rl/trainer/ckpt.py +++ b/src/prime_rl/trainer/ckpt.py @@ -109,6 +109,7 @@ def state_dict(self) -> dict[str, Any]: if isinstance(opt, CPUOffloadOptimizer): opt._move_states("cpu") if has_cpu_offload: + torch.cuda.synchronize() gc.collect() torch.cuda.empty_cache() diff --git a/src/prime_rl/trainer/optim.py b/src/prime_rl/trainer/optim.py index cb0b3d2fb0..5c6ea65567 100644 --- a/src/prime_rl/trainer/optim.py +++ b/src/prime_rl/trainer/optim.py @@ -1,4 +1,5 @@ import copy +import re from typing import Callable import torch @@ -20,58 +21,270 @@ class CPUOffloadOptimizer: Unlike FSDP's CPUOffload which offloads weights too, this keeps weights on GPU. With activation checkpointing, activations and optimizer states are never on GPU at the same time: peak memory becomes max(activations, opt_states) instead of sum. + + When ``named_params`` is provided, the step is performed per-transformer-layer + ("chunked") instead of all-at-once. Each chunk's optimizer states are moved to + GPU, the optimizer step is run for just that chunk, and the states are moved back + to CPU before the next chunk. This bounds peak GPU optimizer-state memory to a + single layer's worth rather than the full model's, preventing OOM when + weight + grad + all_opt_states exceeds available VRAM. + + H2D (state→GPU) and D2H (state→CPU) transfers are issued on dedicated CUDA + streams so they overlap with the optimizer compute of the current chunk: the next + chunk's states are prefetched while the current chunk computes, and the previous + chunk's states are evicted while the next chunk computes. The cost is one extra + layer's worth of optimizer state on GPU during the overlap window, which is + negligible compared to the savings from not loading the full model's states. """ - def __init__(self, optimizer: Optimizer, pin_memory: bool = True): + def __init__( + self, + optimizer: Optimizer, + named_params: list[tuple[str, nn.Parameter]] | None = None, + pin_memory: bool = True, + stream: bool = True, + ): self.optimizer = optimizer self.pin_memory = pin_memory + self.stream = stream self._initialized = False + # Build per-layer chunks from the optimizer's actual param_groups so + # we only include params the optimizer tracks (excludes frozen params + # that Muon / _create_optimizer filtered out). + self._chunks: list[list[nn.Parameter]] | None = None + if named_params is not None: + self._chunks = self._build_chunks(named_params) + + # ------------------------------------------------------------------ # + # Chunk construction + # ------------------------------------------------------------------ # + + @staticmethod + def _extract_layer_idx(name: str) -> int | None: + """Return the transformer-layer index encoded in a parameter name. + + Matches patterns like ``model.layers.0.self_attn.q_proj.weight`` or + ``language_model.layers.12.mlp.experts.3.weight`` and returns the + integer layer index. Returns ``None`` for non-layer params + (embed_tokens, lm_head, norm, …). + """ + m = re.search(r"layers\.(\d+)\.", name) + return int(m.group(1)) if m else None + + def _build_chunks(self, named_params: list[tuple[str, nn.Parameter]]) -> list[list[nn.Parameter]]: + """Partition optimizer parameters into per-layer chunks. + + Parameters from the same transformer layer form one chunk; non-layer + parameters (embeddings, lm_head, final norm) are collected into a + trailing "misc" chunk. Only parameters actually present in the + optimizer's ``param_groups`` are included. + """ + # Map param identity → layer index + param_layer: dict[int, int | None] = {} + for name, param in named_params: + param_layer[id(param)] = self._extract_layer_idx(name) + + # Group optimizer params by layer + by_layer: dict[int, list[nn.Parameter]] = {} + misc: list[nn.Parameter] = [] + for group in self.optimizer.param_groups: + for param in group["params"]: + layer_idx = param_layer.get(id(param)) + if layer_idx is not None: + by_layer.setdefault(layer_idx, []).append(param) + else: + misc.append(param) + + chunks = [by_layer[k] for k in sorted(by_layer)] + if misc: + chunks.append(misc) + return chunks or [misc] # edge-case: no layer params at all + + def _chunk_param_ids(self, chunk_idx: int) -> set[int]: + return {id(p) for p in self._chunks[chunk_idx]} + + # ------------------------------------------------------------------ # + # State movement + # ------------------------------------------------------------------ # + + def _move_tensor(self, v: torch.Tensor, device: str) -> torch.Tensor: + if device == "cpu": + if self.pin_memory: + # Allocate pinned CPU destination first, then async-copy into it. + # This avoids the race where pin_memory() would read a tensor + # whose async D2H copy hasn't completed yet. + dst = torch.empty_like(v, device="cpu").pin_memory() + dst.copy_(v, non_blocking=True) + return dst + return v.to("cpu") + return v.to(device, non_blocking=True) + + def _move_dtensor(self, v: DTensor, device: str) -> DTensor: + new_local = self._move_tensor(v._local_tensor, device) + new_dtensor = copy.copy(v) + new_dtensor._local_tensor = new_local + return new_dtensor + def _move_states(self, device: str): - """Move optimizer states to CPU or back to GPU (matching each parameter's device).""" + """Move **all** optimizer states to *device* (used by state_dict / first step).""" for p in self.optimizer.state: state = self.optimizer.state[p] for k, v in state.items(): if isinstance(v, DTensor): - local_tensor = v._local_tensor - if device == "cpu": - non_blocking = not self.pin_memory - new_local = local_tensor.to("cpu", non_blocking=non_blocking) - if self.pin_memory and not new_local.is_pinned(): - new_local = new_local.pin_memory() - else: - new_local = local_tensor.to(device, non_blocking=True) - new_dtensor = copy.copy(v) - new_dtensor._local_tensor = new_local - state[k] = new_dtensor + state[k] = self._move_dtensor(v, device) elif isinstance(v, torch.Tensor): - if device == "cpu": - non_blocking = not self.pin_memory - cpu_tensor = v.to("cpu", non_blocking=non_blocking) - if self.pin_memory and not cpu_tensor.is_pinned(): - cpu_tensor = cpu_tensor.pin_memory() - state[k] = cpu_tensor - else: - state[k] = v.to(device, non_blocking=True) + state[k] = self._move_tensor(v, device) + + def _move_chunk_states(self, chunk_idx: int, device: str, stream: torch.cuda.Stream | None = None): + """Move optimizer states for a single chunk to/from *device*. + + When *stream* is provided the transfers are issued on that stream so they + can overlap with work on the default stream. + """ + chunk_ids = self._chunk_param_ids(chunk_idx) + + def _do_move(): + for p in self.optimizer.state: + if id(p) not in chunk_ids: + continue + state = self.optimizer.state[p] + for k, v in state.items(): + if isinstance(v, DTensor): + state[k] = self._move_dtensor(v, device) + elif isinstance(v, torch.Tensor): + state[k] = self._move_tensor(v, device) + + if stream is not None: + with torch.cuda.stream(stream): + _do_move() + else: + _do_move() + + # ------------------------------------------------------------------ # + # Per-chunk optimizer step + # ------------------------------------------------------------------ # + + def _build_chunk_param_groups(self, chunk_idx: int) -> list[dict]: + """Create param_groups containing only the current chunk's params. + + Hyperparameters (lr, weight_decay, betas, algorithm, step, …) are + copied from the original groups so the optimizer behaves identically. + """ + chunk_ids = self._chunk_param_ids(chunk_idx) + chunk_groups: list[dict] = [] + for group in self._original_param_groups: + filtered = [p for p in group["params"] if id(p) in chunk_ids] + if not filtered: + continue + new_group = {k: v for k, v in group.items() if k != "params"} + new_group["params"] = filtered + chunk_groups.append(new_group) + return chunk_groups + + def _step_chunk(self, chunk_idx: int, closure=None): + """Run ``optimizer.step()`` for a single chunk by temporarily swapping param_groups.""" + chunk_groups = self._build_chunk_param_groups(chunk_idx) + self.optimizer.param_groups = chunk_groups + result = self.optimizer.step(closure) + self.optimizer.param_groups = self._original_param_groups + return result + + def _sync_step_counters(self, original_steps: list): + """Increment the original param_groups' step counters by one. + + Muon stores a per-group ``step`` integer that ``step()`` increments at the + top of the call. Because we swap in per-chunk copies, the *original* + groups never see the increment. Standard optimizers (AdamW, SGD, + SignSGD) keep the step counter in ``state[p]['step']`` which is + per-parameter and thus already correct (each param is in exactly one + chunk and its state is touched exactly once). + """ + for group, orig_step in zip(self._original_param_groups, original_steps): + group["step"] = orig_step + 1 + + # ------------------------------------------------------------------ # + # Public API + # ------------------------------------------------------------------ # def step(self, closure=None): - # First step initializes states on GPU - offload after - if not self._initialized: + # Non-chunked path (no named_params provided). + if self._chunks is None or len(self._chunks) <= 1: + if not self._initialized: + result = self.optimizer.step(closure) + self._move_states("cpu") + torch.cuda.synchronize() + self._initialized = True + return result + self._move_states("cuda") result = self.optimizer.step(closure) self._move_states("cpu") - self._initialized = True return result - # Move states to GPU - self._move_states("cuda") + # Chunked path (also used for the first step to avoid materializing + # all optimizer states on GPU at once during initialization). + self._original_param_groups = self.optimizer.param_groups + original_steps = [g.get("step", 0) for g in self._original_param_groups] - # Run optimizer step - result = self.optimizer.step(closure) + if self.stream: + self._step_chunked_streamed(closure) + else: + self._step_chunked(closure) - # Move states back to CPU - self._move_states("cpu") + self._sync_step_counters(original_steps) + self._original_param_groups = None - return result + if not self._initialized: + self._initialized = True + + return None + + def _step_chunked(self, closure=None): + """Per-layer optimizer step without stream overlap — simple sequential loop.""" + for i in range(len(self._chunks)): + self._move_chunk_states(i, "cuda") + self._step_chunk(i, closure if i == 0 else None) + self._move_chunk_states(i, "cpu") + torch.cuda.synchronize() + + def _step_chunked_streamed(self, closure=None): + """Per-layer optimizer step with stream-overlapped H2D / D2H. + + H2D (state→GPU) for chunk i+1 is issued on a dedicated stream while + the optimizer compute for chunk i runs on the compute stream. D2H + (state→CPU) for chunk i runs on a second dedicated stream, overlapping + with the compute for chunk i+1. + """ + h2d_stream = torch.cuda.Stream() + d2h_stream = torch.cuda.Stream() + compute_stream = torch.cuda.current_stream() + n = len(self._chunks) + + # Issue H2D for chunk 0 on the h2d stream. + self._move_chunk_states(0, "cuda", h2d_stream) + + for i in range(n): + # Wait for H2D of chunk i to land on GPU. + compute_stream.wait_stream(h2d_stream) + + # Prefetch H2D for chunk i+1 so it overlaps with step i on the + # compute stream. Wait for the previous chunk's D2H to finish + # first so we don't overwrite pinned CPU buffers still in flight. + if i + 1 < n: + h2d_stream.wait_stream(d2h_stream) + self._move_chunk_states(i + 1, "cuda", h2d_stream) + + # Run optimizer step for chunk i. + # Only evaluate the closure once (first chunk) to avoid recomputing + # loss/grads per chunk — the optimizer contract expects one evaluation per step. + self._step_chunk(i, closure if i == 0 else None) + + # Evict chunk i's states back to CPU (overlaps with step i+1). + d2h_stream.wait_stream(compute_stream) + self._move_chunk_states(i, "cpu", d2h_stream) + + # Block the CPU until all async transfers are complete. + torch.cuda.synchronize() def zero_grad(self, set_to_none: bool = True): self.optimizer.zero_grad(set_to_none=set_to_none) @@ -84,6 +297,7 @@ def state_dict(self): sd = self.optimizer.state_dict() if self._initialized: self._move_states("cpu") + torch.cuda.synchronize() return sd def load_state_dict(self, state_dict): @@ -114,6 +328,8 @@ def setup_optimizer( parallel_dims: ParallelDims, lora: bool = False, cpu_offload: bool = False, + cpu_offload_chunked: bool = False, + cpu_offload_stream: bool = False, ) -> Optimizer | CPUOffloadOptimizer: if lora: # Wait for run 0 to be created in the multi run manager @@ -126,7 +342,11 @@ def setup_optimizer( if cpu_offload: get_logger().info("Wrapping optimizer with CPUOffloadOptimizer for optimizer state CPU offloading") - return CPUOffloadOptimizer(optimizer) + return CPUOffloadOptimizer( + optimizer, + named_params=named_params if cpu_offload_chunked else None, + stream=cpu_offload_stream, + ) return optimizer diff --git a/src/prime_rl/trainer/rl/train.py b/src/prime_rl/trainer/rl/train.py index 29a178a76f..39cca469a7 100644 --- a/src/prime_rl/trainer/rl/train.py +++ b/src/prime_rl/trainer/rl/train.py @@ -166,6 +166,8 @@ def train(config: TrainerConfig): parallel_dims, lora=config.model.lora is not None, cpu_offload=config.model.optim_cpu_offload, + cpu_offload_chunked=config.model.optim_cpu_offload_chunked, + cpu_offload_stream=config.model.optim_cpu_offload_stream, ) scheduler = setup_scheduler(optimizer, config.scheduler, config.max_steps, config.optim.lr) else: diff --git a/src/prime_rl/trainer/sft/train.py b/src/prime_rl/trainer/sft/train.py index 1f2784bccf..0d23275209 100644 --- a/src/prime_rl/trainer/sft/train.py +++ b/src/prime_rl/trainer/sft/train.py @@ -178,7 +178,12 @@ def train(config: SFTConfig): # Set up the optimizer logger.info(f"Initializing optimizer ({config.optim})") optimizer = setup_optimizer( - config.optim, list(model.named_parameters()), parallel_dims, cpu_offload=config.model.optim_cpu_offload + config.optim, + list(model.named_parameters()), + parallel_dims, + cpu_offload=config.model.optim_cpu_offload, + cpu_offload_chunked=config.model.optim_cpu_offload_chunked, + cpu_offload_stream=config.model.optim_cpu_offload_stream, ) # Set up the learning rate scheduler