Skip to content

Revert "[Dev] Optimizer State and Master Weight Offloading (#2760)" #2984

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
ko3n1g merged 1 commit into from
Jan 16, 2026
Merged

Revert "[Dev] Optimizer State and Master Weight Offloading (#2760)" #2984

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
315 changes: 0 additions & 315 deletions megatron/core/optimizer/cpu_offloading/optimizer_state_offloader.py
View file
Open in desktop

This file was deleted.

25 changes: 0 additions & 25 deletions megatron/core/optimizer/distrib_optimizer.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -49,7 +49,6 @@
from ..fp8_utils import dequantize_fp8_tensor, is_float8tensor, quantize_param_shard
from ..transformer.fsdp_dtensor_checkpoint import handle_experts_in_state_dict
from ..transformer.module import MegatronModule
from .cpu_offloading.optimizer_state_offloader import OptimizerStateOffloader
from .grad_scaler import MegatronGradScaler
from .optimizer import MixedPrecisionOptimizer, _zero_grad_group_helper, param_group_identifier_keys
from .optimizer_config import OptimizerConfig
Expand Down Expand Up @@ -605,10 +604,6 @@ def __init__(
self.optimizer.param_groups = [g["orig_group"] for g in self.opt_group_ranges]
self.optimizer.load_state_dict(self.optimizer.state_dict())

self._state_offloader: Optional[OptimizerStateOffloader] = None
if self.config.offload_optimizer_states:
self._state_offloader = OptimizerStateOffloader(self)

def _get_model_param_range_map(self, param: torch.nn.Parameter):
"""
Given a model param, get the index sub-range of the param that this
Expand Down Expand Up @@ -2585,8 +2580,6 @@ def step_with_ready_grads(self) -> bool:
Under the hood, either launch synchronous param all-gathers or get ready to launch
asynchorous all-gathers that get overlapped with the next forward pass.
"""
if self._state_offloader is not None:
self._state_offloader.sync_before_step()
update_successful = super().step_with_ready_grads()

timers = self.config.timers
Expand All @@ -2607,22 +2600,4 @@ def step_with_ready_grads(self) -> bool:
if timers is not None:
timers('params-all-gather').stop()

if self._state_offloader is not None:
self._state_offloader.mark_optimizer_states_initialized()

return update_successful

def offload_states(self):
"""Offload states to CPU."""
if self._state_offloader is not None:
self._state_offloader.offload()

def reload_offloaded_states(self):
"""Start async reload of offloaded states."""
if self._state_offloader is not None:
self._state_offloader.reload()

def release_offloaded_gpu_states(self):
"""Release GPU memory after D2H completes. For delayed release case."""
if self._state_offloader is not None:
self._state_offloader.release_gpu_memory()
6 changes: 0 additions & 6 deletions megatron/core/optimizer/optimizer_config.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -266,12 +266,6 @@ class OptimizerConfig:
pin_cpu_params: bool = True
"""If True, pin the optimizer parameters to CPU memory."""

offload_optimizer_states: bool = False
"""
If True, offload optimizer states to CPU after each optimizer step and
reload them before the next optimizer step.
"""

################
# Miscellaneous
################
Expand Down
Loading
Loading