diff --git a/megatron/core/distributed/distributed_data_parallel_config.py b/megatron/core/distributed/distributed_data_parallel_config.py index 56ec9e89539..c01d68e8abd 100644 --- a/megatron/core/distributed/distributed_data_parallel_config.py +++ b/megatron/core/distributed/distributed_data_parallel_config.py @@ -85,6 +85,10 @@ class DistributedDataParallelConfig: """If true, keep the compute param in fp8 (do not use any other intermediate dtype) and perform the param all-gather in fp8.""" + preserve_fp8_columnwise: bool = True + """If true, preserve FP8 columnwise parameter storage across optimizer updates. + This is required when the parameter all-gather is captured by a CUDA graph.""" + fp4_param_gather: bool = False """If true, keep the compute param in fp4 (do not use any other intermediate dtype) and perform the param all-gather in fp4.""" diff --git a/megatron/core/distributed/param_and_grad_buffer.py b/megatron/core/distributed/param_and_grad_buffer.py index 4911b4ce887..125afd511ba 100644 --- a/megatron/core/distributed/param_and_grad_buffer.py +++ b/megatron/core/distributed/param_and_grad_buffer.py @@ -330,7 +330,11 @@ def _post_param_sync(self): quantized_params = [] for bucket in self.buckets: for param in bucket.params: - if _param_uses_quantized_storage(param): + if ( + is_nvfp4tensor(param) + or is_grouped_tensor_with_quantized_storage(param) + or (self.ddp_config.preserve_fp8_columnwise and is_float8tensor(param)) + ): quantized_params.append(param) if len(quantized_params) > 0: post_all_gather_processing(quantized_params) diff --git a/megatron/core/fp8_utils.py b/megatron/core/fp8_utils.py index 1cc61ce2f38..4e9ed230626 100644 --- a/megatron/core/fp8_utils.py +++ b/megatron/core/fp8_utils.py @@ -386,6 +386,7 @@ def _quantize_param_shard_impl( start_offsets: List[int], data_parallel_group: torch.distributed.ProcessGroup, fsdp_shard_model_params: Optional[List[torch.Tensor]] = None, + preserve_columnwise: bool = True, ) -> None: if len(model_params) == 0: return @@ -409,7 +410,7 @@ def _quantize_param_shard_impl( # columnwise data and manually call post_all_gather_processing after all-gather, this # makes fp8 params compatible with CUDA graph. kwargs = {} - if te_post_all_gather_processing is not None: + if preserve_columnwise and te_post_all_gather_processing is not None: kwargs["manual_post_all_gather_processing"] = True cast_master_weights_to_fp8(*args, **kwargs) @@ -437,6 +438,7 @@ def _quantize_param_shard_impl( start_offsets: List[int], data_parallel_group: torch.distributed.ProcessGroup, fsdp_shard_model_params: Optional[List[torch.Tensor]] = None, + preserve_columnwise: bool = True, ) -> None: # Avoid circular import from megatron.core.optimizer.optimizer import _multi_tensor_copy_this_to_that @@ -527,6 +529,7 @@ def _quantize_param_shard_impl( start_offsets: List[int], data_parallel_group: torch.distributed.ProcessGroup, fsdp_shard_model_params: Optional[List[torch.Tensor]] = None, + preserve_columnwise: bool = True, ) -> None: # Avoid circular import from megatron.core.optimizer.optimizer import _multi_tensor_copy_this_to_that @@ -632,11 +635,21 @@ def modify_underlying_storage(tensor: torch.Tensor, new_raw_data: torch.Tensor): # Interface Function def quantize_param_shard( - model_params, main_params, start_offsets, data_parallel_group, fsdp_shard_model_params=None + model_params, + main_params, + start_offsets, + data_parallel_group, + fsdp_shard_model_params=None, + preserve_columnwise=True, ): """Cast shard fp32 main params to fp8 model params.""" _quantize_param_shard_impl( - model_params, main_params, start_offsets, data_parallel_group, fsdp_shard_model_params + model_params, + main_params, + start_offsets, + data_parallel_group, + fsdp_shard_model_params, + preserve_columnwise, ) diff --git a/megatron/core/optimizer/__init__.py b/megatron/core/optimizer/__init__.py index 27b675d1b8d..3c667becf6a 100644 --- a/megatron/core/optimizer/__init__.py +++ b/megatron/core/optimizer/__init__.py @@ -594,7 +594,13 @@ 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) + # TE >= 2.1.0.dev0 (the same versions that accept the + # store_param_remainders kwarg above) requires it as a + # positional arg here as well. + if is_te_min_version("2.1.0.dev0"): + opt.initialize_state(p, config.store_param_remainders) + else: + opt.initialize_state(p) elif config.optimizer == 'lion': if not HAVE_EMERGING_OPTIMIZERS: diff --git a/megatron/core/optimizer/cpu_offloading/optimizer_state_offloader.py b/megatron/core/optimizer/cpu_offloading/optimizer_state_offloader.py index 81fd116c8ba..10b40294a8a 100644 --- a/megatron/core/optimizer/cpu_offloading/optimizer_state_offloader.py +++ b/megatron/core/optimizer/cpu_offloading/optimizer_state_offloader.py @@ -2,7 +2,7 @@ """Optimizer state offloading class.""" -from typing import TYPE_CHECKING, Dict, List, Tuple +from typing import TYPE_CHECKING, Dict, List, Optional, Tuple import torch @@ -15,10 +15,11 @@ 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. + Full-reload mode overlaps D2H/H2D transfers via CUDA streams; the chunked + path is synchronous per chunk. 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 Adam optimizer state (when FusedAdam was constructed with master_weights=True) - In mcore's shard_fp32_from_float16_groups """ @@ -60,16 +61,40 @@ def __init__(self, distrib_optimizer: "DistributedOptimizer"): # List[List[cpu_tensor]] self._shard_fp32_from_float16_cpu_buffers: List[List[torch.Tensor]] = [] + # Lazily built {id(gpu master shard): cpu buffer} map for checkpoint-save reads + self._master_cpu_by_param_id: Optional[Dict[int, torch.Tensor]] = None + # State tracking self._offloaded = False self._offloaded_state_keys: Tuple[str, ...] = () self._offloaded_mcore_master_weights = False + self._d2h_inflight = False + self._h2d_pending_state_keys: Tuple[str, ...] = () + self._h2d_pending_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 initialize_offloaded_mcore_master_weights(self, cpu_buffers: List[List[torch.Tensor]]): + """Register mcore master weights that were constructed directly on CPU.""" + assert not self.optimizer_contains_master_weights + assert len(cpu_buffers) == len(self.dist_optimizer.shard_fp32_from_float16_groups) + for cpu_group, gpu_group in zip( + cpu_buffers, self.dist_optimizer.shard_fp32_from_float16_groups + ): + assert len(cpu_group) == len(gpu_group) + for cpu_tensor, gpu_tensor in zip(cpu_group, gpu_group): + assert cpu_tensor.device.type == "cpu" + assert cpu_tensor.shape == gpu_tensor.shape + assert cpu_tensor.dtype == gpu_tensor.dtype + assert gpu_tensor.is_cuda + assert gpu_tensor.untyped_storage().size() == 0 + self._shard_fp32_from_float16_cpu_buffers = cpu_buffers + self._offloaded_mcore_master_weights = True + self._offloaded = True + def mark_optimizer_states_initialized(self): """ Mark that optimizer states (exp_avg, exp_avg_sq) are now available. @@ -77,6 +102,55 @@ def mark_optimizer_states_initialized(self): """ self._optimizer_states_initialized = True + @staticmethod + def _tensor_storage_bytes(tensor: torch.Tensor) -> int: + if not isinstance(tensor, torch.Tensor): + return 0 + try: + return int(tensor.untyped_storage().size()) + except RuntimeError: + return 0 + + def _collect_memory_state(self) -> Dict[str, int]: + """Collect GPU/CPU residency byte counts. Used by unit tests to assert residency.""" + state_gpu_bytes = {key: 0 for key in (*self.OPTIMIZER_STATE_KEYS, self.MASTER_WEIGHT_KEY)} + state_cpu_bytes = {key: 0 for key in (*self.OPTIMIZER_STATE_KEYS, self.MASTER_WEIGHT_KEY)} + for param, param_state in self.adam_optimizer.state.items(): + for key in state_gpu_bytes: + tensor = param_state.get(key, None) + if isinstance(tensor, torch.Tensor) and tensor.is_cuda: + state_gpu_bytes[key] += self._tensor_storage_bytes(tensor) + for key, tensor in self._opt_state_cpu_buffers.get(param, {}).items(): + if key in state_cpu_bytes: + state_cpu_bytes[key] += self._tensor_storage_bytes(tensor) + + mcore_master_gpu_bytes = 0 + for group in self.dist_optimizer.shard_fp32_from_float16_groups: + for tensor in group: + mcore_master_gpu_bytes += self._tensor_storage_bytes(tensor) + + mcore_master_cpu_bytes = 0 + for group in self._shard_fp32_from_float16_cpu_buffers: + for tensor in group: + mcore_master_cpu_bytes += self._tensor_storage_bytes(tensor) + + allocated = torch.cuda.memory_allocated() + reserved = torch.cuda.memory_reserved() + max_allocated = torch.cuda.max_memory_allocated() + return { + "allocated": allocated, + "reserved": reserved, + "max_allocated": max_allocated, + "state_gpu_exp_avg": state_gpu_bytes['exp_avg'], + "state_gpu_exp_avg_sq": state_gpu_bytes['exp_avg_sq'], + "state_gpu_master": state_gpu_bytes[self.MASTER_WEIGHT_KEY], + "state_cpu_exp_avg": state_cpu_bytes['exp_avg'], + "state_cpu_exp_avg_sq": state_cpu_bytes['exp_avg_sq'], + "state_cpu_master": state_cpu_bytes[self.MASTER_WEIGHT_KEY], + "mcore_master_gpu": mcore_master_gpu_bytes, + "mcore_master_cpu": mcore_master_cpu_bytes, + } + def _get_state_keys_to_offload( self, offload_optimizer_states: bool, offload_master_weights: bool ) -> Tuple[str, ...]: @@ -91,6 +165,17 @@ def _get_state_keys_to_offload( keys.append(self.MASTER_WEIGHT_KEY) return tuple(keys) + def _get_initialized_state_keys_to_offload( + self, offload_optimizer_states: bool, offload_master_weights: bool + ) -> Tuple[str, ...]: + """Get state keys that may already exist for an incrementally stepped param subset.""" + keys = [] + 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: @@ -117,28 +202,44 @@ def _offload_shard_groups( 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) + self._ensure_shard_group_cpu_buffers(shard_groups, cpu_buffers, pin_memory) # Copy D2H for group_idx, group in enumerate(shard_groups): for param_idx, gpu_tensor in enumerate(group): + if ( + not isinstance(gpu_tensor, torch.Tensor) + or not gpu_tensor.is_cuda + or gpu_tensor.untyped_storage().size() == 0 + ): + continue 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 _ensure_shard_group_cpu_buffers( + self, + shard_groups: List[List[torch.Tensor]], + cpu_buffers: List[List[torch.Tensor]], + pin_memory: bool = True, + ): + """Initialize CPU buffers matching shard groups on first use.""" + if len(cpu_buffers) != 0: + return + + 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) + def _offload_states( self, offload_optimizer_states: bool, @@ -158,7 +259,11 @@ def _offload_states( continue gpu_tensor = param_state[state_key] - if not isinstance(gpu_tensor, torch.Tensor) or not gpu_tensor.is_cuda: + if ( + not isinstance(gpu_tensor, torch.Tensor) + or not gpu_tensor.is_cuda + or gpu_tensor.untyped_storage().size() == 0 + ): continue cpu_buffer = self._ensure_state_cpu_buffer( @@ -195,6 +300,40 @@ def _release_states(self): for gpu_tensor in group: gpu_tensor.untyped_storage().resize_(0) + def _release_states_for_params(self, params: List[torch.Tensor], state_keys: Tuple[str, ...]): + """Release selected optimizer state GPU tensors after their CPU copies complete.""" + states = self.adam_optimizer.state + + for param in params: + param_state = states.get(param, None) + if param_state is None: + continue + for state_key in state_keys: + state_tensor = param_state.get(state_key, None) + if isinstance(state_tensor, torch.Tensor) and state_tensor.is_cuda: + state_tensor.untyped_storage().resize_(0) + + def _reload_states_for_params( + self, params: List[torch.Tensor], state_keys: Tuple[str, ...], is_allocate_stage: bool + ): + """Reload selected optimizer state tensors from CPU buffers for a subset of params.""" + states = self.adam_optimizer.state + + for param in params: + param_state = states.get(param, None) + if param_state is None or param not in self._opt_state_cpu_buffers: + continue + for state_key in 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()) + def _reload_shard_groups( self, shard_groups: List[List[torch.Tensor]], @@ -214,6 +353,59 @@ def _reload_shard_groups( cpu_buffer, non_blocking=cpu_buffer.is_pinned() ) + def _iter_shard_group_entries_for_params( + self, + shard_groups: List[List[torch.Tensor]], + cpu_buffers: List[List[torch.Tensor]], + params: List[torch.Tensor], + ): + """Yield shard group entries whose GPU tensor is in params.""" + param_ids = {id(param) for param in params} + for group_idx, group in enumerate(shard_groups): + for param_idx, gpu_tensor in enumerate(group): + if id(gpu_tensor) in param_ids: + yield gpu_tensor, cpu_buffers[group_idx][param_idx] + + def _reload_shard_groups_for_params( + self, + params: List[torch.Tensor], + shard_groups: List[List[torch.Tensor]], + cpu_buffers: List[List[torch.Tensor]], + is_allocate_stage: bool, + ): + """Reload selected shard tensors from CPU to GPU.""" + for gpu_tensor, cpu_buffer in self._iter_shard_group_entries_for_params( + shard_groups, cpu_buffers, params + ): + if is_allocate_stage: + gpu_tensor.untyped_storage().resize_(cpu_buffer.untyped_storage().size()) + else: + gpu_tensor.copy_(cpu_buffer, non_blocking=cpu_buffer.is_pinned()) + + def _offload_shard_groups_for_params( + self, + params: List[torch.Tensor], + shard_groups: List[List[torch.Tensor]], + cpu_buffers: List[List[torch.Tensor]], + ): + """Offload selected shard tensors to their existing CPU buffers.""" + self._ensure_shard_group_cpu_buffers(shard_groups, cpu_buffers) + for gpu_tensor, cpu_buffer in self._iter_shard_group_entries_for_params( + shard_groups, cpu_buffers, params + ): + cpu_buffer.copy_(gpu_tensor, non_blocking=cpu_buffer.is_pinned()) + gpu_tensor.record_stream(self._d2h_stream) + + def _release_shard_groups_for_params( + self, params: List[torch.Tensor], shard_groups: List[List[torch.Tensor]] + ): + """Release selected shard tensors after D2H copy completes.""" + param_ids = {id(param) for param in params} + for group in shard_groups: + for gpu_tensor in group: + if id(gpu_tensor) in param_ids: + gpu_tensor.untyped_storage().resize_(0) + def _reload_states(self, is_allocate_stage: bool): """ Reload optimizer states and/or master weights from CPU to GPU. @@ -250,6 +442,81 @@ def _reload_states(self, is_allocate_stage: bool): is_allocate_stage, ) + def _has_offloaded_work(self) -> bool: + return bool(self._offloaded_state_keys) or self._offloaded_mcore_master_weights + + def _has_h2d_pending_work(self) -> bool: + return bool(self._h2d_pending_state_keys) or self._h2d_pending_mcore_master_weights + + def _mark_h2d_pending(self, state_keys: Tuple[str, ...], reload_mcore_master_weights: bool): + self._h2d_pending_state_keys = tuple( + dict.fromkeys((*self._h2d_pending_state_keys, *state_keys)) + ) + self._h2d_pending_mcore_master_weights = ( + self._h2d_pending_mcore_master_weights or reload_mcore_master_weights + ) + + def _clear_h2d_pending(self): + self._h2d_pending_state_keys = () + self._h2d_pending_mcore_master_weights = False + + def sync_pending_h2d(self): + """Synchronize pending H2D reloads before CPU-side optimizer state reads.""" + if not self._has_h2d_pending_work(): + return + self._h2d_stream.synchronize() + self._clear_h2d_pending() + + def _mcore_master_cpu_buffers_by_param_id(self) -> Dict[int, torch.Tensor]: + if not self._shard_fp32_from_float16_cpu_buffers: + return {} + if self._master_cpu_by_param_id is None: + self._master_cpu_by_param_id = { + id(gpu_tensor): cpu_buffer + for group, cpu_group in zip( + self.dist_optimizer.shard_fp32_from_float16_groups, + self._shard_fp32_from_float16_cpu_buffers, + ) + for gpu_tensor, cpu_buffer in zip(group, cpu_group) + } + return self._master_cpu_by_param_id + + def get_offloaded_states_for_read(self, main_param: torch.Tensor) -> Dict[str, torch.Tensor]: + """Return this param's offloaded states as their CPU copies, for checkpoint save. + + Only states whose GPU storage was actually released are returned (keyed by the + FusedAdam state name, plus MASTER_WEIGHT_KEY for mcore-managed master shards). + CPU bytes are the true values only for unscaled fp32/bf16 states; that is + enforced at config validation. Requires sync save: an async writer would race + the next step's chunk offloads overwriting these same CPU buffers. + """ + # An enqueued reload() flips _offloaded off before its H2D copies finish; + # callers then fall back to raw GPU reads, so those copies must complete + # first (also lets skipped-read paths retire the pending-work flags). + self.sync_pending_h2d() + if not self._offloaded: + return {} + if self._d2h_inflight: + self._d2h_stream.synchronize() + self._d2h_inflight = False + + tensors: Dict[str, torch.Tensor] = {} + param_state = self.adam_optimizer.state.get(main_param, {}) + cpu_state = self._opt_state_cpu_buffers.get(main_param, {}) + for state_key in self._offloaded_state_keys: + gpu_tensor = param_state.get(state_key, None) + if ( + isinstance(gpu_tensor, torch.Tensor) + and gpu_tensor.untyped_storage().size() == 0 + and state_key in cpu_state + ): + tensors[state_key] = cpu_state[state_key] + if self._offloaded_mcore_master_weights and main_param.untyped_storage().size() == 0: + cpu_master = self._mcore_master_cpu_buffers_by_param_id().get(id(main_param), None) + if cpu_master is not None: + tensors[self.MASTER_WEIGHT_KEY] = cpu_master + return tensors + def offload(self, offload_optimizer_states: bool = True, offload_master_weights: bool = True): """ Offload optimizer states and/or master weights to CPU. @@ -262,12 +529,130 @@ def offload(self, offload_optimizer_states: bool = True, offload_master_weights: if not offload_optimizer_states and not offload_master_weights: return + self.sync_pending_h2d() + # 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 = self._has_offloaded_work() + self._d2h_inflight = self._offloaded + + def reload_master_weights_for_params(self, params: List[torch.Tensor]): + """Synchronously reload master weights for the params needed by one chunk.""" + if not self._offloaded: + return + + reload_master_state = self.MASTER_WEIGHT_KEY in self._offloaded_state_keys + reload_mcore_master_weights = self._offloaded_mcore_master_weights + if not reload_master_state and not reload_mcore_master_weights: + return + mcore_params_to_reload = [ + param + for param in params + if isinstance(param, torch.Tensor) and param.untyped_storage().size() == 0 + ] + + self._h2d_stream.wait_stream(self._d2h_stream) + + if reload_master_state: + self._reload_states_for_params( + params, (self.MASTER_WEIGHT_KEY,), is_allocate_stage=True + ) + if reload_mcore_master_weights and mcore_params_to_reload: + self._reload_shard_groups_for_params( + mcore_params_to_reload, + self.dist_optimizer.shard_fp32_from_float16_groups, + self._shard_fp32_from_float16_cpu_buffers, + is_allocate_stage=True, + ) + + self._h2d_stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(self._h2d_stream): + if reload_master_state: + self._reload_states_for_params( + params, (self.MASTER_WEIGHT_KEY,), is_allocate_stage=False + ) + if reload_mcore_master_weights and mcore_params_to_reload: + self._reload_shard_groups_for_params( + mcore_params_to_reload, + self.dist_optimizer.shard_fp32_from_float16_groups, + self._shard_fp32_from_float16_cpu_buffers, + is_allocate_stage=False, + ) + + torch.cuda.current_stream().wait_stream(self._h2d_stream) + + def reload_optimizer_states_for_params(self, params: List[torch.Tensor]): + """Synchronously reload Adam states for the params needed by the next chunk update.""" + state_keys = tuple( + key for key in self.OPTIMIZER_STATE_KEYS if key in self._offloaded_state_keys + ) + if not state_keys: + return + + self._h2d_stream.wait_stream(self._d2h_stream) + + self._reload_states_for_params(params, state_keys, is_allocate_stage=True) + + self._h2d_stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(self._h2d_stream): + self._reload_states_for_params(params, state_keys, is_allocate_stage=False) + + torch.cuda.current_stream().wait_stream(self._h2d_stream) + + def offload_initialized_states_for_params( + self, + params: List[torch.Tensor], + offload_optimizer_states: bool = True, + offload_master_weights: bool = True, + ): + """Offload optimizer states that were just initialized for a subset of params.""" + state_keys = self._get_initialized_state_keys_to_offload( + offload_optimizer_states, offload_master_weights + ) + offload_mcore_master_weights = ( + offload_master_weights and not self.optimizer_contains_master_weights + ) + if not state_keys and not offload_mcore_master_weights: + return + + self._d2h_stream.wait_stream(torch.cuda.current_stream()) + + with torch.cuda.stream(self._d2h_stream): + states = self.adam_optimizer.state + for param in params: + param_state = states.get(param, None) + if param_state is None: + continue + for state_key in state_keys: + gpu_tensor = param_state.get(state_key, None) + 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) + cpu_buffer.copy_(gpu_tensor, non_blocking=True) + gpu_tensor.record_stream(self._d2h_stream) + if offload_mcore_master_weights: + self._offload_shard_groups_for_params( + params, + self.dist_optimizer.shard_fp32_from_float16_groups, + self._shard_fp32_from_float16_cpu_buffers, + ) + + self._d2h_stream.synchronize() + self._release_states_for_params(params, state_keys) + if offload_mcore_master_weights: + self._release_shard_groups_for_params( + params, self.dist_optimizer.shard_fp32_from_float16_groups + ) + self._offloaded_state_keys = tuple( + dict.fromkeys((*self._offloaded_state_keys, *state_keys)) + ) + self._offloaded_mcore_master_weights = ( + self._offloaded_mcore_master_weights or offload_mcore_master_weights + ) self._offloaded = True def release_gpu_memory(self): @@ -281,6 +666,9 @@ def release_gpu_memory(self): if not self._offloaded: return + if self._d2h_inflight: + self._d2h_stream.synchronize() + self._d2h_inflight = False self._release_states() def reload(self): @@ -291,6 +679,9 @@ def reload(self): if not self._offloaded: return + reload_state_keys = self._offloaded_state_keys + reload_mcore_master_weights = self._offloaded_mcore_master_weights + # Allocate GPU memory on the current stream to avoid fragmentation. self._reload_states(is_allocate_stage=True) @@ -301,6 +692,7 @@ def reload(self): with torch.cuda.stream(self._h2d_stream): self._reload_states(is_allocate_stage=False) + self._mark_h2d_pending(reload_state_keys, reload_mcore_master_weights) self._offloaded_state_keys = () self._offloaded_mcore_master_weights = False self._offloaded = False @@ -313,3 +705,4 @@ def sync_before_step(self): 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) + self._clear_h2d_pending() diff --git a/megatron/core/optimizer/distrib_optimizer.py b/megatron/core/optimizer/distrib_optimizer.py index 3f9168d88f3..5f34679f83f 100644 --- a/megatron/core/optimizer/distrib_optimizer.py +++ b/megatron/core/optimizer/distrib_optimizer.py @@ -382,6 +382,12 @@ def _build_model_and_main_param_groups( shard_float16_groups = [] shard_fp32_groups = [] shard_fp32_from_float16_groups = [] + shard_fp32_from_float16_cpu_buffers = [] + init_mcore_master_weights_on_cpu = ( + config.offload_optimizer_states + and config.offload_optimizer_states_chunk_numel > 0 + and not config.use_precision_aware_optimizer_no_fp8_or_ds_fp8 + ) # Allocate (or slice) each group's param shard. for group_range in opt_group_ranges: @@ -392,11 +398,16 @@ def _build_model_and_main_param_groups( shard_float16_params_this_group = [] shard_fp32_params_this_group = [] shard_fp32_from_float16_params_this_group = [] + shard_fp32_from_float16_cpu_buffers_this_group = [] model_float16_groups.append(model_float16_params_this_group) model_fp32_groups.append(model_fp32_params_this_group) shard_float16_groups.append(shard_float16_params_this_group) shard_fp32_groups.append(shard_fp32_params_this_group) shard_fp32_from_float16_groups.append(shard_fp32_from_float16_params_this_group) + if init_mcore_master_weights_on_cpu: + shard_fp32_from_float16_cpu_buffers.append( + shard_fp32_from_float16_cpu_buffers_this_group + ) for model_param in group_range["params"]: @@ -428,29 +439,69 @@ def _build_model_and_main_param_groups( # Generate main param. if not config.use_precision_aware_optimizer_no_fp8_or_ds_fp8: + shard_main_param_cpu = None # If we use FP8 params to initialize FP32 main params (compared to using the # bf16/fp16 params to initialize the main params), there will be a loss of # precision at the beginning of training (this problem will not occur if the # training is long enough or if the main params are loaded from a # checkpoint). - if cls._is_distopt_quantized_param(model_param) or is_nvfp4tensor( - model_param - ): - if hasattr(model_param, 'get_high_precision_init_val'): - shard_main_param = ( - model_param.get_high_precision_init_val() - .view(-1)[param_range.start : param_range.end] - .clone() - .to(model_param.device) - .float() - ) - model_param.clear_high_precision_init_val() + if init_mcore_master_weights_on_cpu: + if cls._is_distopt_quantized_param(model_param) or is_nvfp4tensor( + model_param + ): + if hasattr(model_param, 'get_high_precision_init_val'): + init_val = model_param.get_high_precision_init_val() + source_shard = init_val.view(-1)[ + param_range.start : param_range.end + ] + model_param.clear_high_precision_init_val() + else: + source_shard = model_param.float().view(-1)[ + param_range.start : param_range.end + ] else: - shard_main_param = model_param.float().view(-1)[ - param_range.start : param_range.end - ] + source_shard = shard_model_param + shard_main_param_cpu = torch.empty( + source_shard.size(), + dtype=torch.float32, + layout=source_shard.layout, + device='cpu', + pin_memory=True, + ) + shard_main_param_cpu.copy_(source_shard) + shard_main_param = torch.empty( + source_shard.size(), dtype=torch.float32, device=model_param.device + ) + shard_main_param.untyped_storage().resize_(0) + shard_fp32_from_float16_cpu_buffers_this_group.append( + shard_main_param_cpu + ) else: - shard_main_param = shard_model_param.clone().float() + if cls._is_distopt_quantized_param(model_param) or is_nvfp4tensor( + model_param + ): + if hasattr(model_param, 'get_high_precision_init_val'): + shard_main_param = ( + model_param.get_high_precision_init_val() + .view(-1)[param_range.start : param_range.end] + .clone() + .to(model_param.device) + .float() + ) + model_param.clear_high_precision_init_val() + else: + shard_main_param = model_param.float().view(-1)[ + param_range.start : param_range.end + ] + else: + shard_main_param = shard_model_param.clone().float() + + if shard_main_param_cpu is not None: + tensor_parallel.copy_tensor_model_parallel_attributes( + shard_main_param_cpu, model_param + ) + if hasattr(model_param, 'shared'): + shard_main_param_cpu.shared = model_param.shared tensor_parallel.copy_tensor_model_parallel_attributes( shard_main_param, model_param @@ -506,6 +557,7 @@ def _build_model_and_main_param_groups( shard_float16_groups, shard_fp32_groups, shard_fp32_from_float16_groups, + shard_fp32_from_float16_cpu_buffers, ) @staticmethod @@ -782,6 +834,7 @@ def __init__( self.shard_float16_groups, self.shard_fp32_groups, self.shard_fp32_from_float16_groups, + self.shard_fp32_from_float16_cpu_buffers, ) = self._build_model_and_main_param_groups( self.gbuf_ranges, self.model_param_gbuf_map, self.opt_group_ranges, config ) @@ -796,6 +849,10 @@ def __init__( if self.config.offload_optimizer_states: self._state_offloader = OptimizerStateOffloader(self) + if self.shard_fp32_from_float16_cpu_buffers: + self._state_offloader.initialize_offloaded_mcore_master_weights( + self.shard_fp32_from_float16_cpu_buffers + ) def _get_model_param_range_map(self, param: torch.nn.Parameter): """ @@ -815,6 +872,22 @@ def get_grad_stats_parallel_group(self) -> torch.distributed.ProcessGroup: """ return getattr(self, 'grad_stats_parallel_group', None) + def _inner_state_dict_for_bookkeeping(self): + """Inner optimizer state dict for callers that only consume 'param_groups'. + + TE FusedAdam.state_dict() is torch's base packing plus an unscale pass that + materializes every state tensor via get_unscaled_state(); with optimizer + state offloading the bf16->fp32 unscale kernel runs on offload-released + (zero-storage) states — an asynchronous illegal memory access. Both callers + here (state_dict / load_state_dict) drop 'state' and read only + 'param_groups', so under offloading pack raw references with the torch base + implementation instead. Unscaled state values for checkpoints come from the + per-param path (_get_main_param_and_optimizer_states), not from here. + """ + if self._state_offloader is not None: + return torch.optim.Optimizer.state_dict(self.optimizer) + return self.optimizer.state_dict() + def state_dict(self): """ The state dict contains all non-DP-rank-dependent (i.e., non-parameter- @@ -823,7 +896,7 @@ def state_dict(self): optimizer state (e.g., exp_avg, exp_avg_sq) are stored in a separate checkpoint file by calling 'save_parameter_state()'. """ - inner_state_dict = self.optimizer.state_dict() + inner_state_dict = self._inner_state_dict_for_bookkeeping() state_dict = {} # Extract 'step', for non-Apex/TE support. @@ -917,6 +990,13 @@ def load_state_dict(self, state_dict): self.optimizer.load_state_dict(state_dict) return + # Mid-run load with already-initialized states: the inner load below reads + # and rewrites the existing state tensors, which are zero-storage when + # offload-released. Same reload gate as _set_main_param_and_optimizer_states; + # no-op on the fresh-start path (states not yet offloaded). + if len(self.optimizer.state) != 0: + self._reload_offloaded_optimizer_states_for_state_dict() + if len(self.optimizer.state) == 0: if isinstance(self.optimizer, HybridDeviceOptimizer): self.optimizer.dummy_step() @@ -947,7 +1027,7 @@ def make_needed_groups(param_group): for param_group in state_dict["optimizer"]["param_groups"]: needed_groups = make_needed_groups(param_group) param_groups_map[needed_groups] = param_group - inner_state_dict = self.optimizer.state_dict() + inner_state_dict = self._inner_state_dict_for_bookkeeping() state_dict_param_groups = [] for inner_param_group in inner_state_dict["param_groups"]: needed_groups = make_needed_groups(inner_param_group) @@ -1072,10 +1152,34 @@ def make_needed_groups(param_group): else: raise NotImplementedError(f'Unknown sharding_type: {sharding_type}') + @staticmethod + def _assert_readable_state(tensor, key): + """Fail loud on raw reads of offload-released states. + + A released state keeps its shape but has zero-byte storage; any kernel or + copy launched on it is an asynchronous illegal memory access that surfaces + only at a later sync (and can silently corrupt a checkpoint written in + between). Catch it synchronously at the read site instead. + """ + assert ( + not isinstance(tensor, torch.Tensor) + or not tensor.is_cuda + or tensor.numel() == 0 + or tensor.untyped_storage().size() > 0 + ), ( + f"optimizer state '{key}' was offload-released but is being read outside " + f"the offload read gate (get_offloaded_states_for_read); this is an " + f"offload accounting bug" + ) + def _get_main_param_and_optimizer_states(self, model_param): """Return a dict containing the main param and optimizer states corresponding to the input model_param. + If optimizer states are offloaded to CPU, the offloaded tensors are returned as + their CPU copies directly, so checkpoint save never re-materializes the full + optimizer state on GPU. + The structure of the returned dict: tensors = { "param": torch.Tensor @@ -1086,23 +1190,39 @@ def _get_main_param_and_optimizer_states(self, model_param): group_index, group_order = self.model_param_group_index_map[model_param] if self.config.use_precision_aware_optimizer_no_fp8_or_ds_fp8: sharded_model_param = self.optimizer.param_groups[group_index]["params"][group_order] + offloaded = self._offloaded_states_for_read(sharded_model_param) tensors = {} for k in self.optimizer.state[sharded_model_param]: if not isinstance(self.optimizer.state[sharded_model_param][k], torch.Tensor): continue + if k in offloaded: + cpu_tensor = offloaded[k] + # Mirror get_unscaled_state(): bf16 states are checkpointed as fp32. + tensors[k] = ( + cpu_tensor.float() if cpu_tensor.dtype == torch.bfloat16 else cpu_tensor + ) + continue if isinstance(self.optimizer, HybridDeviceOptimizer): tensors[k] = self.optimizer.state[sharded_model_param][k] continue + self._assert_readable_state(self.optimizer.state[sharded_model_param][k], k) tensors[k] = self.optimizer.get_unscaled_state(sharded_model_param, k) tensors["param"] = tensors.pop("master_param") else: main_param = self.optimizer.param_groups[group_index]["params"][group_order] + offloaded = self._offloaded_states_for_read(main_param) optim_state = self.optimizer.state[main_param] - tensors = {"param": main_param} + if OptimizerStateOffloader.MASTER_WEIGHT_KEY not in offloaded: + self._assert_readable_state(main_param, OptimizerStateOffloader.MASTER_WEIGHT_KEY) + tensors = { + "param": offloaded.get(OptimizerStateOffloader.MASTER_WEIGHT_KEY, main_param) + } for k, v in optim_state.items(): if isinstance(v, torch.Tensor): - tensors[k] = v + if k not in offloaded: + self._assert_readable_state(v, k) + tensors[k] = offloaded.get(k, v) return tensors @staticmethod @@ -1179,6 +1299,18 @@ def _expand_quantized_param_shard_for_cast( return expanded_model_params, expanded_shard_main_params, expanded_start_offsets + def _reload_offloaded_optimizer_states_for_state_dict(self): + if self._state_offloader is None: + return + self._state_offloader.reload() + self._state_offloader.sync_pending_h2d() + + def _offloaded_states_for_read(self, main_param): + """CPU copies of main_param's offloaded states for checkpoint save.""" + if self._state_offloader is None: + return {} + return self._state_offloader.get_offloaded_states_for_read(main_param) + def _set_main_param_and_optimizer_states(self, model_param, tensors): """Set the main param and optimizer states corresponding to the input model_param. @@ -1189,6 +1321,7 @@ def _set_main_param_and_optimizer_states(self, model_param, tensors): "exp_avg_sq": torch.Tensor } """ + self._reload_offloaded_optimizer_states_for_state_dict() group_index, group_order = self.model_param_group_index_map[model_param] if self.config.use_precision_aware_optimizer_no_fp8_or_ds_fp8: sharded_model_param = self.optimizer.param_groups[group_index]["params"][group_order] @@ -2144,6 +2277,9 @@ def load_parameter_state_from_dp_zero_legacy(self, state_dict): method, along with `--ckpt-convert-format` and `--ckpt-convert-save` to update a legacy-format checkpoint to the modern format. """ + # This path writes into optimizer state tensors directly; materialize any + # offload-released (zero-storage) states first. + self._reload_offloaded_optimizer_states_for_state_dict() # Data parallelism variables. assert self.data_parallel_group_gloo is not None @@ -2834,6 +2970,7 @@ def _copy_main_params_to_model_params(self): expanded_shard_fp32_from_fp8, expanded_shard_offsets_in_fp8, self.data_parallel_group, + preserve_columnwise=self.ddp_config.preserve_fp8_columnwise, ) # Utility method for copying group params. @@ -2875,6 +3012,59 @@ def prepare_model_params_for_param_sync(self) -> None: model_chunk.zero_grad_buffer() self._copy_main_params_to_param_buffer() + def _copy_main_params_to_model_params_for_params( + self, params: List[torch.Tensor], copy_fp8_params: bool = True + ): + """Copy selected main params to model params after a chunked optimizer step.""" + if self.is_stub_optimizer or self.ddp_config.use_megatron_fsdp: + return + + if self.config.use_precision_aware_optimizer_no_fp8_or_ds_fp8: + return + + param_ids = {id(param) for param in params} + if copy_fp8_params: + fp8_params, shard_fp32_from_fp8, shard_offsets_in_fp8 = ( + self._get_fp8_params_and_shard_fp32_from_fp8() + ) + fp8_triplets = [ + (model_param, shard_main_param, offset) + for model_param, shard_main_param, offset in zip( + fp8_params, shard_fp32_from_fp8, shard_offsets_in_fp8 + ) + if shard_main_param is not None and id(shard_main_param) in param_ids + ] + if fp8_triplets: + raise RuntimeError( + "Chunked optimizer-state offload does not support per-chunk FP8 " + "param copy when fp8_param_gather is disabled. Enable " + "fp8_param_gather for FP8 primary weights, or use the non-chunked " + "optimizer copy path." + ) + + def copy_group_params(shard_main_groups, model_groups): + for shard_main_group, model_group in zip(shard_main_groups, model_groups): + for shard_main_param, model_param in zip(shard_main_group, model_group): + if id(shard_main_param) not in param_ids or ( + self._is_distopt_quantized_param(model_param) or is_nvfp4tensor(model_param) + ): + continue + + param_range_map = self._get_model_param_range_map(model_param) + world_range = param_range_map["gbuf_world_in_bucket"] + assert world_range.size == shard_main_param.nelement() + + gbuf_index, _, bucket_id = self.model_param_gbuf_map[model_param] + model_param_buffer = self.buffers[gbuf_index].buckets[bucket_id].param_data + + shard_model_param = model_param_buffer.view(-1)[ + world_range.start : world_range.end + ] + shard_model_param.data.copy_(shard_main_param) + + copy_group_params(self.shard_fp32_from_float16_groups, self.model_float16_groups) + copy_group_params(self.shard_fp32_groups, self.model_fp32_groups) + def _copy_main_params_to_param_buffer(self): """ This function is only used for MXFP8 params. @@ -3100,6 +3290,20 @@ def _copy_model_params_to_main_params(self, state_dict=None): state_dict ) + if ( + self._state_offloader is not None + and self._state_offloader._offloaded_mcore_master_weights + ): + assert state_dict is not None and model_param_to_state_dict_param_map is not None, ( + "Cannot initialize offloaded main params from live model params. " + "Use --load-main-params-from-ckpt when optimizer state offload is enabled " + "and optimizer state is not loaded from the checkpoint." + ) + self._copy_model_params_to_offloaded_main_params( + model_param_to_state_dict_param_map=model_param_to_state_dict_param_map + ) + return + # Utility method for copying group params. def copy_group_params(model_groups, shard_main_groups): for model_group, shard_main_group in zip(model_groups, shard_main_groups): @@ -3112,6 +3316,13 @@ def copy_group_params(model_groups, shard_main_groups): if state_dict is not None: # Use param from state_dict to initialize main_param model_param = model_param_to_state_dict_param_map[model_param] + assert not self.ddp_config.fp8_param_gather or not ( + self._is_distopt_quantized_param(model_param) + or is_nvfp4tensor(model_param) + ), ( + "fp8_param_gather must initialize master params from high-precision " + "checkpoint tensors, not quantized checkpoint/model tensors." + ) if self._is_distopt_quantized_param(model_param): if self._is_grouped_quantized_tensor(model_param): @@ -3131,6 +3342,206 @@ 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 _copy_model_params_to_offloaded_main_params(self, model_param_to_state_dict_param_map): + chunk_numel = self.config.offload_optimizer_states_chunk_numel + + def copy_one_param(model_param, shard_main_param): + param_range_map = self._get_model_param_range_map(model_param) + param_range = param_range_map["param"] + assert param_range.size == shard_main_param.nelement() + + source_param = model_param_to_state_dict_param_map[model_param] + + assert not self.ddp_config.fp8_param_gather or not self._is_distopt_quantized_param( + source_param + ), ( + "fp8_param_gather must initialize master params from high-precision " + "checkpoint tensors, not FP8 checkpoint/model tensors." + ) + + if self._is_distopt_quantized_param(source_param): + if self._is_grouped_quantized_tensor(source_param): + dequantized_source_param = source_param.float() + else: + dequantized_source_param = dequantize_fp8_tensor(source_param) + shard_model_param = dequantized_source_param.view(-1)[ + param_range.start : param_range.end + ] + else: + shard_model_param = source_param.view(-1)[param_range.start : param_range.end] + shard_main_param.data.copy_(shard_model_param) + + def flush_chunk(chunk): + if not chunk: + return + shard_main_params = [shard_main_param for _, shard_main_param in chunk] + self._state_offloader.reload_master_weights_for_params(shard_main_params) + for model_param, shard_main_param in chunk: + copy_one_param(model_param, shard_main_param) + self._state_offloader.offload_initialized_states_for_params( + shard_main_params, offload_optimizer_states=False, offload_master_weights=True + ) + + chunk = [] + chunk_size = 0 + for model_group, shard_main_group in zip( + self.model_float16_groups, self.shard_fp32_from_float16_groups + ): + for model_param, shard_main_param in zip(model_group, shard_main_group): + if shard_main_param is None: + continue + if ( + chunk_numel > 0 + and chunk + and chunk_size + shard_main_param.numel() > chunk_numel + ): + flush_chunk(chunk) + chunk = [] + chunk_size = 0 + chunk.append((model_param, shard_main_param)) + chunk_size += shard_main_param.numel() + flush_chunk(chunk) + + for model_group, shard_main_group in zip(self.model_fp32_groups, self.shard_fp32_groups): + for model_param, shard_main_param in zip(model_group, shard_main_group): + copy_one_param(model_param, shard_main_param) + + @staticmethod + def _copy_optimizer_step_value(step): + if isinstance(step, torch.Tensor): + return step.clone() + return step + + @staticmethod + def _iter_chunks_by_numel(items, numel_fn, chunk_numel): + current_chunk = [] + current_numel = 0 + + for item in items: + item_numel = numel_fn(item) + if current_chunk and current_numel + item_numel > chunk_numel: + yield current_chunk + current_chunk = [] + current_numel = 0 + current_chunk.append(item) + current_numel += item_numel + + if current_chunk: + yield current_chunk + + def _iter_optimizer_param_chunks(self, params: List[torch.Tensor]): + yield from self._iter_chunks_by_numel( + params, lambda param: param.numel(), self.config.offload_optimizer_states_chunk_numel + ) + + def _use_chunked_optimizer_state_reload(self) -> bool: + return ( + self._state_offloader is not None + and self.config.offload_optimizer_states_chunk_numel > 0 + and not self.is_stub_optimizer + ) + + def _copy_main_params_after_chunked_optimizer_step( + self, params: List[torch.Tensor], copy_fp8_params: bool = True + ): + assert not self.config.reuse_grad_buf_for_mxfp8_param_ag, ( + "Chunked optimizer-state offload does not support " + "reuse_grad_buf_for_mxfp8_param_ag. This feature is intended for MXFP8 " + "param all-gather and needs a separate chunked param-buffer copy path." + ) + self._copy_main_params_to_model_params_for_params(params, copy_fp8_params=copy_fp8_params) + + def _iter_fp8_param_copy_chunks(self): + fp8_params, shard_fp32_from_fp8, shard_offsets_in_fp8 = ( + self._get_fp8_params_and_shard_fp32_from_fp8() + ) + yield from self._iter_chunks_by_numel( + zip(fp8_params, shard_fp32_from_fp8, shard_offsets_in_fp8), + lambda entry: entry[0].numel(), + self.config.offload_optimizer_states_chunk_numel, + ) + + def _copy_fp8_params_after_chunked_optimizer_step(self): + if self.config.use_precision_aware_optimizer_no_fp8_or_ds_fp8: + # FusedAdam updates the fp8 model params in-kernel during the chunked step, so there + # is nothing to copy here (mirrors the non-chunked _copy_main_params_to_model_params). + return + if not self.ddp_config.fp8_param_gather: + return + + for fp8_chunk in self._iter_fp8_param_copy_chunks(): + shard_main_params = [shard_main_param for _, shard_main_param, _ in fp8_chunk] + params_to_reload = [ + shard_main_param + for shard_main_param in shard_main_params + if isinstance(shard_main_param, torch.Tensor) + ] + self._state_offloader.reload_master_weights_for_params(params_to_reload) + + model_params = [] + expanded_shard_main_params = [] + expanded_offsets = [] + for model_param, shard_main_param, offset in fp8_chunk: + expanded_model_params, expanded_masters, expanded_starts = ( + self._expand_quantized_param_shard_for_cast( + model_param, shard_main_param, offset + ) + ) + model_params.extend(expanded_model_params) + expanded_shard_main_params.extend(expanded_masters) + expanded_offsets.extend(expanded_starts) + + quantize_param_shard( + model_params, + expanded_shard_main_params, + expanded_offsets, + self.data_parallel_group, + preserve_columnwise=self.ddp_config.preserve_fp8_columnwise, + ) + + self._state_offloader.offload_initialized_states_for_params( + params_to_reload, offload_optimizer_states=False, offload_master_weights=True + ) + + def _optimizer_step_with_chunked_state_reload(self, reload_optimizer_states: bool): + """Run the inner optimizer step in chunks while keeping Adam states offloaded.""" + original_param_groups = self.optimizer.param_groups + next_steps = {} + + try: + for group_index, group in enumerate(original_param_groups): + starting_step = self._copy_optimizer_step_value(group.get("step", None)) + group_next_step = None + + for chunk in self._iter_optimizer_param_chunks(group["params"]): + chunk_group = group.copy() + chunk_group["params"] = chunk + if starting_step is None: + chunk_group.pop("step", None) + else: + chunk_group["step"] = self._copy_optimizer_step_value(starting_step) + + self._state_offloader.reload_master_weights_for_params(chunk) + if reload_optimizer_states: + self._state_offloader.reload_optimizer_states_for_params(chunk) + self.optimizer.param_groups = [chunk_group] + self.optimizer.step() + group_next_step = self._copy_optimizer_step_value(chunk_group.get("step", None)) + self._copy_main_params_after_chunked_optimizer_step( + chunk, copy_fp8_params=not self.ddp_config.fp8_param_gather + ) + self._state_offloader.offload_initialized_states_for_params( + chunk, offload_master_weights=True + ) + + if group_next_step is not None: + next_steps[group_index] = group_next_step + finally: + self.optimizer.param_groups = original_param_groups + + for group_index, next_step in next_steps.items(): + original_param_groups[group_index]["step"] = next_step + def start_param_sync_for_bucket_group_subset(self) -> None: """Trigger ``start_param_sync`` on DistOpt-managed bucket groups only. @@ -3167,7 +3578,29 @@ def step_with_ready_grads(self) -> bool: """ if self._state_offloader is not None: self._state_offloader.sync_before_step() - update_successful = super().step_with_ready_grads() + use_incremental_state_init = ( + self._use_chunked_optimizer_state_reload() + and not self._state_offloader._optimizer_states_initialized + ) + use_chunked_state_reload = ( + self._use_chunked_optimizer_state_reload() + and self._state_offloader._optimizer_states_initialized + ) + + if use_incremental_state_init or use_chunked_state_reload: + timers = self.config.timers + if timers is not None: + timers('optimizer-inner-step', log_level=1).start( + barrier=self.config.barrier_with_L1_time + ) + self._optimizer_step_with_chunked_state_reload(use_chunked_state_reload) + self._copy_fp8_params_after_chunked_optimizer_step() + if timers is not None: + timers('optimizer-inner-step').stop() + + update_successful = True + else: + update_successful = super().step_with_ready_grads() should_sync_params = not self.ddp_config.overlap_param_gather and not getattr( self, '_defer_param_sync', False @@ -3204,10 +3637,19 @@ def offload_states(self): self._state_offloader.offload() def reload_offloaded_states(self): - """Start async reload of offloaded states.""" - if self._state_offloader is not None: + """Start async reload of offloaded states. + + In chunked mode the reload is deferred: the optimizer step reloads master weights and + states per chunk instead. + """ + if self._state_offloader is not None and not self._use_chunked_optimizer_state_reload(): self._state_offloader.reload() + def on_step_skipped(self): + """Synchronize any pending offloaded-state reload before skipping an update.""" + if self._state_offloader is not None: + self._state_offloader.sync_pending_h2d() + def release_offloaded_gpu_states(self): """Release GPU memory after D2H completes. For delayed release case.""" if self._state_offloader is not None: diff --git a/megatron/core/optimizer/optimizer.py b/megatron/core/optimizer/optimizer.py index 52166c0cc18..f8f4fbca867 100644 --- a/megatron/core/optimizer/optimizer.py +++ b/megatron/core/optimizer/optimizer.py @@ -158,6 +158,10 @@ def __init__( self.config = config self.init_state_fn = init_state_fn + def on_step_skipped(self): + """Hook for subclasses that need cleanup when a step is skipped.""" + return + def get_parameters(self) -> List[torch.nn.Parameter]: """ Get list of parameters wrapped in optimizer. @@ -747,6 +751,7 @@ def step(self): found_inf_flag = self.prepare_grads() if found_inf_flag: + self.on_step_skipped() return False, None, None # Clip the main gradients. @@ -1145,6 +1150,7 @@ def step(self): found_inf_flag = self.prepare_grads() if found_inf_flag: + self.on_step_skipped() return False, None, None # Clip gradients. @@ -1704,6 +1710,8 @@ def step(self): self.grad_norms_by_group = {} found_inf_flag = self.prepare_grads() if found_inf_flag: + for optimizer in self.chained_optimizers: + optimizer.on_step_skipped() return False, None, None grad_norm = self.get_grad_norm() diff --git a/megatron/core/optimizer/optimizer_config.py b/megatron/core/optimizer/optimizer_config.py index 1048220953a..eff3de9bdbd 100644 --- a/megatron/core/optimizer/optimizer_config.py +++ b/megatron/core/optimizer/optimizer_config.py @@ -372,6 +372,15 @@ class OptimizerConfig: """ If True, offload optimizer states to CPU after each optimizer step and reload them before the next optimizer step. + Checkpoint saves are served from the offloaded CPU copies; the legacy + non-distributed save path reloads states to GPU first. + """ + + offload_optimizer_states_chunk_numel: int = 0 + """ + When > 0, run each offloaded optimizer step in chunks of at most this many + local parameter elements (reload, update, offload per chunk; a larger + parameter forms its own chunk). When 0, fully reload states for each step. """ ################ @@ -489,6 +498,25 @@ def __post_init__(self): self.exp_avg_sq_dtype == torch.float32 ), "exp_avg_sq_dtype can only be fp32 when not using precision-aware optimizer" + if self.offload_optimizer_states: + assert not self.optimizer_cuda_graph, ( + "offload_optimizer_states releases optimizer-state storage between " + "steps and cannot be captured by an optimizer CUDA graph" + ) + # Checkpoint save reads offloaded states straight from the CPU buffers, which + # is byte-exact only for unscaled dtypes: TE FusedAdam stores fp16/fp8 states + # scaled and store_param_remainders masters as int16 remainders. + assert self.exp_avg_dtype in (torch.float32, torch.bfloat16) and ( + self.exp_avg_sq_dtype in (torch.float32, torch.bfloat16) + ), "offload_optimizer_states requires fp32/bf16 exp_avg/exp_avg_sq dtypes" + if self.use_precision_aware_optimizer_no_fp8_or_ds_fp8: + assert ( + self.main_params_dtype == torch.float32 and not self.store_param_remainders + ), ( + "offload_optimizer_states with TE-managed master weights requires " + "main_params_dtype=fp32 and store_param_remainders=False" + ) + # Backward-compatible aliases (deprecated; use OptimizerConfig directly). AdamOptimizerConfig = OptimizerConfig diff --git a/megatron/training/arguments.py b/megatron/training/arguments.py index 577f375fcdc..d0c18f804a5 100644 --- a/megatron/training/arguments.py +++ b/megatron/training/arguments.py @@ -1979,7 +1979,9 @@ def validate_args(args, defaults={}): ) if args.load_main_params_from_ckpt: - assert args.no_load_optim, '--load-main-params-from-ckpt must be used with --no-load-optim.' + assert ( + args.no_load_optim or args.finetune + ), '--load-main-params-from-ckpt must be used with --no-load-optim or --finetune.' if args.use_dist_ckpt and args.async_save: if not args.use_persistent_ckpt_worker: @@ -2049,6 +2051,45 @@ def validate_args(args, defaults={}): assert ( not args.use_megatron_fsdp ), "offload_optimizer_states does not support Megatron-FSDP for now." + assert not args.optimizer_cuda_graph, ( + "offload_optimizer_states releases and reallocates optimizer-state " + "storage between steps; a CUDA-graph-captured optimizer step would " + "replay on freed pointers." + ) + assert not args.rl_offload_optimizer_during_inference, ( + "offload_optimizer_states is incompatible with " + "rl_offload_optimizer_during_inference, which reads optimizer state " + "directly and would touch offload-released storage." + ) + assert not args.optimizer_cpu_offload, ( + "offload_optimizer_states is incompatible with optimizer_cpu_offload " + "(HybridDeviceOptimizer); pick one." + ) + assert not args.async_save, ( + "offload_optimizer_states serves checkpoint saves from the offloaded CPU " + "copies; an async writer would race the next step's offloads overwriting " + "those same CPU buffers." + ) + if (args.no_load_optim or args.finetune) and ( + args.load is not None or args.pretrained_checkpoint is not None + ): + if not args.load_main_params_from_ckpt: + args.load_main_params_from_ckpt = True + print( + 'Warning: enabling --load-main-params-from-ckpt because ' + '--offload-optimizer-states is used without loading optimizer state.' + ) + assert ( + args.offload_optimizer_states_chunk_numel >= 0 + ), "offload_optimizer_states_chunk_numel must be non-negative" + if args.offload_optimizer_states_chunk_numel > 0: + assert not args.reuse_grad_buf_for_mxfp8_param_ag, ( + "chunked offload_optimizer_states does not support " + "reuse_grad_buf_for_mxfp8_param_ag" + ) + assert not args.fp4_param_gather, ( + "chunked offload_optimizer_states does not support fp4_param_gather" + ) if args.non_persistent_ckpt_type == "local": assert ( @@ -3752,6 +3793,14 @@ def _add_training_args(parser): 'Note that this still uses pure GPU optimizer instead of ' 'HybridDeviceOptimizer for --optimizer-cpu-offload.', ) + group.add_argument( + '--offload-optimizer-states-chunk-numel', + type=int, + default=0, + help='If positive, run each offloaded optimizer step in chunks of at ' + 'most this many local parameter elements (a larger parameter forms its ' + 'own chunk). If 0, fully reload states for each step.', + ) group.add_argument( '--dataloader-type', type=str, diff --git a/megatron/training/checkpointing.py b/megatron/training/checkpointing.py index e6965dd74f7..b806b018eb2 100644 --- a/megatron/training/checkpointing.py +++ b/megatron/training/checkpointing.py @@ -24,6 +24,7 @@ from torch.distributed.checkpoint import FileSystemReader, default_planner from megatron.core import dist_checkpointing, mpu, tensor_parallel +from megatron.core.dist_checkpointing.dict_utils import nested_values from megatron.core.dist_checkpointing.mapping import ShardedObject from megatron.core.dist_checkpointing.strategies.async_utils import _disable_gc from megatron.core.dist_checkpointing.strategies.fully_parallel import ( @@ -35,6 +36,7 @@ TorchDistSaveShardedStrategy, get_async_strategy, ) +from megatron.core.fp8_utils import dequantize_fp8_tensor, is_float8tensor from megatron.core.msc_utils import MultiStorageClientFeature, open_file from megatron.core.num_microbatches_calculator import update_num_microbatches from megatron.core.optimizer import DistributedOptimizer @@ -1483,6 +1485,8 @@ def _load_global_dist_base_checkpoint( ) if checkpointing_context is not None: checkpointing_context["load_strategy"] = load_strategy + if _should_load_fp8_skeleton_on_cpu(args): + _replace_fp8_skeleton_data_with_cpu_placeholders(sharded_state_dict) state_dict = dist_checkpointing.load( sharded_state_dict, checkpoint_name, @@ -1494,6 +1498,27 @@ def _load_global_dist_base_checkpoint( return state_dict, checkpoint_name, release, CheckpointType.GLOBAL +def _should_load_fp8_skeleton_on_cpu(args): + return ( + getattr(args, 'offload_optimizer_states', False) + and getattr(args, 'fp8_param_gather', False) + and getattr(args, 'load_main_params_from_ckpt', False) + and (getattr(args, 'no_load_optim', False) or getattr(args, 'finetune', False)) + ) + + +def _replace_fp8_skeleton_data_with_cpu_placeholders(sharded_state_dict): + """Move FP8 load targets to CPU to avoid a full-model GPU dequantization peak.""" + for value in nested_values(sharded_state_dict): + if hasattr(value, "data") and is_float8tensor(value.data): + dtype = getattr(value, "dtype", value.data.dtype) + placeholder = torch.empty(tuple(value.data.shape), dtype=dtype, device="cpu") + placeholder.copy_(dequantize_fp8_tensor(value.data)) + value.data = placeholder + if hasattr(value, "dtype"): + value.dtype = placeholder.dtype + + def _get_checkpoint_format(checkpoint_name, args): """Get the format of an existing checkpoint.""" if MultiStorageClientFeature.is_enabled(): @@ -2259,6 +2284,16 @@ def load_model_state_dict(module, state_dict, strict: bool): update_legacy_format=args.ckpt_convert_update_legacy_dist_opt_format, ) + if getattr(args, 'offload_optimizer_states', False) and optimizer is not None: + for optim_instance in getattr(optimizer, 'chained_optimizers', []): + state_offloader = getattr(optim_instance, '_state_offloader', None) + if state_offloader is not None and any( + 'exp_avg' in state for state in optim_instance.optimizer.state.values() + ): + state_offloader.mark_optimizer_states_initialized() + optim_instance.offload_states() + optim_instance.release_offloaded_gpu_states() + # Load scheduler. if opt_param_scheduler is not None: if 'lr_scheduler' in state_dict: # backward compatbility @@ -2275,6 +2310,13 @@ def load_model_state_dict(module, state_dict, strict: bool): raise e else: if (args.fp16 or args.bf16) and optimizer is not None: + if getattr(args, 'offload_optimizer_states', False) and ( + args.no_load_optim or args.finetune + ): + assert args.load_main_params_from_ckpt, ( + "--offload-optimizer-states without optimizer state requires " + "--load-main-params-from-ckpt" + ) if args.load_main_params_from_ckpt: optimizer.reload_model_params(state_dict=state_dict) else: diff --git a/megatron/training/training.py b/megatron/training/training.py index ea74b708daa..dcd799a18ac 100644 --- a/megatron/training/training.py +++ b/megatron/training/training.py @@ -2343,6 +2343,9 @@ def get_megatron_ddp_config(args: argparse.Namespace) -> DistributedDataParallel args.ddp_param_name_patterns_for_fp32_local_accumulation ) kwargs["average_in_collective"] = args.ddp_average_in_collective + kwargs["preserve_fp8_columnwise"] = ( + args.cuda_graph_impl != "none" or args.optimizer_cuda_graph + ) # Megatron-FSDP arguments. kwargs["megatron_fsdp_main_params_dtype"] = args.megatron_fsdp_main_params_dtype kwargs["megatron_fsdp_main_grads_dtype"] = args.megatron_fsdp_main_grads_dtype @@ -2486,6 +2489,12 @@ def setup_model_and_optimizer( if ( args.load is not None or args.pretrained_checkpoint is not None ) and not args.moe_use_upcycling: + if args.offload_optimizer_states and (args.no_load_optim or args.finetune): + for optim_instance in optimizer.chained_optimizers: + if isinstance(optim_instance, DistributedOptimizer): + optim_instance.offload_states() + optim_instance.release_offloaded_gpu_states() + one_logger and one_logger.log_metrics( {'load_checkpoint_start_time': one_logger_utils.get_timestamp_in_ms()} ) @@ -2664,6 +2673,12 @@ def train_step( args.save_dgrads_interval is not None and (iteration + 1) % args.save_dgrads_interval == 0 ) while rerun_state_machine.should_run_forward_backward(data_iterator): + # Offload optimizer states to CPU if enabled. + if args.offload_optimizer_states: + for optim_instance in optimizer.chained_optimizers: + if isinstance(optim_instance, DistributedOptimizer): + optim_instance.offload_states() + # Set grad to zero. for model_chunk in model: model_chunk.zero_grad_buffer() @@ -2705,6 +2720,14 @@ def train_step( if isinstance(optim_instance, DistributedOptimizer): optim_instance._copy_main_params_to_param_buffer() + # Release GPU memory for offloaded optimizer states. + # This needs to be done after _copy_main_params_to_param_buffer(). + # Separate offload and release to allow early D2H transfer to overlap with other operations. + if args.offload_optimizer_states: + for optim_instance in optimizer.chained_optimizers: + if isinstance(optim_instance, DistributedOptimizer): + optim_instance.release_offloaded_gpu_states() + # Forward pass. if save_activations_in_this_iteration: enable_activation_logging(model, args.save) @@ -3826,7 +3849,21 @@ def train( config.param_sync_func = [model_chunk.start_param_sync for model_chunk in model] if len(model) == 1: config.param_sync_func = config.param_sync_func[0] - config.finalize_model_grads_func = finalize_model_grads + + # Wrap finalize_model_grads to reload offloaded optimizer states before grad finalization. + # This allows H2D transfer to overlap with grad all-reduce. + if args.offload_optimizer_states: + + def finalize_model_grads_with_state_reload(*fmg_args, **fmg_kwargs): + # Reload offloaded states for all DistributedOptimizer instances + for optim_instance in optimizer.chained_optimizers: + if isinstance(optim_instance, DistributedOptimizer): + optim_instance.reload_offloaded_states() + return finalize_model_grads(*fmg_args, **fmg_kwargs) + + config.finalize_model_grads_func = finalize_model_grads_with_state_reload + else: + config.finalize_model_grads_func = finalize_model_grads if args.log_energy: energy_monitor.setup() diff --git a/tests/unit_tests/test_optimizer_state_offloading.py b/tests/unit_tests/test_optimizer_state_offloading.py index baaab355182..dbdeecda5b3 100644 --- a/tests/unit_tests/test_optimizer_state_offloading.py +++ b/tests/unit_tests/test_optimizer_state_offloading.py @@ -2,22 +2,65 @@ """Unit tests for OptimizerStateOffloader.""" +import os +import sys +from functools import partial +from types import SimpleNamespace +from unittest import mock + import pytest import torch import torch.nn as nn +from megatron.core import parallel_state +from megatron.core.dist_checkpointing.dict_utils import diff from megatron.core.distributed import DistributedDataParallel, DistributedDataParallelConfig -from megatron.core.optimizer import OptimizerConfig, get_megatron_optimizer +from megatron.core.enums import ModelType +from megatron.core.fp8_utils import dequantize_fp8_tensor, is_float8tensor +from megatron.core.models.gpt.gpt_layer_specs import get_gpt_layer_with_transformer_engine_spec +from megatron.core.models.gpt.gpt_model import GPTModel +from megatron.core.num_microbatches_calculator import destroy_num_microbatches_calculator +from megatron.core.optimizer import ChainedOptimizer, OptimizerConfig, get_megatron_optimizer +from megatron.core.optimizer.distrib_optimizer import DistributedOptimizer +from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed from megatron.core.transformer import TransformerConfig +from megatron.core.utils import is_te_min_version, unwrap_model +from megatron.training.arguments import core_transformer_config_from_args, parse_args, validate_args +from megatron.training.checkpointing import ( + _replace_fp8_skeleton_data_with_cpu_placeholders, + _should_load_fp8_skeleton_on_cpu, + load_checkpoint, + save_checkpoint, +) +from megatron.training.global_vars import ( + destroy_global_vars, + get_args, + set_args, + set_global_variables, +) +from megatron.training.training import get_model, setup_model_and_optimizer +from tests.unit_tests.dist_checkpointing import ( + TempNamedDir, + init_basic_mock_args, + init_checkpointing_mock_args, + initialize_gpt_model, +) from tests.unit_tests.test_utilities import Utils try: + from transformer_engine.pytorch.fp8 import check_fp8_support from transformer_engine.pytorch.optimizers import FusedAdam # noqa: F401 + FP8_AVAILABLE, REASON_FOR_NO_FP8 = check_fp8_support() TE_FUSED_ADAM_AVAILABLE = True except ImportError: + FP8_AVAILABLE = False + REASON_FOR_NO_FP8 = "Transformer Engine FP8 support is not available" TE_FUSED_ADAM_AVAILABLE = False +_SEED = 1234 +RUN_FP8_OFFLOAD_INTEGRATION = os.getenv("MEGATRON_RUN_FP8_OFFLOAD_INTEGRATION") == "1" + class SimpleModel(nn.Module): """Simple model for testing.""" @@ -31,9 +74,35 @@ def forward(self, x): return self.fc2(torch.relu(self.fc1(x))) -def create_model_and_optimizer(hidden_size=256, offload_optimizer_states=True, **optimizer_kwargs): +class FakeFloat8Tensor: + def __init__(self, shape, dtype): + self.shape = shape + self.dtype = dtype + + +class MockGroupedLinear(nn.Module): + def __init__(self, num_gemms=2, single_grouped_weight=True, single_grouped_bias=True): + super().__init__() + self.num_gemms = num_gemms + self.single_grouped_weight = single_grouped_weight + self.single_grouped_bias = single_grouped_bias + self.use_bias = True + + def _split_grouped_checkpoint_tensor(self, tensor, _key): + return list(tensor.unbind(dim=0)) + + +class MockGroupedModel(nn.Module): + def __init__(self, grouped_module): + super().__init__() + self.experts = grouped_module + + +def create_model_and_optimizer( + hidden_size=256, offload_optimizer_states=True, model_dtype=torch.bfloat16, **optimizer_kwargs +): """Helper to create model and optimizer for tests.""" - model = SimpleModel(hidden_size=hidden_size).bfloat16().cuda() + model = SimpleModel(hidden_size=hidden_size).to(dtype=model_dtype, device="cuda") ddp_config = DistributedDataParallelConfig(use_distributed_optimizer=True) model = DistributedDataParallel( TransformerConfig(num_attention_heads=1, num_layers=1), ddp_config, model @@ -41,7 +110,8 @@ def create_model_and_optimizer(hidden_size=256, offload_optimizer_states=True, * default_config = dict( optimizer='adam', - bf16=True, + bf16=model_dtype == torch.bfloat16, + fp16=model_dtype == torch.float16, lr=0.001, use_distributed_optimizer=True, offload_optimizer_states=offload_optimizer_states, @@ -53,15 +123,981 @@ def create_model_and_optimizer(hidden_size=256, offload_optimizer_states=True, * return model, optim -def run_forward_backward_step(model, optim, hidden_size=256): +def get_single_distributed_optimizer(optim): + """Return the distributed optimizer under the test wrapper.""" + return optim.chained_optimizers[0] + + +def run_forward_backward_step(model, optim, hidden_size=256, input_dtype=torch.bfloat16): """Run a single forward-backward-step cycle.""" - input_tensor = torch.randn(8, hidden_size, dtype=torch.bfloat16, device='cuda') + input_tensor = torch.randn(8, hidden_size, dtype=input_dtype, device='cuda') output = model(input_tensor) output.sum().backward() optim.step() optim.zero_grad() +def get_optimizer_memory_state(dist_optim): + """Return deterministic optimizer residency bytes.""" + offloader = dist_optim._state_offloader + if offloader is not None: + return offloader._collect_memory_state() + + state_gpu_bytes = {"exp_avg": 0, "exp_avg_sq": 0, "master_param": 0} + for param_state in dist_optim.optimizer.state.values(): + for key in state_gpu_bytes: + tensor = param_state.get(key, None) + if isinstance(tensor, torch.Tensor) and tensor.is_cuda: + state_gpu_bytes[key] += tensor.untyped_storage().size() + + mcore_master_gpu_bytes = 0 + for group in dist_optim.shard_fp32_from_float16_groups: + for tensor in group: + mcore_master_gpu_bytes += tensor.untyped_storage().size() + + return { + "state_gpu_exp_avg": state_gpu_bytes["exp_avg"], + "state_gpu_exp_avg_sq": state_gpu_bytes["exp_avg_sq"], + "state_gpu_master": state_gpu_bytes["master_param"], + "state_cpu_exp_avg": 0, + "state_cpu_exp_avg_sq": 0, + "state_cpu_master": 0, + "mcore_master_gpu": mcore_master_gpu_bytes, + "mcore_master_cpu": 0, + } + + +def get_model_param_snapshot(model): + return {name: get_tensor_snapshot(param) for name, param in model.named_parameters()} + + +def get_tensor_snapshot(tensor): + if is_float8tensor(tensor): + return dequantize_fp8_tensor(tensor).detach().float().clone() + return tensor.detach().float().clone() + + +def get_optimizer_state_snapshot(dist_optim): + # Offload-released states have zero-storage GPU tensors; read those through the + # offloader's checkpoint read gate (their CPU copies), like checkpoint save does. + offloader = dist_optim._state_offloader + + def snapshot_tensor(tensor): + return tensor.detach().float().cpu().clone() + + def snapshot_master(tensor): + if offloader is not None and tensor.untyped_storage().size() == 0: + cpu_copy = offloader.get_offloaded_states_for_read(tensor).get("master_param") + if cpu_copy is not None: + tensor = cpu_copy + return snapshot_tensor(tensor) + + snapshot = [] + for group in dist_optim.optimizer.param_groups: + group_state = {"step": group.get("step", None), "params": []} + if isinstance(group_state["step"], torch.Tensor): + group_state["step"] = group_state["step"].detach().cpu().clone() + for param in group["params"]: + param_state = dist_optim.optimizer.state.get(param, {}) + released = ( + offloader.get_offloaded_states_for_read(param) if offloader is not None else {} + ) + group_state["params"].append( + { + key: snapshot_tensor(released.get(key, value)) + for key, value in param_state.items() + if key in ("exp_avg", "exp_avg_sq", "master_param") + and isinstance(value, torch.Tensor) + } + ) + snapshot.append(group_state) + snapshot.append( + { + "mcore_master": [ + snapshot_master(tensor) + for group in dist_optim.shard_fp32_from_float16_groups + for tensor in group + ] + } + ) + return snapshot + + +def fp8_model_provider(pre_process=True, post_process=True, config=None, **_): + model_parallel_cuda_manual_seed(_SEED) + args = get_args() + if config is None: + config = core_transformer_config_from_args(args) + return GPTModel( + config=config, + transformer_layer_spec=get_gpt_layer_with_transformer_engine_spec(), + vocab_size=args.vocal_size, + max_sequence_length=args.max_position_embeddings, + pre_process=pre_process, + post_process=post_process, + fp16_lm_cross_entropy=args.fp16_lm_cross_entropy, + parallel_output=True, + share_embeddings_and_output_weights=not args.untie_embeddings_and_output_weights, + position_embedding_type=args.position_embedding_type, + rotary_percent=args.rotary_percent, + ) + + +def destroy_fp8_test_state(): + Utils.destroy_model_parallel() + destroy_global_vars() + destroy_num_microbatches_calculator() + + +def create_fp8_test_args( + mode, + train_iters, + load_main_params_from_ckpt=False, + load=None, + no_load_optim=None, + fp8_param_gather=True, +): + destroy_global_vars() + destroy_num_microbatches_calculator() + + sys.argv = ["test_optimizer_state_offloading.py"] + args = parse_args() + args.num_layers = 1 + args.vocal_size = 256 + args.hidden_size = 64 + args.num_attention_heads = 4 + args.max_position_embeddings = 64 + args.micro_batch_size = 1 + args.create_attention_mask_in_dataloader = True + args.seq_length = 64 + args.tensor_model_parallel_size = 1 + args.sequence_parallel = False + args.pipeline_model_parallel_size = 1 + args.context_parallel_size = 1 + args.train_iters = train_iters + args.lr = 3e-5 + args.optimizer = "adam" + args.neptune_project = "" + args.wandb_project = "" + args.tensorboard_dir = None + args.bf16 = True + args.add_bias_linear = False + args.swiglu = True + args.use_distributed_optimizer = True + args.fp8 = "e4m3" + args.fp8_recipe = "blockwise" + args.fp8_param_gather = fp8_param_gather + args.no_load_optim = load_main_params_from_ckpt if no_load_optim is None else no_load_optim + args.load_main_params_from_ckpt = load_main_params_from_ckpt + args.load = load + args.use_precision_aware_optimizer = True + args.main_grads_dtype = "fp32" + args.main_params_dtype = "fp32" + args.exp_avg_dtype = "bf16" + args.exp_avg_sq_dtype = "bf16" + args.offload_optimizer_states = mode != "baseline" + if mode == "chunked_state_offload": + args.offload_optimizer_states_chunk_numel = 1 + + validate_args(args) + set_global_variables(args, False) + return args + + +def get_fp8_batch(seq_length, micro_batch_size): + data = torch.arange(seq_length, dtype=torch.int64, device="cuda") + input_ids = data.repeat((micro_batch_size, 1)) + labels = (data + 1).repeat((micro_batch_size, 1)) + position_ids = data.repeat((micro_batch_size, 1)) + attention_mask = torch.ones( + (micro_batch_size, 1, seq_length, seq_length), dtype=bool, device="cuda" + ) + loss_mask = torch.ones((micro_batch_size, seq_length), device="cuda") + return input_ids, labels, position_ids, attention_mask, loss_mask + + +def get_high_precision_model_state_dict(model): + return { + name: get_tensor_snapshot(param).to(device=param.device) + for name, param in model.named_parameters() + } + + +def run_fp8_training_variant(mode, train_iters=1, load_main_params_from_ckpt=False): + destroy_fp8_test_state() + args = create_fp8_test_args(mode, train_iters, load_main_params_from_ckpt) + set_args(args) + torch.manual_seed(_SEED) + Utils.initialize_model_parallel(tensor_model_parallel_size=1) + + gpt_model, optim, _ = setup_model_and_optimizer( + fp8_model_provider, ModelType.encoder_or_decoder + ) + assert len(gpt_model) == 1 + model = gpt_model[0] + dist_optim = get_single_distributed_optimizer(optim) + offloader = dist_optim._state_offloader + + num_fp8_params = sum(1 for _, param in model.named_parameters() if is_float8tensor(param)) + assert num_fp8_params > 0 + + if load_main_params_from_ckpt: + optim.reload_model_params(state_dict={"model": get_high_precision_model_state_dict(model)}) + torch.cuda.synchronize() + + input_ids, labels, position_ids, attention_mask, loss_mask = get_fp8_batch( + args.seq_length, args.micro_batch_size + ) + losses = [] + grad_norms = [] + + for _ in range(train_iters): + model.zero_grad_buffer() + optim.zero_grad() + model.set_is_first_microbatch() + output = model( + input_ids=input_ids, + position_ids=position_ids, + attention_mask=attention_mask, + labels=labels, + loss_mask=loss_mask, + ) + loss = output.float().mean() + loss.backward() + + if mode == "state_offload" and offloader._offloaded: + dist_optim.reload_offloaded_states() + + update_successful, grad_norm, _ = optim.step() + assert update_successful + losses.append(loss.detach().float().clone()) + if isinstance(grad_norm, torch.Tensor): + grad_norms.append(grad_norm.detach().float().clone()) + else: + grad_norms.append(torch.tensor(grad_norm, dtype=torch.float32, device="cuda")) + + if mode == "state_offload": + dist_optim.offload_states() + dist_optim.release_offloaded_gpu_states() + + torch.cuda.synchronize() + residency_after_step = get_optimizer_memory_state(dist_optim) + moment_dtypes = { + state[key].dtype + for state in dist_optim.optimizer.state.values() + for key in ("exp_avg", "exp_avg_sq") + if key in state + } + offloaded_moment_dtypes = ( + { + state[key].dtype + for state in offloader._opt_state_cpu_buffers.values() + for key in ("exp_avg", "exp_avg_sq") + if key in state + } + if offloader is not None + else set() + ) + + if offloader is not None and offloader._offloaded: + dist_optim.get_parameter_state_dp_reshardable() + + torch.cuda.synchronize() + return { + "losses": losses, + "grad_norms": grad_norms, + "model": get_model_param_snapshot(model), + "optimizer": get_optimizer_state_snapshot(dist_optim), + "residency_after_step": residency_after_step, + "moment_dtypes": moment_dtypes, + "offloaded_moment_dtypes": offloaded_moment_dtypes, + } + + +def assert_optimizer_snapshots_close(actual, expected): + assert len(actual) == len(expected) + for actual_group, expected_group in zip(actual[:-1], expected[:-1]): + if isinstance(expected_group["step"], torch.Tensor): + torch.testing.assert_close(actual_group["step"], expected_group["step"]) + else: + assert actual_group["step"] == expected_group["step"] + assert len(actual_group["params"]) == len(expected_group["params"]) + for actual_param, expected_param in zip(actual_group["params"], expected_group["params"]): + assert actual_param.keys() == expected_param.keys() + for key in expected_param: + torch.testing.assert_close(actual_param[key], expected_param[key]) + + actual_mcore = actual[-1]["mcore_master"] + expected_mcore = expected[-1]["mcore_master"] + assert len(actual_mcore) == len(expected_mcore) + for actual_tensor, expected_tensor in zip(actual_mcore, expected_mcore): + torch.testing.assert_close(actual_tensor, expected_tensor) + + +def assert_residency_offloaded(state): + assert state["state_gpu_exp_avg"] == 0 + assert state["state_gpu_exp_avg_sq"] == 0 + assert state["mcore_master_gpu"] == 0 + assert state["state_cpu_exp_avg"] > 0 + assert state["state_cpu_exp_avg_sq"] > 0 + assert state["mcore_master_cpu"] > 0 + + +def test_load_fp8_skeleton_on_cpu_gate(): + args = SimpleNamespace( + offload_optimizer_states=True, + fp8_param_gather=True, + load_main_params_from_ckpt=True, + no_load_optim=True, + finetune=False, + ) + assert _should_load_fp8_skeleton_on_cpu(args) + + args.no_load_optim = False + assert not _should_load_fp8_skeleton_on_cpu(args) + + args.finetune = True + assert _should_load_fp8_skeleton_on_cpu(args) + + +def test_replace_fp8_skeleton_data_with_cpu_placeholders(monkeypatch): + fp8_data = FakeFloat8Tensor((2, 3), torch.bfloat16) + sharded_state_dict = { + "fp8": SimpleNamespace(data=fp8_data, dtype=fp8_data.dtype), + "factory": SimpleNamespace(data=FakeFloat8Tensor((4,), torch.float32)), + "plain": SimpleNamespace(data=torch.ones(1, device="cpu"), dtype=torch.float32), + } + + monkeypatch.setattr( + "megatron.training.checkpointing.is_float8tensor", + lambda data: isinstance(data, FakeFloat8Tensor), + ) + monkeypatch.setattr( + "megatron.training.checkpointing.dequantize_fp8_tensor", + lambda data: torch.full(tuple(data.shape), 3.0, dtype=data.dtype), + ) + + _replace_fp8_skeleton_data_with_cpu_placeholders(sharded_state_dict) + + fp8_placeholder = sharded_state_dict["fp8"].data + assert isinstance(fp8_placeholder, torch.Tensor) + assert fp8_placeholder.device.type == "cpu" + assert fp8_placeholder.shape == fp8_data.shape + assert fp8_placeholder.dtype == fp8_data.dtype + assert sharded_state_dict["fp8"].dtype == fp8_data.dtype + # The placeholder must hold the dequantized current weights, not + # uninitialized memory: under non-strict loading, keys the checkpoint does + # not overwrite keep these values, and torch.empty garbage here previously + # caused loss spikes after finetune loads. + torch.testing.assert_close(fp8_placeholder, torch.full((2, 3), 3.0, dtype=torch.bfloat16)) + factory_placeholder = sharded_state_dict["factory"].data + assert factory_placeholder.device.type == "cpu" + assert factory_placeholder.dtype == torch.float32 + torch.testing.assert_close(factory_placeholder, torch.full((4,), 3.0, dtype=torch.float32)) + torch.testing.assert_close(sharded_state_dict["plain"].data, torch.ones(1)) + + +def test_chained_optimizer_split_state_dict_single_model_chunk(): + class MockOptimizer: + def __init__(self, model_chunk): + self.config = None + self.model_chunks = [model_chunk] + self.reload_state_dict = None + + def reload_model_params(self, state_dict=None): + self.reload_state_dict = state_dict + + model_chunk = object() + optimizer_1 = MockOptimizer(model_chunk) + optimizer_2 = MockOptimizer(model_chunk) + chained_optimizer = ChainedOptimizer([optimizer_1, optimizer_2]) + + state_dict = {"weight": torch.ones(1)} + chained_optimizer.reload_model_params(state_dict) + + assert optimizer_1.reload_state_dict is state_dict + assert optimizer_2.reload_state_dict is state_dict + + +def test_normalize_grouped_state_dict_stacks_indexed_tensors(): + model = MockGroupedModel( + MockGroupedLinear(single_grouped_weight=True, single_grouped_bias=True) + ) + state_dict = { + "decoder.experts.weight0": torch.ones(2, 3), + "decoder.experts.weight1": torch.full((2, 3), 2.0), + "decoder.experts.bias0": torch.ones(2), + "decoder.experts.bias1": torch.full((2,), 2.0), + } + + DistributedOptimizer._normalize_state_dict_for_grouped_params(state_dict, model) + + assert set(state_dict) == {"decoder.experts.weight", "decoder.experts.bias"} + torch.testing.assert_close( + state_dict["decoder.experts.weight"], + torch.stack([torch.ones(2, 3), torch.full((2, 3), 2.0)], dim=0), + ) + torch.testing.assert_close( + state_dict["decoder.experts.bias"], + torch.stack([torch.ones(2), torch.full((2,), 2.0)], dim=0), + ) + + +def test_normalize_grouped_state_dict_splits_grouped_tensors(): + model = MockGroupedModel( + MockGroupedLinear(single_grouped_weight=False, single_grouped_bias=False) + ) + state_dict = { + "decoder.experts.weight": torch.stack([torch.ones(2, 3), torch.full((2, 3), 2.0)], dim=0), + "decoder.experts.bias": torch.stack([torch.ones(2), torch.full((2,), 2.0)], dim=0), + } + + DistributedOptimizer._normalize_state_dict_for_grouped_params(state_dict, model) + + assert set(state_dict) == { + "decoder.experts.weight0", + "decoder.experts.weight1", + "decoder.experts.bias0", + "decoder.experts.bias1", + } + torch.testing.assert_close(state_dict["decoder.experts.weight0"], torch.ones(2, 3)) + torch.testing.assert_close(state_dict["decoder.experts.weight1"], torch.full((2, 3), 2.0)) + torch.testing.assert_close(state_dict["decoder.experts.bias0"], torch.ones(2)) + torch.testing.assert_close(state_dict["decoder.experts.bias1"], torch.full((2,), 2.0)) + + +def _load_checkpoint_no_arg_checks(*args, **kwargs): + with mock.patch('megatron.training.checkpointing.check_checkpoint_args'): + with mock.patch('megatron.training.checkpointing.update_num_microbatches'): + return load_checkpoint(*args, **kwargs) + + +def _check_equal_dp_zero_state(state_a, state_b): + if parallel_state.get_data_parallel_rank(with_context_parallel=True) == 0: + diffs = diff(state_a, state_b) + is_equal = not any(map(bool, diffs)) + else: + diffs = None + is_equal = True + + all_equal = torch.tensor(int(is_equal), device='cuda') + torch.distributed.all_reduce(all_equal, op=torch.distributed.ReduceOp.MIN) + if not bool(all_equal.item()): + raise RuntimeError(f'[{Utils.rank}] {diffs}') + + +@pytest.mark.skipif(not TE_FUSED_ADAM_AVAILABLE, reason="Requires TE FusedAdam") +def test_chunked_offload_optimizer_save_load(tmp_path_dist_ckpt): + """Offloaded optimizer states must survive an fs_model_space save/load round trip.""" + Utils.initialize_model_parallel() + + def assert_masters_offloaded(dist_optim): + for group in dist_optim.shard_fp32_from_float16_groups: + for tensor in group: + assert tensor.untyped_storage().size() == 0 + + def setup_offloaded_model_and_optimizer(seed, with_initialized_states): + builder_args = parse_args(ignore_unknown_args=True) + with mock.patch('megatron.training.training.get_args', new=lambda: builder_args): + init_basic_mock_args(builder_args, tp=1, pp=1, bf16=True) + model = get_model( + partial( + initialize_gpt_model, + seed=seed, + tensor_model_parallel_size=1, + pipeline_model_parallel_size=1, + pipeline_dtype=torch.bfloat16, + bf16=True, + ) + ) + config = OptimizerConfig( + bf16=True, + params_dtype=torch.bfloat16, + use_distributed_optimizer=True, + offload_optimizer_states=True, + offload_optimizer_states_chunk_numel=1, + ) + optimizer = get_megatron_optimizer(config, model) + + if with_initialized_states: + # Mimic mid-training state on the save side: Adam states exist and + # everything is offloaded. The load side keeps the real startup state + # (lazy Adam states not initialized yet, only masters offloaded). + torch.manual_seed(seed + 1) + model_parallel_cuda_manual_seed(seed + 1) + for group in optimizer.optimizer.param_groups: + for param in group['params']: + if len(optimizer.optimizer.state[param]) == 0: + optimizer.optimizer.state[param]['exp_avg'] = torch.rand_like(param.data) + optimizer.optimizer.state[param]['exp_avg_sq'] = torch.rand_like(param.data) + + dist_optim = optimizer.chained_optimizers[0] + dist_optim._state_offloader.mark_optimizer_states_initialized() + dist_optim.offload_states() + dist_optim.release_offloaded_gpu_states() + torch.cuda.synchronize() + return unwrap_model(model), optimizer + + with TempNamedDir( + tmp_path_dist_ckpt / 'test_chunked_offload_optimizer_save_load', sync=True + ) as ckpt_dir: + mock_args = parse_args(ignore_unknown_args=True) + with mock.patch('megatron.training.checkpointing.get_args', new=lambda: mock_args): + init_basic_mock_args(mock_args, tp=1, pp=1) + init_checkpointing_mock_args(mock_args, ckpt_dir) + mock_args.offload_optimizer_states = True + + model_a, optimizer_a = setup_offloaded_model_and_optimizer( + seed=2, with_initialized_states=True + ) + dist_optim_a = optimizer_a.chained_optimizers[0] + assert_masters_offloaded(dist_optim_a) + + from megatron.training.training import preprocess_common_state_dict + + save_checkpoint( + 10, + model_a, + optimizer_a, + None, + 0, + preprocess_common_state_dict_fn=preprocess_common_state_dict, + ) + state_a = dist_optim_a.get_parameter_state_dp_zero(use_gloo_comm=False) + + # Both save consumers above must be served from the offloader CPU copies: + # nothing gets re-materialized on GPU by the save itself. + assert_masters_offloaded(dist_optim_a) + for param_state in dist_optim_a.optimizer.state.values(): + for key in ('exp_avg', 'exp_avg_sq'): + moment = param_state.get(key) + if isinstance(moment, torch.Tensor): + assert moment.untyped_storage().size() == 0 + + model_b, optimizer_b = setup_offloaded_model_and_optimizer( + seed=3, with_initialized_states=False + ) + dist_optim_b = optimizer_b.chained_optimizers[0] + assert_masters_offloaded(dist_optim_b) + _load_checkpoint_no_arg_checks(model_b, optimizer_b, None) + + state_b = dist_optim_b.get_parameter_state_dp_zero(use_gloo_comm=False) + _check_equal_dp_zero_state(state_a, state_b) + + # The loaded values must also survive a fresh offload/reload cycle, i.e. they + # must land in the offloader CPU buffers, not only in the live GPU tensors. + dist_optim_b.offload_states() + dist_optim_b.release_offloaded_gpu_states() + torch.cuda.synchronize() + assert_masters_offloaded(dist_optim_b) + state_b_after_offload = dist_optim_b.get_parameter_state_dp_zero(use_gloo_comm=False) + _check_equal_dp_zero_state(state_a, state_b_after_offload) + + Utils.destroy_model_parallel() + + +def run_training_variant(mode, inputs, hidden_size=64): + torch.manual_seed(1234) + offload_optimizer_states = mode != "baseline" + optimizer_kwargs = {"lr": 0.01} + if mode == "chunked_state_offload": + optimizer_kwargs["offload_optimizer_states_chunk_numel"] = 1 + model, optim = create_model_and_optimizer( + hidden_size=hidden_size, + offload_optimizer_states=offload_optimizer_states, + **optimizer_kwargs, + ) + dist_optim = get_single_distributed_optimizer(optim) + losses = [] + + for input_tensor in inputs: + output = model(input_tensor.clone()) + loss = output.float().sum() + loss.backward() + + if mode == "state_offload" and dist_optim._state_offloader._offloaded: + dist_optim.reload_offloaded_states() + + optim.step() + optim.zero_grad() + losses.append(loss.detach().float().clone()) + + if mode == "state_offload": + dist_optim.offload_states() + dist_optim.release_offloaded_gpu_states() + + torch.cuda.synchronize() + residency_after_step = get_optimizer_memory_state(dist_optim) + + if dist_optim._state_offloader is not None and dist_optim._state_offloader._offloaded: + dist_optim.get_parameter_state_dp_reshardable() + + torch.cuda.synchronize() + return { + "losses": losses, + "model": get_model_param_snapshot(model), + "optimizer": get_optimizer_state_snapshot(dist_optim), + "residency_after_step": residency_after_step, + } + + +@pytest.mark.skipif(not TE_FUSED_ADAM_AVAILABLE, reason="Requires TE FusedAdam") +def test_baseline_full_offload_and_chunked_offload_e2e_equivalence(): + """Compare baseline, full offload, and chunked offload training semantics.""" + Utils.initialize_model_parallel() + try: + hidden_size = 64 + torch.manual_seed(5678) + inputs = [ + torch.randn(8, hidden_size, dtype=torch.bfloat16, device='cuda') for _ in range(3) + ] + + baseline = run_training_variant("baseline", inputs, hidden_size) + state_offload = run_training_variant("state_offload", inputs, hidden_size) + chunked = run_training_variant("chunked_state_offload", inputs, hidden_size) + + for variant in (state_offload, chunked): + for actual_loss, expected_loss in zip(variant["losses"], baseline["losses"]): + torch.testing.assert_close(actual_loss, expected_loss, rtol=1e-4, atol=1e-4) + assert variant["model"].keys() == baseline["model"].keys() + for name in baseline["model"]: + torch.testing.assert_close( + variant["model"][name], + baseline["model"][name], + rtol=1e-4, + atol=1e-3, + msg=f"model parameter {name} mismatch", + ) + assert_optimizer_snapshots_close(variant["optimizer"], baseline["optimizer"]) + + baseline_state = baseline["residency_after_step"] + chunked_state = chunked["residency_after_step"] + full_offload_state = state_offload["residency_after_step"] + assert baseline_state["state_gpu_exp_avg"] > 0 + assert baseline_state["state_gpu_exp_avg_sq"] > 0 + assert baseline_state["mcore_master_gpu"] > 0 + assert_residency_offloaded(full_offload_state) + assert_residency_offloaded(chunked_state) + finally: + Utils.destroy_model_parallel() + + +@pytest.mark.skipif(not TE_FUSED_ADAM_AVAILABLE, reason="Requires TE FusedAdam") +@pytest.mark.skipif(not FP8_AVAILABLE, reason=REASON_FOR_NO_FP8) +@pytest.mark.skipif(not is_te_min_version("2.4.0.dev0"), reason="TE 2.4.0.dev0 is required") +@pytest.mark.timeout(180) +def test_fp8_param_gather_chunked_offload_load_main_params_smoke(): + try: + chunked = run_fp8_training_variant("chunked_state_offload", load_main_params_from_ckpt=True) + + assert len(chunked["losses"]) == 1 + assert len(chunked["grad_norms"]) == 1 + assert_residency_offloaded(chunked["residency_after_step"]) + finally: + destroy_fp8_test_state() + + +@pytest.mark.skipif(not TE_FUSED_ADAM_AVAILABLE, reason="Requires TE FusedAdam") +@pytest.mark.skipif(not FP8_AVAILABLE, reason=REASON_FOR_NO_FP8) +@pytest.mark.skipif(not is_te_min_version("2.4.0.dev0"), reason="TE 2.4.0.dev0 is required") +@pytest.mark.skipif( + not RUN_FP8_OFFLOAD_INTEGRATION, + reason="Set MEGATRON_RUN_FP8_OFFLOAD_INTEGRATION=1 to run this FP8 integration test", +) +@pytest.mark.timeout(300) +def test_fp8_blockwise_full_offload_and_chunked_offload_integration(): + """Compare FP8 blockwise baseline, full offload, and chunked offload semantics.""" + try: + baseline = run_fp8_training_variant("baseline") + state_offload = run_fp8_training_variant("state_offload") + chunked = run_fp8_training_variant("chunked_state_offload") + + for variant in (state_offload, chunked): + for actual_loss, expected_loss in zip(variant["losses"], baseline["losses"]): + torch.testing.assert_close(actual_loss, expected_loss, rtol=1e-4, atol=1e-4) + for actual_grad_norm, expected_grad_norm in zip( + variant["grad_norms"], baseline["grad_norms"] + ): + torch.testing.assert_close( + actual_grad_norm, expected_grad_norm, rtol=1e-4, atol=1e-4 + ) + assert variant["model"].keys() == baseline["model"].keys() + for name in baseline["model"]: + torch.testing.assert_close( + variant["model"][name], + baseline["model"][name], + rtol=1e-3, + atol=1e-2, + msg=f"FP8 model parameter {name} mismatch", + ) + assert_optimizer_snapshots_close(variant["optimizer"], baseline["optimizer"]) + + baseline_state = baseline["residency_after_step"] + full_offload_state = state_offload["residency_after_step"] + chunked_state = chunked["residency_after_step"] + assert baseline_state["state_gpu_exp_avg"] > 0 + assert baseline_state["state_gpu_exp_avg_sq"] > 0 + assert baseline_state["mcore_master_gpu"] > 0 + assert_residency_offloaded(full_offload_state) + assert_residency_offloaded(chunked_state) + finally: + destroy_fp8_test_state() + + +@pytest.mark.skipif(not TE_FUSED_ADAM_AVAILABLE, reason="Requires TE FusedAdam") +def test_offload_optimizer_states_checkpoint_without_optimizer_forces_main_params_from_ckpt(): + try: + args = create_fp8_test_args( + "state_offload", + train_iters=1, + load="/tmp/checkpoint", + no_load_optim=True, + fp8_param_gather=False, + ) + assert args.load_main_params_from_ckpt + finally: + destroy_fp8_test_state() + + +@pytest.mark.skipif(not TE_FUSED_ADAM_AVAILABLE, reason="Requires TE FusedAdam") +def test_offload_optimizer_states_checkpoint_with_optimizer_keeps_main_params_reload_off(): + try: + args = create_fp8_test_args( + "state_offload", + train_iters=1, + load="/tmp/checkpoint", + no_load_optim=False, + fp8_param_gather=False, + ) + assert not args.load_main_params_from_ckpt + finally: + destroy_fp8_test_state() + + +@pytest.mark.skipif(not TE_FUSED_ADAM_AVAILABLE, reason="Requires TE FusedAdam") +def test_chunked_offload_initializes_mcore_master_weights_on_cpu(): + """Check chunked offload starts with mcore master weights already offloaded.""" + Utils.initialize_model_parallel() + try: + model, optim = create_model_and_optimizer( + hidden_size=64, offload_optimizer_states_chunk_numel=1 + ) + dist_optim = get_single_distributed_optimizer(optim) + offloader = dist_optim._state_offloader + + state = offloader._collect_memory_state() + assert offloader._offloaded is True + assert offloader._offloaded_mcore_master_weights is True + assert offloader._optimizer_states_initialized is False + assert state["mcore_master_gpu"] == 0 + assert state["mcore_master_cpu"] > 0 + assert state["state_cpu_exp_avg"] == 0 + assert state["state_cpu_exp_avg_sq"] == 0 + + original_cpu_buffers = [ + [tensor.clone() for tensor in group] + for group in offloader._shard_fp32_from_float16_cpu_buffers + ] + for group in dist_optim.shard_fp32_from_float16_groups: + for tensor in group: + assert tensor.untyped_storage().size() == 0 + + offloader.reload() + offloader.sync_before_step() + + state = offloader._collect_memory_state() + assert offloader._offloaded is False + assert state["mcore_master_gpu"] == state["mcore_master_cpu"] + for group_idx, group in enumerate(dist_optim.shard_fp32_from_float16_groups): + for param_idx, tensor in enumerate(group): + assert tensor.untyped_storage().size() > 0 + torch.testing.assert_close(tensor.cpu(), original_cpu_buffers[group_idx][param_idx]) + finally: + Utils.destroy_model_parallel() + + +@pytest.mark.skipif(not TE_FUSED_ADAM_AVAILABLE, reason="Requires TE FusedAdam") +def test_chunked_offload_releases_mcore_master_weights_and_saves_from_cpu(): + """Check chunked residency, offload idempotence, and that checkpoint reads + are served from the CPU copies without restoring GPU residency.""" + Utils.initialize_model_parallel() + try: + model, optim = create_model_and_optimizer( + hidden_size=64, offload_optimizer_states_chunk_numel=1 + ) + dist_optim = get_single_distributed_optimizer(optim) + offloader = dist_optim._state_offloader + + run_forward_backward_step(model, optim, hidden_size=64) + + state = offloader._collect_memory_state() + assert offloader._offloaded is True + assert offloader._offloaded_mcore_master_weights is True + assert state["state_gpu_exp_avg"] == 0 + assert state["state_gpu_exp_avg_sq"] == 0 + assert state["mcore_master_gpu"] == 0 + assert state["state_cpu_exp_avg"] > 0 + assert state["state_cpu_exp_avg_sq"] > 0 + assert state["mcore_master_cpu"] > 0 + for group in dist_optim.shard_fp32_from_float16_groups: + for tensor in group: + assert tensor.untyped_storage().size() == 0 + + offloader.offload() + offloader.release_gpu_memory() + torch.cuda.synchronize() + + state = offloader._collect_memory_state() + assert state["state_gpu_exp_avg"] == 0 + assert state["state_gpu_exp_avg_sq"] == 0 + assert state["mcore_master_gpu"] == 0 + + dist_optim.get_parameter_state_dp_reshardable() + + state = offloader._collect_memory_state() + assert offloader._offloaded is True + assert state["state_gpu_exp_avg"] == 0 + assert state["state_gpu_exp_avg_sq"] == 0 + assert state["mcore_master_gpu"] == 0 + assert state["state_cpu_exp_avg"] > 0 + assert state["state_cpu_exp_avg_sq"] > 0 + assert state["mcore_master_cpu"] > 0 + for group in dist_optim.shard_fp32_from_float16_groups: + for tensor in group: + assert tensor.untyped_storage().size() == 0 + finally: + Utils.destroy_model_parallel() + + +@pytest.mark.skipif(not TE_FUSED_ADAM_AVAILABLE, reason="Requires TE FusedAdam") +def test_chunked_reload_offloaded_states_defers_mcore_master_weights(): + """Check grad-finalize reload hook does not full-load chunked master weights.""" + Utils.initialize_model_parallel() + try: + model, optim = create_model_and_optimizer( + hidden_size=64, offload_optimizer_states_chunk_numel=1 + ) + dist_optim = get_single_distributed_optimizer(optim) + offloader = dist_optim._state_offloader + + run_forward_backward_step(model, optim, hidden_size=64) + + state = offloader._collect_memory_state() + assert state["mcore_master_gpu"] == 0 + assert state["mcore_master_cpu"] > 0 + assert offloader._offloaded_mcore_master_weights is True + + dist_optim.reload_offloaded_states() + torch.cuda.synchronize() + + state = offloader._collect_memory_state() + assert state["mcore_master_gpu"] == 0 + assert state["mcore_master_cpu"] > 0 + assert offloader._offloaded_mcore_master_weights is True + finally: + Utils.destroy_model_parallel() + + +@pytest.mark.skipif(not TE_FUSED_ADAM_AVAILABLE, reason="Requires TE FusedAdam") +def test_incremental_first_step_state_init_offloads_chunks(): + """Test that first-step lazy optimizer states can be initialized and offloaded in chunks.""" + Utils.initialize_model_parallel() + model, optim = create_model_and_optimizer(offload_optimizer_states_chunk_numel=1) + dist_optim = get_single_distributed_optimizer(optim) + offloader = dist_optim._state_offloader + + assert offloader._optimizer_states_initialized is False + + run_forward_backward_step(model, optim) + + assert offloader._optimizer_states_initialized is True + assert offloader._offloaded is True + assert offloader._opt_state_cpu_buffers + + for state in offloader.adam_optimizer.state.values(): + for state_name in offloader.OPTIMIZER_STATE_KEYS: + if state_name in state: + assert state[state_name].untyped_storage().size() == 0 + + dist_optim.get_parameter_state_dp_reshardable() + + assert offloader._offloaded is True + + for state in offloader.adam_optimizer.state.values(): + for state_name in offloader.OPTIMIZER_STATE_KEYS: + if state_name in state: + assert state[state_name].untyped_storage().size() == 0 + + Utils.destroy_model_parallel() + + +@pytest.mark.skipif(not TE_FUSED_ADAM_AVAILABLE, reason="Requires TE FusedAdam") +def test_state_dict_syncs_pending_reload(): + """Test that checkpoint state reads synchronize an already-enqueued H2D reload.""" + Utils.initialize_model_parallel() + model, optim = create_model_and_optimizer() + dist_optim = get_single_distributed_optimizer(optim) + offloader = dist_optim._state_offloader + + run_forward_backward_step(model, optim) + + offloader.offload() + assert offloader._d2h_inflight is True + offloader.release_gpu_memory() + assert offloader._d2h_inflight is False + + offloader.reload() + assert offloader._has_h2d_pending_work() + + dist_optim.get_parameter_state_dp_reshardable() + + assert not offloader._has_h2d_pending_work() + Utils.destroy_model_parallel() + + +@pytest.mark.skipif(not TE_FUSED_ADAM_AVAILABLE, reason="Requires TE FusedAdam") +def test_fp16_overflow_skip_syncs_pending_reload(): + """Check skipped fp16 steps do not leave offload H2D reloads in flight.""" + Utils.initialize_model_parallel() + try: + hidden_size = 64 + model, optim = create_model_and_optimizer( + hidden_size=hidden_size, model_dtype=torch.float16, initial_loss_scale=128.0 + ) + dist_optim = get_single_distributed_optimizer(optim) + offloader = dist_optim._state_offloader + + run_forward_backward_step(model, optim, hidden_size=hidden_size, input_dtype=torch.float16) + + dist_optim.offload_states() + dist_optim.release_offloaded_gpu_states() + + input_tensor = torch.randn(8, hidden_size, dtype=torch.float16, device='cuda') + output = model(input_tensor) + output.sum().backward() + + injected_inf = False + for param in model.parameters(): + main_grad = getattr(param, "main_grad", None) + if isinstance(main_grad, torch.Tensor): + main_grad.fill_(float("inf")) + injected_inf = True + assert injected_inf + + dist_optim.reload_offloaded_states() + assert offloader._has_h2d_pending_work() + + update_successful, _, _ = optim.step() + + assert update_successful is False + assert not offloader._has_h2d_pending_work() + + dist_optim.offload_states() + dist_optim.release_offloaded_gpu_states() + finally: + Utils.destroy_model_parallel() + + # ============================================================================= # Test 1: Basic OptimizerStateOffloader Initialization # =============================================================================