Add chunked GPU optimizer-state offloading

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."""

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)

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,
     )
 
 

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: