fix(fsdp): support Megatron-Bridge checkpoints under fsdp_dtensor

Three small fixes that together let FSDP-DTensor resume from a checkpoint converted with Megatron-Bridge: 1. `split_swiglu_linear_fc1` / `split_gdn_fused`: accept a plain `Tensor` (not only `DTensor`) for `data.to_local()`. Bridge writes some already-local tensors that previously crashed the splitter. 2. `split_gdn_fused`: when the GDN tensor is already TP-local (`data.shape[split_dim] == sum(split_sizes)`), use `split_dtensor(..., update_uneven_dtensor_chunk_meta=True)` instead of the TP-mesh-based split that assumes the full unsharded shape. 3. `megatron/training/checkpointing.py::_load_base_checkpoint`: stash top-level state-dict keys that DCP storage metadata doesn't know about (`args`, `iteration`, ...), then restore them after `dcp.load`. Megatron-saved checkpoints contain these as `BytesStorageMetadata`; Bridge-converted ones don't, and DCP errors without this stash. 4. `DistributedOptimizer._copy_model_params_to_main_params`: under `use_megatron_fsdp`, return instead of raising. DCP already loads directly into the fp32 main buffer, and a post-load hook copies main → model weights — both are correct at this call site. Also wires `handle_gdn_in_state_dict` into `preprocess_fsdp_dtensor_state_dict`. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com> Co-Authored-By: xuwchen <79835960+xuwchen@users.noreply.github.com> Co-Authored-By: conver334 <56124251+conver334@users.noreply.github.com> Co-Authored-By: BestJuly <19769279+BestJuly@users.noreply.github.com>
NVIDIA · yaox12 · May 19, 2026 · May 12, 2026 · Apr 22, 2026 · May 13, 2026
commit c9dfa74dd25f5e37cf23a990becc47f7abd19506
@@ -2928,9 +2928,10 @@ def _copy_model_params_to_main_params(self, state_dict=None):
             return
 
         if self.ddp_config.use_megatron_fsdp:
-            raise NotImplementedError(
-                "Megatron-FSDP does not implement a model-to-main parameter update."
-            )
+            # For FSDP, DCP loads checkpoint data directly into main_weight_buffer
+            # (fp32 main params), and a post-load hook copies main → model weights.
+            # Both main params and model params are already correct at this point.
+            return
 
         # When using precision-aware optimizer, main params are held by self.optimizer. It will also
         # do the work of copying data from main params to model params.

@@ -31,6 +31,7 @@
         make_fsdp_dtensor,
     )
     from megatron.core.distributed.fsdp.src.megatron_fsdp.uneven_dtensor import (
+        split_dtensor,
         uneven_dtensor_to_full_tensor,
     )
     from megatron.core.distributed.fsdp.src.megatron_fsdp.utils import (
@@ -291,7 +292,7 @@ def split_swiglu_linear_fc1(data, dist_param, swiglu_shard_axis, is_expert_param
 
         view_shape = list(data.shape)
         view_shape[swiglu_shard_axis] = -1
-        local_tensor = data.to_local()
+        local_tensor = data.to_local() if isinstance(data, DTensor) else data
         weight_w = local_tensor.view(-1)[
             offset_slice(intersection(fsdp_slice, w_slice), -fsdp_slice.start)
         ]
@@ -463,15 +464,26 @@ def offset_slice(s, offset):
 
     def split_gdn_fused(data, dist_param, split_sizes, split_dim):
         """Split a fused GDN projection DTensor into per-component DTensors."""
+        total_split = sum(split_sizes)
+        if isinstance(data, DTensor) and data.shape[split_dim] == total_split:
+            # GDN tensors are already TP-local here.
+            return list(
+                split_dtensor(
+                    data,
+                    split_sizes,
+                    dim=split_dim,
+                    update_uneven_dtensor_chunk_meta=True,
+                )
+            )
+
         fsdp_slice = dist_param.megatron_fsdp_slice
         dist_index = dist_param.megatron_fsdp_dist_index
         tp_mesh = dist_index.get_submesh([dist_index.tp_dim], is_expert_parallel=False)
 
         data_size = data.numel() // tp_mesh.mesh.numel()
-        total_split = sum(split_sizes)
         elems_per_unit = data_size // total_split
 
-        local_tensor = data.to_local()
+        local_tensor = data.to_local() if isinstance(data, DTensor) else data
         view_shape = list(data.shape)
 
         per_tp_rank_shape = list(data.shape)

@@ -55,6 +55,7 @@
     from megatron.core.transformer.fsdp_dtensor_checkpoint import (
         handle_experts_in_state_dict,
         handle_fp8_extra_state_case,
+        handle_gdn_in_state_dict,
         handle_swiglu_in_state_dict,
         print_diff_in_state_dicts,
     )
@@ -1275,6 +1276,15 @@ def preprocess_fsdp_dtensor_state_dict(args, raw_state_dict, model):
         else:
             model_state_dict, _ = handle_swiglu_in_state_dict(model, state_dict["model"], None)
             state_dict["model"] = model_state_dict
+    if "optimizer" in state_dict:
+        model_state_dict, optimizer_state_dict = handle_gdn_in_state_dict(
+            model, state_dict["model"], state_dict["optimizer"]
+        )
+        state_dict["model"] = model_state_dict
+        state_dict["optimizer"] = optimizer_state_dict
+    else:
+        model_state_dict, _ = handle_gdn_in_state_dict(model, state_dict["model"], None)
+        state_dict["model"] = model_state_dict
     if args.num_experts:
         state_dict["model"] = handle_experts_in_state_dict(state_dict["model"], args.num_experts)
     preprocess_state_dict_for_uneven_dtensor(state_dict)
@@ -1645,8 +1655,26 @@ def _load_base_checkpoint(
         ckpt_type = CheckpointType.FSDP_DTENSOR
         fs_storage_reader = torch.distributed.checkpoint.FileSystemReader(checkpoint_name)
         allow_partial_load = not getattr(args, 'strict_fsdp_dtensor_load', False)
+
+        # Stash state_dict keys that have no presence in the checkpoint.
+        # For Megatron-saved checkpoints, DCP stores all entries (tensors
+        # as TensorStorageMetadata, non-tensors as BytesStorageMetadata),
+        # so nothing needs to be stashed.
+        # For Bridge-converted checkpoints, metadata keys like 'args',
+        # 'iteration', etc. don't exist — stash them so DCP doesn't
+        # error, then restore placeholder values after load.
+        state_dict_metadata = fs_storage_reader.read_metadata().state_dict_metadata
+        _stashed_keys = {}
+        for key in list(state_dict.keys()):
+            # Keep the key if it exists as a top-level metadata entry
+            # (e.g. 'args', 'iteration') OR if any flattened sub-key
+            # exists (e.g. 'model.layer.weight' for top-level 'model').
+            if key not in state_dict_metadata and not any(
+                k.startswith(key + '.') for k in state_dict_metadata
+            ):
+                _stashed_keys[key] = state_dict.pop(key)
+
         if allow_partial_load:
-            state_dict_metadata = fs_storage_reader.read_metadata().state_dict_metadata
             rank = torch.distributed.get_rank()
             import time as _time
 
@@ -1658,6 +1686,9 @@ def _load_base_checkpoint(
             state_dict=state_dict, storage_reader=fs_storage_reader, planner=planner
         )
 
+        # Restore stashed keys (Bridge fallback placeholders).
+        state_dict.update(_stashed_keys)
+
         if raw_optimizer_state_dict is not None:
             state_dict["optimizer"] = raw_optimizer_state_dict