From 35317ce0c1f567502a8771a49e0a174321efb239 Mon Sep 17 00:00:00 2001 From: yuxinshan Date: Thu, 23 Jul 2026 10:43:56 +0800 Subject: [PATCH] [fix] save dist checkpoint for MoE models Signed-off-by: yuxinshan --- docker/npu_patch/megatron.patch | 31 ++++++++++++++++++++++ vime/backends/megatron_utils/checkpoint.py | 14 +++++++++- 2 files changed, 44 insertions(+), 1 deletion(-) diff --git a/docker/npu_patch/megatron.patch b/docker/npu_patch/megatron.patch index 0c60a64b5..296947a7e 100644 --- a/docker/npu_patch/megatron.patch +++ b/docker/npu_patch/megatron.patch @@ -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 +@@ -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, diff --git a/vime/backends/megatron_utils/checkpoint.py b/vime/backends/megatron_utils/checkpoint.py index d196ad24d..1e6f9193a 100644 --- a/vime/backends/megatron_utils/checkpoint.py +++ b/vime/backends/megatron_utils/checkpoint.py @@ -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 @@ -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 @@ -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 + except ImportError: pass -logger = logging.getLogger(__name__) __all__ = ["save_checkpoint"]