diff --git a/examples/moe_recipes/deepseek_v3/gb200/mxfp8_256GPU_TP1PP4EP64.yaml b/examples/moe_recipes/deepseek_v3/gb200/mxfp8_256GPU_TP1PP4EP64.yaml index 85bf55a693a..7c391012e89 100644 --- a/examples/moe_recipes/deepseek_v3/gb200/mxfp8_256GPU_TP1PP4EP64.yaml +++ b/examples/moe_recipes/deepseek_v3/gb200/mxfp8_256GPU_TP1PP4EP64.yaml @@ -202,7 +202,8 @@ ARGS: recompute_granularity: selective recompute_modules: - mla_up_proj - offload_optimizer_states: true + chunked_optimizer_state_offload: true + optimizer_state_offload_chunk_size_mb: 256 fine_grained_activation_offloading: true offload_modules: - expert_fc1 diff --git a/megatron/core/optimizer/__init__.py b/megatron/core/optimizer/__init__.py index 27b675d1b8d..3e1dde9981d 100644 --- a/megatron/core/optimizer/__init__.py +++ b/megatron/core/optimizer/__init__.py @@ -594,7 +594,7 @@ def init_state_fn(opt, config=None): opt.state[p]['exp_avg'] = torch.zeros_like(p.data) opt.state[p]['exp_avg_sq'] = torch.zeros_like(p.data) else: - opt.initialize_state(p) + opt.initialize_state(p, config.store_param_remainders) elif config.optimizer == 'lion': if not HAVE_EMERGING_OPTIMIZERS: diff --git a/megatron/core/optimizer/cpu_offloading/README.md b/megatron/core/optimizer/cpu_offloading/README.md index 68bfba54ba2..95101c4e00c 100644 --- a/megatron/core/optimizer/cpu_offloading/README.md +++ b/megatron/core/optimizer/cpu_offloading/README.md @@ -11,3 +11,91 @@ Add these flags to enable optimizer cpu offload in MCore. ## Configuration Recommendations Gradient copy from GPU to CPU, CPU optimizer step, and subsequent parameter copy from CPU to GPU can be time-consuming operations, and it is recommended to use the flag `--overlap-cpu-optimizer-d2h-h2d` to execute them concurrently. + +## Chunked GPU optimizer-state offload + +`--chunked-optimizer-state-offload` keeps a configurable fraction of optimizer tensor state and +master weights in pinned CPU memory between updates, while the optimizer update itself still runs +on GPU. Tensor state is restored and updated in bounded parameter chunks. All selected master +weights use one full restore window and remain on CUDA across every state chunk during an update or +distributed-checkpoint state initialization. + +```bash +--chunked-optimizer-state-offload +--optimizer-state-offload-chunk-size-mb 256 +--optimizer-state-offload-fraction 1.0 +--ckpt-format torch_dist +``` + +A chunk size of `0` restores all selected tensor state for one update and therefore does not bound +the temporary tensor-state peak; argument validation emits a warning for this setting. A positive +size bounds that tensor-state window, but never the full selected master-weight window described +above. The fraction is an approximate byte fraction because a parameter and all of its optimizer +state are an atomic bundle. Partial selection prioritizes bundles that actually include a separate +master before state-only bundles; precision-aware FusedAdam can own a separate master for native +FP32 DistributedOptimizer shards too. +Supported paths are Adam with `DistributedOptimizer`, and BF16 Muon with compact-layout +`LayerWiseDistributedOptimizer`; BF16 and FP8 parameter gather are supported. Full-iteration and +optimizer CUDA graphs are not supported while the fraction is nonzero. When optimizer state is +saved, async distributed-checkpoint save is rejected: its background writer can retain references +to the same pinned CPU buffers that the next optimizer update reuses. Synchronous distributed +optimizer-state save is supported. + +Chunked execution also requires optimizer parameter-group metadata updates to be independent of +the parameter subset passed to each chunk. CPU tensor-valued group fields are compared by value; +CUDA tensor-valued fields are only checked for matching shape, dtype, and device to avoid a host +synchronization per chunk. Optimizers whose CUDA group metadata depends on the subset are unsupported. + +Steady-state updates use two reusable GPU state windows per active offload manager: H2D for the +next chunk, optimizer compute for the current chunk, and D2H for the previous chunk can overlap. +The normal tensor-state staging bound is therefore about +`2 * optimizer_state_offload_chunk_size_mb` per manager that is actively stepping, not one +chunk; a single parameter's atomic tensor-state payload can exceed the target and produces a +runtime warning. LayerWise has at most one non-empty Muon child, so its internal state manager uses +the same two-slot bound. In the outer split Muon + Adam chain, however, a sibling +`DistributedOptimizer` can retain its first prefetched chunk while the LayerWise manager owns two +slots. The aggregate across those two managers can therefore still approach +`3 * optimizer_state_offload_chunk_size_mb`; the 256 MiB example can hold about 768 MiB during +that overlap window. Temporary state and master windows are allocated on the current compute +stream so their reserved memory returns to the same allocator pool used by forward/backward; the +H2D stream waits for both that compute stream and the D2H stream before writing them. D2H +completion does not block the host at the forward boundary. + +Every `ChainedOptimizer` recursively shares one copy-stream pair across its LayerWise and sibling +`DistributedOptimizer` managers, including ordinary Adam dense/expert chains. Stream sharing +serializes transfer order but does not merge managers' two-slot pools. LayerWise starts its single +managed child's first state chunk during gradient finalization, then pipelines that child's later +chunks with optimizer compute. A host synchronization is used only for synchronous +distributed-checkpoint access and one-time lazy state initialization. + +Selected master bindings and supported full-size state bindings become CPU tensors as soon as their +asynchronous D2H copies are enqueued. CUDA readers must restore them through the optimizer offload +lifecycle, and host readers must first call the checkpoint synchronization hook. The training loop +enters this state only at the optimizer-to-forward boundary, where no other master reader is +allowed. +MCore's model-to-main, main-to-model, and MXFP8 param-buffer copy entry points restore a +CPU-bound selected master asynchronously, order the current compute stream after that H2D, and then +validate residency. This makes external reload and param-staging entry points self-healing. + +Setting `optimizer_state_offload_fraction` to zero is a full disable even when the feature flag is +present: no manager, transfer stream, training hook, or checkpoint-format restriction is installed. +Checkpoint format and async-save restrictions are also skipped when optimizer state is explicitly +excluded from both load and save; excluding only saves permits async model-only checkpoints in +formats that otherwise support async save. + +Distributed checkpointing is required whenever optimizer state is loaded or saved because its +sharded-state construction lets the offloader initialize selected state in bounded chunks, and its +optimizer load handoff can preserve CPU canonical tensors instead of device-casting the full state +to CUDA. A loader is not guaranteed to fill every pinned destination in place: fully-parallel +loading may create transport temporaries or replace a destination with a non-pinned CPU tensor. +The offloader adopts and pins selected CPU state afterward, while moving an unselected parameter's +full-size tensor state back to that parameter's CUDA device. +The chunk size therefore bounds the tensor-state portion of optimizer initialization, steady-state +updates, and the offloader's own load conversions. It does not bound the full selected master +window or all temporary memory internal to a distributed checkpoint strategy. The legacy torch +checkpoint path can reconstruct the full state on CUDA and is not supported. + +The deprecated `--offload-optimizer-states` argument remains as a parser compatibility spelling. +Configuration validation emits a visible `FutureWarning`, enables the new mode, and consumes the +legacy field so copying the normalized configuration does not warn again. Its default-equivalent +settings are chunk size `0` and fraction `1.0`, with the restrictions above. diff --git a/megatron/core/optimizer/cpu_offloading/chunked_optimizer_state_offload.py b/megatron/core/optimizer/cpu_offloading/chunked_optimizer_state_offload.py new file mode 100644 index 00000000000..930bf7f8486 --- /dev/null +++ b/megatron/core/optimizer/cpu_offloading/chunked_optimizer_state_offload.py @@ -0,0 +1,1013 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +"""Chunked GPU optimizer-state execution backed by CPU canonical storage. + +The wrapped optimizer continues to own its normal param_groups and state mappings. +Parameters selected for offload keep their master weights and tensor states in pinned CPU +memory between updates. Immediately before an update, master weights are restored as one +window and optimizer tensor states are restored one chunk at a time. The external optimizer +therefore only ever observes regular CUDA tensors and does not need an offload-specific API. +""" + +from __future__ import annotations + +import collections +import dataclasses +import logging +import math +from typing import Callable, Dict, Iterable, List, Mapping, Sequence + +import torch + +from megatron.core.utils import log_single_rank + +_MASTER_PARAM_KEY = "master_param" +_NON_OFFLOADABLE_STATE_KEYS = frozenset({_MASTER_PARAM_KEY, "step", "found_inf"}) + +logger = logging.getLogger(__name__) + + +@dataclasses.dataclass(frozen=True) +class OptimizerStateChunk: + """A group of optimizer parameters updated as one GPU-resident window. + + Attributes: + params: Parameters whose optimizer tensor states are staged together. + """ + + params: tuple[torch.Tensor, ...] + + +@dataclasses.dataclass +class _StatePrefetch: + """Prefetch token whose existence also records a scheduled no-copy chunk.""" + + event: torch.cuda.Event | None + + +@dataclasses.dataclass +class _StateStagingSlot: + """Reusable GPU storage for one optimizer-state pipeline window.""" + + buffers: Dict[tuple[torch.device, torch.dtype], torch.Tensor] = dataclasses.field( + default_factory=dict + ) + + +class ChunkedOptimizerStateOffloader: + """Manage CPU canonical optimizer state and chunked GPU optimizer updates. + + chunk_size_bytes is a soft state-size limit. Parameters are atomic units: a single + parameter whose state exceeds the limit forms an oversized chunk. This is required by + matrix optimizers such as Muon and also avoids fabricating optimizer parameters for slices + of a DistributedOptimizer shard. + + Master weights are always offloaded for selected mixed-precision parameters. They use a + full-window fallback: all selected masters are restored before the first state chunk and + copied back after the optimizer/model-parameter staging phase. This keeps the external + optimizer and FP8 parameter-gather paths unchanged while still removing masters from the + forward/backward residency set. + + Args: + optimizer: The external torch-compatible optimizer to execute. + master_params: MCore-owned FP32 master parameters. These objects also appear in + optimizer.param_groups and retain their identity while their data moves between + CPU and CUDA. + chunk_size_bytes: Target bytes per optimizer-state chunk. Zero means one full chunk. + offload_fraction: Fraction of optimizer parameter bundles selected for offload. + state_dtypes: Dtypes of the supported full-size optimizer tensor states. For example, + Adam supplies its two moment dtypes and Muon supplies its momentum dtype. These are + used only for deterministic byte planning; the external optimizer remains the + authority for the actual state schema. + optimizer_owned_master_dtypes: Dtype of each optimizer-owned ``master_param`` entry. + Passing the exact parameters that materialize this entry avoids guessing and + accounts for compact int16 parameter remainders. + d2h_stream: Optional transfer stream shared by related optimizer wrappers. + h2d_stream: Optional transfer stream shared by related optimizer wrappers. + """ + + def __init__( + self, + optimizer: torch.optim.Optimizer, + master_params: Sequence[torch.Tensor], + chunk_size_bytes: int, + offload_fraction: float, + state_dtypes: Sequence[torch.dtype], + optimizer_owned_master_dtypes: Mapping[torch.Tensor, torch.dtype] | None = None, + d2h_stream: torch.cuda.Stream | None = None, + h2d_stream: torch.cuda.Stream | None = None, + ) -> None: + if chunk_size_bytes < 0: + raise ValueError(f"chunk_size_bytes must be non-negative, got {chunk_size_bytes}") + if not 0.0 <= offload_fraction <= 1.0: + raise ValueError(f"offload_fraction must be in [0, 1], got {offload_fraction}") + if not state_dtypes: + raise ValueError("state_dtypes must contain at least one optimizer-state dtype") + + self.optimizer = optimizer + self.chunk_size_bytes = chunk_size_bytes + self.offload_fraction = offload_fraction + self.state_dtypes = tuple(state_dtypes) + self._state_bytes_per_param = sum(dtype.itemsize for dtype in self.state_dtypes) + + self._params = self._unique_optimizer_params(optimizer.param_groups) + # Parameter-group membership is fixed for the optimizer lifetime. Cache it once so each + # state chunk only walks its own parameters instead of rescanning every optimizer group. + self._param_group_index_by_id = { + id(param): group_index + for group_index, group in enumerate(optimizer.param_groups) + for param in group["params"] + } + self._param_devices = {param: param.device for param in self._params} + self._explicit_master_param_ids = {id(param) for param in master_params} + self._master_in_optimizer_state = bool(getattr(optimizer, "master_weights", False)) + if self._explicit_master_param_ids and self._master_in_optimizer_state: + raise ValueError( + "chunked optimizer state offload cannot manage both MCore-owned master " + "parameters and optimizer state['master_param'] entries in one optimizer" + ) + if optimizer_owned_master_dtypes is None: + # Keep direct construction backward compatible. Production DistributedOptimizer + # callers pass the exact optimizer-owned subset and its real storage dtype. + self._optimizer_owned_master_dtypes = ( + {id(param): torch.float32 for param in self._params} + if self._master_in_optimizer_state + else {} + ) + else: + self._optimizer_owned_master_dtypes = { + id(param): dtype for param, dtype in optimizer_owned_master_dtypes.items() + } + unknown_param_ids = set(self._optimizer_owned_master_dtypes) - { + id(param) for param in self._params + } + if unknown_param_ids: + raise ValueError( + "optimizer_owned_master_dtypes contains parameters outside " + "optimizer.param_groups" + ) + if self._optimizer_owned_master_dtypes and not self._master_in_optimizer_state: + raise ValueError( + "optimizer_owned_master_dtypes requires an optimizer with master_weights" + ) + + self._selected_params = self._select_params_for_offload() + self._selected_param_ids = {id(param) for param in self._selected_params} + self._resident_params = tuple( + param for param in self._params if id(param) not in self._selected_param_ids + ) + self._chunks = self._build_chunks(self._selected_params) + + self._cpu_state: Dict[torch.Tensor, Dict[str, torch.Tensor]] = collections.defaultdict(dict) + # This map can be keyed by parameter alone because the constructor rejects combining + # MCore-owned ``param.data`` masters with optimizer-owned ``state['master_param']`` in + # one manager, and ``_validate_master_storage`` rechecks that invariant after lazy init. + self._cpu_master: Dict[torch.Tensor, torch.Tensor] = {} + + self._d2h_stream = d2h_stream if d2h_stream is not None else torch.cuda.Stream() + self._h2d_stream = h2d_stream if h2d_stream is not None else torch.cuda.Stream() + self._master_h2d_event: torch.cuda.Event | None = None + self._master_weights_resident = True + # Two reusable state windows allow H2D(N+1), step(N), and D2H(N-1) to overlap + # without letting host run-ahead allocate one CUDA tensor set per chunk. + self._state_staging_slots = (_StateStagingSlot(), _StateStagingSlot()) + self._next_state_staging_slot = 0 + self._first_prefetch: _StatePrefetch | None = None + + @staticmethod + def _unique_optimizer_params(param_groups: Sequence[dict]) -> tuple[torch.Tensor, ...]: + params = [] + seen = set() + for group in param_groups: + for param in group["params"]: + if id(param) in seen: + continue + seen.add(id(param)) + params.append(param) + return tuple(params) + + def _estimated_bundle_bytes(self, param: torch.Tensor) -> int: + state_bytes = self._estimated_state_bytes(param) + if id(param) in self._explicit_master_param_ids: + master_bytes = param.numel() * param.element_size() + elif id(param) in self._optimizer_owned_master_dtypes: + master_bytes = param.numel() * self._optimizer_owned_master_dtypes[id(param)].itemsize + else: + master_bytes = 0 + return state_bytes + master_bytes + + def _param_has_master(self, param: torch.Tensor) -> bool: + return ( + id(param) in self._explicit_master_param_ids + or id(param) in self._optimizer_owned_master_dtypes + ) + + def _select_params_for_offload(self) -> tuple[torch.Tensor, ...]: + if not self._params or self.offload_fraction == 0.0: + return () + if self.offload_fraction == 1.0: + return self._params + + target_bytes = math.ceil( + sum(self._estimated_bundle_bytes(param) for param in self._params) + * self.offload_fraction + ) + # Prefer bundles that actually have a separate master, so a partial byte budget removes + # both state and master residency before selecting state-only bundles. Python's sort is + # stable, preserving optimizer order within the two classes. + candidates = sorted(self._params, key=lambda param: not self._param_has_master(param)) + selected = [] + selected_bytes = 0 + for param in candidates: + if selected_bytes >= target_bytes: + break + selected.append(param) + selected_bytes += self._estimated_bundle_bytes(param) + return tuple(selected) + + def _estimated_state_bytes(self, param: torch.Tensor) -> int: + return param.numel() * self._state_bytes_per_param + + def _build_chunks(self, params: Sequence[torch.Tensor]) -> tuple[OptimizerStateChunk, ...]: + if not params: + return () + if self.chunk_size_bytes == 0: + return (OptimizerStateChunk(tuple(params)),) + + chunks = [] + current_params = [] + current_bytes = 0 + oversized_param_count = 0 + largest_oversized_param_bytes = 0 + for param in params: + param_bytes = self._estimated_state_bytes(param) + if param_bytes > self.chunk_size_bytes: + oversized_param_count += 1 + largest_oversized_param_bytes = max(largest_oversized_param_bytes, param_bytes) + if current_params and current_bytes + param_bytes > self.chunk_size_bytes: + chunks.append(OptimizerStateChunk(tuple(current_params))) + current_params = [] + current_bytes = 0 + current_params.append(param) + current_bytes += param_bytes + if current_params: + chunks.append(OptimizerStateChunk(tuple(current_params))) + if oversized_param_count: + mib = 1024**2 + log_single_rank( + logger, + logging.WARNING, + "Optimizer-state offload target is %.2f MiB, but %d selected parameter(s) " + "have an atomic tensor-state payload larger than the target; the largest is " + "%.2f MiB for %s. These parameters cannot be split, so this manager's " + "two-slot GPU tensor-state window can exceed roughly twice the configured " + "target.", + self.chunk_size_bytes / mib, + oversized_param_count, + largest_oversized_param_bytes / mib, + type(self.optimizer).__name__, + ) + return tuple(chunks) + + @property + def chunks(self) -> tuple[OptimizerStateChunk, ...]: + """Return the immutable optimizer-state chunk plan.""" + + return self._chunks + + @property + def selected_params(self) -> tuple[torch.Tensor, ...]: + """Return optimizer parameters whose state/master bundle is offloaded.""" + + return self._selected_params + + def is_param_offloaded(self, param: torch.Tensor) -> bool: + """Return whether param belongs to the offloaded fraction.""" + + return id(param) in self._selected_param_ids + + @property + def transfer_streams(self) -> tuple[torch.cuda.Stream, torch.cuda.Stream]: + """Return the D2H/H2D stream pair used by this manager.""" + + return self._d2h_stream, self._h2d_stream + + def use_transfer_streams( + self, d2h_stream: torch.cuda.Stream, h2d_stream: torch.cuda.Stream + ) -> None: + """Adopt streams shared by sibling optimizers before any transfer is scheduled.""" + + if d2h_stream is self._d2h_stream and h2d_stream is self._h2d_stream: + return + if ( + self._master_h2d_event is not None + or self._first_prefetch is not None + or self._cpu_state + or self._cpu_master + ): + raise RuntimeError("optimizer-state transfer streams cannot change after first use") + self._d2h_stream = d2h_stream + self._h2d_stream = h2d_stream + + @staticmethod + def _new_cpu_buffer(tensor: torch.Tensor) -> torch.Tensor: + return torch.empty( + tensor.size(), dtype=tensor.dtype, layout=tensor.layout, device="cpu", pin_memory=True + ) + + def _cpu_buffer_for_state( + self, param: torch.Tensor, key: str, tensor: torch.Tensor + ) -> torch.Tensor: + current = self._cpu_state[param].get(key) + if current is None or current.shape != tensor.shape or current.dtype != tensor.dtype: + current = self._new_cpu_buffer(tensor) + self._cpu_state[param][key] = current + return current + + def _cpu_buffer_for_master(self, param: torch.Tensor, tensor: torch.Tensor) -> torch.Tensor: + current = self._cpu_master.get(param) + if current is None or current.shape != tensor.shape or current.dtype != tensor.dtype: + current = self._new_cpu_buffer(tensor) + self._cpu_master[param] = current + return current + + @staticmethod + def _is_offloadable_state(param: torch.Tensor, key: str, value: object) -> bool: + """Return whether a tensor belongs to the supported chunked state schema.""" + + # Shape is sufficient only for the currently admitted fp32/bf16 Adam moments and the + # exact Muon momentum schema. FP8 moments carry one-element scale metadata that can match + # a one-element DistOpt shard; supporting them requires explicit state-key classification. + return ( + key not in _NON_OFFLOADABLE_STATE_KEYS + and isinstance(value, torch.Tensor) + and value.shape == param.shape + ) + + def _state_items_to_offload(self, param: torch.Tensor) -> Iterable[tuple[str, torch.Tensor]]: + for key, value in self.optimizer.state.get(param, {}).items(): + if self._is_offloadable_state(param, key, value) and value.is_cuda: + yield key, value + + def _has_unregistered_cuda_state(self, params: Sequence[torch.Tensor]) -> bool: + """Return whether a step lazily created tensor state outside the staging pool.""" + + return any( + key not in self._cpu_state.get(param, {}) + for param in params + for key, _ in self._state_items_to_offload(param) + ) + + def _schedule_state_d2h(self, params: Sequence[torch.Tensor]) -> None: + transfers = [ + (param, key, gpu_tensor) + for param in params + for key, gpu_tensor in self._state_items_to_offload(param) + ] + if not transfers: + return + + self._d2h_stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(self._d2h_stream): + for param, key, gpu_tensor in transfers: + cpu_tensor = self._cpu_buffer_for_state(param, key, gpu_tensor) + cpu_tensor.copy_(gpu_tensor, non_blocking=True) + gpu_tensor.record_stream(self._d2h_stream) + self.optimizer.state[param][key] = cpu_tensor + + def adopt_cpu_optimizer_state(self, params: Sequence[torch.Tensor] | None = None) -> None: + """Register CPU tensor states created or loaded outside the manager. + + Naturally CPU scalar state is left resident. Tensor state matching an optimizer + parameter shape is treated as offloaded state; this covers Adam moments and Muon + momentum while avoiding ordinary CPU step counters. Tracked buffers that are no longer + part of the optimizer's canonical state are released so checkpoint schema changes do + not retain stale pinned host memory. + """ + + params = self._selected_params if params is None else params + for param in params: + if not self.is_param_offloaded(param): + continue + adopted_keys = set() + for key, value in tuple(self.optimizer.state.get(param, {}).items()): + if self._is_offloadable_state(param, key, value) and value.device.type == "cpu": + if not value.is_pinned(): + pinned_value = self._new_cpu_buffer(value) + pinned_value.copy_(value) + self.optimizer.state[param][key] = pinned_value + value = pinned_value + self._cpu_state[param][key] = value + adopted_keys.add(key) + + tracked_state = self._cpu_state.get(param) + if tracked_state is None: + continue + for key in tuple(tracked_state): + if key not in adopted_keys: + del tracked_state[key] + if not tracked_state: + del self._cpu_state[param] + + def initialize_state_for_loading(self, init_state_fn: Callable | None, config: object) -> None: + """Initialize checkpoint state without materializing all selected state on CUDA. + + The optimizer's own initializer remains authoritative. This matters for precision-aware + FusedAdam, whose state can include scaling metadata and parameter remainders, and for + Muon, whose state schema is owned by Emerging Optimizers. Resident parameters initialize + normally; selected parameters initialize one state chunk at a time and are immediately + returned to CPU canonical storage. ``init_state_fn`` must only initialize per-parameter + optimizer state; parameter-group metadata writes target temporary subset-group clones and + are intentionally not committed to the original groups. + """ + + if init_state_fn is None: + raise RuntimeError( + "chunked optimizer state offload requires an optimizer state initializer " + "when preparing a distributed-checkpoint load" + ) + + # Checkpoint preparation may have put explicit MCore masters on CPU. Restore the + # selected master window so third-party initializers always observe CUDA parameters. + self._schedule_master_h2d() + self._wait_master_h2d() + + original_groups = list(self.optimizer.param_groups) + base_group_metadata = self._snapshot_group_metadata(original_groups) + + def initialize_subset(params: Sequence[torch.Tensor]) -> None: + groups, _ = self._make_subset_groups(params, base_group_metadata) + if not groups: + return + self.optimizer.param_groups = groups + init_state_fn(self.optimizer, config) + + try: + if self._resident_params: + initialize_subset(self._resident_params) + + for chunk in self._chunks: + initialize_subset(chunk.params) + self._schedule_state_d2h(chunk.params) + # Initialization is an infrequent checkpoint operation. Synchronizing one + # chunk here bounds its CUDA allocation and makes the CPU destination stable + # for distributed-checkpoint sharded-state construction. + self._d2h_stream.synchronize() + self.adopt_cpu_optimizer_state(chunk.params) + + # Precision-aware FusedAdam creates masters in optimizer.state during init. + # Scan and offload the full master window once, after all chunks exist. + self._schedule_master_d2h() + self._d2h_stream.synchronize() + self.adopt_cpu_optimizer_state() + finally: + self.optimizer.param_groups = original_groups + + def load_state_dict_without_device_cast(self, state_dict: dict) -> None: + """Load optimizer metadata while preserving each state's staging device. + + ``torch.optim.Optimizer.load_state_dict`` casts every tensor state to the + corresponding parameter device. The distributed-checkpoint template initializes + selected states in pinned CPU memory, although a loading strategy may preserve or + replace those destinations. Casting the resulting CPU state would temporarily + reconstruct the full state on CUDA. Reproduce the id remapping and parameter-group + restoration without that cast. + """ + + saved_groups = state_dict["param_groups"] + current_groups = self.optimizer.param_groups + if len(saved_groups) != len(current_groups): + raise ValueError( + "loaded optimizer has a different number of parameter groups: " + f"{len(saved_groups)} != {len(current_groups)}" + ) + + saved_param_ids = [param_id for group in saved_groups for param_id in group["params"]] + current_params = [param for group in current_groups for param in group["params"]] + if len(saved_param_ids) != len(current_params): + raise ValueError( + "loaded optimizer has a different number of parameters: " + f"{len(saved_param_ids)} != {len(current_params)}" + ) + id_map = dict(zip(saved_param_ids, current_params)) + current_param_ids = {id(param) for param in current_params} + + restored_groups = [] + for saved_group, current_group in zip(saved_groups, current_groups): + if len(saved_group["params"]) != len(current_group["params"]): + raise ValueError("loaded optimizer parameter group has a different size") + restored_group = dict(saved_group) + restored_group["params"] = current_group["params"] + restored_groups.append(restored_group) + + set_scaled_state = getattr(self.optimizer, "set_scaled_state", None) + state_dtype_map = getattr(self.optimizer, "name_to_dtype_map", None) + if callable(set_scaled_state) and isinstance(state_dtype_map, dict): + # TE FusedAdam serializes precision-aware state in its unscaled representation. + # Invoke TE's public scaling API on CUDA in bounded state chunks, then return selected + # destinations to CPU instead of installing serialized tensors verbatim. + # ``initialize_state_for_loading`` has already populated every state slot in the + # supported schema, so preserving ``current_state`` gives TE valid setter + # destinations. Unlike torch Optimizer.load_state_dict, a checkpoint that omits a + # parameter state retains that initialized value; supported checkpoints are complete. + current_state = self.optimizer.state + self.optimizer.__setstate__({"state": current_state, "param_groups": restored_groups}) + mapped_states = {} + for key, saved_state in state_dict["state"].items(): + param = id_map.get(key, key) + if id(param) not in current_param_ids: + self.optimizer.state[param] = saved_state + continue + + mapped_states[param] = saved_state + + def restore_scaled_state(param: torch.Tensor, saved_state: dict) -> None: + for state_name, value in saved_state.items(): + if value is None: + continue + if isinstance(value, torch.Tensor) and state_name in state_dtype_map: + store_remainder = ( + bool(getattr(self.optimizer, "store_param_remainders", False)) + and state_name == _MASTER_PARAM_KEY + and param.dtype == torch.bfloat16 + ) + target_dtype = value.dtype if store_remainder else torch.float32 + scaled_value = value.to( + device=self._param_devices[param], + dtype=target_dtype, + non_blocking=value.device.type == "cpu" and value.is_pinned(), + ) + set_scaled_state(param, state_name, scaled_value) + else: + self.optimizer.state[param][state_name] = value + + for param, saved_state in mapped_states.items(): + if not self.is_param_offloaded(param): + restore_scaled_state(param, saved_state) + + self._schedule_master_h2d() + self._wait_master_h2d() + for chunk in self._chunks: + for param in chunk.params: + saved_state = mapped_states.get(param) + if saved_state is not None: + restore_scaled_state(param, saved_state) + # ``initialize_state_for_loading`` already installed CPU destinations. TE's + # setter scales on CUDA and copies into those existing tensors, so there is no + # CUDA optimizer state to schedule back through the D2H staging path here. + self._release_state_staging_slots() + self._schedule_master_d2h() + # The TE setter's CUDA-to-CPU copies run on the current stream. This explicit edge + # also covers precision-aware optimizers without an MCore-owned master D2H transfer. + self._d2h_stream.wait_stream(torch.cuda.current_stream()) + self._d2h_stream.synchronize() + else: + restored_state = collections.defaultdict(dict) + for key, value in state_dict["state"].items(): + param = id_map.get(key, key) + if id(param) in current_param_ids and isinstance(value, dict): + value = dict(value) + if not self.is_param_offloaded(param): + for state_name, state_value in tuple(value.items()): + if ( + self._is_offloadable_state(param, state_name, state_value) + and state_value.device.type == "cpu" + ): + value[state_name] = state_value.to( + device=self._param_devices[param], + non_blocking=state_value.is_pinned(), + ) + restored_state[param] = value + self.optimizer.__setstate__({"state": restored_state, "param_groups": restored_groups}) + self.adopt_cpu_optimizer_state() + + def _validate_master_storage(self) -> None: + for param in self._selected_params: + if id(param) not in self._explicit_master_param_ids: + continue + state_master = self.optimizer.state.get(param, {}).get(_MASTER_PARAM_KEY) + if isinstance(state_master, torch.Tensor): + raise RuntimeError( + "optimizer created state['master_param'] for a parameter that already uses " + "an MCore-owned master parameter" + ) + + def _master_bindings_are_cuda(self) -> bool: + """Return whether every materialized selected master is CUDA-bound.""" + + for param in self._selected_params: + if id(param) in self._explicit_master_param_ids and not param.data.is_cuda: + return False + state_master = self.optimizer.state.get(param, {}).get(_MASTER_PARAM_KEY) + if isinstance(state_master, torch.Tensor) and not state_master.is_cuda: + return False + return True + + def assert_master_weights_resident(self, operation: str) -> None: + """Reject an external master reader while D2H/H2D changes its canonical binding.""" + + if self._master_weights_resident and self._master_bindings_are_cuda(): + return + raise RuntimeError( + f"{operation} requires CUDA-resident optimizer master weights, but chunked " + "optimizer state offload has rebound at least one selected master to CPU or an " + "H2D restore is still pending" + ) + + def _schedule_master_d2h(self) -> None: + if not self._selected_params: + return + self._validate_master_storage() + # A rerun or checkpoint can request offload after prefetch but before step(). + # Make the current stream depend on the H2D copy before the D2H stream follows it. + self._wait_master_h2d() + + transfers = [] + for param in self._selected_params: + if id(param) in self._explicit_master_param_ids and param.data.is_cuda: + transfers.append((param, None, param.data)) + state_master = self.optimizer.state.get(param, {}).get(_MASTER_PARAM_KEY) + if isinstance(state_master, torch.Tensor) and state_master.is_cuda: + transfers.append((param, _MASTER_PARAM_KEY, state_master)) + if not transfers: + return + + self._master_weights_resident = False + self._d2h_stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(self._d2h_stream): + for param, state_key, gpu_tensor in transfers: + cpu_tensor = self._cpu_buffer_for_master(param, gpu_tensor) + cpu_tensor.copy_(gpu_tensor, non_blocking=True) + gpu_tensor.record_stream(self._d2h_stream) + # CPU is the canonical binding as soon as the copy is enqueued. CUDA's + # allocator keeps the source alive through record_stream(), and any H2D + # consumer waits on the D2H stream before reading this buffer. + if state_key is None: + param.data = cpu_tensor + else: + self.optimizer.state[param][state_key] = cpu_tensor + + def _schedule_master_h2d(self) -> None: + if not self._selected_params or self._master_h2d_event is not None: + return + self._validate_master_storage() + + transfers = [] + for param in self._selected_params: + if id(param) in self._explicit_master_param_ids and not param.data.is_cuda: + transfers.append((param, None, param.data)) + state_master = self.optimizer.state.get(param, {}).get(_MASTER_PARAM_KEY) + if isinstance(state_master, torch.Tensor) and not state_master.is_cuda: + transfers.append((param, _MASTER_PARAM_KEY, state_master)) + if not transfers: + self._master_weights_resident = self._master_bindings_are_cuda() + return + + self._master_weights_resident = False + gpu_transfers = [ + ( + param, + state_key, + cpu_tensor, + torch.empty_like(cpu_tensor, device=self._param_devices[param]), + ) + for param, state_key, cpu_tensor in transfers + ] + + self._order_h2d_after_source_streams() + with torch.cuda.stream(self._h2d_stream): + for param, state_key, cpu_tensor, gpu_tensor in gpu_transfers: + gpu_tensor.copy_(cpu_tensor, non_blocking=cpu_tensor.is_pinned()) + gpu_tensor.record_stream(self._h2d_stream) + if state_key is None: + param.data = gpu_tensor + else: + self.optimizer.state[param][state_key] = gpu_tensor + self._cpu_master[param] = cpu_tensor + + event = torch.cuda.Event() + event.record(self._h2d_stream) + + self._master_h2d_event = event + + def _wait_master_h2d(self) -> None: + if self._master_h2d_event is None: + return + torch.cuda.current_stream().wait_event(self._master_h2d_event) + self._master_h2d_event = None + self._master_weights_resident = self._master_bindings_are_cuda() + + def _order_h2d_after_source_streams(self) -> None: + """Order H2D after CPU production and compute-stream GPU allocation/use. + + CPU canonical buffers may still be D2H destinations. GPU staging buffers are + deliberately allocated on the current compute stream so the caching allocator can + recycle them into later forward/backward allocations. The H2D stream must therefore + wait for both producer/owner streams before writing either side. + """ + + self._h2d_stream.wait_stream(self._d2h_stream) + self._h2d_stream.wait_stream(torch.cuda.current_stream()) + + def _state_staging_views( + self, chunk: OptimizerStateChunk + ) -> List[tuple[torch.Tensor, str, torch.Tensor, torch.Tensor]]: + """Return state views backed by the next reusable CUDA staging slot.""" + + entries = [] + required_numel = collections.defaultdict(int) + for param in chunk.params: + state = self.optimizer.state.get(param, {}) + for key, cpu_tensor in self._cpu_state.get(param, {}).items(): + current = state.get(key) + if not isinstance(current, torch.Tensor) or current.is_cuda: + continue + pool_key = (self._param_devices[param], cpu_tensor.dtype) + entries.append((param, key, cpu_tensor, pool_key)) + alignment_numel = max(1, 256 // cpu_tensor.element_size()) + required_numel[pool_key] = ( + (required_numel[pool_key] + alignment_numel - 1) // alignment_numel + ) * alignment_numel + required_numel[pool_key] += cpu_tensor.numel() + + if not entries: + return [] + + slot = self._state_staging_slots[self._next_state_staging_slot] + self._next_state_staging_slot = (self._next_state_staging_slot + 1) % len( + self._state_staging_slots + ) + for (device, dtype), numel in required_numel.items(): + buffer = slot.buffers.get((device, dtype)) + if buffer is None or buffer.numel() < numel: + slot.buffers[(device, dtype)] = torch.empty(numel, dtype=dtype, device=device) + + offsets = collections.defaultdict(int) + views = [] + for param, key, cpu_tensor, pool_key in entries: + alignment_numel = max(1, 256 // cpu_tensor.element_size()) + offset = ( + (offsets[pool_key] + alignment_numel - 1) // alignment_numel + ) * alignment_numel + buffer = slot.buffers[pool_key] + gpu_tensor = buffer.narrow(0, offset, cpu_tensor.numel()).view(cpu_tensor.shape) + offsets[pool_key] = offset + cpu_tensor.numel() + views.append((param, key, cpu_tensor, gpu_tensor)) + return views + + def _release_state_staging_slots(self) -> None: + """Drop GPU pool ownership after all live views have been queued for D2H.""" + + for slot in self._state_staging_slots: + slot.buffers.clear() + self._next_state_staging_slot = 0 + + def _prefetch_state(self, chunk: OptimizerStateChunk) -> _StatePrefetch: + staging_views = self._state_staging_views(chunk) + if not staging_views: + return _StatePrefetch(None) + + self._order_h2d_after_source_streams() + with torch.cuda.stream(self._h2d_stream): + for param, key, cpu_tensor, gpu_tensor in staging_views: + gpu_tensor.copy_(cpu_tensor, non_blocking=cpu_tensor.is_pinned()) + gpu_tensor.record_stream(self._h2d_stream) + self.optimizer.state[param][key] = gpu_tensor + event = torch.cuda.Event() + event.record(self._h2d_stream) + return _StatePrefetch(event) + + @staticmethod + def _wait_state_prefetch(prefetch: _StatePrefetch) -> None: + if prefetch.event is not None: + torch.cuda.current_stream().wait_event(prefetch.event) + + def prefetch_for_step(self) -> None: + """Asynchronously restore masters and the first state chunk for the next update.""" + + if not self._selected_params: + return + self._schedule_master_h2d() + if self._chunks and self._first_prefetch is None: + self._first_prefetch = self._prefetch_state(self._chunks[0]) + + def prefetch_master_for_step(self) -> None: + """Asynchronously restore only masters, leaving tensor state CPU-resident.""" + + if self._selected_params: + self._schedule_master_h2d() + + def ensure_master_for_param_sync(self) -> None: + """Ensure selected master weights are available for pre-forward FP8 staging.""" + + if self._master_weights_resident and self._master_h2d_event is None: + return + self._schedule_master_h2d() + self._wait_master_h2d() + + def offload_for_forward(self, offload_master: bool = True) -> None: + """Start D2H for tensors still resident after an update or load. + + Selected master bindings become CPU tensors as soon as their asynchronous D2H copy is + enqueued. CUDA consumers must restore them through this manager, and host consumers must + first call :meth:`synchronize_for_checkpoint`. The training loop deliberately calls this + only at the optimizer-to-forward lifecycle boundary, where no other master reader runs. + + Args: + offload_master: If false, stage only optimizer state. The training loop uses this + while an MXFP8 parameter buffer still needs to read updated masters. + """ + + if not self._selected_params: + return + if self._first_prefetch is not None: + self._wait_state_prefetch(self._first_prefetch) + self._first_prefetch = None + for chunk in self._chunks: + self._schedule_state_d2h(chunk.params) + self._release_state_staging_slots() + if offload_master: + self._schedule_master_d2h() + + def synchronize_for_checkpoint(self) -> None: + """Make CPU canonical state stable and visible to distributed checkpointing.""" + + self.offload_for_forward() + self._d2h_stream.synchronize() + self.adopt_cpu_optimizer_state() + + @classmethod + def _clone_group_value(cls, value: object) -> object: + """Clone mutable group metadata so each chunk starts from one logical step.""" + + if isinstance(value, torch.Tensor): + return value.clone() + if isinstance(value, dict): + return {key: cls._clone_group_value(item) for key, item in value.items()} + if isinstance(value, list): + return [cls._clone_group_value(item) for item in value] + if isinstance(value, tuple): + return tuple(cls._clone_group_value(item) for item in value) + if isinstance(value, set): + return {cls._clone_group_value(item) for item in value} + return value + + @classmethod + def _group_values_equal(cls, left: object, right: object) -> bool: + if isinstance(left, torch.Tensor) or isinstance(right, torch.Tensor): + compatible = ( + isinstance(left, torch.Tensor) + and isinstance(right, torch.Tensor) + and left.shape == right.shape + and left.dtype == right.dtype + and left.device == right.device + ) + if not compatible: + return False + # CUDA value comparison would synchronize the host once per chunk. Each subset + # starts from its own clone of the same metadata; commit the first result and only + # validate structural compatibility for CUDA tensor-valued fields. Snapshot/commit + # replaces tensor-valued group fields with clones, so identity-sensitive consumers + # must not retain references to those tensors; optimizer CUDA graphs are rejected. + return left.is_cuda or torch.equal(left, right) + if isinstance(left, dict) or isinstance(right, dict): + return ( + isinstance(left, dict) + and isinstance(right, dict) + and left.keys() == right.keys() + and all(cls._group_values_equal(left[key], right[key]) for key in left) + ) + if isinstance(left, (list, tuple)) or isinstance(right, (list, tuple)): + return ( + isinstance(left, (list, tuple)) + and isinstance(right, (list, tuple)) + and len(left) == len(right) + and all(cls._group_values_equal(x, y) for x, y in zip(left, right)) + ) + try: + result = left == right + except (TypeError, ValueError): + return left is right + try: + return bool(result) + except (TypeError, ValueError): + return left is right + + @classmethod + def _snapshot_group_metadata(cls, groups: Sequence[dict]) -> List[dict]: + return [ + {key: cls._clone_group_value(value) for key, value in group.items() if key != "params"} + for group in groups + ] + + def _make_subset_groups( + self, params: Sequence[torch.Tensor], base_group_metadata: Sequence[dict] + ) -> tuple[List[dict], List[tuple[int, dict]]]: + params_by_group = collections.defaultdict(list) + for param in params: + group_index = self._param_group_index_by_id.get(id(param)) + if group_index is None: + raise RuntimeError( + "optimizer parameter-group membership changed after offload setup" + ) + params_by_group[group_index].append(param) + + groups = [] + indexed_groups = [] + for group_index in sorted(params_by_group): + group = { + key: self._clone_group_value(value) + for key, value in base_group_metadata[group_index].items() + } + group["params"] = params_by_group[group_index] + groups.append(group) + indexed_groups.append((group_index, group)) + return groups, indexed_groups + + def _step_subset( + self, + params: Sequence[torch.Tensor], + original_groups: Sequence[dict], + base_group_metadata: Sequence[dict], + resulting_group_metadata: Dict[int, dict], + ) -> None: + groups, indexed_groups = self._make_subset_groups(params, base_group_metadata) + if not groups: + return + self.optimizer.param_groups = groups + try: + self.optimizer.step() + for group_index, group in indexed_groups: + metadata = self._snapshot_group_metadata([group])[0] + previous = resulting_group_metadata.get(group_index) + if previous is not None and not self._group_values_equal(previous, metadata): + raise RuntimeError( + "optimizer parameter-group updates differ across state chunks for " + "metadata comparable without CUDA value synchronization; chunked " + "optimizer state offload requires group-level updates to be independent " + "of the parameter subset. CUDA tensor-valued fields are validated only " + "by shape, dtype, and device" + ) + if previous is None: + resulting_group_metadata[group_index] = metadata + finally: + self.optimizer.param_groups = original_groups + + @torch.no_grad() + def step(self) -> None: + """Run the external optimizer over resident parameters and staged state chunks.""" + + if not self._selected_params: + self.optimizer.step() + return + + self.prefetch_for_step() + self._wait_master_h2d() + + original_groups = list(self.optimizer.param_groups) + base_group_metadata = self._snapshot_group_metadata(original_groups) + resulting_group_metadata: Dict[int, dict] = {} + + try: + # Resident state needs no transfers and provides useful compute while the first + # selected chunk is being prefetched. + if self._resident_params: + self._step_subset( + self._resident_params, + original_groups, + base_group_metadata, + resulting_group_metadata, + ) + + current_prefetch = self._first_prefetch + self._first_prefetch = None + for chunk_index, chunk in enumerate(self._chunks): + if current_prefetch is None: + current_prefetch = self._prefetch_state(chunk) + self._wait_state_prefetch(current_prefetch) + + next_prefetch = None + if chunk_index + 1 < len(self._chunks): + next_prefetch = self._prefetch_state(self._chunks[chunk_index + 1]) + + self._step_subset( + chunk.params, original_groups, base_group_metadata, resulting_group_metadata + ) + created_state = self._has_unregistered_cuda_state(chunk.params) + self._schedule_state_d2h(chunk.params) + if created_state: + # External optimizers commonly allocate moments lazily on their first step. + # That storage is not backed by the reusable pool yet, so drain this one-time + # initialization copy before the next chunk to preserve the peak-memory bound. + self._d2h_stream.synchronize() + current_prefetch = next_prefetch + finally: + self.optimizer.param_groups = original_groups + self._release_state_staging_slots() + + for group_index, metadata in resulting_group_metadata.items(): + original_group = original_groups[group_index] + for key in tuple(original_group): + if key != "params" and key not in metadata: + del original_group[key] + original_group.update(metadata) diff --git a/megatron/core/optimizer/cpu_offloading/optimizer_state_offloader.py b/megatron/core/optimizer/cpu_offloading/optimizer_state_offloader.py deleted file mode 100644 index 81fd116c8ba..00000000000 --- a/megatron/core/optimizer/cpu_offloading/optimizer_state_offloader.py +++ /dev/null @@ -1,315 +0,0 @@ -# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. - -"""Optimizer state offloading class.""" - -from typing import TYPE_CHECKING, Dict, List, Tuple - -import torch - -if TYPE_CHECKING: - from megatron.core.optimizer.distrib_optimizer import DistributedOptimizer - - -class OptimizerStateOffloader: - """ - Manages offloading of optimizer states and master weights to CPU. - Used with DistributedOptimizer to reduce GPU memory usage. - - Supports overlapped D2H/H2D transfers using CUDA streams. - - Master weights can be stored in two locations: - - In adam optimizer state (when use_precision_aware_optimizer_no_fp8_or_ds_fp8 is True) - - In mcore's shard_fp32_from_float16_groups - """ - - OPTIMIZER_STATE_KEYS = ('exp_avg', 'exp_avg_sq') - MASTER_WEIGHT_KEY = 'master_param' - - def __init__(self, distrib_optimizer: "DistributedOptimizer"): - """ - Args: - distrib_optimizer: The DistributedOptimizer to offload states and master weights from. - """ - self.dist_optimizer = distrib_optimizer - self.adam_optimizer = distrib_optimizer.optimizer - - # Only support TE FusedAdam optimizer for now. - try: - from transformer_engine.pytorch.optimizers import FusedAdam - - assert isinstance(self.adam_optimizer, FusedAdam), ( - f"OptimizerStateOffloader requires TE FusedAdam optimizer, " - f"but got {type(self.adam_optimizer).__name__}" - ) - except ImportError: - raise ImportError( - "OptimizerStateOffloader requires transformer_engine.pytorch.optimizers.FusedAdam" - ) - - # Check if master weights are stored in adam optimizer state - self.optimizer_contains_master_weights = self.adam_optimizer.master_weights - - # CUDA streams for async transfers - self._d2h_stream = torch.cuda.Stream() - self._h2d_stream = torch.cuda.Stream() - - # CPU buffers for optimizer states: {param: {key: cpu_tensor}} - self._opt_state_cpu_buffers: Dict[torch.Tensor, Dict[str, torch.Tensor]] = {} - - # CPU buffers for mcore master weights, matching the structure of source groups - # List[List[cpu_tensor]] - self._shard_fp32_from_float16_cpu_buffers: List[List[torch.Tensor]] = [] - - # State tracking - self._offloaded = False - self._offloaded_state_keys: Tuple[str, ...] = () - self._offloaded_mcore_master_weights = False - - # Track whether optimizer states (exp_avg, exp_avg_sq) have been initialized. - # These are lazily initialized by FusedAdam during the first optimizer.step(). - # Master weights (shard_fp32_from_float16_groups) are available from the start. - self._optimizer_states_initialized = False - - def mark_optimizer_states_initialized(self): - """ - Mark that optimizer states (exp_avg, exp_avg_sq) are now available. - Should be called after the first optimizer.step() completes. - """ - self._optimizer_states_initialized = True - - def _get_state_keys_to_offload( - self, offload_optimizer_states: bool, offload_master_weights: bool - ) -> Tuple[str, ...]: - """Get the state keys in FusedAdam to offload based on configuration.""" - keys = [] - # Skip optimizer states offloading if they haven't been initialized yet. - # Optimizer states are lazily initialized by FusedAdam during the first optimizer.step(). - if self._optimizer_states_initialized: - if offload_optimizer_states: - keys.extend(self.OPTIMIZER_STATE_KEYS) - if offload_master_weights and self.optimizer_contains_master_weights: - keys.append(self.MASTER_WEIGHT_KEY) - return tuple(keys) - - def _ensure_state_cpu_buffer( - self, param: torch.Tensor, state_key: str, gpu_tensor: torch.Tensor, pin_memory: bool = True - ) -> torch.Tensor: - """Get or create a CPU buffer for a state tensor.""" - if param not in self._opt_state_cpu_buffers: - self._opt_state_cpu_buffers[param] = {} - - if state_key not in self._opt_state_cpu_buffers[param]: - cpu_buffer = torch.empty( - gpu_tensor.size(), - dtype=gpu_tensor.dtype, - layout=gpu_tensor.layout, - device='cpu', - pin_memory=pin_memory, - ) - self._opt_state_cpu_buffers[param][state_key] = cpu_buffer - - return self._opt_state_cpu_buffers[param][state_key] - - def _offload_shard_groups( - self, - shard_groups: List[List[torch.Tensor]], - cpu_buffers: List[List[torch.Tensor]], - pin_memory: bool = True, - ): - """Offload a shard group to CPU buffers.""" - # Initialize CPU buffers on first call - if len(cpu_buffers) == 0: - for group in shard_groups: - group_buffers = [] - for gpu_tensor in group: - cpu_buffer = torch.empty( - gpu_tensor.size(), - dtype=gpu_tensor.dtype, - layout=gpu_tensor.layout, - device='cpu', - pin_memory=pin_memory, - ) - group_buffers.append(cpu_buffer) - cpu_buffers.append(group_buffers) - - # Copy D2H - for group_idx, group in enumerate(shard_groups): - for param_idx, gpu_tensor in enumerate(group): - cpu_buffer = cpu_buffers[group_idx][param_idx] - cpu_buffer.copy_(gpu_tensor, non_blocking=pin_memory) - gpu_tensor.record_stream(self._d2h_stream) - - def _offload_states( - self, - offload_optimizer_states: bool, - offload_master_weights: bool, - use_pin_memory: bool = True, - ): - """Offload optimizer states and/or master weights to CPU.""" - # Offload states from adam optimizer - self._offloaded_state_keys = self._get_state_keys_to_offload( - offload_optimizer_states, offload_master_weights - ) - states = self.adam_optimizer.state - - for param, param_state in states.items(): - for state_key in self._offloaded_state_keys: - if state_key not in param_state: - continue - - gpu_tensor = param_state[state_key] - if not isinstance(gpu_tensor, torch.Tensor) or not gpu_tensor.is_cuda: - continue - - cpu_buffer = self._ensure_state_cpu_buffer( - param, state_key, gpu_tensor, use_pin_memory - ) - cpu_buffer.copy_(gpu_tensor, non_blocking=use_pin_memory) - gpu_tensor.record_stream(self._d2h_stream) - - # Offload mcore master weights if not in optimizer state - if offload_master_weights and not self.optimizer_contains_master_weights: - self._offload_shard_groups( - self.dist_optimizer.shard_fp32_from_float16_groups, - self._shard_fp32_from_float16_cpu_buffers, - use_pin_memory, - ) - self._offloaded_mcore_master_weights = True - - def _release_states(self): - """Replace optimizer state GPU tensors with CPU tensors to free GPU memory.""" - states = self.adam_optimizer.state - - for param, param_state in states.items(): - if param not in self._opt_state_cpu_buffers: - continue - - for state_key in self._offloaded_state_keys: - if state_key not in self._opt_state_cpu_buffers[param]: - continue - - param_state[state_key].untyped_storage().resize_(0) - - if self._offloaded_mcore_master_weights: - for group in self.dist_optimizer.shard_fp32_from_float16_groups: - for gpu_tensor in group: - gpu_tensor.untyped_storage().resize_(0) - - def _reload_shard_groups( - self, - shard_groups: List[List[torch.Tensor]], - cpu_buffers: List[List[torch.Tensor]], - is_allocate_stage: bool, - ): - """Reload shard groups from CPU to GPU.""" - for group_idx, group in enumerate(shard_groups): - for param_idx, _ in enumerate(group): - cpu_buffer = cpu_buffers[group_idx][param_idx] - if is_allocate_stage: - shard_groups[group_idx][param_idx].untyped_storage().resize_( - cpu_buffer.untyped_storage().size() - ) - else: - shard_groups[group_idx][param_idx].copy_( - cpu_buffer, non_blocking=cpu_buffer.is_pinned() - ) - - def _reload_states(self, is_allocate_stage: bool): - """ - Reload optimizer states and/or master weights from CPU to GPU. - - If is_allocate_stage is True, only allocate GPU memory for the states and master weights, - but do not copy the data from CPU to GPU. Otherwise, copy the data from CPU to GPU. - The two processes are separated to make sure that the GPU memory is allocated on the - default stream to avoid fragmentation. - """ - # Reload states to adam optimizer - states = self.adam_optimizer.state - - for param, param_state in states.items(): - if param not in self._opt_state_cpu_buffers: - continue - - for state_key in self._offloaded_state_keys: - if state_key not in self._opt_state_cpu_buffers[param]: - continue - - cpu_buffer = self._opt_state_cpu_buffers[param][state_key] - if is_allocate_stage: - param_state[state_key].untyped_storage().resize_( - cpu_buffer.untyped_storage().size() - ) - else: - param_state[state_key].copy_(cpu_buffer, non_blocking=cpu_buffer.is_pinned()) - - # Reload mcore master weights if not in optimizer state - if self._offloaded_mcore_master_weights: - self._reload_shard_groups( - self.dist_optimizer.shard_fp32_from_float16_groups, - self._shard_fp32_from_float16_cpu_buffers, - is_allocate_stage, - ) - - def offload(self, offload_optimizer_states: bool = True, offload_master_weights: bool = True): - """ - Offload optimizer states and/or master weights to CPU. - Starts async D2H transfer that can overlap with other operations. - - Args: - offload_optimizer_states: Whether to offload exp_avg, exp_avg_sq. - offload_master_weights: Whether to offload master weights. - """ - if not offload_optimizer_states and not offload_master_weights: - return - - # Wait for current stream finishing updating the optimizer states. - self._d2h_stream.wait_stream(torch.cuda.current_stream()) - - with torch.cuda.stream(self._d2h_stream): - self._offload_states(offload_optimizer_states, offload_master_weights) - - self._offloaded = True - - def release_gpu_memory(self): - """ - Release GPU memory for optimizer states and master weights after D2H copy completes. - - This is separated from offload() to allow delayed GPU memory release, - which is needed for mxfp8 + overlap_param_gather case where master weights - must remain on GPU until after _copy_main_params_to_param_buffer() is called. - """ - if not self._offloaded: - return - - self._release_states() - - def reload(self): - """ - Reload optimizer states and/or master weights from CPU to GPU. - Call before optimizer.step() to ensure states are on GPU. - """ - if not self._offloaded: - return - - # Allocate GPU memory on the current stream to avoid fragmentation. - self._reload_states(is_allocate_stage=True) - - self._h2d_stream.wait_stream(self._d2h_stream) - self._h2d_stream.wait_stream(torch.cuda.current_stream()) - - # Reload states on the h2d stream to overlap with other operations. - with torch.cuda.stream(self._h2d_stream): - self._reload_states(is_allocate_stage=False) - - self._offloaded_state_keys = () - self._offloaded_mcore_master_weights = False - self._offloaded = False - - def sync_before_step(self): - """ - Wait for H2D reload to complete before optimizer.step(). - Must be called to ensure states are on GPU before optimizer uses them. - - This is separated from reload() to make it possible to move the reload ahead of time. - """ - torch.cuda.current_stream().wait_stream(self._h2d_stream) diff --git a/megatron/core/optimizer/distrib_optimizer.py b/megatron/core/optimizer/distrib_optimizer.py index 9baed75a11d..5f5a9c716e5 100644 --- a/megatron/core/optimizer/distrib_optimizer.py +++ b/megatron/core/optimizer/distrib_optimizer.py @@ -62,7 +62,6 @@ ) from ..transformer.fsdp_dtensor_checkpoint import handle_experts_in_state_dict from ..transformer.module import MegatronModule -from .cpu_offloading.optimizer_state_offloader import OptimizerStateOffloader from .grad_scaler import MegatronGradScaler from .optimizer import ( MixedPrecisionOptimizer, @@ -702,8 +701,6 @@ def __init__( "due to checkpointing requirements." ) - self._state_offloader: Optional[OptimizerStateOffloader] = None - # when freezing sub-models we have no real optimizer # but still need a stub DistributedOptimizer class if optimizer is None: @@ -812,8 +809,42 @@ def __init__( self.optimizer.param_groups = [g["orig_group"] for g in self.opt_group_ranges] self.optimizer.load_state_dict(self.optimizer.state_dict()) - if self.config.offload_optimizer_states: - self._state_offloader = OptimizerStateOffloader(self) + if ( + self.config.chunked_optimizer_state_offload + and self.config.optimizer_state_offload_fraction > 0.0 + ): + separate_master_params = [] + optimizer_owned_master_dtypes = {} + if bool(getattr(self.optimizer, "master_weights", False)): + # Precision-aware FusedAdam creates optimizer-owned masters for every + # optimizer parameter, including native-fp32 DistributedOptimizer shards. + for group in itertools.chain(self.shard_float16_groups, self.shard_fp32_groups): + for param in group: + if not isinstance(param, torch.Tensor): + continue + master_dtype = ( + torch.int16 + if self.config.store_param_remainders and param.dtype == torch.bfloat16 + else self.config.main_params_dtype + ) + optimizer_owned_master_dtypes[param] = master_dtype + else: + separate_master_params = [ + param + for group in self.shard_fp32_from_float16_groups + for param in group + if isinstance(param, torch.Tensor) + ] + state_dtypes = ( + (self.config.exp_avg_dtype, self.config.exp_avg_sq_dtype) + if self.config.use_precision_aware_optimizer + else (torch.float32, torch.float32) + ) + self.enable_chunked_optimizer_state_offload( + master_params=separate_master_params, + state_dtypes=state_dtypes, + optimizer_owned_master_dtypes=optimizer_owned_master_dtypes, + ) def _get_model_param_range_map(self, param: torch.nn.Parameter): """ @@ -1046,9 +1077,16 @@ def make_needed_groups(param_group): v["step"] = step.detach().clone() # Optimizer. - self.optimizer.load_state_dict( - {"state": state_dict_state, "param_groups": state_dict_param_groups} - ) + optimizer_state_dict = {"state": state_dict_state, "param_groups": state_dict_param_groups} + if self._optimizer_state_offloader is None: + self.optimizer.load_state_dict(optimizer_state_dict) + else: + # The distributed checkpoint path has already allocated selected tensor + # states on CPU. Avoid Optimizer.load_state_dict's automatic cast to the + # parameter device, which would recreate the complete state on CUDA. + self._optimizer_state_offloader.load_state_dict_without_device_cast( + optimizer_state_dict + ) # Grad scaler. if 'grad_scaler' not in state_dict: @@ -1513,6 +1551,14 @@ def sharded_state_dict( Regular state dict parameters are saved on DP rank 0 and loaded on all ranks. """ + if self._optimizer_state_offloader is not None: + if is_loading: + self._optimizer_state_offloader.initialize_state_for_loading( + self.init_state_fn, self.config + ) + else: + self._optimizer_state_offloader.synchronize_for_checkpoint() + if sharding_type is not None: log_single_rank( logger, @@ -1563,7 +1609,7 @@ def sharded_state_dict( for k, v in state_dict.items() } - if is_loading: + if is_loading and self._optimizer_state_offloader is None: # Call the distributed optimizer's specialized load_state_dict(), # which conditionally skips re-allocating the optimizer's state if # already initialized, which in turn reduces memory fragmentation. @@ -2828,6 +2874,9 @@ def _copy_main_params_to_model_params(self): if self.config.use_precision_aware_optimizer_no_fp8_or_ds_fp8: return + self.ensure_master_weights_for_param_sync() + self.assert_master_weights_resident("_copy_main_params_to_model_params") + if self.ddp_config.fp4_param_gather: # Quantize FP32 master shards back to NVFP4 model params (rowwise only) quantize_nvfp4_param_shard( @@ -2909,6 +2958,20 @@ def _copy_main_params_to_param_buffer(self): raise NotImplementedError( "_copy_main_params_to_param_buffer not supported for Megatron-FSDP." ) + self.ensure_master_weights_for_param_sync() + self.assert_master_weights_resident("_copy_main_params_to_param_buffer") + + # A speculative next-bucket gather can still own the shared MXFP8 parameter + # buffer here. Finish only outstanding gathers before overwriting that buffer. + # force_sync consumes an existing handle without dispatching a new gather and + # is safe even if gradient finalization already reset the logical dispatch flag. + for model_chunk in self.model_chunks: + for bucket_group in ( + model_chunk.bucket_groups + model_chunk.expert_parallel_bucket_groups + ): + if bucket_group.param_gather_handle is not None: + model_chunk._start_bucket_group_param_sync(bucket_group, force_sync=True) + for shard_main_group, model_group in zip( self.shard_fp32_from_float16_groups, self.model_float16_groups ): @@ -3079,6 +3142,9 @@ def _copy_model_params_to_main_params(self, state_dict=None): the model params. This copy does not make use of the grad buffer as an intermediary. """ + self.ensure_master_weights_for_param_sync() + self.assert_master_weights_resident("_copy_model_params_to_main_params") + if isinstance(self.optimizer, HybridDeviceOptimizer): self.optimizer.update_fp32_param_by_new_param() return @@ -3154,7 +3220,7 @@ def copy_group_params(model_groups, shard_main_groups): copy_group_params(self.model_float16_groups, self.shard_fp32_from_float16_groups) copy_group_params(self.model_fp32_groups, self.shard_fp32_groups) - def start_param_sync_for_bucket_group_subset(self) -> None: + def start_param_sync_for_bucket_group_subset(self, force_sync: bool = False) -> None: """Trigger ``start_param_sync`` on DistOpt-managed bucket groups only. Walks each model chunk's DDP bucket groups and skips those tagged @@ -3180,7 +3246,7 @@ def start_param_sync_for_bucket_group_subset(self) -> None: bucket_group.buckets[0], default_for_untagged=False ): continue - model_chunk._start_bucket_group_param_sync(bucket_group, force_sync=False) + model_chunk._start_bucket_group_param_sync(bucket_group, force_sync=force_sync) @torch.no_grad() def step_with_ready_grads(self) -> bool: @@ -3188,8 +3254,6 @@ def step_with_ready_grads(self) -> bool: Under the hood, either launch synchronous param all-gathers or get ready to launch asynchorous all-gathers that get overlapped with the next forward pass. """ - if self._state_offloader is not None: - self._state_offloader.sync_before_step() update_successful = super().step_with_ready_grads() should_sync_params = not self.ddp_config.overlap_param_gather and not getattr( @@ -3216,22 +3280,4 @@ def step_with_ready_grads(self) -> bool: if timers is not None and (self.ddp_config.use_megatron_fsdp or should_sync_params): timers('params-all-gather').stop() - if self._state_offloader is not None: - self._state_offloader.mark_optimizer_states_initialized() - return update_successful - - def offload_states(self): - """Offload states to CPU.""" - if self._state_offloader is not None: - self._state_offloader.offload() - - def reload_offloaded_states(self): - """Start async reload of offloaded states.""" - if self._state_offloader is not None: - self._state_offloader.reload() - - def release_offloaded_gpu_states(self): - """Release GPU memory after D2H completes. For delayed release case.""" - if self._state_offloader is not None: - self._state_offloader.release_gpu_memory() diff --git a/megatron/core/optimizer/layer_wise_optimizer.py b/megatron/core/optimizer/layer_wise_optimizer.py index 27d2bb8e132..97628ddfef0 100644 --- a/megatron/core/optimizer/layer_wise_optimizer.py +++ b/megatron/core/optimizer/layer_wise_optimizer.py @@ -480,6 +480,58 @@ def __init__( # ``_copy_main_params_to_param_buffer``); the gather staging reuses the grad buffer. optimizers[i]._layer_wise_non_distopt_child = True + if ( + config.chunked_optimizer_state_offload + and config.optimizer_state_offload_fraction > 0.0 + ): + from .emerging_optimizers import TensorParallelMuon + + offload_children = [] + state_dtypes = (torch.float32,) + for optimizer in optimizers: + if optimizer.optimizer is None: + continue + # The chunk planner assumes exactly the state schema of TensorParallelMuon: + # one fp32 full-size momentum. Reject subclasses such as AdaptiveMuon rather + # than silently underestimating their tensor-state window. + raw_optimizer = optimizer.optimizer + assert type(raw_optimizer) is TensorParallelMuon, ( + "LayerWise chunked optimizer state offload currently supports exactly " + "TensorParallelMuon with one fp32 full-size momentum tensor; got " + f"{type(raw_optimizer).__name__}. Only Muon-managed parameter groups may " + "be LayerWise optimizer children" + ) + # Whole-parameter sharding can leave this rank with an optimizer object but + # no locally owned Muon parameters. Preserve the no-offloader path for it. + if not any(group["params"] for group in raw_optimizer.param_groups): + continue + master_params = [ + param for group in optimizer.fp32_from_float16_groups for param in group + ] + offload_children.append((optimizer, master_params)) + + assert len(offload_children) <= 1, ( + "LayerWise chunked optimizer state offload assumes at most one non-empty " + f"TensorParallelMuon child; got {len(offload_children)}" + ) + for optimizer, master_params in offload_children: + optimizer.enable_chunked_optimizer_state_offload( + master_params=master_params, state_dtypes=state_dtypes + ) + # The compact FP8 gather reads fp32 main_param directly. With overlap + # enabled, force that gather at the next iteration boundary before + # master D2H so masters are still absent during forward/backward. + offloader = optimizer._optimizer_state_offloader + optimizer.set_optimizer_state_offload_deferred_lifecycle( + state_prefetch_to_step=offloader is not None, + master_offload_for_param_sync=( + offloader is not None + and bool(offloader.selected_params) + and self.use_fp8_param_sync + and self.overlap_param_gather + ), + ) + # shard_params() removed non-owned params from the local optimizer groups, so the # Float16 wrapping above only clears the TE high-precision init copy (a full-size CPU # tensor per fp8 param) for locally owned params. Without this sweep every DP rank @@ -498,6 +550,12 @@ def __init__( super().__init__(optimizers) + self._managed_optimizer_state_offload_indices = tuple( + index + for index, optimizer in enumerate(self.chained_optimizers) + if getattr(optimizer, '_optimizer_state_offloader', None) is not None + ) + # Assign self.model_chunks AFTER super().__init__: ChainedOptimizer.__init__ # resets self.model_chunks to [] and then repopulates only from chained # children that have a model_chunks attribute (DistOpt does, Float16-wrapped @@ -972,7 +1030,34 @@ def count_zeros(self): use_decoupled_grad=self.config.use_precision_aware_optimizer_no_fp8_or_ds_fp8, ) - def start_param_sync_for_bucket_group_subset(self) -> None: + def _managed_optimizer_state_offload_child_indices(self) -> tuple[int, ...]: + """Return child indices that execute through a chunked state offloader.""" + + if ( + not self.config.chunked_optimizer_state_offload + or self.config.optimizer_state_offload_fraction == 0.0 + ): + return () + return self._managed_optimizer_state_offload_indices + + def prefetch_optimizer_state_for_gradient_finalization(self) -> None: + """Prefetch all masters and the managed Muon child's first state chunk.""" + + self.prefetch_optimizer_master_weights_for_step() + managed_indices = self._managed_optimizer_state_offload_child_indices() + if managed_indices: + self.chained_optimizers[managed_indices[0]].prefetch_optimizer_state_for_step() + + def _before_child_step(self, optimizer_idx: int) -> None: + """Ensure the managed Muon child's state is prefetched before child steps.""" + + managed_indices = self._managed_optimizer_state_offload_child_indices() + if not managed_indices: + return + if optimizer_idx == 0: + self.chained_optimizers[managed_indices[0]].prefetch_optimizer_state_for_step() + + def start_param_sync_for_bucket_group_subset(self, force_sync: bool = False) -> None: """Trigger ``start_param_sync`` on LayerWise-managed bucket groups only. Walks each model chunk's dense + expert-parallel bucket groups and @@ -989,7 +1074,7 @@ def start_param_sync_for_bucket_group_subset(self) -> None: if bucket_group.buckets and _bucket_is_managed_by_layer_wise_optimizer( bucket_group.buckets[0] ): - model_chunk._start_bucket_group_param_sync(bucket_group, force_sync=False) + model_chunk._start_bucket_group_param_sync(bucket_group, force_sync=force_sync) @torch.no_grad() def step_with_ready_grads(self) -> bool: diff --git a/megatron/core/optimizer/optimizer.py b/megatron/core/optimizer/optimizer.py index d045f64dcc3..d22f6f9f4a5 100644 --- a/megatron/core/optimizer/optimizer.py +++ b/megatron/core/optimizer/optimizer.py @@ -49,6 +49,7 @@ from ..transformer.module import param_is_not_shared from ..utils import log_single_rank from .clip_grads import clip_grad_by_total_norm_fp32, count_zeros_fp32, get_grad_norm_fp32 +from .cpu_offloading.chunked_optimizer_state_offload import ChunkedOptimizerStateOffloader from .grad_scaler import MegatronGradScaler from .optimizer_config import OptimizerConfig @@ -143,6 +144,12 @@ class MegatronOptimizer(ABC): init_state_fn (Callable, optional): Function to initialize optimizer state. """ + # ChainedOptimizer intentionally does not call this base class's constructor. + # Keep inherited offload helpers safe when they are invoked on that wrapper. + _optimizer_state_offloader: ChunkedOptimizerStateOffloader | None = None + _defer_optimizer_state_prefetch_to_step = False + _defer_optimizer_master_offload_for_param_sync = False + def __init__( self, optimizer: torch.optim.Optimizer, @@ -158,6 +165,127 @@ def __init__( ) self.config = config self.init_state_fn = init_state_fn + self._optimizer_state_offloader: ChunkedOptimizerStateOffloader | None = None + self._defer_optimizer_state_prefetch_to_step = False + self._defer_optimizer_master_offload_for_param_sync = False + + def enable_chunked_optimizer_state_offload( + self, + master_params: List[torch.Tensor], + state_dtypes: Tuple[torch.dtype, ...], + optimizer_owned_master_dtypes: Dict[torch.Tensor, torch.dtype] | None = None, + d2h_stream: torch.cuda.Stream | None = None, + h2d_stream: torch.cuda.Stream | None = None, + ) -> None: + """Enable chunked state execution for this wrapped optimizer. + + Args: + master_params: Separate MCore-owned master parameters to offload. + state_dtypes: Dtypes of full-size tensor states used by the byte planner. + optimizer_owned_master_dtypes: Exact storage dtype for each optimizer-owned + master parameter, including native FP32 parameters when the optimizer creates + a separate ``master_param`` entry for them. + d2h_stream: Optional D2H stream shared by related optimizer wrappers. + h2d_stream: Optional H2D stream shared by related optimizer wrappers. + """ + + if ( + self.optimizer is None + or not self.config.chunked_optimizer_state_offload + or self.config.optimizer_state_offload_fraction == 0.0 + ): + return + self._optimizer_state_offloader = ChunkedOptimizerStateOffloader( + optimizer=self.optimizer, + master_params=master_params, + chunk_size_bytes=self.config.optimizer_state_offload_chunk_size_mb * 1024 * 1024, + offload_fraction=self.config.optimizer_state_offload_fraction, + state_dtypes=state_dtypes, + optimizer_owned_master_dtypes=optimizer_owned_master_dtypes, + d2h_stream=d2h_stream, + h2d_stream=h2d_stream, + ) + + def set_optimizer_state_offload_deferred_lifecycle( + self, *, state_prefetch_to_step: bool, master_offload_for_param_sync: bool + ) -> None: + """Configure deferred work for an optimizer managed by an outer chain. + + Args: + state_prefetch_to_step: Restore tensor state from the outer child-step pipeline + instead of during gradient finalization. + master_offload_for_param_sync: Keep masters resident until the outer optimizer has + staged parameters for an explicit pre-forward synchronization. + """ + + if ( + state_prefetch_to_step or master_offload_for_param_sync + ) and self._optimizer_state_offloader is None: + raise RuntimeError("deferred optimizer-state offload requires an enabled offloader") + self._defer_optimizer_state_prefetch_to_step = state_prefetch_to_step + self._defer_optimizer_master_offload_for_param_sync = master_offload_for_param_sync + + def offload_optimizer_state_for_forward(self, offload_master: bool = True) -> None: + """Begin moving optimizer tensor state and, optionally, master weights to CPU.""" + + if self._optimizer_state_offloader is not None: + self._optimizer_state_offloader.offload_for_forward(offload_master=offload_master) + + def prefetch_optimizer_state_for_step(self) -> None: + """Prefetch master weights and the first optimizer-state chunk.""" + + if self._optimizer_state_offloader is not None: + self._optimizer_state_offloader.prefetch_for_step() + + def prefetch_optimizer_state_for_gradient_finalization(self) -> None: + """Start the optimizer prefetch that may overlap gradient finalization.""" + + self.prefetch_optimizer_state_for_step() + + def prefetch_optimizer_master_weights_for_step(self) -> None: + """Prefetch master weights without restoring an optimizer-state chunk.""" + + if self._optimizer_state_offloader is not None: + self._optimizer_state_offloader.prefetch_master_for_step() + + def ensure_master_weights_for_param_sync(self) -> None: + """Restore all master weights needed by an explicit parameter staging pass.""" + + if self._optimizer_state_offloader is not None: + self._optimizer_state_offloader.ensure_master_for_param_sync() + + def ensure_master_weights_for_pre_forward_param_sync(self) -> None: + """Restore masters only when this optimizer requires pre-forward parameter sync.""" + + if self.optimizer_state_offload_requires_pre_forward_param_sync(): + self.ensure_master_weights_for_param_sync() + + def assert_master_weights_resident(self, operation: str) -> None: + """Assert that an external MCore copy can safely read or write master weights.""" + + if self._optimizer_state_offloader is not None: + self._optimizer_state_offloader.assert_master_weights_resident(operation) + + def synchronize_optimizer_state_for_checkpoint(self) -> None: + """Make CPU canonical optimizer tensors stable for checkpointing.""" + + if self._optimizer_state_offloader is not None: + self._optimizer_state_offloader.synchronize_for_checkpoint() + + def optimizer_state_offload_requires_pre_forward_param_sync(self) -> bool: + """Return whether a parameter gather must consume masters before their D2H copy.""" + + return self._defer_optimizer_master_offload_for_param_sync + + def start_param_sync_for_bucket_group_subset(self, force_sync: bool = False) -> None: + """Synchronize an optimizer-owned DDP bucket subset, if the optimizer has one.""" + + if self.optimizer_state_offload_requires_pre_forward_param_sync(): + raise NotImplementedError( + f"{type(self).__name__} must override " + "start_param_sync_for_bucket_group_subset when it requires pre-forward " + "optimizer parameter synchronization" + ) def get_parameters(self) -> List[torch.nn.Parameter]: """ @@ -680,6 +808,15 @@ def prepare_grads(self) -> bool: """Pre-processing gradients before the optimizer step, returns whether inf/nan is found.""" timers = self.config.timers + # The training loop starts this earlier to overlap H2D with gradient finalization. + # Keep this idempotent fallback so direct optimizer.step() callers restore the master + # parameter device before assigning CUDA gradients to it. + if self._optimizer_state_offloader is not None: + if self._defer_optimizer_state_prefetch_to_step: + self._optimizer_state_offloader.prefetch_master_for_step() + else: + self._optimizer_state_offloader.prefetch_for_step() + # Copy gradients from model params to main params. if timers is not None: timers('optimizer-copy-to-main-grad', log_level=1).start( @@ -721,7 +858,10 @@ def step_with_ready_grads(self) -> bool: barrier=self.config.barrier_with_L1_time ) if not self.is_stub_optimizer: - self.optimizer.step() + if self._optimizer_state_offloader is None: + self.optimizer.step() + else: + self._optimizer_state_offloader.step() if timers is not None: timers('optimizer-inner-step').stop() @@ -946,6 +1086,9 @@ def _copy_model_grads_to_main_grads(self): model_param.grad = model_param.main_grad def _copy_main_params_to_model_params(self): + self.ensure_master_weights_for_param_sync() + self.assert_master_weights_resident("_copy_main_params_to_model_params") + # Non-DistOpt LayerWise fp8: route master->model through bf16 (Q(bf16(master))) to match the # fp8-param-gather-OFF baseline (a direct fp32->fp8 copy would write Q(fp32 master)). This # also covers MoE expert weights at expt_dp==1, which are not gathered. @@ -991,6 +1134,8 @@ def _copy_main_params_to_param_buffer(self): def _copy_model_params_to_main_params(self, state_dict=None): assert state_dict is None, "Initialize main params from state dict is not supported" + self.ensure_master_weights_for_param_sync() + self.assert_master_weights_resident("_copy_model_params_to_main_params") # Only needed for the float16 params. model_data, main_data = self._get_model_and_main_params_data_float16() _multi_tensor_copy_this_to_that( @@ -1016,7 +1161,14 @@ def sharded_state_dict( ): if is_loading: - self.init_state_fn(self.optimizer, self.config) + if self._optimizer_state_offloader is not None: + self._optimizer_state_offloader.initialize_state_for_loading( + self.init_state_fn, self.config + ) + else: + self.init_state_fn(self.optimizer, self.config) + else: + self.synchronize_optimizer_state_for_checkpoint() state_dict = self.state_dict() @@ -1099,7 +1251,12 @@ def load_state_dict(self, state_dict): state_dict[optimizer_key]['param_groups'] = self._filter_and_reorder_param_groups( self.optimizer.param_groups, state_dict[optimizer_key]['param_groups'] ) - self.optimizer.load_state_dict(state_dict[optimizer_key]) + if self._optimizer_state_offloader is None: + self.optimizer.load_state_dict(state_dict[optimizer_key]) + else: + self._optimizer_state_offloader.load_state_dict_without_device_cast( + state_dict[optimizer_key] + ) # Grad scaler. if 'grad_scaler' not in state_dict: @@ -1183,6 +1340,11 @@ def step_with_ready_grads(self) -> bool: """Step the optimizer with ready gradients, return successful.""" if self.is_stub_optimizer: return True + if self._optimizer_state_offloader is not None: + raise RuntimeError( + "FP32Optimizer.step_with_ready_grads does not support chunked optimizer state " + "offload" + ) timers = self.config.timers # Update parameters. @@ -1343,6 +1505,27 @@ def __init__(self, chained_optimizers: List[MegatronOptimizer]): else: self.is_stub_optimizer = True self.chained_optimizers = chained_optimizers + self._share_optimizer_state_offload_streams() + + def _share_optimizer_state_offload_streams(self) -> None: + """Serialize all nested offload managers on one D2H/H2D stream pair.""" + + managers = [] + + def collect(current_optimizer: MegatronOptimizer) -> None: + manager = getattr(current_optimizer, '_optimizer_state_offloader', None) + if manager is not None: + managers.append(manager) + for child in getattr(current_optimizer, 'chained_optimizers', ()): + collect(child) + + collect(self) + if len(managers) < 2: + return + + d2h_stream, h2d_stream = managers[0].transfer_streams + for manager in managers[1:]: + manager.use_transfer_streams(d2h_stream, h2d_stream) @property def optimizer(self): @@ -1382,6 +1565,74 @@ def zero_grad(self, set_to_none=True): for optimizer in self.chained_optimizers: optimizer.zero_grad(set_to_none) + def offload_optimizer_state_for_forward(self, offload_master: bool = True) -> None: + """Begin state and optional master D2H on every optimizer in the chain.""" + + for optimizer in self.chained_optimizers: + optimizer.offload_optimizer_state_for_forward(offload_master=offload_master) + + def prefetch_optimizer_state_for_step(self) -> None: + """Eagerly prefetch state/master data for every optimizer in the chain. + + This general-purpose entry can exceed a configured chunk staging bound because every + child may hold a prefetched window concurrently. LayerWise uses a dedicated + gradient-finalization policy for its single managed Muon child. + """ + + for optimizer in self.chained_optimizers: + optimizer.prefetch_optimizer_state_for_step() + + def prefetch_optimizer_state_for_gradient_finalization(self) -> None: + """Start each child's gradient-finalization prefetch policy.""" + + for optimizer in self.chained_optimizers: + optimizer.prefetch_optimizer_state_for_gradient_finalization() + + def prefetch_optimizer_master_weights_for_step(self) -> None: + """Prefetch only master weights for every optimizer in the chain.""" + + for optimizer in self.chained_optimizers: + optimizer.prefetch_optimizer_master_weights_for_step() + + def ensure_master_weights_for_param_sync(self) -> None: + """Restore all child master weights for a general parameter staging pass.""" + + for optimizer in self.chained_optimizers: + optimizer.ensure_master_weights_for_param_sync() + + def ensure_master_weights_for_pre_forward_param_sync(self) -> None: + """Restore only children that require pre-forward parameter synchronization.""" + + for optimizer in self.chained_optimizers: + optimizer.ensure_master_weights_for_pre_forward_param_sync() + + def assert_master_weights_resident(self, operation: str) -> None: + """Assert that every child can safely access its master weights.""" + + for optimizer in self.chained_optimizers: + optimizer.assert_master_weights_resident(operation) + + def synchronize_optimizer_state_for_checkpoint(self) -> None: + """Make every child optimizer's CPU canonical state checkpoint-safe.""" + + for optimizer in self.chained_optimizers: + optimizer.synchronize_optimizer_state_for_checkpoint() + + def optimizer_state_offload_requires_pre_forward_param_sync(self) -> bool: + """Return whether any child must gather parameters before master offload.""" + + return any( + optimizer.optimizer_state_offload_requires_pre_forward_param_sync() + for optimizer in self.chained_optimizers + ) + + def start_param_sync_for_bucket_group_subset(self, force_sync: bool = False) -> None: + """Synchronize only child subsets that require masters before offload.""" + + for optimizer in self.chained_optimizers: + if optimizer.optimizer_state_offload_requires_pre_forward_param_sync(): + optimizer.start_param_sync_for_bucket_group_subset(force_sync=force_sync) + def get_loss_scale(self): if self.chained_optimizers: return self.chained_optimizers[0].get_loss_scale() @@ -1537,10 +1788,14 @@ def prepare_grads(self) -> bool: return found_inf_flag + def _before_child_step(self, optimizer_idx: int) -> None: + """Hook for subclasses to enqueue work immediately before a child optimizer step.""" + def _step(self) -> bool: """Step all optimizers in this chain.""" success = True for optimizer_idx, optimizer in enumerate(self.chained_optimizers): + self._before_child_step(optimizer_idx) success &= optimizer.step_with_ready_grads() if self.config.overlap_param_gather_with_optimizer_step and optimizer_idx == 0: assert success diff --git a/megatron/core/optimizer/optimizer_config.py b/megatron/core/optimizer/optimizer_config.py index 1048220953a..ecb47110015 100644 --- a/megatron/core/optimizer/optimizer_config.py +++ b/megatron/core/optimizer/optimizer_config.py @@ -1,6 +1,7 @@ # Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. import fnmatch +import warnings from dataclasses import dataclass, field from typing import Callable, Optional, Tuple, Union @@ -368,12 +369,22 @@ class OptimizerConfig: pin_cpu_params: bool = True """If True, pin the optimizer parameters to CPU memory.""" - offload_optimizer_states: bool = False - """ - If True, offload optimizer states to CPU after each optimizer step and - reload them before the next optimizer step. + chunked_optimizer_state_offload: bool = False + """Keep selected optimizer tensor states and master weights on CPU between updates.""" + + optimizer_state_offload_chunk_size_mb: int = 0 + """Target GPU staging-window size for optimizer tensor state, in MiB. + + Zero restores all selected tensor state together. Master weights are restored in one full + window and are not bounded by this value. """ + optimizer_state_offload_fraction: float = 1.0 + """Fraction of optimizer parameter bundles whose state and master weights are offloaded.""" + + offload_optimizer_states: bool = False + """Deprecated alias for :attr:`chunked_optimizer_state_offload`.""" + ################ # Miscellaneous ################ @@ -404,6 +415,77 @@ class OptimizerConfig: def __post_init__(self): """Check the validity of the config.""" + used_deprecated_optimizer_state_offload = self.offload_optimizer_states + if self.offload_optimizer_states: + warnings.warn( + "offload_optimizer_states is deprecated; use " + "chunked_optimizer_state_offload instead. The deprecated option now enables " + "the replacement with its legacy-equivalent defaults: chunk size 0 and " + "offload fraction 1.0 unless those settings were explicitly changed.", + FutureWarning, + stacklevel=2, + ) + self.chunked_optimizer_state_offload = True + self.offload_optimizer_states = False + + assert ( + self.optimizer_state_offload_chunk_size_mb >= 0 + ), "optimizer_state_offload_chunk_size_mb must be non-negative" + assert ( + 0.0 <= self.optimizer_state_offload_fraction <= 1.0 + ), "optimizer_state_offload_fraction must be in [0.0, 1.0]" + + if self.chunked_optimizer_state_offload and self.optimizer_state_offload_fraction > 0.0: + if used_deprecated_optimizer_state_offload: + assert self.optimizer == 'adam', ( + "the deprecated --offload-optimizer-states alias currently supports only " + "Adam; use --chunked-optimizer-state-offload for supported Muon offload" + ) + else: + assert self.optimizer in ( + 'adam', + 'muon', + ), "chunked optimizer state offload currently supports Adam and Muon" + + assert ( + not self.optimizer_cpu_offload + ), "chunked_optimizer_state_offload and optimizer_cpu_offload are mutually exclusive" + assert ( + not self.optimizer_cuda_graph + ), "chunked optimizer state offload does not support optimizer CUDA graphs" + + if self.use_precision_aware_optimizer: + assert self.main_params_dtype == torch.float32, ( + "chunked optimizer state offload with precision-aware optimizer requires " + "main_params_dtype=torch.float32" + ) + supported_state_dtypes = (torch.float32, torch.bfloat16) + assert self.exp_avg_dtype in supported_state_dtypes, ( + "chunked optimizer state offload with precision-aware optimizer requires " + "exp_avg_dtype to be torch.float32 or torch.bfloat16" + ) + assert self.exp_avg_sq_dtype in supported_state_dtypes, ( + "chunked optimizer state offload with precision-aware optimizer requires " + "exp_avg_sq_dtype to be torch.float32 or torch.bfloat16" + ) + + if self.optimizer == 'muon': + assert ( + self.use_layer_wise_distributed_optimizer + ), "Muon optimizer state offload requires LayerWiseDistributedOptimizer" + assert self.bf16, "Muon optimizer state offload currently requires bf16" + assert not self.use_layer_wise_param_layout, ( + "Muon optimizer state offload requires --use-layer-wise-param-layout to be " + "disabled so the decoupled compact LayerWise DDP layout is used" + ) + else: + assert not self.use_layer_wise_distributed_optimizer, ( + "Adam optimizer state offload does not support " "LayerWiseDistributedOptimizer" + ) + assert ( + self.use_distributed_optimizer + ), "Adam optimizer state offload requires use_distributed_optimizer" + # The following condition is used to avoid repetition in distrib_optimizer.py. # This is because in distrib_optimizer.py, the process to handle parameters are # different for different training precision settings. FP8 cases require different @@ -421,8 +503,6 @@ def __post_init__(self): if self.fp8_recipe == "mxfp8": if not self.reuse_grad_buf_for_mxfp8_param_ag: - import warnings - warnings.warn( "mxfp8 without using reuse_grad_buf_for_mxfp8_param_ag and fp8_param_gather" "will use significant amount additional GPU memory." diff --git a/megatron/training/arguments.py b/megatron/training/arguments.py index a168381eaa7..98add4da0e4 100644 --- a/megatron/training/arguments.py +++ b/megatron/training/arguments.py @@ -1859,6 +1859,17 @@ def validate_args(args, defaults={}): args.use_layer_wise_distributed_optimizer = True args.use_distributed_optimizer = False + if ( + args.optimizer == 'muon' + and (args.chunked_optimizer_state_offload or args.offload_optimizer_states) + and args.optimizer_state_offload_fraction > 0.0 + ): + assert args.use_layer_wise_distributed_optimizer, ( + "Muon optimizer state offload requires the LayerWise distributed optimizer " + "(--use-distributed-optimizer) so non-Muon parameter groups are routed through " + "a sibling DistributedOptimizer" + ) + assert not args.use_torch_fsdp2, "Emerging optimizer does not support Torch-FSDP2 for now." assert ( not args.use_megatron_fsdp @@ -2025,16 +2036,44 @@ def validate_args(args, defaults={}): "must be used in conjunction with `--fp8-recipe delayed`." ) - if args.offload_optimizer_states: - assert ( - args.use_distributed_optimizer - ), "offload_optimizer_states is only supported with distributed optimizer" - assert ( - args.optimizer == 'adam' - ), "offload_optimizer_states is only supported with adam optimizer" + # Optimizer configuration normalizes the deprecated spelling. Include it here so + # cross-config validation still runs before that configuration is constructed. + if ( + args.chunked_optimizer_state_offload or args.offload_optimizer_states + ) and args.optimizer_state_offload_fraction > 0.0: + if args.optimizer_state_offload_chunk_size_mb == 0: + warn_rank_0( + "optimizer state offload is enabled with chunk size 0, so all " + "selected optimizer tensor state is temporarily restored together. Set " + "--optimizer-state-offload-chunk-size-mb to a positive value to bound the " + "tensor-state GPU window. Selected master weights always use one full " + "restore window regardless of chunk size." + ) assert ( not args.use_megatron_fsdp - ), "offload_optimizer_states does not support Megatron-FSDP for now." + ), "chunked optimizer state offload does not support Megatron-FSDP" + if not args.no_save_optim or not args.no_load_optim: + assert args.ckpt_format == 'torch_dist', ( + "chunked optimizer state offload requires --ckpt-format torch_dist when " + "optimizer state is saved or loaded because its sharded load hooks preserve " + "CPU canonical storage" + ) + if not args.no_save_optim: + assert not args.async_save, ( + "chunked optimizer state offload does not support --async-save when optimizer " + "state is saved because the background writer can retain tensors backed by " + "reusable pinned CPU buffers" + ) + assert ( + args.cuda_graph_impl != 'full_iteration' + ), "chunked optimizer state offload does not support full-iteration CUDA graphs" + if args.optimizer == 'muon' and args.fp8_param_gather and args.overlap_param_gather: + warn_rank_0( + "Muon compact FP8 parameter gather with chunked optimizer state offload " + "forces LayerWise-owned pre-forward parameter buckets to complete before " + "masters are offloaded; sibling DistributedOptimizer gathers and " + "optimizer-state transfers remain overlapped." + ) if args.non_persistent_ckpt_type == "local": assert ( @@ -3714,7 +3753,8 @@ def _add_training_args(parser): '--optimizer-offload-fraction', type=float, default=1.0, - help='Ratio of optimizer state to offload to CPU', + help='Fraction used by --optimizer-cpu-offload. This is distinct from ' + '--optimizer-state-offload-fraction, and the two offload modes are mutually exclusive.', ) group.add_argument( '--use-torch-optimizer-for-cpu-offload', @@ -3747,15 +3787,43 @@ def _add_training_args(parser): dest='pin_cpu_params', help='Disable pinning of CPU memory for parameters.', ) + group.add_argument( + '--chunked-optimizer-state-offload', + action='store_true', + help='Keep selected optimizer tensor states and master weights in CPU memory between ' + 'updates. GPU tensor-state updates run in bounded chunks, while selected master ' + 'weights use one full restore window.', + ) + group.add_argument( + '--optimizer-state-offload-chunk-size-mb', + type=int, + default=0, + help='Target size of each optimizer tensor-state staging buffer in MiB. Two buffers ' + 'can be live for overlap (roughly 2x this value per active offload manager). Zero ' + 'temporarily restores all selected tensor state for one full GPU update. A single ' + 'atomic parameter state may exceed this soft target and is reported at runtime. This ' + 'value does not bound the selected master-weight window, which is always restored in ' + 'full.', + ) + group.add_argument( + '--optimizer-state-offload-fraction', + type=float, + default=1.0, + help='Fraction of optimizer parameter bundles to offload. Selected bundles always ' + 'include both tensor state and any separate master weight. This is distinct from ' + '--optimizer-offload-fraction, which belongs to --optimizer-cpu-offload; the modes are ' + 'mutually exclusive. Zero disables optimizer state offload even when its flag is set.', + ) group.add_argument( '--offload-optimizer-states', action='store_true', dest='offload_optimizer_states', - help='Offload optimizer states to CPU after each optimizer step and ' - 'reload them before the next optimizer step. ' - 'Only support TE FusedAdam optimizer.' - 'Note that this still uses pure GPU optimizer instead of ' - 'HybridDeviceOptimizer for --optimizer-cpu-offload.', + help='Deprecated spelling for --chunked-optimizer-state-offload. It automatically ' + 'enables the replacement; the legacy-equivalent settings are chunk size 0 and offload ' + 'fraction 1.0 unless new tuning arguments are supplied. ' + 'With a nonzero fraction, optimizer-state checkpoint I/O requires --ckpt-format ' + 'torch_dist, optimizer-state saves must be synchronous, and optimizer or full-iteration ' + 'CUDA graphs are unsupported.', ) group.add_argument( '--dataloader-type', diff --git a/megatron/training/checkpointing.py b/megatron/training/checkpointing.py index bfa3564c6cd..395d5a57eca 100644 --- a/megatron/training/checkpointing.py +++ b/megatron/training/checkpointing.py @@ -604,6 +604,27 @@ def save_checkpoint( start_ckpt = time() args = get_args() + # Keep the safety invariant even when arguments came through a validation path that does not + # perform the flat CLI cross-check, or were mutated after startup. + optimizer_config = getattr(optimizer, "config", None) if optimizer is not None else None + chunked_optimizer_state_offload = bool( + optimizer_config is not None + and getattr(optimizer_config, "chunked_optimizer_state_offload", False) + and getattr(optimizer_config, "optimizer_state_offload_fraction", 1.0) > 0.0 + ) + if chunked_optimizer_state_offload and not args.no_save_optim: + if args.async_save: + raise RuntimeError( + "chunked optimizer state offload does not support --async-save when optimizer " + "state is saved because reusable pinned CPU buffers cannot be handed to a " + "background writer" + ) + if getattr(args, "ckpt_format", None) != "torch_dist": + raise RuntimeError( + "chunked optimizer state offload requires --ckpt-format torch_dist; checkpoint " + "arguments may have overwritten the startup setting during resume" + ) + if args.async_save and not is_empty_async_queue(): print_rank_0( 'WARNING: Starting a checkpoint save before previous has finished. Consider increasing the checkpoint interval.' @@ -699,7 +720,11 @@ def save_checkpoint( optimizer.save_parameter_state(optim_checkpoint_name) # LayerWiseDistributedOptimizer save optimizer state to file on different ranks - if getattr(args, "use_layer_wise_distributed_optimizer", False) and args.ckpt_format == 'torch': + if ( + getattr(args, "use_layer_wise_distributed_optimizer", False) + and not args.no_save_optim + and args.ckpt_format == 'torch' + ): dp_rank = mpu.get_data_parallel_rank() optim_checkpoint_name = os.path.join( os.path.dirname(checkpoint_name), f"layer_wise_optimizer_{dp_rank}.pt" diff --git a/megatron/training/training.py b/megatron/training/training.py index 6f958a8f6d3..5a8aab667d7 100644 --- a/megatron/training/training.py +++ b/megatron/training/training.py @@ -2744,29 +2744,55 @@ def train_step( timers = get_timers() num_microbatches = get_num_microbatches() - offload_optimizer_states = getattr(args, 'offload_optimizer_states', False) - if offload_optimizer_states: - # Reload optimizer states as late as possible so the H2D transfer can overlap - # with gradient finalization. Preserve custom finalize hooks installed by a - # model builder, and avoid wrapping the hook again on every training step. + optimizer_config = getattr(optimizer, 'config', None) + chunked_optimizer_state_offload = bool( + optimizer_config is not None + and getattr(optimizer_config, 'chunked_optimizer_state_offload', False) + and getattr(optimizer_config, 'optimizer_state_offload_fraction', 1.0) > 0.0 + ) + pre_forward_param_sync_before_master_offload = ( + chunked_optimizer_state_offload + and optimizer.optimizer_state_offload_requires_pre_forward_param_sync() + ) + delay_master_offload_for_param_buffer = ( + chunked_optimizer_state_offload + and args.reuse_grad_buf_for_mxfp8_param_ag + and args.overlap_param_gather + ) + if chunked_optimizer_state_offload: + # Prefetch masters and the first state chunk as late as possible so H2D overlaps + # gradient finalization. Preserve custom hooks and avoid wrapping more than once. finalize_model_grads_func = getattr(config, 'finalize_model_grads_func', None) if ( - getattr(finalize_model_grads_func, '_optimizer_state_offload_wrapped_optimizer', None) + getattr( + finalize_model_grads_func, + '_chunked_optimizer_state_offload_wrapped_optimizer', + None, + ) is not optimizer ): - base_finalize_model_grads_func = finalize_model_grads_func or finalize_model_grads + base_finalize_model_grads_func = getattr( + finalize_model_grads_func, + '_chunked_optimizer_state_offload_base_finalize_model_grads_func', + None, + ) + if base_finalize_model_grads_func is None: + base_finalize_model_grads_func = finalize_model_grads_func or finalize_model_grads def finalize_model_grads_with_state_reload(*fmg_args, **fmg_kwargs): - for optim_instance in optimizer.chained_optimizers: - if isinstance(optim_instance, DistributedOptimizer): - optim_instance.reload_offloaded_states() + optimizer.prefetch_optimizer_state_for_gradient_finalization() return base_finalize_model_grads_func(*fmg_args, **fmg_kwargs) setattr( finalize_model_grads_with_state_reload, - '_optimizer_state_offload_wrapped_optimizer', + '_chunked_optimizer_state_offload_wrapped_optimizer', optimizer, ) + setattr( + finalize_model_grads_with_state_reload, + '_chunked_optimizer_state_offload_base_finalize_model_grads_func', + base_finalize_model_grads_func, + ) config.finalize_model_grads_func = finalize_model_grads_with_state_reload rerun_state_machine = get_rerun_state_machine() @@ -2791,11 +2817,12 @@ def finalize_model_grads_with_state_reload(*fmg_args, **fmg_kwargs): args.save_dgrads_interval is not None and (iteration + 1) % args.save_dgrads_interval == 0 ) while rerun_state_machine.should_run_forward_backward(rerun_data_iterator): - # Start the D2H transfer before zeroing gradients to maximize overlap. - if offload_optimizer_states: - for optim_instance in optimizer.chained_optimizers: - if isinstance(optim_instance, DistributedOptimizer): - optim_instance.offload_states() + # Start D2H before zeroing gradients. In the MXFP8 staging path, updated masters + # remain readable until their delayed D2H after the main-param copy below. + if chunked_optimizer_state_offload and not pre_forward_param_sync_before_master_offload: + optimizer.offload_optimizer_state_for_forward( + offload_master=not delay_master_offload_for_param_buffer + ) # Set grad to zero. for model_chunk in model: @@ -2804,6 +2831,16 @@ def finalize_model_grads_with_state_reload(*fmg_args, **fmg_kwargs): model_chunk.force_all_reduce = save_wgrads_in_this_iteration optimizer.zero_grad() + if pre_forward_param_sync_before_master_offload: + # Compact LayerWise FP8 gather consumes fp32 masters. Finish the normally + # overlapped gather after grad-buffer reset. If a sibling DistOpt still needs + # MXFP8 staging, keep masters until that copy below completes. + optimizer.ensure_master_weights_for_pre_forward_param_sync() + optimizer.start_param_sync_for_bucket_group_subset(force_sync=True) + optimizer.offload_optimizer_state_for_forward( + offload_master=not delay_master_offload_for_param_buffer + ) + if has_nvidia_modelopt and getattr(args, "modelopt_enabled", False): # Distillation shape-adjust reads parallel_state; only for modelopt-enabled runs. adjust_tensor_shapes_fn = get_tensor_shapes_adjust_fn_for_distillation( @@ -2836,14 +2873,15 @@ def finalize_model_grads_with_state_reload(*fmg_args, **fmg_kwargs): if forward_pre_hook_enabled or full_cg_captured: for optim_instance in optimizer.chained_optimizers: if isinstance(optim_instance, DistributedOptimizer): + # Only this DistOpt sibling consumes masters in the MXFP8 param-buffer + # staging pass. The staging entry restores its own offloaded masters; + # do not perturb the LayerWise/Muon master lifecycle here. optim_instance._copy_main_params_to_param_buffer() - # Master weights must remain resident until any main-param copy above is - # complete. Releasing here keeps optimizer memory out of forward/backward. - if offload_optimizer_states: - for optim_instance in optimizer.chained_optimizers: - if isinstance(optim_instance, DistributedOptimizer): - optim_instance.release_offloaded_gpu_states() + # In the delayed MXFP8 path, optimizer state was offloaded above while masters + # remained resident for main-param staging. Start master D2H after that copy. + if delay_master_offload_for_param_buffer: + optimizer.offload_optimizer_state_for_forward() # Forward pass. if save_activations_in_this_iteration: diff --git a/tests/unit_tests/dist_checkpointing/test_layer_wise_optimizer.py b/tests/unit_tests/dist_checkpointing/test_layer_wise_optimizer.py index 42ef0a401ee..f728a080b16 100644 --- a/tests/unit_tests/dist_checkpointing/test_layer_wise_optimizer.py +++ b/tests/unit_tests/dist_checkpointing/test_layer_wise_optimizer.py @@ -130,6 +130,108 @@ def setup_method(self, method): def teardown_method(self, method): Utils.destroy_model_parallel() + @pytest.mark.parametrize( + ("source_offload", "destination_offload"), [(False, True), (True, False)] + ) + def test_chunked_offload_distributed_checkpoint_compatibility( + self, tmp_path_dist_ckpt, source_offload, destination_offload + ): + """Compact Muon checkpoints remain compatible across the offload switch.""" + + Utils.initialize_model_parallel(1, 1) + + def clone_to_cpu(value): + if isinstance(value, torch.Tensor): + return value.detach().cpu().clone() + if isinstance(value, dict): + return {key: clone_to_cpu(item) for key, item in value.items()} + if isinstance(value, list): + return [clone_to_cpu(item) for item in value] + if isinstance(value, tuple): + return tuple(clone_to_cpu(item) for item in value) + return deepcopy(value) + + def snapshot_optimizer(optimizer): + children = getattr(optimizer, 'chained_optimizers', ()) + if children: + return [snapshot_optimizer(child) for child in children] + + snapshot = {'common': clone_to_cpu(optimizer.state_dict())} + if hasattr(optimizer, 'model_param_group_index_map'): + parameter_state = {} + for model_param in optimizer.model_param_group_index_map: + tensors = optimizer._get_main_param_and_optimizer_states(model_param) + parameter_state[optimizer._param_name(model_param)] = { + key: tensor.detach().cpu().clone() + for key, tensor in tensors.items() + if key != 'step' + } + snapshot['parameter_state'] = parameter_state + return snapshot + + with ( + TempNamedDir(tmp_path_dist_ckpt / 'layerwise_offload_A', sync=True) as ckpt_dir_A, + TempNamedDir(tmp_path_dist_ckpt / 'layerwise_offload_B', sync=True) as ckpt_dir_B, + ): + metadata = {'distrib_optim_sharding_type': 'dp_reshardable'} + model_A, optimizer_A = setup_model_and_optimizer( + seed=2, + tp=1, + pp=1, + bf16=True, + dist_opt=True, + optimizer='dist_muon', + use_param_layout=True, + chunked_optimizer_state_offload=source_offload, + optimizer_state_offload_chunk_size_mb=1, + ) + model_sharded_sd_A = model_A[0].sharded_state_dict() + save(optimizer_A.sharded_state_dict(model_sharded_sd_A, metadata=metadata), ckpt_dir_A) + source_state = snapshot_optimizer(optimizer_A) + + model_B, optimizer_B = setup_model_and_optimizer( + seed=3, + tp=1, + pp=1, + bf16=True, + dist_opt=True, + optimizer='dist_muon', + use_param_layout=True, + chunked_optimizer_state_offload=destination_offload, + optimizer_state_offload_chunk_size_mb=1, + initialize_optimizer_state=False, + ) + model_sharded_sd_B = model_B[0].sharded_state_dict() + load_template = optimizer_B.sharded_state_dict( + model_sharded_sd_B, is_loading=True, metadata=metadata + ) + optimizer_B.load_state_dict(load(load_template, ckpt_dir_A)) + check_equal(source_state, snapshot_optimizer(optimizer_B)) + + if destination_offload: + + def collect_managers(optimizer): + managers = [] + manager = getattr(optimizer, '_optimizer_state_offloader', None) + if manager is not None: + managers.append(manager) + for child in getattr(optimizer, 'chained_optimizers', ()): + managers.extend(collect_managers(child)) + return managers + + managers = collect_managers(optimizer_B) + # Compact Muon state and its sibling DistributedOptimizer state must both + # preserve CPU canonical storage after the distributed load. + assert len(managers) >= 2 + for manager in managers: + for param in manager.selected_params: + assert param.device.type == 'cpu' + for key, value in manager.optimizer.state[param].items(): + if manager._is_offloadable_state(param, key, value): + assert value.device.type == 'cpu' + + save(optimizer_B.sharded_state_dict(model_sharded_sd_B, metadata=metadata), ckpt_dir_B) + def test_parameter_sharding(self): """Test that parameters are correctly sharded across DP ranks.""" Utils.initialize_model_parallel(1, 1) diff --git a/tests/unit_tests/dist_checkpointing/test_optimizer.py b/tests/unit_tests/dist_checkpointing/test_optimizer.py index 3849e044545..cfe46aafad6 100644 --- a/tests/unit_tests/dist_checkpointing/test_optimizer.py +++ b/tests/unit_tests/dist_checkpointing/test_optimizer.py @@ -29,7 +29,7 @@ from megatron.core.tensor_parallel import model_parallel_cuda_manual_seed from megatron.core.transformer import MLATransformerConfig, TransformerConfig from megatron.core.transformer.mlp import apply_swiglu_sharded_factory -from megatron.core.utils import is_torch_min_version +from megatron.core.utils import is_te_min_version, is_torch_min_version from megatron.training.arguments import parse_args from megatron.training.checkpointing import load_checkpoint, save_checkpoint from tests.unit_tests.dist_checkpointing import ( @@ -606,6 +606,100 @@ def setup_method(self, method): def teardown_method(self, method): Utils.destroy_model_parallel() + @pytest.mark.parametrize( + ("source_offload", "destination_offload"), [(False, True), (True, False)] + ) + @pytest.mark.parametrize( + "use_precision_aware_optimizer", + [ + False, + pytest.param( + True, + marks=pytest.mark.skipif( + not is_te_min_version("2.1.0.dev0"), + reason="Precision-aware FusedAdam requires TE 2.1.0.dev0", + ), + ), + ], + ) + def test_chunked_offload_distributed_checkpoint_compatibility( + self, tmp_path_dist_ckpt, source_offload, destination_offload, use_precision_aware_optimizer + ): + """Offloaded and legacy Adam DistOpt checkpoints load in both directions.""" + + Utils.initialize_model_parallel(1, 1) + with ( + TempNamedDir(tmp_path_dist_ckpt / 'chunked_offload_A', sync=True) as ckpt_dir_A, + TempNamedDir(tmp_path_dist_ckpt / 'chunked_offload_B', sync=True) as ckpt_dir_B, + ): + metadata = {'distrib_optim_sharding_type': 'dp_reshardable'} + model_A, optimizer_A = setup_model_and_optimizer( + seed=2, + tp=1, + pp=1, + chunked_optimizer_state_offload=source_offload, + optimizer_state_offload_chunk_size_mb=1, + use_precision_aware_optimizer=use_precision_aware_optimizer, + ) + model_sharded_sd_A = model_A[0].sharded_state_dict() + save(optimizer_A.sharded_state_dict(model_sharded_sd_A, metadata=metadata), ckpt_dir_A) + + dist_optimizer_A = optimizer_A.chained_optimizers[0] + source_states = {} + for model_param in dist_optimizer_A.model_param_group_index_map: + tensors = dist_optimizer_A._get_main_param_and_optimizer_states(model_param) + source_states[dist_optimizer_A._param_name(model_param)] = { + key: tensor.detach().cpu().clone() + for key, tensor in tensors.items() + if key != 'step' + } + + model_B, optimizer_B = setup_model_and_optimizer( + seed=3, + tp=1, + pp=1, + chunked_optimizer_state_offload=destination_offload, + optimizer_state_offload_chunk_size_mb=1, + use_precision_aware_optimizer=use_precision_aware_optimizer, + initialize_optimizer_state=False, + ) + model_sharded_sd_B = model_B[0].sharded_state_dict() + load_template = optimizer_B.sharded_state_dict( + model_sharded_sd_B, is_loading=True, metadata=metadata + ) + optimizer_B.load_state_dict(load(load_template, ckpt_dir_A)) + + dist_optimizer_B = optimizer_B.chained_optimizers[0] + destination_param_names = { + dist_optimizer_B._param_name(param) + for param in dist_optimizer_B.model_param_group_index_map + } + assert destination_param_names == source_states.keys() + for model_param in dist_optimizer_B.model_param_group_index_map: + param_name = dist_optimizer_B._param_name(model_param) + restored = dist_optimizer_B._get_main_param_and_optimizer_states(model_param) + for key, expected in source_states[param_name].items(): + assert torch.equal(restored[key].detach().cpu(), expected), ( + f"optimizer state '{key}' of parameter '{param_name}' " + "was not restored from the checkpoint" + ) + + if destination_offload: + manager = dist_optimizer_B._optimizer_state_offloader + assert manager is not None + for param in manager.selected_params: + for value in manager.optimizer.state[param].values(): + if isinstance(value, torch.Tensor) and value.numel() == param.numel(): + assert value.device.type == 'cpu' + if manager._master_in_optimizer_state: + master = manager.optimizer.state[param].get('master_param') + assert isinstance(master, torch.Tensor) + assert master.device.type == 'cpu' + elif id(param) in manager._explicit_master_param_ids: + assert param.device.type == 'cpu' + + save(optimizer_B.sharded_state_dict(model_sharded_sd_B, metadata=metadata), ckpt_dir_B) + @pytest.mark.parametrize("fully_parallel", [False, True]) @pytest.mark.parametrize( ("tp_pp_ep", "is_moe", "is_mla", "test_step", "kwargs"), diff --git a/tests/unit_tests/dist_checkpointing/utils.py b/tests/unit_tests/dist_checkpointing/utils.py index 81851f7f80e..142ee5f0928 100644 --- a/tests/unit_tests/dist_checkpointing/utils.py +++ b/tests/unit_tests/dist_checkpointing/utils.py @@ -185,6 +185,11 @@ def setup_model_and_optimizer( dist_opt=True, optimizer='adam', use_param_layout=False, + chunked_optimizer_state_offload=False, + optimizer_state_offload_chunk_size_mb=0, + optimizer_state_offload_fraction=1.0, + use_precision_aware_optimizer=False, + initialize_optimizer_state=True, ): optimizer_type = optimizer use_layer_wise = False @@ -226,6 +231,10 @@ def setup_model_and_optimizer( use_distributed_optimizer=ddp_use_dist_opt, use_layer_wise_distributed_optimizer=use_layer_wise, optimizer=optimizer, + chunked_optimizer_state_offload=chunked_optimizer_state_offload, + optimizer_state_offload_chunk_size_mb=optimizer_state_offload_chunk_size_mb, + optimizer_state_offload_fraction=optimizer_state_offload_fraction, + use_precision_aware_optimizer=use_precision_aware_optimizer, ) if optimizer_type in ('muon', 'dist_muon'): @@ -250,18 +259,21 @@ def _init_states(optimizer): if not hasattr(optimizer, 'optimizer'): optimizer.init_state_fn(optimizer) else: - optimizer.init_state_fn(optimizer.optimizer) + optimizer.init_state_fn(optimizer.optimizer, optimizer.config) - if isinstance(optimizer, ChainedOptimizer): - _init_states(optimizer) - else: - for group in optimizer.optimizer.param_groups: - for p in group['params']: - if len(optimizer.optimizer.state[p]) == 0: - optimizer.optimizer.state[p]['exp_avg'] = torch.rand_like(p.data) - optimizer.optimizer.state[p]['exp_avg_sq'] = torch.rand_like(p.data) + if initialize_optimizer_state: + if isinstance(optimizer, ChainedOptimizer): + _init_states(optimizer) + else: + for group in optimizer.optimizer.param_groups: + for p in group['params']: + if len(optimizer.optimizer.state[p]) == 0: + optimizer.optimizer.state[p]['exp_avg'] = torch.rand_like(p.data) + optimizer.optimizer.state[p]['exp_avg_sq'] = torch.rand_like(p.data) optimizer.reload_model_params() + if chunked_optimizer_state_offload: + optimizer.offload_optimizer_state_for_forward() CachedMetadataFileSystemReader.clear_metadata_cache() return unwrap_model(model), optimizer diff --git a/tests/unit_tests/test_checkpointing.py b/tests/unit_tests/test_checkpointing.py index 9b62c26d674..5e35d3bda9f 100644 --- a/tests/unit_tests/test_checkpointing.py +++ b/tests/unit_tests/test_checkpointing.py @@ -123,6 +123,87 @@ def create_args(): yield args +def test_async_save_rejects_reusable_chunked_optimizer_buffers(create_args): + """Checkpoint entry must preserve the CLI guard when startup validation was bypassed.""" + + create_args.async_save = True + create_args.ckpt_format = 'torch_dist' + optimizer = SimpleNamespace(config=SimpleNamespace(chunked_optimizer_state_offload=True)) + with ( + mock.patch("megatron.training.checkpointing.get_args", return_value=create_args), + pytest.raises(RuntimeError, match="does not support --async-save"), + ): + save_checkpoint(0, [], optimizer, None, 0) + + +def test_checkpoint_save_rejects_mutated_format_with_chunked_optimizer_offload(create_args): + """Checkpoint entry revalidates the format after resume-time argument restoration.""" + + create_args.ckpt_format = 'torch' + optimizer = SimpleNamespace(config=SimpleNamespace(chunked_optimizer_state_offload=True)) + with ( + mock.patch("megatron.training.checkpointing.get_args", return_value=create_args), + pytest.raises(RuntimeError, match="requires --ckpt-format torch_dist"), + ): + save_checkpoint(0, [], optimizer, None, 0) + + +def test_checkpoint_save_treats_zero_offload_fraction_as_disabled(create_args): + """A zero fraction must not impose async-save or distributed-format restrictions.""" + + create_args.async_save = True + create_args.ckpt_format = 'torch' + optimizer = SimpleNamespace( + config=SimpleNamespace( + chunked_optimizer_state_offload=True, optimizer_state_offload_fraction=0.0 + ) + ) + with ( + mock.patch("megatron.training.checkpointing.get_args", return_value=create_args), + mock.patch( + "megatron.training.checkpointing.is_empty_async_queue", + side_effect=RuntimeError("continued past optimizer offload guard"), + ), + pytest.raises(RuntimeError, match="continued past optimizer offload guard"), + ): + save_checkpoint(0, [], optimizer, None, 0) + + +def test_checkpoint_save_skips_offload_guards_without_optimizer_state(create_args): + """Model-only saves do not expose reusable optimizer buffers to the async writer.""" + + create_args.async_save = True + create_args.ckpt_format = 'torch_dist' + create_args.no_save_optim = True + optimizer = SimpleNamespace(config=SimpleNamespace(chunked_optimizer_state_offload=True)) + with ( + mock.patch("megatron.training.checkpointing.get_args", return_value=create_args), + mock.patch( + "megatron.training.checkpointing.is_empty_async_queue", + side_effect=RuntimeError("continued past optimizer offload guard"), + ), + pytest.raises(RuntimeError, match="continued past optimizer offload guard"), + ): + save_checkpoint(0, [], optimizer, None, 0) + + +def test_checkpoint_save_allows_legacy_format_without_optimizer_state(create_args): + """A model-only save is not forced through optimizer-specific distributed hooks.""" + + create_args.ckpt_format = 'torch' + create_args.no_save_optim = True + optimizer = SimpleNamespace(config=SimpleNamespace(chunked_optimizer_state_offload=True)) + with ( + mock.patch("megatron.training.checkpointing.get_args", return_value=create_args), + mock.patch( + "megatron.training.checkpointing.on_save_checkpoint_start", + side_effect=RuntimeError("continued past optimizer offload guard"), + ), + pytest.raises(RuntimeError, match="continued past optimizer offload guard"), + ): + save_checkpoint(0, [], optimizer, None, 0) + + @pytest.fixture def create_ckpt_load_args(create_args): """Setup dummy args allowing checkpoint load.""" diff --git a/tests/unit_tests/test_fp8_param.py b/tests/unit_tests/test_fp8_param.py index 69265906a4c..6eb69ba5932 100644 --- a/tests/unit_tests/test_fp8_param.py +++ b/tests/unit_tests/test_fp8_param.py @@ -89,8 +89,8 @@ def model_provider( layer_spec_fn=get_gpt_layer_with_transformer_engine_spec, **config_kwargs, ): - model_parallel_cuda_manual_seed(_SEED) args = get_args() + model_parallel_cuda_manual_seed(_SEED, te_rng_tracker=args.te_rng_tracker) config = core_transformer_config_from_args(args) transformer_layer_spec = layer_spec_fn( num_experts=args.num_experts, moe_grouped_gemm=args.moe_grouped_gemm @@ -220,6 +220,7 @@ def _run_test_helper( fp8_param_gather: bool = True, use_cuda_graph: bool = False, eval_transition: bool = False, + num_steps: int = 100, **kwargs, ): """Test fp8_param with a small GPT model.""" @@ -249,7 +250,9 @@ def _run_test_helper( input_ids, labels, position_ids, attention_mask, loss_mask = self.get_batch( self.seq_length, self.micro_batch_size ) - model_parallel_cuda_manual_seed(_SEED) + model_parallel_cuda_manual_seed( + _SEED, te_rng_tracker=args.te_rng_tracker, force_reset_rng=True + ) cfg_container = Utils.pretrain_config_from_global_args(args, "gpt") pg_collection = ProcessGroupCollection.use_mpu_process_groups() if inference: @@ -331,9 +334,19 @@ def _run_test_helper( loss_list = [] eval_loss_list = [] + delay_master_offload = ( + not inference + and args.chunked_optimizer_state_offload + and args.reuse_grad_buf_for_mxfp8_param_ag + and args.overlap_param_gather + ) - for i in range(100): + for i in range(num_steps): if not inference: + if args.chunked_optimizer_state_offload: + optimizer.offload_optimizer_state_for_forward( + offload_master=not delay_master_offload + ) gpt_model[0].zero_grad_buffer() optimizer.zero_grad() @@ -352,8 +365,12 @@ def _run_test_helper( # we need to call the _copy_main_params_to_param_buffer() after the grad buffer # is zeroed by zero_grad_buffer() because param and grad buffer are shared. if args.reuse_grad_buf_for_mxfp8_param_ag and args.overlap_param_gather: + # Each DistOpt staging entry restores its own offloaded master window. self.copy_main_params_to_param_buffer(gpt_model, optimizer) + if delay_master_offload: + optimizer.offload_optimizer_state_for_forward() + gpt_model[0].set_is_first_microbatch() output = gpt_model[0].forward( input_ids=input_ids, @@ -380,6 +397,8 @@ def _run_test_helper( for name, param in gpt_model[0].named_parameters(): assert param.main_grad is not None + if args.chunked_optimizer_state_offload: + optimizer.prefetch_optimizer_state_for_gradient_finalization() update_successful, _, _ = optimizer.step() assert update_successful @@ -395,6 +414,29 @@ def _run_test_helper( ) ) + if not inference and args.chunked_optimizer_state_offload: + optimizer.offload_optimizer_state_for_forward() + managers = [ + child._optimizer_state_offloader + for child in optimizer.chained_optimizers + if child._optimizer_state_offloader is not None + ] + assert managers + for manager in managers: + for param in manager.selected_params: + for value in manager.optimizer.state[param].values(): + if isinstance(value, torch.Tensor) and value.numel() == param.numel(): + assert value.device.type == 'cpu' + assert any( + param.device.type == 'cpu' + or ( + isinstance(manager.optimizer.state[param].get('master_param'), torch.Tensor) + and manager.optimizer.state[param]['master_param'].device.type == 'cpu' + ) + for manager in managers + for param in manager.selected_params + ) + if self.cuda_graph_helper is not None and self.cuda_graph_helper.graphs_created(): self.cuda_graph_helper.delete_cuda_graphs() self.cuda_graph_helper = None @@ -438,6 +480,61 @@ def run_test_with_eval_transition(self, tp_size, recipe, **kwargs): torch.testing.assert_close(loss, loss_ref, atol=1e-4, rtol=1e-4) torch.testing.assert_close(eval_loss, eval_loss_ref, atol=1e-4, rtol=1e-4) + @pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8) + @pytest.mark.skipif(not is_te_min_version("2.3.0.dev0"), reason="TE 2.3.0.dev0 is required") + @pytest.mark.parametrize("recipe", ["blockwise", "mxfp8"]) + def test_chunked_optimizer_state_offload(self, recipe): + """Adam DistOpt keeps FP8-gather updates equivalent while state/master are offloaded.""" + + arch = get_device_arch_version() + if recipe == "blockwise" and arch != 9: + pytest.skip("blockwise FP8 is Hopper-only") + if recipe == "mxfp8" and arch < 10: + pytest.skip("mxfp8 requires Blackwell architecture or newer") + + common = dict(overlap_param_gather=True, overlap_grad_reduce=True, num_steps=3) + offloaded = self._run_test_helper( + 1, + recipe, + fp8_param_gather=True, + chunked_optimizer_state_offload=True, + optimizer_state_offload_chunk_size_mb=1, + optimizer_state_offload_fraction=1.0, + ckpt_format='torch_dist', + **common, + ) + reference = self._run_test_helper(1, recipe, fp8_param_gather=True, **common) + torch.testing.assert_close(offloaded, reference, atol=1e-4, rtol=1e-4) + + @pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8) + @pytest.mark.skipif(not is_te_min_version("2.3.0.dev0"), reason="TE 2.3.0.dev0 is required") + def test_precision_aware_mxfp8_chunked_optimizer_state_offload(self): + """Precision-aware FusedAdam masters and moments remain CPU canonical with MXFP8.""" + + if get_device_arch_version() < 10: + pytest.skip("mxfp8 requires Blackwell architecture or newer") + + common = dict( + overlap_param_gather=True, + overlap_grad_reduce=True, + num_steps=3, + use_precision_aware_optimizer=True, + exp_avg_dtype='bf16', + exp_avg_sq_dtype='bf16', + ) + offloaded = self._run_test_helper( + 1, + "mxfp8", + fp8_param_gather=True, + chunked_optimizer_state_offload=True, + optimizer_state_offload_chunk_size_mb=1, + optimizer_state_offload_fraction=1.0, + ckpt_format='torch_dist', + **common, + ) + reference = self._run_test_helper(1, "mxfp8", fp8_param_gather=True, **common) + torch.testing.assert_close(offloaded, reference, atol=1e-4, rtol=1e-4) + @pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8) @pytest.mark.parametrize("tp_size", [2]) @pytest.mark.parametrize("dp_overlap", [(True, True)]) diff --git a/tests/unit_tests/test_layer_wise_optimizer.py b/tests/unit_tests/test_layer_wise_optimizer.py index aaa06b2ddbe..6550945676e 100644 --- a/tests/unit_tests/test_layer_wise_optimizer.py +++ b/tests/unit_tests/test_layer_wise_optimizer.py @@ -1,6 +1,7 @@ # Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. import os +from types import SimpleNamespace import pytest import torch @@ -10,7 +11,7 @@ from megatron.core import parallel_state from megatron.core.distributed import DistributedDataParallel, DistributedDataParallelConfig -from megatron.core.optimizer import OptimizerConfig, get_megatron_optimizer +from megatron.core.optimizer import ChainedOptimizer, OptimizerConfig, get_megatron_optimizer from megatron.core.optimizer.layer_wise_optimizer import LayerWiseDistributedOptimizer from megatron.core.optimizer.optimizer import Float16OptimizerWithFloat16Params from megatron.core.process_groups_config import ProcessGroupCollection @@ -59,6 +60,159 @@ def forward(self, x): return self.fc1(x) +def test_layerwise_step_hook_is_noop_without_chunked_offload(): + """Disabled offload must not inspect child-only manager attributes or rebuild pipeline state.""" + + optimizer = object.__new__(LayerWiseDistributedOptimizer) + optimizer.config = SimpleNamespace(chunked_optimizer_state_offload=False) + optimizer.chained_optimizers = [object()] + + assert optimizer._managed_optimizer_state_offload_child_indices() == () + optimizer._before_child_step(0) + + +def test_layerwise_fraction_zero_skips_offload_specific_child_inspection(monkeypatch): + """A zero fraction must bypass Muon schema inspection and offloader construction.""" + + class FakeWrappedOptimizer: + def __init__(self, optimizer, config, grad_scaler, init_state_fn): + self.optimizer = optimizer + self.config = config + self._optimizer_state_offloader = None + + def fake_shard_params(self, optimizers, full_param_layouts=None, model_chunks=None): + self.dp_cp_params_list = [] + self.expt_dp_params_list = [] + + monkeypatch.setattr( + "megatron.core.optimizer.layer_wise_optimizer.Float16OptimizerWithFloat16Params", + FakeWrappedOptimizer, + ) + monkeypatch.setattr(LayerWiseDistributedOptimizer, "shard_params", fake_shard_params) + config = SimpleNamespace( + bf16=True, + chunked_optimizer_state_offload=True, + optimizer_state_offload_fraction=0.0, + overlap_param_gather=False, + use_layer_wise_param_layout=False, + ) + + optimizer = LayerWiseDistributedOptimizer([object()], config) + + assert optimizer._managed_optimizer_state_offload_child_indices() == () + LayerWiseDistributedOptimizer._before_child_step(optimizer, 0) + + +def test_layerwise_offload_rejects_multiple_nonempty_muon_children(monkeypatch): + """The supported LayerWise offload schema has at most one local Muon child.""" + + from megatron.core.optimizer.emerging_optimizers import TensorParallelMuon + + class FakeWrappedOptimizer: + def __init__(self, optimizer, config, grad_scaler, init_state_fn): + self.optimizer = optimizer + self.config = config + self.fp32_from_float16_groups = [[]] + self._optimizer_state_offloader = None + + def fake_shard_params(self, optimizers, full_param_layouts=None, model_chunks=None): + self.dp_cp_params_list = [] + self.expt_dp_params_list = [] + + raw_optimizers = [] + for _ in range(2): + raw_optimizer = object.__new__(TensorParallelMuon) + raw_optimizer.param_groups = [{"params": [object()]}] + raw_optimizers.append(raw_optimizer) + + monkeypatch.setattr( + "megatron.core.optimizer.layer_wise_optimizer.Float16OptimizerWithFloat16Params", + FakeWrappedOptimizer, + ) + monkeypatch.setattr(LayerWiseDistributedOptimizer, "shard_params", fake_shard_params) + config = SimpleNamespace( + bf16=True, + chunked_optimizer_state_offload=True, + optimizer_state_offload_fraction=1.0, + overlap_param_gather=False, + use_layer_wise_param_layout=False, + ) + + with pytest.raises(AssertionError, match="at most one non-empty TensorParallelMuon child"): + LayerWiseDistributedOptimizer(raw_optimizers, config) + + +def test_layerwise_param_sync_subset_excludes_distopt_buckets(): + """Forced pre-forward sync must leave sibling DistributedOptimizer buckets untouched.""" + + layerwise_param = SimpleNamespace(is_managed_by_layer_wise_optimizer=True) + distopt_param = SimpleNamespace(is_managed_by_layer_wise_optimizer=False) + layerwise_group = SimpleNamespace(buckets=[SimpleNamespace(params_list=[layerwise_param])]) + distopt_group = SimpleNamespace(buckets=[SimpleNamespace(params_list=[distopt_param])]) + calls = [] + model_chunk = SimpleNamespace( + bucket_groups=[layerwise_group, distopt_group], + expert_parallel_bucket_groups=[], + _start_bucket_group_param_sync=lambda group, force_sync: calls.append((group, force_sync)), + ) + optimizer = SimpleNamespace(model_chunks=[model_chunk]) + + LayerWiseDistributedOptimizer.start_param_sync_for_bucket_group_subset( + optimizer, force_sync=True + ) + + assert calls == [(layerwise_group, True)] + + +def test_outer_chain_only_syncs_children_requiring_master_before_offload(): + calls = [] + + def child(name, requires_sync): + return SimpleNamespace( + optimizer_state_offload_requires_pre_forward_param_sync=lambda: requires_sync, + start_param_sync_for_bucket_group_subset=lambda force_sync: calls.append( + (name, force_sync) + ), + ) + + optimizer = SimpleNamespace( + chained_optimizers=[child("layerwise", True), child("distopt", False)] + ) + ChainedOptimizer.start_param_sync_for_bucket_group_subset(optimizer, force_sync=True) + assert calls == [("layerwise", True)] + + +def test_layerwise_prefetches_single_managed_child_state(): + """LayerWise prefetches its single managed child before the chained step.""" + + events = [] + + def child(name, managed): + return SimpleNamespace( + _optimizer_state_offloader=object() if managed else None, + prefetch_optimizer_master_weights_for_step=lambda: events.append(f"master:{name}"), + prefetch_optimizer_state_for_step=lambda: events.append(f"state:{name}"), + step_with_ready_grads=lambda: events.append(f"step:{name}") or True, + ) + + children = [child("a", False), child("b", True), child("c", False)] + optimizer = object.__new__(LayerWiseDistributedOptimizer) + optimizer.chained_optimizers = children + optimizer.config = SimpleNamespace( + chunked_optimizer_state_offload=True, + optimizer_state_offload_fraction=1.0, + overlap_param_gather_with_optimizer_step=False, + ) + optimizer._managed_optimizer_state_offload_indices = (1,) + + optimizer.prefetch_optimizer_state_for_gradient_finalization() + assert events == ["master:a", "master:b", "master:c", "state:b"] + + events.clear() + assert optimizer._step() + assert events == ["state:b", "step:a", "step:b", "step:c"] + + @pytest.mark.skipif( int(os.getenv('WORLD_SIZE', '1')) == 1, reason="Multi-rank test requires WORLD_SIZE > 1" ) @@ -82,6 +236,9 @@ def create_model_and_optimizer( use_layer_wise=True, copy_from=None, use_param_layout=False, + chunked_optimizer_state_offload=False, + optimizer_state_offload_chunk_size_mb=0, + optimizer_state_offload_fraction=1.0, ): """Create model, DDP wrapper, and optimizer. @@ -135,6 +292,9 @@ def create_model_and_optimizer( clip_grad=clip_grad, muon_tp_mode="duplicated", use_layer_wise_distributed_optimizer=use_layer_wise, + chunked_optimizer_state_offload=chunked_optimizer_state_offload, + optimizer_state_offload_chunk_size_mb=optimizer_state_offload_chunk_size_mb, + optimizer_state_offload_fraction=optimizer_state_offload_fraction, ) pg_collection = ProcessGroupCollection.use_mpu_process_groups() @@ -299,6 +459,113 @@ def test_basic(self, use_param_layout): f"Parameter {name} differs between rank 0 and rank {i}. {str(e)}" ) from None + def test_compact_muon_chunked_state_and_master_offload(self): + """Compact-layout Muon updates whole tensors while state and masters live on CPU.""" + + model, optimizer, _ = self.create_model_and_optimizer( + use_param_layout=True, + chunked_optimizer_state_offload=True, + optimizer_state_offload_chunk_size_mb=1, + ) + reference_model, reference_optimizer, _ = self.create_model_and_optimizer( + use_param_layout=True, copy_from=model + ) + + layerwise_optimizer = next( + child + for child in optimizer.chained_optimizers + if isinstance(child, LayerWiseDistributedOptimizer) + ) + managed_children = [ + child + for child in layerwise_optimizer.chained_optimizers + if child._optimizer_state_offloader is not None + ] + # Whole-parameter compact layout can leave some DP ranks without a local Muon + # parameter. Those ranks legitimately have no LayerWise state offloader, but must + # continue through the distributed optimizer step with the ranks that do own one. + max_managed_children = torch.tensor(len(managed_children), dtype=torch.int, device='cuda') + torch.distributed.all_reduce(max_managed_children, op=torch.distributed.ReduceOp.MAX) + assert max_managed_children.item() > 0 + if managed_children: + assert ( + len( + {id(child._optimizer_state_offloader._d2h_stream) for child in managed_children} + ) + == 1 + ) + assert ( + len( + {id(child._optimizer_state_offloader._h2d_stream) for child in managed_children} + ) + == 1 + ) + + optimizer.offload_optimizer_state_for_forward() + optimizer.zero_grad() + + for child in managed_children: + manager = child._optimizer_state_offloader + for param in manager.selected_params: + if any( + id(param) == id(master) + for group in child.fp32_from_float16_groups + for master in group + ): + assert param.device.type == 'cpu' + + inputs = torch.randn(16, 80, dtype=torch.bfloat16, device='cuda') + model(inputs).sum().backward() + optimizer.prefetch_optimizer_state_for_gradient_finalization() + assert optimizer.step()[0] + + reference_model(inputs).sum().backward() + assert reference_optimizer.step()[0] + + optimizer.offload_optimizer_state_for_forward() + for child in managed_children: + manager = child._optimizer_state_offloader + for param in manager.selected_params: + for value in manager.optimizer.state[param].values(): + if isinstance(value, torch.Tensor) and value.numel() == param.numel(): + assert value.device.type == 'cpu' + if any( + id(param) == id(master) + for group in child.fp32_from_float16_groups + for master in group + ): + assert param.device.type == 'cpu' + + for offloaded_param, reference_param in zip( + model.parameters(), reference_model.parameters() + ): + torch.testing.assert_close(offloaded_param, reference_param, rtol=1e-4, atol=1e-5) + + def test_split_layerwise_and_distopt_share_offload_streams(self): + """Muon children and sibling Adam DistOpt serialize copies on one stream pair.""" + + _, optimizer, _ = self.create_model_and_optimizer( + use_param_layout=True, + chunked_optimizer_state_offload=True, + optimizer_state_offload_chunk_size_mb=1, + ) + managers = [] + + def collect(current_optimizer): + manager = getattr(current_optimizer, '_optimizer_state_offloader', None) + if manager is not None: + managers.append(manager) + for child in getattr(current_optimizer, 'chained_optimizers', ()): + collect(child) + + collect(optimizer) + max_manager_count = torch.tensor(len(managers), dtype=torch.int, device='cuda') + torch.distributed.all_reduce(max_manager_count, op=torch.distributed.ReduceOp.MAX) + assert max_manager_count.item() >= 2 + if len(managers) >= 2: + assert len({id(manager.transfer_streams[0]) for manager in managers}) == 1 + assert len({id(manager.transfer_streams[1]) for manager in managers}) == 1 + def test_get_grad_norm(self): """Test LayerWiseDistributedOptimizer gradient norm computation.""" model, optimizer, pg_collection = self.create_model_and_optimizer() diff --git a/tests/unit_tests/test_muon_decouple_fp8_param_gather.py b/tests/unit_tests/test_muon_decouple_fp8_param_gather.py index 62677ef32ef..773fb02b085 100644 --- a/tests/unit_tests/test_muon_decouple_fp8_param_gather.py +++ b/tests/unit_tests/test_muon_decouple_fp8_param_gather.py @@ -145,7 +145,13 @@ def model_provider(self, pre_process=True, post_process=True, **kw): ) def _create_args( - self, fp8_param_gather, fp8_recipe, overlap, num_experts=0, expert_model_parallel_size=1 + self, + fp8_param_gather, + fp8_recipe, + overlap, + num_experts=0, + expert_model_parallel_size=1, + chunked_optimizer_state_offload=False, ): destroy_global_vars() destroy_num_microbatches_calculator() @@ -195,6 +201,11 @@ def _create_args( args.fp8 = "e4m3" args.fp8_recipe = fp8_recipe args.fp8_param_gather = fp8_param_gather + args.chunked_optimizer_state_offload = chunked_optimizer_state_offload + if chunked_optimizer_state_offload: + args.optimizer_state_offload_chunk_size_mb = 1 + args.optimizer_state_offload_fraction = 1.0 + args.ckpt_format = 'torch_dist' if fp8_param_gather and fp8_recipe == "mxfp8": args.reuse_grad_buf_for_mxfp8_param_ag = ( True # mxfp8 columnwise needs the bf16 round-trip @@ -230,7 +241,13 @@ def _batch(self): return ids, labels, pos, mask, loss_mask def _build( - self, fp8_param_gather, fp8_recipe, overlap, num_experts=0, expert_model_parallel_size=1 + self, + fp8_param_gather, + fp8_recipe, + overlap, + num_experts=0, + expert_model_parallel_size=1, + chunked_optimizer_state_offload=False, ): args = self._create_args( fp8_param_gather, @@ -238,6 +255,7 @@ def _build( overlap, num_experts=num_experts, expert_model_parallel_size=expert_model_parallel_size, + chunked_optimizer_state_offload=chunked_optimizer_state_offload, ) set_args(args) torch.manual_seed(_SEED) @@ -257,8 +275,27 @@ def _run_steps(self, args, model, optimizer, n): ids, labels, pos, mask, loss_mask = self._batch() losses, outs, grads, masters, params = [], [], [], [], [] for _ in range(n): + delay_master_offload = ( + args.chunked_optimizer_state_offload + and args.reuse_grad_buf_for_mxfp8_param_ag + and args.overlap_param_gather + ) + pre_forward_sync = ( + args.chunked_optimizer_state_offload + and optimizer.optimizer_state_offload_requires_pre_forward_param_sync() + ) + if args.chunked_optimizer_state_offload and not pre_forward_sync: + optimizer.offload_optimizer_state_for_forward( + offload_master=not delay_master_offload + ) model[0].zero_grad_buffer() optimizer.zero_grad() + if pre_forward_sync: + optimizer.ensure_master_weights_for_pre_forward_param_sync() + optimizer.start_param_sync_for_bucket_group_subset(force_sync=True) + optimizer.offload_optimizer_state_for_forward( + offload_master=not delay_master_offload + ) # reuse_grad_buf aliases the bf16 staging buffer onto the just-zeroed grad buffer, so # the DistOpt sibling buffers must re-stage masters before the deferred (overlap) param # all-gather (mirrors the training loop). The non-DistOpt LayerWise children stage via @@ -266,7 +303,10 @@ def _run_steps(self, args, model, optimizer, n): if args.reuse_grad_buf_for_mxfp8_param_ag and args.overlap_param_gather: for opt in optimizer.chained_optimizers: if isinstance(opt, DistributedOptimizer): + # The staging entry restores its own offloaded masters, matching train_step. opt._copy_main_params_to_param_buffer() + if delay_master_offload: + optimizer.offload_optimizer_state_for_forward() model[0].set_is_first_microbatch() out = model[0].forward( input_ids=ids, @@ -284,6 +324,8 @@ def _run_steps(self, args, model, optimizer, n): for name, p in model[0].named_parameters() if p.main_grad is not None } + if args.chunked_optimizer_state_offload: + optimizer.prefetch_optimizer_state_for_gradient_finalization() ok, _, _ = optimizer.step() assert ok params.append(_snapshot_params(model[0])) @@ -349,6 +391,111 @@ def test_on_vs_off_bitwise_identical(self, fp8_recipe, overlap): # iterations (PR #5470 review), so a handful of steps can miss a real divergence. self._check_on_vs_off(fp8_recipe, overlap, n=30) + @pytest.mark.parametrize("fp8_recipe", ["blockwise", "mxfp8"]) + @pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8) + @pytest.mark.skipif(not is_te_min_version("2.3.0.dev0"), reason="TE 2.3.0.dev0 is required") + def test_chunked_optimizer_state_offload(self, fp8_recipe): + """Compact Muon FP8 gather remains equivalent with chunked state/master offload.""" + + arch = get_device_arch_version() + if fp8_recipe == "blockwise" and arch != 9: + pytest.skip("blockwise FP8 is Hopper-only") + if fp8_recipe == "mxfp8" and arch < 10: + pytest.skip("mxfp8 requires Blackwell architecture or newer") + + with deterministic_mode(): + ref_args, ref_model, ref_opt = self._build(True, fp8_recipe, True) + params0 = _snapshot_params(ref_model[0]) + masters0 = _snapshot_masters(ref_model[0]) + off_args, off_model, off_opt = self._build( + True, fp8_recipe, True, chunked_optimizer_state_offload=True + ) + _restore_initial_state(off_model[0], off_opt, params0, masters0) + + reference = self._run_steps(ref_args, ref_model, ref_opt, 30) + offloaded = self._run_steps(off_args, off_model, off_opt, 30) + + for actual_steps, expected_steps in zip(offloaded, reference): + assert len(actual_steps) == len(expected_steps) + for step, (actual, expected) in enumerate(zip(actual_steps, expected_steps)): + if isinstance(actual, dict): + assert actual.keys() == expected.keys() + for name in actual: + _assert_equal(actual[name], expected[name], f"offload step {step} {name}") + else: + _assert_equal(actual, expected, f"offload step {step}") + + off_opt.offload_optimizer_state_for_forward() + layerwise = off_opt.chained_optimizers[0] + managers = [ + child._optimizer_state_offloader + for child in layerwise.chained_optimizers + if child._optimizer_state_offloader is not None + ] + assert managers + for manager in managers: + for param in manager.selected_params: + assert param.device.type == 'cpu' + for value in manager.optimizer.state[param].values(): + if isinstance(value, torch.Tensor) and value.numel() == param.numel(): + assert value.device.type == 'cpu' + + @pytest.mark.launch_on_gb200 + @pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8) + @pytest.mark.skipif(not is_te_min_version("2.3.0.dev0"), reason="TE 2.3.0.dev0 is required") + def test_pending_distopt_gather_does_not_break_next_offload_iteration(self): + """A speculative DistOpt gather must be safely consumed by the next iteration.""" + + if torch.distributed.get_world_size() < 2: + pytest.skip("Requires at least two data-parallel ranks") + if get_device_arch_version() < 10: + pytest.skip("MXFP8 requires Blackwell architecture or newer") + + args, model, optimizer = self._build( + True, "mxfp8", True, chunked_optimizer_state_offload=True + ) + self._run_steps(args, model, optimizer, 1) + + ddp = model[0] + distopt = next( + child + for child in optimizer.chained_optimizers + if isinstance(child, DistributedOptimizer) + ) + + # Stage a valid DistOpt payload, then model the speculative successor gather + # that can survive the end of forward. + distopt._copy_main_params_to_param_buffer() + distopt_groups = [] + for group in ddp.bucket_groups + ddp.expert_parallel_bucket_groups: + if not group.buckets or not group.buckets[0].params_list: + continue + first_param = group.buckets[0].params_list[0] + if not getattr(first_param, "is_managed_by_layer_wise_optimizer", False): + distopt_groups.append(group) + assert distopt_groups, "test precondition: expected a DistOpt-owned bucket group" + + pending_group = distopt_groups[0] + ddp._start_bucket_group_param_sync(pending_group, force_sync=False) + assert pending_group.param_gather_handle is not None + assert pending_group.param_gather_dispatched + + # Wait for the GPU work so the test does not intentionally race buffer writes, + # but retain the handle object to reproduce the pending ownership state. + stale_handle = pending_group.param_gather_handle + stale_handle.wait() + pending_group.finish_grad_sync() + + try: + second_step = self._run_steps(args, model, optimizer, 1) + finally: + # Preserve the real failure while preventing the intentionally retained + # handle from leaking into distributed teardown. + if pending_group.param_gather_handle is stale_handle: + pending_group.param_gather_handle = None + + assert all(len(values) == 1 for values in second_step) + @pytest.mark.parametrize("overlap", [False, True]) @pytest.mark.parametrize("expt_dp_gt_1", [False, True]) @pytest.mark.parametrize("fp8_recipe", ["blockwise", "mxfp8"]) diff --git a/tests/unit_tests/test_optimizer_state_offloading.py b/tests/unit_tests/test_optimizer_state_offloading.py index baaab355182..8eb7445bf79 100644 --- a/tests/unit_tests/test_optimizer_state_offloading.py +++ b/tests/unit_tests/test_optimizer_state_offloading.py @@ -1,14 +1,31 @@ -# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -"""Unit tests for OptimizerStateOffloader.""" +"""Unit tests for chunked optimizer state and master-weight offload.""" + +import dataclasses +import logging +import sys +from contextlib import nullcontext +from types import SimpleNamespace +from unittest import mock import pytest import torch import torch.nn as nn from megatron.core.distributed import DistributedDataParallel, DistributedDataParallelConfig -from megatron.core.optimizer import OptimizerConfig, get_megatron_optimizer +from megatron.core.optimizer import ( + ChainedOptimizer, + FP32Optimizer, + OptimizerConfig, + get_megatron_optimizer, +) +from megatron.core.optimizer.cpu_offloading.chunked_optimizer_state_offload import ( + ChunkedOptimizerStateOffloader, +) from megatron.core.transformer import TransformerConfig +from megatron.core.utils import is_te_min_version +from megatron.training.arguments import parse_args, validate_args from tests.unit_tests.test_utilities import Utils try: @@ -20,20 +37,899 @@ class SimpleModel(nn.Module): - """Simple model for testing.""" + """Small model with enough parameters to exercise multiple state chunks.""" - def __init__(self, hidden_size=256): + def __init__(self, hidden_size: int = 256, num_layers: int = 2) -> None: super().__init__() + assert num_layers >= 2 self.fc1 = nn.Linear(hidden_size, hidden_size) + self.hidden_layers = nn.ModuleList( + nn.Linear(hidden_size, hidden_size) for _ in range(num_layers - 2) + ) self.fc2 = nn.Linear(hidden_size, hidden_size) - def forward(self, x): - return self.fc2(torch.relu(self.fc1(x))) + def forward(self, inputs: torch.Tensor) -> torch.Tensor: + """Run the test MLP.""" + + output = torch.relu(self.fc1(inputs)) + for layer in self.hidden_layers: + output = torch.relu(layer(output)) + return self.fc2(output) + + +@pytest.fixture +def model_parallel_context(): + """Always tear down model-parallel state, including on assertion failure or skip.""" + + try: + Utils.initialize_model_parallel() + yield + finally: + Utils.destroy_model_parallel() + + +def test_legacy_flag_warns_and_enables_chunked_offload(): + """OptimizerConfig retains and normalizes the deprecated compatibility field.""" + + with pytest.warns(FutureWarning, match='offload_optimizer_states is deprecated'): + config = OptimizerConfig( + optimizer='adam', use_distributed_optimizer=True, offload_optimizer_states=True + ) + + assert not config.offload_optimizer_states + assert config.chunked_optimizer_state_offload + + with mock.patch("megatron.core.optimizer.optimizer_config.warnings.warn") as warn: + copied_config = dataclasses.replace(config) + + warn.assert_not_called() + assert not copied_config.offload_optimizer_states + assert copied_config.chunked_optimizer_state_offload + + +def test_legacy_flag_uses_user_facing_optimizer_error(): + """An unsupported legacy invocation names the deprecated CLI alias in its failure.""" + + with ( + pytest.warns(FutureWarning, match='offload_optimizer_states is deprecated'), + pytest.raises(AssertionError, match='deprecated --offload-optimizer-states alias'), + ): + OptimizerConfig(optimizer='sgd', offload_optimizer_states=True) + + +def test_chunked_optimizer_state_offload_config_accepts_supported_optimizers(): + """OptimizerConfig accepts the supported Adam and compact LayerWise Muon modes.""" + + OptimizerConfig( + optimizer='adam', use_distributed_optimizer=True, chunked_optimizer_state_offload=True + ) + OptimizerConfig( + optimizer='muon', + bf16=True, + use_layer_wise_distributed_optimizer=True, + use_layer_wise_param_layout=False, + chunked_optimizer_state_offload=True, + ) + + +def test_chunked_optimizer_state_offload_fraction_zero_is_a_noop(): + """The feature flag may remain set while a zero fraction disables mode restrictions.""" + + config = OptimizerConfig( + optimizer='sgd', chunked_optimizer_state_offload=True, optimizer_state_offload_fraction=0.0 + ) + + assert config.chunked_optimizer_state_offload + assert config.optimizer_state_offload_fraction == 0.0 + + +@pytest.mark.parametrize( + ('overrides', 'error'), + [ + ({'optimizer': 'sgd'}, 'currently supports Adam and Muon'), + ({'use_distributed_optimizer': False}, 'requires use_distributed_optimizer'), + ( + {'optimizer': 'muon', 'use_layer_wise_distributed_optimizer': True}, + 'currently requires bf16', + ), + ( + {'use_layer_wise_distributed_optimizer': True}, + 'Adam optimizer state offload does not support', + ), + ({'optimizer': 'muon', 'bf16': True}, 'requires LayerWiseDistributedOptimizer'), + ( + { + 'optimizer': 'muon', + 'bf16': True, + 'use_layer_wise_distributed_optimizer': True, + 'use_layer_wise_param_layout': True, + }, + '--use-layer-wise-param-layout to be disabled', + ), + ({'optimizer_cpu_offload': True}, 'optimizer_cpu_offload are mutually exclusive'), + ({'optimizer_state_offload_chunk_size_mb': -1}, 'must be non-negative'), + ({'optimizer_state_offload_fraction': 1.1}, 'must be in'), + ({'optimizer_cuda_graph': True}, 'does not support optimizer CUDA graphs'), + ], +) +def test_chunked_optimizer_state_offload_config_rejects_unsupported_modes(overrides, error): + """Optimizer-owned offload constraints are enforced by OptimizerConfig.""" + + config_kwargs = { + 'optimizer': 'adam', + 'use_distributed_optimizer': True, + 'chunked_optimizer_state_offload': True, + } + config_kwargs.update(overrides) + with pytest.raises(AssertionError, match=error): + OptimizerConfig(**config_kwargs) + + +@pytest.mark.parametrize( + ("flag_name", "chunk_size_mb", "expects_unbounded_window_warning"), + [("chunked_optimizer_state_offload", 1, False), ("offload_optimizer_states", 0, True)], +) +def test_training_args_reject_async_save_with_chunked_state_offload( + monkeypatch, flag_name, chunk_size_mb, expects_unbounded_window_warning +): + """Reusable canonical CPU buffers cannot be handed to a background checkpoint writer.""" + + monkeypatch.setattr(sys, 'argv', ['test_optimizer_state_offloading.py']) + args = parse_args() + args.num_layers = 2 + args.vocab_size = 256 + args.hidden_size = 64 + args.num_attention_heads = 4 + args.max_position_embeddings = 128 + args.seq_length = 128 + args.micro_batch_size = 1 + args.train_iters = 1 + args.lr = 1e-4 + args.tokenizer_type = 'NullTokenizer' + args.optimizer = 'adam' + args.use_distributed_optimizer = True + setattr(args, flag_name, True) + args.optimizer_state_offload_chunk_size_mb = chunk_size_mb + args.ckpt_format = 'torch_dist' + args.async_save = True + args.use_persistent_ckpt_worker = True + + with mock.patch("megatron.training.arguments.warn_rank_0") as warn: + with pytest.raises(AssertionError, match="does not support --async-save"): + validate_args(args) + + warning_messages = [call.args[0] for call in warn.call_args_list if call.args] + assert any("enabled with chunk size 0" in message for message in warning_messages) is ( + expects_unbounded_window_warning + ) + + +def test_training_args_treat_zero_offload_fraction_as_disabled(monkeypatch): + """A zero fraction bypasses checkpoint restrictions during startup validation.""" + + monkeypatch.setattr(sys, 'argv', ['test_optimizer_state_offloading.py']) + args = parse_args() + args.num_layers = 2 + args.vocab_size = 256 + args.hidden_size = 64 + args.num_attention_heads = 4 + args.max_position_embeddings = 128 + args.seq_length = 128 + args.micro_batch_size = 1 + args.train_iters = 1 + args.lr = 1e-4 + args.tokenizer_type = 'NullTokenizer' + args.optimizer = 'adam' + args.use_distributed_optimizer = True + args.chunked_optimizer_state_offload = True + args.optimizer_state_offload_fraction = 0.0 + args.ckpt_format = 'torch' + args.use_dist_ckpt = False + args.async_save = True + args.use_persistent_ckpt_worker = True + + validated_args = validate_args(args) + + assert validated_args.async_save + + +def test_training_args_reject_muon_offload_without_distributed_optimizer(monkeypatch): + """Muon offload must route non-Muon groups to a sibling DistributedOptimizer.""" + + monkeypatch.setattr(sys, 'argv', ['test_optimizer_state_offloading.py']) + args = parse_args() + args.num_layers = 2 + args.vocab_size = 256 + args.hidden_size = 64 + args.num_attention_heads = 4 + args.max_position_embeddings = 128 + args.seq_length = 128 + args.micro_batch_size = 1 + args.train_iters = 1 + args.lr = 1e-4 + args.tokenizer_type = 'NullTokenizer' + args.optimizer = 'muon' + args.use_distributed_optimizer = False + args.chunked_optimizer_state_offload = True + args.ckpt_format = 'torch_dist' + + with pytest.raises(AssertionError, match="requires the LayerWise distributed optimizer"): + validate_args(args) + + +def test_training_args_accept_deprecated_dist_muon_offload(monkeypatch): + """The deprecated dist_muon spelling normalizes to the supported LayerWise mode.""" + + monkeypatch.setattr(sys, 'argv', ['test_optimizer_state_offloading.py']) + args = parse_args() + args.num_layers = 2 + args.vocab_size = 256 + args.hidden_size = 64 + args.num_attention_heads = 4 + args.max_position_embeddings = 128 + args.seq_length = 128 + args.micro_batch_size = 1 + args.train_iters = 1 + args.lr = 1e-4 + args.tokenizer_type = 'NullTokenizer' + args.optimizer = 'dist_muon' + args.use_distributed_optimizer = False + args.chunked_optimizer_state_offload = True + args.ckpt_format = 'torch_dist' + + validated_args = validate_args(args) + + assert validated_args.optimizer == 'muon' + assert validated_args.use_layer_wise_distributed_optimizer + assert not validated_args.use_distributed_optimizer + + +@pytest.mark.parametrize( + ("ckpt_format", "async_save", "no_load_optim"), + [("torch", False, True), ("torch_dist", True, False)], +) +def test_training_args_narrow_checkpoint_restrictions_to_optimizer_io( + monkeypatch, ckpt_format, async_save, no_load_optim +): + """Model-only saves skip offload restrictions not needed by optimizer checkpoint I/O.""" + + monkeypatch.setattr(sys, 'argv', ['test_optimizer_state_offloading.py']) + args = parse_args() + args.num_layers = 2 + args.vocab_size = 256 + args.hidden_size = 64 + args.num_attention_heads = 4 + args.max_position_embeddings = 128 + args.seq_length = 128 + args.micro_batch_size = 1 + args.train_iters = 1 + args.lr = 1e-4 + args.tokenizer_type = 'NullTokenizer' + args.optimizer = 'adam' + args.use_distributed_optimizer = True + args.chunked_optimizer_state_offload = True + args.ckpt_format = ckpt_format + args.use_dist_ckpt = ckpt_format == "torch_dist" + args.async_save = async_save + args.use_persistent_ckpt_worker = True + args.no_save_optim = True + args.no_load_optim = no_load_optim + + validated_args = validate_args(args) + + assert validated_args.ckpt_format == ckpt_format + assert validated_args.async_save is async_save + + +def test_standard_chained_optimizer_shares_state_offload_streams(): + """Ordinary chained Adam-style siblings share transfer ordering, not only LayerWise.""" + + class FakeManager: + def __init__(self): + self.transfer_streams = (object(), object()) + + def use_transfer_streams(self, d2h_stream, h2d_stream): + self.transfer_streams = (d2h_stream, h2d_stream) + + class FakeOptimizer: + def __init__(self, config, manager): + self.config = config + self._optimizer_state_offloader = manager + + config = object() + managers = [FakeManager(), FakeManager()] + ChainedOptimizer([FakeOptimizer(config, manager) for manager in managers]) + + assert managers[1].transfer_streams[0] is managers[0].transfer_streams[0] + assert managers[1].transfer_streams[1] is managers[0].transfer_streams[1] + + +def test_chained_optimizer_forwards_master_residency_assertion(): + """Chain assertions reach every child while inherited lifecycle defaults remain safe.""" + + calls = [] + + class FakeOptimizer: + def __init__(self, config, name): + self.config = config + + self.name = name + + def assert_master_weights_resident(self, operation): + calls.append((self.name, operation)) + + config = object() + optimizer = ChainedOptimizer( + [FakeOptimizer(config, 'layerwise'), FakeOptimizer(config, 'distributed')] + ) + + assert optimizer._optimizer_state_offloader is None + optimizer.assert_master_weights_resident("test") + assert calls == [('layerwise', 'test'), ('distributed', 'test')] + optimizer.set_optimizer_state_offload_deferred_lifecycle( + state_prefetch_to_step=False, master_offload_for_param_sync=False + ) + assert not optimizer._defer_optimizer_state_prefetch_to_step + assert not optimizer._defer_optimizer_master_offload_for_param_sync + + +def test_chained_master_restore_targets_only_pre_forward_sync_children(): + """A compact Muon pre-forward restore must not pull sibling DistOpt masters back.""" + + events = [] + + class FakeOptimizer: + def __init__(self, config, name, requires_sync): + self.config = config + self.name = name + self.requires_sync = requires_sync + + def optimizer_state_offload_requires_pre_forward_param_sync(self): + checks.append(self.name) + return self.requires_sync + + def ensure_master_weights_for_pre_forward_param_sync(self): + if self.optimizer_state_offload_requires_pre_forward_param_sync(): + events.append(self.name) + + config = object() + checks = [] + optimizer = ChainedOptimizer( + [FakeOptimizer(config, "layerwise", True), FakeOptimizer(config, "distributed", False)] + ) + + optimizer.ensure_master_weights_for_pre_forward_param_sync() + + assert events == ["layerwise"] + assert checks == ["layerwise", "distributed"] + + +def test_chained_general_master_restore_targets_all_children(): + """The general chain API preserves the leaf contract by restoring every child.""" + + events = [] + + class FakeOptimizer: + def __init__(self, config, name): + self.config = config + self.name = name + + def ensure_master_weights_for_param_sync(self): + events.append(self.name) + + config = object() + optimizer = ChainedOptimizer( + [FakeOptimizer(config, "layerwise"), FakeOptimizer(config, "distributed")] + ) + + optimizer.ensure_master_weights_for_param_sync() + + assert events == ["layerwise", "distributed"] + + +def test_pre_forward_sync_requirement_requires_subset_implementation(): + """A leaf declaring pre-forward sync must implement its bucket-subset dispatch.""" + + optimizer = object.__new__(FP32Optimizer) + optimizer._defer_optimizer_master_offload_for_param_sync = True + + with pytest.raises(NotImplementedError, match="must override"): + optimizer.start_param_sync_for_bucket_group_subset(force_sync=True) + + +def test_fp32_ready_grad_step_rejects_chunked_state_offloader(): + """The unsupported common-base FP32 entry point must not bypass its manager.""" + + optimizer = object.__new__(FP32Optimizer) + optimizer.is_stub_optimizer = False + optimizer._optimizer_state_offloader = object() + + with pytest.raises(RuntimeError, match="FP32Optimizer.step_with_ready_grads"): + optimizer.step_with_ready_grads() + + +def test_offloader_rejects_two_master_storage_schemes(): + """One wrapped optimizer cannot alias explicit and optimizer-owned master slots.""" + + param = nn.Parameter(torch.ones(2)) + optimizer = torch.optim.SGD([param], lr=0.1) + optimizer.master_weights = True + with pytest.raises(ValueError, match="both MCore-owned master parameters"): + ChunkedOptimizerStateOffloader( + optimizer=optimizer, + master_params=[param], + chunk_size_bytes=0, + offload_fraction=1.0, + state_dtypes=(torch.float32,), + ) + + +def test_oversized_atomic_state_chunks_log_once(): + """A soft chunk target overflow is visible without splitting optimizer parameters.""" + + manager = object.__new__(ChunkedOptimizerStateOffloader) + manager.chunk_size_bytes = 1024**2 + manager._state_bytes_per_param = torch.float32.itemsize + manager.optimizer = SimpleNamespace() + params = tuple(nn.Parameter(torch.empty(numel)) for numel in (300_000, 400_000, 10)) + + with mock.patch( + "megatron.core.optimizer.cpu_offloading.chunked_optimizer_state_offload.log_single_rank" + ) as log: + chunks = manager._build_chunks(params) + + assert [[id(param) for param in chunk.params] for chunk in chunks] == [ + [id(param)] for param in params + ] + log.assert_called_once() + _, level, _, target_mib, oversized_count, largest_mib, optimizer_name = log.call_args.args + assert level == logging.WARNING + assert target_mib == 1.0 + assert oversized_count == 2 + assert largest_mib == pytest.approx(400_000 * torch.float32.itemsize / 1024**2) + assert optimizer_name == "SimpleNamespace" + + +def test_h2d_targets_record_the_transfer_stream(): + """State and master allocations remain live if an update exits before its H2D completes.""" + + class FakeTensor: + def __init__(self, *, is_cuda=False, pinned=False): + self.is_cuda = is_cuda + self.pinned = pinned + self.recorded_streams = [] + + def is_pinned(self): + return self.pinned + + def copy_(self, source, non_blocking=False): + return self + + def record_stream(self, stream): + self.recorded_streams.append(stream) + + class FakeEvent: + def __init__(self): + self.recorded_stream = None + + def record(self, stream): + self.recorded_stream = stream + + h2d_stream = object() + state_param = object() + state_cpu = FakeTensor(pinned=True) + state_gpu = FakeTensor(is_cuda=True) + state_manager = object.__new__(ChunkedOptimizerStateOffloader) + state_manager._h2d_stream = h2d_stream + state_manager._state_staging_views = lambda chunk: [ + (state_param, "moment", state_cpu, state_gpu) + ] + state_manager._order_h2d_after_source_streams = lambda: None + state_manager.optimizer = SimpleNamespace(state={state_param: {}}) + + class FakeParam: + pass + + master_param = FakeParam() + master_cpu = FakeTensor(pinned=True) + master_gpu = FakeTensor(is_cuda=True) + master_param.data = master_cpu + master_manager = object.__new__(ChunkedOptimizerStateOffloader) + master_manager._selected_params = (master_param,) + master_manager._master_h2d_event = None + master_manager._explicit_master_param_ids = {id(master_param)} + master_manager._master_weights_resident = True + master_manager._param_devices = {master_param: object()} + master_manager._cpu_master = {} + master_manager._h2d_stream = h2d_stream + master_manager._validate_master_storage = lambda: None + master_manager._order_h2d_after_source_streams = lambda: None + master_manager.optimizer = SimpleNamespace(state={master_param: {}}) + + with ( + mock.patch.object(torch.cuda, "stream", side_effect=lambda stream: nullcontext()), + mock.patch.object(torch.cuda, "Event", side_effect=FakeEvent), + mock.patch.object(torch, "empty_like", return_value=master_gpu), + ): + state_manager._prefetch_state(object()) + master_manager._schedule_master_h2d() + + assert state_gpu.recorded_streams == [h2d_stream] + assert master_gpu.recorded_streams == [h2d_stream] + + +def test_initialize_state_for_loading_requires_an_initializer(): + """A missing external optimizer initializer must fail instead of silently skipping state.""" + + manager = object.__new__(ChunkedOptimizerStateOffloader) + + with pytest.raises(RuntimeError, match="requires an optimizer state initializer"): + manager.initialize_state_for_loading(None, object()) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA") +def test_transfer_stream_adoption_is_idempotent_after_first_use(): + """Re-adopting the current stream pair is safe even after canonical CPU state exists.""" + + param = nn.Parameter(torch.zeros(8, device='cuda')) + optimizer = torch.optim.SGD([param], lr=0.1) + manager = ChunkedOptimizerStateOffloader( + optimizer=optimizer, + master_params=[], + chunk_size_bytes=0, + offload_fraction=1.0, + state_dtypes=(torch.float32,), + ) + d2h_stream, h2d_stream = manager.transfer_streams + manager._cpu_state[param]['moment'] = torch.empty(param.shape, device='cpu', pin_memory=True) + + manager.use_transfer_streams(d2h_stream, h2d_stream) + + with pytest.raises(RuntimeError, match="cannot change after first use"): + manager.use_transfer_streams(torch.cuda.Stream(), h2d_stream) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA") +def test_chunked_step_releases_staging_slots_when_optimizer_raises(): + """An optimizer failure must not leave reusable staging buffers owned by the manager.""" + + class FailingOptimizer(torch.optim.Optimizer): + def __init__(self, params): + super().__init__(params, defaults={}) + + def step(self, closure=None): + raise RuntimeError("optimizer failure") + + param = nn.Parameter(torch.zeros(8, device='cuda')) + optimizer = FailingOptimizer([param]) + manager = ChunkedOptimizerStateOffloader( + optimizer=optimizer, + master_params=[], + chunk_size_bytes=0, + offload_fraction=1.0, + state_dtypes=(torch.float32,), + ) + cpu_state = torch.empty(param.shape, device='cpu', pin_memory=True) + optimizer.state[param]['moment'] = cpu_state + manager._cpu_state[param]['moment'] = cpu_state + + with pytest.raises(RuntimeError, match="optimizer failure"): + manager.step() + + assert all(not slot.buffers for slot in manager._state_staging_slots) + manager.synchronize_for_checkpoint() + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA") +def test_chunked_step_propagates_group_writes_and_clones_tensor_step(): + """Group metadata is committed once without rescanning original group parameters.""" -def create_model_and_optimizer(hidden_size=256, offload_optimizer_states=True, **optimizer_kwargs): - """Helper to create model and optimizer for tests.""" - model = SimpleModel(hidden_size=hidden_size).bfloat16().cuda() + class ParamsThatMustNotBeIterated(list): + def __iter__(self): + raise AssertionError("chunk execution rescanned the full optimizer parameter group") + + class GroupStateOptimizer(torch.optim.Optimizer): + def __init__(self, params): + super().__init__( + params, + defaults={"step": torch.tensor(0), "updates": 0, "removed_during_step": True}, + ) + + @torch.no_grad() + def step(self, closure=None): + for group in self.param_groups: + group["step"].add_(1) + group["updates"] += 1 + group["new_field"] = "preserved" + group.pop("removed_during_step") + for param in group["params"]: + param.add_(1) + + params = [nn.Parameter(torch.zeros(256, device="cuda")) for _ in range(3)] + optimizer = GroupStateOptimizer(params) + manager = ChunkedOptimizerStateOffloader( + optimizer=optimizer, + master_params=[], + chunk_size_bytes=1024, + offload_fraction=1.0, + state_dtypes=(torch.float32,), + ) + assert len(manager.chunks) == 3 + optimizer.param_groups[0]["params"] = ParamsThatMustNotBeIterated( + optimizer.param_groups[0]["params"] + ) + + manager.step() + + group = optimizer.param_groups[0] + assert group["step"].item() == 1 + assert group["updates"] == 1 + assert group["new_field"] == "preserved" + assert "removed_during_step" not in group + assert all(torch.equal(param, torch.ones_like(param)) for param in params) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA") +def test_state_staging_views_are_aligned_and_non_overlapping(): + """Flat staging pools must preserve a distinct aligned region for every state tensor.""" + + params = [nn.Parameter(torch.zeros(size, device="cuda")) for size in (3, 5, 7)] + optimizer = torch.optim.SGD(params, lr=0.1) + manager = ChunkedOptimizerStateOffloader( + optimizer=optimizer, + master_params=[], + chunk_size_bytes=0, + offload_fraction=1.0, + state_dtypes=(torch.float32,), + ) + for param in params: + cpu_state = torch.empty(param.shape, dtype=torch.float32, device="cpu", pin_memory=True) + optimizer.state[param]["moment"] = cpu_state + manager._cpu_state[param]["moment"] = cpu_state + + staging_views = manager._state_staging_views(manager.chunks[0]) + byte_ranges = [] + for _, _, _, gpu_tensor in staging_views: + start = gpu_tensor.data_ptr() + assert start % 256 == 0 + byte_ranges.append((start, start + gpu_tensor.numel() * gpu_tensor.element_size())) + + byte_ranges.sort() + assert all(left[1] <= right[0] for left, right in zip(byte_ranges, byte_ranges[1:])) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA") +def test_state_h2d_waits_for_compute_stream_before_reusing_staging_buffer(): + """H2D must not overwrite compute-stream storage while prior work still uses it.""" + + param = nn.Parameter(torch.zeros(4096, device="cuda")) + optimizer = torch.optim.SGD([param], lr=0.1) + manager = ChunkedOptimizerStateOffloader( + optimizer=optimizer, + master_params=[], + chunk_size_bytes=0, + offload_fraction=1.0, + state_dtypes=(torch.float32,), + ) + cpu_state = torch.zeros(param.shape, dtype=torch.float32, device="cpu", pin_memory=True) + optimizer.state[param]["moment"] = cpu_state + manager._cpu_state[param]["moment"] = cpu_state + + compute_stream = torch.cuda.Stream() + with torch.cuda.stream(compute_stream): + # Allocate the reusable slot on the compute stream, then make its last compute use + # finish well after the tiny CPU-to-GPU copy would finish without stream ordering. + staging_views = manager._state_staging_views(manager.chunks[0]) + manager._next_state_staging_slot = 0 + torch.cuda._sleep(50_000_000) + staging_views[0][3].fill_(7) + prefetch = manager._prefetch_state(manager.chunks[0]) + + assert prefetch.event is not None + prefetch.event.synchronize() + torch.cuda.synchronize() + torch.testing.assert_close( + optimizer.state[param]["moment"], torch.zeros_like(optimizer.state[param]["moment"]) + ) + manager.synchronize_for_checkpoint() + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA") +def test_master_h2d_allocates_on_compute_stream(): + """A restored master window must return to the compute stream's allocator pool.""" + + param = nn.Parameter(torch.ones(4096, device="cuda")) + optimizer = torch.optim.SGD([param], lr=0.1) + manager = ChunkedOptimizerStateOffloader( + optimizer=optimizer, + master_params=[param], + chunk_size_bytes=0, + offload_fraction=1.0, + state_dtypes=(torch.float32,), + ) + manager.synchronize_for_checkpoint() + assert param.device.type == "cpu" + + allocation_streams = [] + compute_stream = torch.cuda.Stream() + empty_like = torch.empty_like + + def record_allocation_stream(*args, **kwargs): + allocation_streams.append(torch.cuda.current_stream().cuda_stream) + tensor = empty_like(*args, **kwargs) + # If master H2D does not wait for this allocation/owner stream, the delayed + # sentinel write can race after the CPU master copy and remain visible. + torch.cuda._sleep(50_000_000) + tensor.fill_(7) + return tensor + + with ( + torch.cuda.stream(compute_stream), + mock.patch.object(torch, "empty_like", side_effect=record_allocation_stream), + ): + manager.prefetch_master_for_step() + + manager.ensure_master_for_param_sync() + torch.cuda.synchronize() + assert allocation_streams == [compute_stream.cuda_stream] + assert param.is_cuda + torch.testing.assert_close(param, torch.ones_like(param)) + manager.synchronize_for_checkpoint() + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA") +def test_state_eligibility_is_identical_for_d2h_and_cpu_adoption(): + """A scalar tensor is not state for a one-element parameter unless its shape matches.""" + + param = nn.Parameter(torch.zeros(1, device="cuda")) + optimizer = torch.optim.SGD([param], lr=0.1) + manager = ChunkedOptimizerStateOffloader( + optimizer=optimizer, + master_params=[], + chunk_size_bytes=0, + offload_fraction=1.0, + state_dtypes=(torch.float32,), + ) + optimizer.state[param]["moment"] = torch.ones_like(param) + optimizer.state[param]["step"] = torch.ones((), device="cuda") + optimizer.state[param]["found_inf"] = torch.ones_like(param) + + manager.synchronize_for_checkpoint() + + assert optimizer.state[param]["moment"].device.type == "cpu" + assert optimizer.state[param]["step"].is_cuda + assert optimizer.state[param]["found_inf"].is_cuda + assert "moment" in manager._cpu_state.get(param, {}) + assert "step" not in manager._cpu_state.get(param, {}) + assert "found_inf" not in manager._cpu_state.get(param, {}) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA") +def test_load_releases_cpu_buffers_from_an_old_state_schema(): + """Replacing checkpoint state must not retain obsolete pinned canonical buffers.""" + + param = nn.Parameter(torch.zeros(8, device="cuda")) + optimizer = torch.optim.SGD([param], lr=0.1) + manager = ChunkedOptimizerStateOffloader( + optimizer=optimizer, + master_params=[], + chunk_size_bytes=0, + offload_fraction=1.0, + state_dtypes=(torch.float32,), + ) + old_state = torch.empty_like(param, device="cpu", pin_memory=True) + new_state = torch.empty_like(param, device="cpu", pin_memory=True) + optimizer.state[param]["old_moment"] = old_state + manager._cpu_state[param]["old_moment"] = old_state + state_dict = optimizer.state_dict() + state_dict["state"][0] = {"new_moment": new_state} + + manager.load_state_dict_without_device_cast(state_dict) + + assert "old_moment" not in manager._cpu_state[param] + assert manager._cpu_state[param]["new_moment"] is new_state + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA") +def test_partial_load_moves_resident_tensor_state_back_to_cuda(): + """Fully-parallel CPU loads keep selected state canonical but restore resident moments.""" + + params = [nn.Parameter(torch.zeros(8, device="cuda")) for _ in range(2)] + optimizer = torch.optim.SGD(params, lr=0.1) + manager = ChunkedOptimizerStateOffloader( + optimizer=optimizer, + master_params=[], + chunk_size_bytes=0, + offload_fraction=0.5, + state_dtypes=(torch.float32,), + ) + state_dict = optimizer.state_dict() + state_dict["state"] = { + param_id: { + "moment": torch.full(param.shape, index + 1.0, dtype=param.dtype, device="cpu"), + "step": torch.tensor(index + 1), + } + for index, (param_id, param) in enumerate( + zip(state_dict["param_groups"][0]["params"], params) + ) + } + + manager.load_state_dict_without_device_cast(state_dict) + + selected_param = manager.selected_params[0] + resident_param = manager._resident_params[0] + assert optimizer.state[selected_param]["moment"].device.type == "cpu" + assert optimizer.state[selected_param]["moment"].is_pinned() + assert optimizer.state[resident_param]["moment"].is_cuda + assert optimizer.state[resident_param]["step"].device.type == "cpu" + torch.testing.assert_close( + optimizer.state[resident_param]["moment"], torch.full_like(resident_param, 2.0) + ) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA") +def test_byte_planning_uses_state_dtypes_and_prioritizes_master_bundles(): + """A partial budget uses real moment widths and spends bytes on removable masters first.""" + + native_param = nn.Parameter(torch.zeros(128, device="cuda")) + master_param = nn.Parameter(torch.zeros(128, device="cuda")) + optimizer = torch.optim.SGD([native_param, master_param], lr=0.1) + manager = ChunkedOptimizerStateOffloader( + optimizer=optimizer, + master_params=[master_param], + chunk_size_bytes=0, + offload_fraction=0.5, + state_dtypes=(torch.bfloat16, torch.bfloat16), + ) + + assert manager._estimated_state_bytes(native_param) == 128 * 2 * torch.bfloat16.itemsize + assert len(manager.selected_params) == 1 + assert manager.selected_params[0] is master_param + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA") +def test_precision_aware_master_planning_uses_real_storage(): + """Precision-aware planner accounts for every declared master using its real dtype.""" + + low_precision_param = nn.Parameter(torch.zeros(128, dtype=torch.bfloat16, device="cuda")) + native_fp32_param = nn.Parameter(torch.zeros(128, dtype=torch.float32, device="cuda")) + optimizer = torch.optim.SGD([native_fp32_param, low_precision_param], lr=0.1) + optimizer.master_weights = True + manager = ChunkedOptimizerStateOffloader( + optimizer=optimizer, + master_params=[], + chunk_size_bytes=0, + offload_fraction=1.0, + state_dtypes=(torch.bfloat16, torch.bfloat16), + optimizer_owned_master_dtypes={ + native_fp32_param: torch.float32, + low_precision_param: torch.int16, + }, + ) + + state_bytes = 128 * 2 * torch.bfloat16.itemsize + assert manager._estimated_bundle_bytes(native_fp32_param) == ( + state_bytes + 128 * torch.float32.itemsize + ) + assert manager._estimated_bundle_bytes(low_precision_param) == ( + state_bytes + 128 * torch.int16.itemsize + ) + + +def create_model_and_optimizer( + hidden_size: int = 256, + num_layers: int = 2, + chunked_optimizer_state_offload: bool = True, + chunk_size_mb: int = 1, + offload_fraction: float = 1.0, + include_native_fp32_param: bool = False, + **optimizer_kwargs, +): + """Create a bf16 DDP model and DistributedOptimizer.""" + + model = SimpleModel(hidden_size=hidden_size, num_layers=num_layers).bfloat16().cuda() + if include_native_fp32_param: + model.native_fp32_param = nn.Parameter(torch.zeros(hidden_size, device='cuda')) ddp_config = DistributedDataParallelConfig(use_distributed_optimizer=True) model = DistributedDataParallel( TransformerConfig(num_attention_heads=1, num_layers=1), ddp_config, model @@ -44,294 +940,319 @@ def create_model_and_optimizer(hidden_size=256, offload_optimizer_states=True, * bf16=True, lr=0.001, use_distributed_optimizer=True, - offload_optimizer_states=offload_optimizer_states, + chunked_optimizer_state_offload=chunked_optimizer_state_offload, + optimizer_state_offload_chunk_size_mb=chunk_size_mb, + optimizer_state_offload_fraction=offload_fraction, ) default_config.update(optimizer_kwargs) + optimizer = get_megatron_optimizer(OptimizerConfig(**default_config), [model]) + return model, optimizer - optimizer_config = OptimizerConfig(**default_config) - optim = get_megatron_optimizer(optimizer_config, [model]) - return model, optim +@pytest.mark.skipif( + not TE_FUSED_ADAM_AVAILABLE or not is_te_min_version("2.1.0.dev0"), + reason="Requires TE FusedAdam with precision-aware optimizer support", +) +@pytest.mark.usefixtures("model_parallel_context") +def test_distributed_optimizer_plans_native_fp32_optimizer_owned_master(): + """DistOpt tells the planner that TE owns a master for native-FP32 shards too.""" -def run_forward_backward_step(model, optim, hidden_size=256): - """Run a single forward-backward-step cycle.""" - input_tensor = torch.randn(8, hidden_size, dtype=torch.bfloat16, device='cuda') - output = model(input_tensor) - output.sum().backward() - optim.step() - optim.zero_grad() + model, optimizer = create_model_and_optimizer( + use_precision_aware_optimizer=True, include_native_fp32_param=True, offload_fraction=0.5 + ) + dist_optimizer = optimizer.chained_optimizers[0] + manager = dist_optimizer._optimizer_state_offloader + native_fp32_shards = [param for group in dist_optimizer.shard_fp32_groups for param in group] + + assert native_fp32_shards + assert manager is not None + assert all(manager._param_has_master(param) for param in native_fp32_shards) + assert all( + manager._optimizer_owned_master_dtypes[id(param)] == dist_optimizer.config.main_params_dtype + for param in native_fp32_shards + ) + del model, optimizer -# ============================================================================= -# Test 1: Basic OptimizerStateOffloader Initialization -# ============================================================================= -@pytest.mark.skipif(not TE_FUSED_ADAM_AVAILABLE, reason="Requires TE FusedAdam") -def test_offloader_initialization(): - """Test that OptimizerStateOffloader initializes correctly.""" - Utils.initialize_model_parallel() - model, optim = create_model_and_optimizer() - dist_optim = optim.chained_optimizers[0] +def test_group_metadata_comparison_accepts_bool_convertible_equality(): + """Non-bool scalar equality results are interpreted through their truth value.""" + + class TruthValue: + def __bool__(self): + return True + + class Metadata: + def __eq__(self, other): + return TruthValue() + + assert ChunkedOptimizerStateOffloader._group_values_equal(Metadata(), Metadata()) - # Offloader is created in __init__ when offload_optimizer_states=True - assert dist_optim._state_offloader is not None - offloader = dist_optim._state_offloader - # Verify offloader properties - assert offloader.adam_optimizer is not None - assert offloader._d2h_stream is not None - assert offloader._h2d_stream is not None - assert offloader._offloaded is False +def offload_before_forward(optimizer) -> None: + """Move selected state/master bundles to CPU and release their CUDA tensors.""" - # Before first step, optimizer states are not initialized yet - assert offloader._optimizer_states_initialized is False + optimizer.offload_optimizer_state_for_forward() + optimizer.zero_grad() - # Run one step to initialize optimizer states - run_forward_backward_step(model, optim) - # After first step, optimizer states should be marked as initialized - assert offloader._optimizer_states_initialized is True - Utils.destroy_model_parallel() +def run_forward_backward_step( + model, optimizer, hidden_size: int = 256, input_tensor: torch.Tensor | None = None +) -> None: + """Run forward/backward and one chunked optimizer update.""" + + if input_tensor is None: + input_tensor = torch.randn(8, hidden_size, dtype=torch.bfloat16, device='cuda') + output = model(input_tensor) + output.sum().backward() + optimizer.prefetch_optimizer_state_for_step() + optimizer.step() + optimizer.zero_grad() -# ============================================================================= -# Test 2: Early Master Weight Offloading Before First Step -# ============================================================================= @pytest.mark.skipif(not TE_FUSED_ADAM_AVAILABLE, reason="Requires TE FusedAdam") -def test_early_master_weight_offloading(): - """Test that master weights can be offloaded before the first optimizer step.""" - Utils.initialize_model_parallel() - model, optim = create_model_and_optimizer() - dist_optim = optim.chained_optimizers[0] - - # Offloader is created in __init__ - assert dist_optim._state_offloader is not None - offloader = dist_optim._state_offloader - - # Before first step, optimizer states are not initialized - assert offloader._optimizer_states_initialized is False - - # Capture original master weights before offload - original_master_weights = [] - for group in dist_optim.shard_fp32_from_float16_groups: - group_weights = [tensor.clone() for tensor in group] - original_master_weights.append(group_weights) - - # Offload before first step - should only offload master weights - offloader.offload() - offloader.release_gpu_memory() - torch.cuda.synchronize() +@pytest.mark.usefixtures("model_parallel_context") +def test_chunk_plan_and_initial_master_offload(): + """The manager builds bounded chunks and moves masters to CPU before the first forward.""" + + model, optimizer = create_model_and_optimizer(hidden_size=512, chunk_size_mb=1) + dist_optimizer = optimizer.chained_optimizers[0] + manager = dist_optimizer._optimizer_state_offloader + + assert manager is not None + # Direct manager tests cover exact chunk splitting independently of DP sharding. This + # integration test only needs a valid plan before checking the initial master offload. + assert manager.chunks + assert manager.selected_params + + expected_masters = { + id(param): param.detach().clone() + for group in dist_optimizer.shard_fp32_from_float16_groups + for param in group + if manager.is_param_offloaded(param) + } + offload_before_forward(optimizer) + optimizer.synchronize_optimizer_state_for_checkpoint() + + for group in dist_optimizer.shard_fp32_from_float16_groups: + for param in group: + if manager.is_param_offloaded(param): + assert param.device.type == 'cpu' + torch.testing.assert_close(param, expected_masters[id(param)].cpu()) + + # Public reload entry points restore an offloaded master window before writing it. + dist_optimizer.reload_model_params() + for group in dist_optimizer.shard_fp32_from_float16_groups: + for param in group: + if manager.is_param_offloaded(param): + assert param.is_cuda + torch.testing.assert_close(param, expected_masters[id(param)]) + + del model, optimizer - # Verify master weights were offloaded (storage resized to 0) - for group in dist_optim.shard_fp32_from_float16_groups: - for tensor in group: - assert tensor.untyped_storage().size() == 0, "Master weight should be offloaded" - - # Reload master weights - offloader.reload() - offloader.sync_before_step() - - # Verify master weights match after reload - for group_idx, group in enumerate(dist_optim.shard_fp32_from_float16_groups): - for param_idx, tensor in enumerate(group): - original = original_master_weights[group_idx][param_idx] - torch.testing.assert_close( - tensor, - original, - msg=f"Master weight [{group_idx}][{param_idx}] mismatch after offload/reload", - ) - # Now run a step and verify optimizer states can be offloaded after - run_forward_backward_step(model, optim) - assert offloader._optimizer_states_initialized is True +@pytest.mark.skipif(not TE_FUSED_ADAM_AVAILABLE, reason="Requires TE FusedAdam") +@pytest.mark.usefixtures("model_parallel_context") +def test_chunk_size_zero_uses_one_full_state_window(): + """Chunk size zero preserves full temporary GPU restore semantics.""" - Utils.destroy_model_parallel() + model, optimizer = create_model_and_optimizer(chunk_size_mb=0) + manager = optimizer.chained_optimizers[0]._optimizer_state_offloader + assert manager is not None + assert len(manager.chunks) == 1 + del model, optimizer -# ============================================================================= -# Test 3: Offload and Reload Correctness -# ============================================================================= @pytest.mark.skipif(not TE_FUSED_ADAM_AVAILABLE, reason="Requires TE FusedAdam") -@pytest.mark.parametrize("offload_optimizer_states", [True, False]) -@pytest.mark.parametrize("offload_master_weights", [True, False]) -def test_offload_reload_correctness(offload_optimizer_states, offload_master_weights): - """Test that offload/reload preserves optimizer state values.""" - if not offload_optimizer_states and not offload_master_weights: - pytest.skip("At least one offload type required") - - Utils.initialize_model_parallel() - model, optim = create_model_and_optimizer() - dist_optim = optim.chained_optimizers[0] - - # Run steps to build up optimizer state - for _ in range(3): - run_forward_backward_step(model, optim) +@pytest.mark.usefixtures("model_parallel_context") +def test_step_leaves_state_and_master_cpu_canonical(): + """A completed update leaves selected moments and masters CPU-resident before forward.""" - offloader = dist_optim._state_offloader + model, optimizer = create_model_and_optimizer(hidden_size=512, chunk_size_mb=1) + dist_optimizer = optimizer.chained_optimizers[0] + manager = dist_optimizer._optimizer_state_offloader - # Capture original states before offload - original_states = {} - for param, state in offloader.adam_optimizer.state.items(): - original_states[param] = { - k: v.clone() for k, v in state.items() if isinstance(v, torch.Tensor) - } + offload_before_forward(optimizer) + run_forward_backward_step(model, optimizer, hidden_size=512) + optimizer.offload_optimizer_state_for_forward() - # Offload - offloader.offload( - offload_optimizer_states=offload_optimizer_states, - offload_master_weights=offload_master_weights, - ) + for param in manager.selected_params: + for key, value in manager.optimizer.state[param].items(): + if isinstance(value, torch.Tensor) and value.numel() == param.numel(): + assert value.device.type == 'cpu', f"{key} was not offloaded" + + for group in dist_optimizer.shard_fp32_from_float16_groups: + for param in group: + if manager.is_param_offloaded(param): + assert param.device.type == 'cpu' + + del model, optimizer - # Release GPU memory - offloader.release_gpu_memory() - torch.cuda.synchronize() - # Reload - offloader.reload() - offloader.sync_before_step() - - # Verify states match after reload - for param, state in offloader.adam_optimizer.state.items(): - if param in original_states: - for key, original_tensor in original_states[param].items(): - if key in state and isinstance(state[key], torch.Tensor): - reloaded_tensor = state[key] - assert reloaded_tensor.device.type == 'cuda', f"State {key} should be on GPU" - torch.testing.assert_close( - reloaded_tensor, - original_tensor, - msg=f"State {key} mismatch after offload/reload", - ) - Utils.destroy_model_parallel() - - -# ============================================================================= -# Test 4: GPU Memory Release Verification -# ============================================================================= @pytest.mark.skipif(not TE_FUSED_ADAM_AVAILABLE, reason="Requires TE FusedAdam") -def test_gpu_memory_release(): - """Test that GPU memory is actually freed after release_gpu_memory().""" - Utils.initialize_model_parallel() - # Use larger model for measurable memory impact - model, optim = create_model_and_optimizer(hidden_size=1024) - dist_optim = optim.chained_optimizers[0] +@pytest.mark.usefixtures("model_parallel_context") +def test_state_can_offload_before_master_for_mxfp8_staging(): + """The delayed-master path keeps masters on CUDA until parameter staging completes.""" + + model, optimizer = create_model_and_optimizer(hidden_size=512, chunk_size_mb=1) + dist_optimizer = optimizer.chained_optimizers[0] + manager = dist_optimizer._optimizer_state_offloader + + offload_before_forward(optimizer) + run_forward_backward_step(model, optimizer, hidden_size=512) + optimizer.offload_optimizer_state_for_forward(offload_master=False) + + for param in manager.selected_params: + for key, value in manager.optimizer.state[param].items(): + if key != "master_param" and isinstance(value, torch.Tensor): + if value.numel() == param.numel(): + assert value.device.type == "cpu" + for group in dist_optimizer.shard_fp32_from_float16_groups: + for param in group: + if manager.is_param_offloaded(param): + assert param.is_cuda + + optimizer.offload_optimizer_state_for_forward() + for group in dist_optimizer.shard_fp32_from_float16_groups: + for param in group: + if manager.is_param_offloaded(param): + assert param.device.type == "cpu" + + optimizer.synchronize_optimizer_state_for_checkpoint() + del model, optimizer - # Initialize optimizer states - run_forward_backward_step(model, optim, hidden_size=1024) - offloader = dist_optim._state_offloader +@pytest.mark.skipif(not TE_FUSED_ADAM_AVAILABLE, reason="Requires TE FusedAdam") +@pytest.mark.usefixtures("model_parallel_context") +def test_offload_fraction_applies_to_state_master_bundles(): + """Partial offload selects matching state/master bundles deterministically.""" - # Measure memory before offload - torch.cuda.synchronize() - torch.cuda.empty_cache() - memory_before = torch.cuda.memory_allocated() + model, optimizer = create_model_and_optimizer( + hidden_size=512, chunk_size_mb=1, offload_fraction=0.5 + ) + dist_optimizer = optimizer.chained_optimizers[0] + manager = dist_optimizer._optimizer_state_offloader + assert manager is not None + total_bytes = sum(manager._estimated_bundle_bytes(param) for param in manager._params) + selected_bytes = sum( + manager._estimated_bundle_bytes(param) for param in manager.selected_params + ) + assert manager.selected_params + assert selected_bytes >= (total_bytes + 1) // 2 - # Offload and release - offloader.offload() - offloader.release_gpu_memory() + selected_ids = {id(param) for param in manager.selected_params} + resident_master = any( + id(param) in manager._explicit_master_param_ids and id(param) not in selected_ids + for param in manager._params + ) + selected_without_master = any( + id(param) not in manager._explicit_master_param_ids for param in manager.selected_params + ) + assert not (resident_master and selected_without_master) - # Wait for async operations - torch.cuda.synchronize() - torch.cuda.empty_cache() - memory_after = torch.cuda.memory_allocated() + offload_before_forward(optimizer) + for group in dist_optimizer.shard_fp32_from_float16_groups: + for param in group: + if id(param) in selected_ids: + assert param.device.type == 'cpu' + else: + assert param.is_cuda - # Memory should decrease - memory_freed = memory_before - memory_after - assert memory_freed > 0, f"Expected memory to be freed, but got {memory_freed} bytes difference" - Utils.destroy_model_parallel() + del model, optimizer -# ============================================================================= -# Test 5: Multiple Offload/Reload Cycles -# ============================================================================= @pytest.mark.skipif(not TE_FUSED_ADAM_AVAILABLE, reason="Requires TE FusedAdam") -def test_multiple_offload_reload_cycles(): - """Test that multiple offload/reload cycles work correctly.""" - Utils.initialize_model_parallel() - model, optim = create_model_and_optimizer() - dist_optim = optim.chained_optimizers[0] +@pytest.mark.usefixtures("model_parallel_context") +def test_partial_offload_keeps_data_parallel_ranks_in_sync(): + """Rank-local partial selections must still produce identical gathered parameters.""" - # Initialize - run_forward_backward_step(model, optim) + if torch.distributed.get_world_size() < 2: + pytest.skip("Requires multiple data-parallel ranks") - offloader = dist_optim._state_offloader + torch.manual_seed(42) + model, optimizer = create_model_and_optimizer( + hidden_size=512, chunk_size_mb=1, offload_fraction=0.5 + ) + torch.manual_seed(123) + for _ in range(2): + inputs = torch.randn(8, 512, dtype=torch.bfloat16, device="cuda") + offload_before_forward(optimizer) + run_forward_backward_step(model, optimizer, hidden_size=512, input_tensor=inputs) - # Run multiple cycles - for cycle in range(5): - # Offload - offloader.offload() - offloader.release_gpu_memory() + optimizer.offload_optimizer_state_for_forward() + for param in model.parameters(): + gathered = [torch.empty_like(param) for _ in range(torch.distributed.get_world_size())] + torch.distributed.all_gather(gathered, param) + for other in gathered[1:]: + torch.testing.assert_close(param, other, atol=0, rtol=0) - # Reload - offloader.reload() - offloader.sync_before_step() + optimizer.synchronize_optimizer_state_for_checkpoint() + del model, optimizer - # Run optimizer step - run_forward_backward_step(model, optim) - # Verify model can still produce valid outputs - input_tensor = torch.randn(8, 256, dtype=torch.bfloat16, device='cuda') - output = model(input_tensor) - assert not output.isnan().any(), "Model output contains NaN after multiple cycles" - Utils.destroy_model_parallel() +@pytest.mark.skipif(not TE_FUSED_ADAM_AVAILABLE, reason="Requires TE FusedAdam") +@pytest.mark.usefixtures("model_parallel_context") +def test_chunked_step_advances_group_step_once(): + """Calling the external optimizer for several chunks advances logical step only once.""" + + # DistOpt shards the flat parameter buffer across DP ranks. Give every rank about three + # small layers so each local shard contains several bundles and must span multiple chunks. + num_layers = 3 * torch.distributed.get_world_size() + model, optimizer = create_model_and_optimizer(num_layers=num_layers, chunk_size_mb=1) + dist_optimizer = optimizer.chained_optimizers[0] + manager = dist_optimizer._optimizer_state_offloader + assert manager is not None and len(manager.chunks) > 1 + for _ in range(3): + offload_before_forward(optimizer) + run_forward_backward_step(model, optimizer) -# ============================================================================= -# Test 6: Training Correctness with Offloading -# ============================================================================= -@pytest.mark.skipif(not TE_FUSED_ADAM_AVAILABLE, reason="Requires TE FusedAdam") -def test_training_correctness_with_offloading(): - """Test that training with offloading produces same results as without.""" - Utils.initialize_model_parallel() - torch.manual_seed(42) + steps = [ + group['step'] + for group in manager.optimizer.param_groups + if group['params'] and 'step' in group + ] + assert steps and all(step == 3 for step in steps) + + del model, optimizer - # Model 1: with offloading - model1, optim1 = create_model_and_optimizer(offload_optimizer_states=True, lr=0.01) - # Model 2: without offloading (reference) +@pytest.mark.skipif(not TE_FUSED_ADAM_AVAILABLE, reason="Requires TE FusedAdam") +@pytest.mark.usefixtures("model_parallel_context") +def test_training_matches_non_offloaded_optimizer(): + """Chunked updates remain numerically equivalent to a full FusedAdam update.""" + torch.manual_seed(42) - model2, optim2 = create_model_and_optimizer(offload_optimizer_states=False, lr=0.01) + model_offload, optimizer_offload = create_model_and_optimizer( + chunked_optimizer_state_offload=True, chunk_size_mb=1, lr=0.01 + ) + torch.manual_seed(42) + model_reference, optimizer_reference = create_model_and_optimizer( + chunked_optimizer_state_offload=False, lr=0.01 + ) - # Train both models - n_steps = 10 torch.manual_seed(123) - dist_optim1 = optim1.chained_optimizers[0] - - # Offloader is created in __init__ when offload_optimizer_states=True - assert dist_optim1._state_offloader is not None - offloader = dist_optim1._state_offloader - - for step in range(n_steps): - input_tensor = torch.randn(8, 256, dtype=torch.bfloat16, device='cuda') - - # Model 1 with offloading - # Offload states (master weights can be offloaded from the start, - # optimizer states will be skipped until after first step) - offloader.offload() - offloader.release_gpu_memory() - - output1 = model1(input_tensor) - loss1 = output1.sum() - loss1.backward() - - offloader.reload() - offloader.sync_before_step() - optim1.step() - optim1.zero_grad() - - # Model 2 without offloading - output2 = model2(input_tensor) - loss2 = output2.sum() - loss2.backward() - optim2.step() - optim2.zero_grad() - - # Compare final model weights - for (n1, p1), (n2, p2) in zip(model1.named_parameters(), model2.named_parameters()): + for _ in range(5): + inputs = torch.randn(8, 256, dtype=torch.bfloat16, device='cuda') + + offload_before_forward(optimizer_offload) + run_forward_backward_step(model_offload, optimizer_offload, input_tensor=inputs) + + output = model_reference(inputs) + output.sum().backward() + optimizer_reference.step() + optimizer_reference.zero_grad() + + optimizer_offload.offload_optimizer_state_for_forward() + for (name_offload, param_offload), (name_reference, param_reference) in zip( + model_offload.named_parameters(), model_reference.named_parameters() + ): + assert name_offload == name_reference torch.testing.assert_close( - p1.data, - p2.data, + param_offload, + param_reference, atol=1e-5, rtol=1e-4, - msg=f"Parameter {n1} mismatch between offloaded and non-offloaded training", + msg=f"Parameter {name_offload} differs from the full-step reference", ) - Utils.destroy_model_parallel() + + del model_offload, optimizer_offload, model_reference, optimizer_reference diff --git a/tests/unit_tests/training/test_train_step_schedule_plumbing.py b/tests/unit_tests/training/test_train_step_schedule_plumbing.py index a88db1aa3e0..a3bdefd278e 100644 --- a/tests/unit_tests/training/test_train_step_schedule_plumbing.py +++ b/tests/unit_tests/training/test_train_step_schedule_plumbing.py @@ -72,19 +72,291 @@ def test_train_step_defaults_to_none(): assert captured["p2p_communicator"] is None and captured["pg_collection"] is None -def test_train_step_drives_optimizer_state_offload_lifecycle(): - """The standard train path must offload, release, and reload optimizer states.""" +def test_train_step_drives_chunked_optimizer_state_offload_lifecycle(): + """MXFP8 staging must offload state first, then masters, through one API.""" + events = [] + + class FakeOptimizer: + config = SimpleNamespace(chunked_optimizer_state_offload=True) + + def offload_optimizer_state_for_forward(self, offload_master=True): + events.append(f"offload:{offload_master}") + + def prefetch_optimizer_state_for_gradient_finalization(self): + events.append("prefetch") + + def optimizer_state_offload_requires_pre_forward_param_sync(self): + return False + + def zero_grad(self): + events.append("zero_grad") + + args = SimpleNamespace( + save_params_interval=None, + save_activations_interval=None, + save_tokens_per_expert_interval=None, + save_wgrads_interval=None, + save_dgrads_interval=None, + reuse_grad_buf_for_mxfp8_param_ag=True, + overlap_param_gather=True, + seq_length=8, + global_batch_size=1, + micro_batch_size=1, + decoder_seq_length=None, + empty_unused_memory_level=0, + ) + optimizer = FakeOptimizer() + model = [ + SimpleNamespace( + force_all_reduce=False, + zero_grad_buffer=lambda: events.append("zero_grad_buffer"), + remove_forward_pre_hook_handles=[], + ) + ] + config = SimpleNamespace( + sequence_packing_scheduler=None, + finalize_model_grads_func=lambda *args, **kwargs: events.append("finalize"), + ) + + def forward_backward(**kwargs): + events.append("forward_backward") + config.finalize_model_grads_func() + return [] + + with ( + mock.patch.object(training_mod, "get_args", return_value=args), + mock.patch.object(training_mod, "get_timers", return_value=mock.MagicMock()), + mock.patch.object(training_mod, "get_rerun_state_machine", return_value=_Rerun()), + mock.patch.object(training_mod, "get_num_microbatches", return_value=1), + mock.patch.object(training_mod, "has_nvidia_modelopt", False), + mock.patch.object(training_mod, "get_moe_router_tracer", return_value=None), + ): + training_mod.train_step( + forward_step_func=lambda *args, **kwargs: None, + data_iterator=iter([]), + model=model, + optimizer=optimizer, + opt_param_scheduler=None, + config=config, + forward_backward_func=forward_backward, + iteration=0, + ) + + assert events == [ + "offload:False", + "zero_grad_buffer", + "zero_grad", + "offload:True", + "forward_backward", + "prefetch", + "finalize", + ] + + +def test_train_step_rebinds_finalize_prefetch_without_nesting_optimizers(): + """A new optimizer must replace the prefetch binding while preserving the base hook.""" + + events = [] + + class FakeOptimizer: + config = SimpleNamespace(chunked_optimizer_state_offload=True) + + def __init__(self, name): + self.name = name + + def offload_optimizer_state_for_forward(self, offload_master=True): + pass + + def prefetch_optimizer_state_for_gradient_finalization(self): + events.append(f"prefetch:{self.name}") + + def optimizer_state_offload_requires_pre_forward_param_sync(self): + return False + + def zero_grad(self): + pass + + args = SimpleNamespace( + save_params_interval=None, + save_activations_interval=None, + save_tokens_per_expert_interval=None, + save_wgrads_interval=None, + save_dgrads_interval=None, + reuse_grad_buf_for_mxfp8_param_ag=False, + overlap_param_gather=False, + seq_length=8, + global_batch_size=1, + micro_batch_size=1, + decoder_seq_length=None, + empty_unused_memory_level=0, + ) + model = [SimpleNamespace(force_all_reduce=False, zero_grad_buffer=lambda: None)] + + def base_finalize_model_grads(*args, **kwargs): + events.append("finalize") + + config = SimpleNamespace( + sequence_packing_scheduler=None, finalize_model_grads_func=base_finalize_model_grads + ) + + def forward_backward(**kwargs): + config.finalize_model_grads_func() + return [] + + optimizers = [FakeOptimizer("a"), FakeOptimizer("b")] + with ( + mock.patch.object(training_mod, "get_args", return_value=args), + mock.patch.object(training_mod, "get_timers", return_value=mock.MagicMock()), + mock.patch.object(training_mod, "get_rerun_state_machine", side_effect=_Rerun), + mock.patch.object(training_mod, "get_num_microbatches", return_value=1), + mock.patch.object(training_mod, "has_nvidia_modelopt", False), + mock.patch.object(training_mod, "get_moe_router_tracer", return_value=None), + ): + for optimizer in optimizers: + training_mod.train_step( + forward_step_func=lambda *args, **kwargs: None, + data_iterator=iter([]), + model=model, + optimizer=optimizer, + opt_param_scheduler=None, + config=config, + forward_backward_func=forward_backward, + iteration=0, + ) + + assert events == ["prefetch:a", "finalize", "prefetch:b", "finalize"] + wrapper = config.finalize_model_grads_func + assert wrapper._chunked_optimizer_state_offload_wrapped_optimizer is optimizers[-1] + assert ( + wrapper._chunked_optimizer_state_offload_base_finalize_model_grads_func + is base_finalize_model_grads + ) + + +def test_mxfp8_staging_delegates_master_restore_to_distributed_optimizer(): + """The second MXFP8 staging pass delegates master restoration to its DistOpt entry.""" + events = [] class FakeDistributedOptimizer: - def offload_states(self): + def ensure_master_weights_for_param_sync(self): + raise AssertionError("train_step must not restore DistOpt masters separately") + + def _copy_main_params_to_param_buffer(self): + events.append("distopt_stage") + + class FakeOptimizer: + config = SimpleNamespace(chunked_optimizer_state_offload=True) + + def __init__(self): + self.chained_optimizers = [SimpleNamespace(), FakeDistributedOptimizer()] + + def offload_optimizer_state_for_forward(self, offload_master=True): + events.append(f"offload:{offload_master}") + + def prefetch_optimizer_state_for_gradient_finalization(self): + events.append("prefetch") + + def optimizer_state_offload_requires_pre_forward_param_sync(self): + return False + + def ensure_master_weights_for_param_sync(self): + raise AssertionError("must not restore the whole optimizer chain") + + def zero_grad(self): + events.append("zero_grad") + + args = SimpleNamespace( + save_params_interval=None, + save_activations_interval=None, + save_tokens_per_expert_interval=None, + save_wgrads_interval=None, + save_dgrads_interval=None, + reuse_grad_buf_for_mxfp8_param_ag=True, + overlap_param_gather=True, + seq_length=8, + global_batch_size=1, + micro_batch_size=1, + decoder_seq_length=None, + empty_unused_memory_level=0, + ) + model = [ + SimpleNamespace( + force_all_reduce=False, + zero_grad_buffer=lambda: events.append("zero_grad_buffer"), + remove_forward_pre_hook_handles=[object()], + ) + ] + config = SimpleNamespace( + sequence_packing_scheduler=None, + finalize_model_grads_func=lambda *args, **kwargs: events.append("finalize"), + ) + + def forward_backward(**kwargs): + events.append("forward_backward") + config.finalize_model_grads_func() + return [] + + with ( + mock.patch.object(training_mod, "DistributedOptimizer", FakeDistributedOptimizer), + mock.patch.object(training_mod, "get_args", return_value=args), + mock.patch.object(training_mod, "get_timers", return_value=mock.MagicMock()), + mock.patch.object(training_mod, "get_rerun_state_machine", return_value=_Rerun()), + mock.patch.object(training_mod, "get_num_microbatches", return_value=1), + mock.patch.object(training_mod, "has_nvidia_modelopt", False), + mock.patch.object(training_mod, "get_moe_router_tracer", return_value=None), + ): + training_mod.train_step( + forward_step_func=lambda *args, **kwargs: None, + data_iterator=iter([]), + model=model, + optimizer=FakeOptimizer(), + opt_param_scheduler=None, + config=config, + forward_backward_func=forward_backward, + iteration=0, + ) + + assert events == [ + "offload:False", + "zero_grad_buffer", + "zero_grad", + "distopt_stage", + "offload:True", + "forward_backward", + "prefetch", + "finalize", + ] + + +def test_train_step_pre_forward_sync_only_uses_layerwise_bucket_subset(): + """Master-dependent FP8 sync must not dispatch sibling DistOpt bucket groups.""" + events = [] + + class FakeOptimizer: + config = SimpleNamespace(chunked_optimizer_state_offload=True) + + def optimizer_state_offload_requires_pre_forward_param_sync(self): + return True + + def ensure_master_weights_for_pre_forward_param_sync(self): + events.append("ensure_master") + + def ensure_master_weights_for_param_sync(self): + raise AssertionError("train_step must use the pre-forward subset restore API") + + def start_param_sync_for_bucket_group_subset(self, force_sync=False): + assert force_sync + events.append("layerwise_param_sync") + + def offload_optimizer_state_for_forward(self, offload_master=True): events.append("offload") - def release_offloaded_gpu_states(self): - events.append("release") + def prefetch_optimizer_state_for_gradient_finalization(self): + events.append("prefetch") - def reload_offloaded_states(self): - events.append("reload") + def zero_grad(self): + events.append("zero_grad") args = SimpleNamespace( save_params_interval=None, @@ -94,20 +366,18 @@ def reload_offloaded_states(self): save_dgrads_interval=None, reuse_grad_buf_for_mxfp8_param_ag=False, overlap_param_gather=False, - offload_optimizer_states=True, seq_length=8, global_batch_size=1, micro_batch_size=1, decoder_seq_length=None, empty_unused_memory_level=0, ) - dist_optimizer = FakeDistributedOptimizer() - optimizer = SimpleNamespace( - chained_optimizers=[dist_optimizer, object()], zero_grad=lambda: events.append("zero_grad") - ) + whole_model_sync = mock.Mock(side_effect=AssertionError("must not sync all DDP buckets")) model = [ SimpleNamespace( - force_all_reduce=False, zero_grad_buffer=lambda: events.append("zero_grad_buffer") + force_all_reduce=False, + zero_grad_buffer=lambda: events.append("zero_grad_buffer"), + start_param_sync=whole_model_sync, ) ] config = SimpleNamespace( @@ -127,26 +397,27 @@ def forward_backward(**kwargs): mock.patch.object(training_mod, "get_num_microbatches", return_value=1), mock.patch.object(training_mod, "has_nvidia_modelopt", False), mock.patch.object(training_mod, "get_moe_router_tracer", return_value=None), - mock.patch.object(training_mod, "DistributedOptimizer", FakeDistributedOptimizer), ): training_mod.train_step( forward_step_func=lambda *args, **kwargs: None, data_iterator=iter([]), model=model, - optimizer=optimizer, + optimizer=FakeOptimizer(), opt_param_scheduler=None, config=config, forward_backward_func=forward_backward, iteration=0, ) + whole_model_sync.assert_not_called() assert events == [ - "offload", "zero_grad_buffer", "zero_grad", - "release", + "ensure_master", + "layerwise_param_sync", + "offload", "forward_backward", - "reload", + "prefetch", "finalize", ]