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

This reverts commit 6b157e0.

Author: ko3n1g

megatron/core/optimizer/cpu_offloading/optimizer_state_offloader.py

@@ -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
@@ -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
@@ -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
@@ -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()

megatron/core/optimizer/distrib_optimizer.py

@@ -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
     ################