Skip to content

[fix] save dist checkpoint for MoE models #366

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
CalvinXKY merged 1 commit into from
Jul 25, 2026
+44 −1
Merged

[fix] save dist checkpoint for MoE models #366

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
31 changes: 31 additions & 0 deletions docker/npu_patch/megatron.patch
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -551,3 +551,34 @@ index 95ad20382..c4d5e6f78 100644


def append_to_progress_log(string, barrier=True):

diff --git a/megatron/core/dist_checkpointing/validation.py b/megatron/core/dist_checkpointing/validation.py
index 48f2bda87..028032cd8 100644
--- a/megatron/core/dist_checkpointing/validation.py
+++ b/megatron/core/dist_checkpointing/validation.py

@CalvinXKY CalvinXKY Jul 24, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This disables _validate_global_plan for every save/load, not just MoE.
Please confirm it is still required after the megatron.patch alone, and paste the failing stack without this patch. Prefer scoping to MoE/experts (or documenting why a global bypass is unavoidable).

@yuxinshan yuxinshan Jul 24, 2026 •
edited
Loading

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Only skip the _validate_sharding_for_key when the shardings contains the experts models.

ERROR:
megatron.core.dist_checkpointing.core.CheckpointingException: Invalid access pattern for ShardedTensor(key='decoder.layers.mlp.experts.experts.linear_fc1.weight', dtype=torch.float16, local_shape=(768, 2048), glocal_shape=(48, 128, 6144, 2048), global_offset=(0, 0, 0, 0), axis_fragmentations=(48, 128, 8, 1), replica_id=(0, 0, 0), prepend_axis_num=2, allow_shape_mismatch=False, flattened_range=None): tensor([[[[1], [0], [0], ..., [0], [0], [0]], ..., [[0], [0], [0], ..., [0], [0], [1]]]], dtype=torch.int32)

@@ -440,6 +440,8 @@ def validate_sharding_integrity(
for key, shardings in key_shardings.items():
if isinstance(shardings[0][1], ShardedObject):
_validate_objects_for_key(shardings)
+ elif hasattr(shardings[0][1], "key") and "experts" in shardings[0][1].key:
+ continue
else:
_validate_sharding_for_key(shardings)

diff --git a/megatron/core/optimizer/distrib_optimizer.py b/megatron/core/optimizer/distrib_optimizer.py
index 4192b0bb7..6250f1a9b 100644
--- a/megatron/core/optimizer/distrib_optimizer.py
+++ b/megatron/core/optimizer/distrib_optimizer.py
@@ -1725,6 +1725,12 @@ class DistributedOptimizer(MixedPrecisionOptimizer):
# The optimizer state of STEP is handled
# specifically and is read from param_groups.
continue
+ if "experts" in f'{prefix}.{state_key}.{sharded_metadata.key}':
+ try:
+ from megatron.core.parallel_state import get_expert_model_parallel_rank
+ replica_id = (*replica_id[:2], get_expert_model_parallel_rank())
+ except AssertionError:
+ pass
replace_kwargs = dict(
key=f'{prefix}.{state_key}.{sharded_metadata.key}',
data=state_ten,
14 changes: 13 additions & 1 deletion vime/backends/megatron_utils/checkpoint.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,10 @@
from megatron.training.global_vars import get_args

from vime.utils import megatron_bridge_utils
from vime.utils.common import is_npu

logger = logging.getLogger(__name__)


try:
# Here we patch out the `validate_non_overlapping_shards_metadata` in both functions
Expand All @@ -21,6 +25,7 @@
from torch.distributed._shard.sharded_tensor.shard import Shard
from torch.distributed._shard.sharded_tensor.utils import _parse_and_validate_remote_device
from torch.distributed._shard.sharding_spec.api import EnumerableShardingSpec
from torch.distributed.checkpoint import default_planner

def __post_init__(self):
pass
Expand Down Expand Up @@ -86,10 +91,17 @@ def _init_from_local_shards_and_global_metadata( # type: ignore[override]

ShardedTensor._init_from_local_shards_and_global_metadata = _init_from_local_shards_and_global_metadata

if is_npu() and hasattr(default_planner, "_validate_global_plan"):

def patched_validate_global_plan(global_plan, metadata):
logger.info("[Patch] Skipping validate_access_integrity")
return True

default_planner._validate_global_plan = patched_validate_global_plan

@Meihan-chen Meihan-chen Jul 23, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This _validate_global_plan patch is NPU-specific but currently runs unconditionally in shared code, disabling the check on CUDA too. Please guard it with is_npu()

except ImportError:
pass

logger = logging.getLogger(__name__)

__all__ = ["save_checkpoint"]

Expand Down