Support DCP checkpointing for the Megatron-FSDP v2 optimizer

megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/__init__.py

@@ -17,7 +17,7 @@
 from .checkpoint import load_checkpoint, save_checkpoint
 from .dbuffer import DBuffer
 from .fully_shard import Placements, fully_shard, fully_shard_context, microbatch
-from .optimizer import fully_shard_optimizer
+from .optimizer import fully_shard_optimizer, init_optimizer_state
 from .schedule import SchedulePolicy
 
 __all__ = [
@@ -27,6 +27,7 @@
     "fully_shard",
     "fully_shard_context",
     "fully_shard_optimizer",
+    "init_optimizer_state",
     "load_checkpoint",
     "microbatch",
     "save_checkpoint",

megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/checkpoint.py

@@ -37,37 +37,13 @@
     set_optimizer_state_dict,
 )
 
+from .optimizer import init_optimizer_state
 from .parameter_group import sync_model_weights_from_main_weights
 from .uneven_dtensor import attach_uneven_dtensor_metadata
 
 __all__ = ["save_checkpoint", "load_checkpoint"]
 
 
-def _init_optimizer_state(optimizer: torch.optim.Optimizer) -> None:
-    """Allocate optimizer state so a DCP load has DTensors to fill.
-
-    :func:`get_optimizer_state_dict` initializes empty optimizer state via torch's
-    ``_init_optim_state``, but that assigns a parameter-dtype gradient. A Megatron-FSDP sharded
-    parameter advertises the FSDP gradient dtype through ``grad_dtype``, which differs from the
-    (main-weight) parameter dtype under mixed precision, and rejects a mismatched gradient. So
-    initialize the state here with a ``grad_dtype``-matched zero gradient; the subsequent load
-    overwrites it. This is a no-op once the state exists (for example after a training step).
-
-    TODO: this function becomes unnecessary once torch's ``_init_optim_state`` honors a parameter's
-    ``grad_dtype`` when it allocates the placeholder gradient (``torch.zeros_like(param)`` in
-    ``torch/distributed/checkpoint/state_dict.py``); an upstream issue is being filed.
-    """
-    if optimizer.state:
-        return
-    for group in optimizer.param_groups:
-        for param in group["params"]:
-            if param.grad is None:
-                grad_dtype = getattr(param, "grad_dtype", None) or param.dtype
-                param.grad = torch.zeros_like(param, dtype=grad_dtype)
-    optimizer.step()
-    optimizer.zero_grad()
-
-
 def save_checkpoint(
     model: torch.nn.Module, optimizer: torch.optim.Optimizer, checkpoint_dir: str | os.PathLike
 ) -> None:
@@ -107,7 +83,7 @@ def load_checkpoint(
         checkpoint_dir: Source directory of the DCP checkpoint.
         sync_model_weights: Refresh compute weights from the loaded main weights afterwards.
     """
-    _init_optimizer_state(optimizer)
+    init_optimizer_state(optimizer)
     model_state_dict = get_model_state_dict(model)
     optimizer_state_dict = get_optimizer_state_dict(model, optimizer)
     attach_uneven_dtensor_metadata(model, model_state_dict, optimizer_state_dict)

megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/optimizer.py

@@ -22,6 +22,37 @@
 
 from .parameter_group import get_containing_parameter_group, sync_model_weights_from_main_weights
 
+__all__ = ["fully_shard_optimizer", "init_optimizer_state"]
+
+
+def init_optimizer_state(optimizer: torch.optim.Optimizer) -> None:
+    """Allocate optimizer state so a DCP load has DTensors to fill.
+
+    :func:`~torch.distributed.checkpoint.state_dict.get_optimizer_state_dict` initializes empty
+    optimizer state via torch's ``_init_optim_state``, but that assigns a parameter-dtype gradient.
+    A Megatron-FSDP sharded parameter advertises the FSDP gradient dtype through ``grad_dtype``,
+    which differs from the (main-weight) parameter dtype under mixed precision, and rejects a
+    mismatched gradient. So initialize the state here with a ``grad_dtype``-matched zero gradient;
+    the subsequent load overwrites it. This is a no-op once the state exists (for example after a
+    training step).
+
+    TODO: this function becomes unnecessary once torch's ``_init_optim_state`` honors a parameter's
+    ``grad_dtype`` when it allocates the placeholder gradient (``torch.zeros_like(param)`` in
+    ``torch/distributed/checkpoint/state_dict.py``); an upstream issue is being filed.
+
+    Args:
+        optimizer: Optimizer whose (possibly empty) state should be materialized.
+    """
+    if optimizer.state:
+        return
+    for group in optimizer.param_groups:
+        for param in group["params"]:
+            if param.grad is None:
+                grad_dtype = getattr(param, "grad_dtype", None) or param.dtype
+                param.grad = torch.zeros_like(param, dtype=grad_dtype)
+    optimizer.step()
+    optimizer.zero_grad()
+
 
 def fully_shard_optimizer(
     optimizer: torch.optim.Optimizer, *, precision_aware: bool = False