Fix DeepSeek-V3 checkpoint export compatibility (flagos-ai#1202)

## Summary

- Add an explicit `--skip-mtp` conversion option so
DeepSeek-V3/Moonlight checkpoints can be exported to HF implementations
that only contain the main LM layers.
- Keep loader and saver MTP handling consistent when MTP is skipped,
avoiding unconsumed queue messages and missing HF MTP layer lookups.
- Align DeepSeek-V3 HF config export with current Megatron argument
names and true vocabulary size handling.
- Guard Engram-only checkpoint fields for non-Engram checkpoints and
remove unsupported legacy Megatron conversion arguments.

## Validation

- `PYTHONPYCACHEPREFIX=/private/tmp/flagscale_pycache python3 -m
py_compile tools/checkpoint/convert.py tools/checkpoint/utils.py
tools/checkpoint/loader_mcore.py tools/checkpoint/loader_transformers.py
tools/checkpoint/saver_mcore.py tools/checkpoint/saver_transformers.py
tools/checkpoint/deepseek_v3/args.py`
- `git diff --check`

Author: lihongyang1990

tools/checkpoint/convert.py

@@ -64,6 +64,14 @@ def main():
     parser.add_argument(
         "--max-queue-size", type=int, default=50, help="Maximum number of tensors in the queue"
     )
+    parser.add_argument(
+        "--skip-mtp",
+        action="store_true",
+        help=(
+            "Skip Multi-Token Prediction (MTP) modules during conversion. "
+            "Use this when the target implementation only contains the main LM layers."
+        ),
+    )
 
     extend_cases = [["mistral", "mixtral"]]
 

tools/checkpoint/deepseek_v3/args.py

@@ -23,7 +23,7 @@ def load_args_hf2mg(args):
     args.swiglu = True if hidden_act == "silu" else False
     args.max_position_embeddings = deepseek_v3_args["max_position_embeddings"]
     args.init_method_std = deepseek_v3_args["initializer_range"]
-    args.norm_epsilon = deepseek_v3_args["rms_norm_eps"]
+    args.layernorm_epsilon = deepseek_v3_args["rms_norm_eps"]
     args.untie_embeddings_and_output_weights = not deepseek_v3_args["tie_word_embeddings"]
     args.rotary_base = deepseek_v3_args["rope_theta"]
     args.disable_bias_linear = not deepseek_v3_args["attention_bias"]
@@ -93,13 +93,14 @@ def load_args_hf2mg(args):
 def save_args_mg2hf(args):
     first_k_dense_replace = args.moe_layer_freq.index(1)
     seq_aux = True if args.moe_router_load_balancing_type == "seq_aux_loss" else False
+    mtp_num_layers = getattr(args, "mtp_num_layers", 0) or 0
     config = DeepseekV3Config(
         vocab_size=args.vocab_size,
         hidden_size=args.hidden_size,
         intermediate_size=args.ffn_hidden_size,
         moe_intermediate_size=args.moe_ffn_hidden_size,
         num_hidden_layers=args.num_layers,
-        num_nextn_predict_layers=args.mtp_num_layers,
+        num_nextn_predict_layers=mtp_num_layers,
         num_attention_heads=args.num_attention_heads,
         num_key_value_heads=args.num_query_groups,
         n_shared_experts=args.moe_shared_expert_intermediate_size // args.moe_ffn_hidden_size,
@@ -118,7 +119,7 @@ def save_args_mg2hf(args):
         seq_aux=seq_aux,
         max_position_embeddings=args.max_position_embeddings,
         initializer_range=args.init_method_std,
-        rms_norm_eps=args.norm_epsilon,
+        rms_norm_eps=args.layernorm_epsilon,
         tie_word_embeddings=not args.untie_embeddings_and_output_weights,
         rope_theta=args.rotary_base,
         attention_dropout=args.attention_dropout,

tools/checkpoint/deepseek_v3/ckpt.py

@@ -46,12 +46,18 @@ def _load_checkpoint(queue, args):
         from megatron.training.checkpointing import load_args_from_checkpoint, load_checkpoint
         from megatron.legacy.model import module
         from megatron.core import mpu
-        from megatron.legacy import fused_kernels
         from megatron.core.tensor_parallel.random import (
                 get_cuda_rng_tracker, _DATA_PARALLEL_RNG_TRACKER_NAME,
                 _EXPERT_PARALLEL_RNG_TRACKER_NAME, _MODEL_PARALLEL_RNG_TRACKER_NAME
             )
-        from tools.checkpoint.utils import _ConverterFakeProcessGroup
+        from tools.checkpoint.utils import (
+            _ConverterFakeProcessGroup,
+            get_expert_model_parallel_rank,
+            get_expert_tensor_parallel_rank,
+            get_mcore_model_parallel_size,
+            get_tensor_model_parallel_rank,
+            validate_mcore_parallel_size,
+        )
     except ModuleNotFoundError:
         print("Unable to import Megatron, please specify the path to Megatron using --megatron-path. Exiting.")
         queue.put("exit")
@@ -78,7 +84,6 @@ def queue_put(name, msg):
         '--no-masked-softmax-fusion',
         '--no-bias-gelu-fusion',
         '--no-bias-dropout-fusion',
-        '--no-async-tensor-model-parallel-allreduce',
         '--use-cpu-initialization',
         '--micro-batch-size', '1',
         '--no-load-optim',
@@ -125,20 +130,27 @@ def _set_arg(arg_name):
     _set_arg("hetero_pipeline_layer_split")
 
     # for engram
-    _set_arg("use_engram")
-    _set_arg("engram_layer_ids")
-    _set_arg("engram_hc_mult")
-    _set_arg("engram_kernel_size")
-    _set_arg("engram_pad_id")
-    _set_arg("engram_seed")
-    _set_arg("engram_vocab_size")
-    _set_arg("engram_tokenizer_name_or_path")
-    _set_arg("max_ngram_size")
-    _set_arg("n_embed_per_ngram")
-    _set_arg("n_head_per_ngram")
-    setattr(margs, "vocab_size", args.true_vocab_size)
-    engram_tokenizer_path_ckpt = getattr(checkpoint_args, "engram_tokenizer_name_or_path", None)
-    setattr(margs, "engram_tokenizer_name_or_path", os.path.join(root_path, engram_tokenizer_path_ckpt))
+    if getattr(checkpoint_args, "use_engram", False):
+        _set_arg("use_engram")
+        _set_arg("engram_layer_ids")
+        _set_arg("engram_hc_mult")
+        _set_arg("engram_kernel_size")
+        _set_arg("engram_pad_id")
+        _set_arg("engram_seed")
+        _set_arg("engram_vocab_size")
+        _set_arg("engram_tokenizer_name_or_path")
+        _set_arg("max_ngram_size")
+        _set_arg("n_embed_per_ngram")
+        _set_arg("n_head_per_ngram")
+        if args.true_vocab_size is not None:
+            setattr(margs, "vocab_size", args.true_vocab_size)
+        engram_tokenizer_path_ckpt = getattr(checkpoint_args, "engram_tokenizer_name_or_path", None)
+        if engram_tokenizer_path_ckpt and not os.path.isabs(engram_tokenizer_path_ckpt):
+            engram_tokenizer_path_ckpt = os.path.join(root_path, engram_tokenizer_path_ckpt)
+        setattr(margs, "engram_tokenizer_name_or_path", engram_tokenizer_path_ckpt)
+    else:
+        setattr(margs, "use_engram", False)
+        setattr(margs, "engram_layer_ids", [])
 
     # for hetero
     if margs.hetero_process_meshes is not None:
@@ -148,7 +160,7 @@ def _set_arg(arg_name):
 
     # Arguments do sanity checks on the world size, but we don't care,
     # so trick it into thinking we are plenty of processes
-    margs.world_size = margs.tensor_model_parallel_size * margs.pipeline_model_parallel_size * margs.expert_model_parallel_size
+    margs.world_size = get_mcore_model_parallel_size(margs) * margs.pipeline_model_parallel_size
 
     # Explicitly copy data types from checkpoint.
     margs.fp16 = checkpoint_args.fp16
@@ -170,6 +182,7 @@ def _set_arg(arg_name):
 
     print("*"*20 + "validate loader arguments" + "*"*20)
     margs = validate_args(margs)
+    validate_mcore_parallel_size(margs)
 
     def check_for_arg(arg_name, default=None):
         if getattr(margs, arg_name, None) is None:
@@ -224,14 +237,20 @@ def check_for_arg(arg_name, default=None):
     tp_size = margs.tensor_model_parallel_size
     pp_size = margs.pipeline_model_parallel_size
     ep_size = margs.expert_model_parallel_size
+    etp_size = margs.expert_tensor_parallel_size
+    mcore_model_parallel_size = get_mcore_model_parallel_size(margs)
     vp_size = margs.virtual_pipeline_model_parallel_size or 1
     mpu.set_tensor_model_parallel_world_size(tp_size)
     mpu.set_pipeline_model_parallel_world_size(pp_size)
     mpu.set_expert_model_parallel_world_size(ep_size)
+    if hasattr(mpu, "set_expert_tensor_parallel_world_size"):
+        mpu.set_expert_tensor_parallel_world_size(etp_size)
     mpu.set_virtual_pipeline_model_parallel_world_size(margs.virtual_pipeline_model_parallel_size)
     mpu.set_tensor_model_parallel_rank(0)
     mpu.set_pipeline_model_parallel_rank(0)
     mpu.set_expert_model_parallel_rank(0)
+    if hasattr(mpu, "set_expert_tensor_parallel_rank"):
+        mpu.set_expert_tensor_parallel_rank(0)
     mpu.set_virtual_pipeline_model_parallel_rank(0)
     # For backward compatibility during local parallel states refactoring
     fake_tp_group = _ConverterFakeProcessGroup(size=tp_size)
@@ -242,10 +261,10 @@ def check_for_arg(arg_name, default=None):
     fake_pp_group = _ConverterFakeProcessGroup(size=margs.pipeline_model_parallel_size)
     fake_cp_group = _ConverterFakeProcessGroup(size=margs.context_parallel_size)
     fake_dp_group = _ConverterFakeProcessGroup(size=margs.data_parallel_size)
-    fake_etp_group = _ConverterFakeProcessGroup(size=margs.expert_tensor_parallel_size)
-    edp_parallel_size = margs.tensor_model_parallel_size * margs.context_parallel_size // (margs.expert_tensor_parallel_size * margs.expert_model_parallel_size)
+    fake_etp_group = _ConverterFakeProcessGroup(size=etp_size)
+    edp_parallel_size = mcore_model_parallel_size // (etp_size * ep_size)
     fake_edp_group = _ConverterFakeProcessGroup(size=edp_parallel_size)
-    fake_etp_ep_group = _ConverterFakeProcessGroup(size=margs.expert_tensor_parallel_size*margs.expert_model_parallel_size)
+    fake_etp_ep_group = _ConverterFakeProcessGroup(size=etp_size * ep_size)
     fake_tcp_group = _ConverterFakeProcessGroup(size=margs.tensor_model_parallel_size*margs.context_parallel_size)
     mpu._PIPELINE_MODEL_PARALLEL_GROUP = fake_pp_group
     mpu._CONTEXT_PARALLEL_GROUP = fake_cp_group
@@ -254,14 +273,11 @@ def check_for_arg(arg_name, default=None):
     mpu._EXPERT_DATA_PARALLEL_GROUP = fake_edp_group
     mpu._EXPERT_TENSOR_AND_MODEL_PARALLEL_GROUP = fake_etp_ep_group
     mpu._TENSOR_AND_CONTEXT_PARALLEL_GROUP = fake_tcp_group
-    mpu._EXPERT_TENSOR_PARALLEL_GROUP = fake_tp_group
+    mpu._EXPERT_TENSOR_PARALLEL_GROUP = fake_etp_group
     mpu._DATA_PARALLEL_GROUP_WITH_CP = fake_dp_group
     mpu._INTRA_PARTIAL_DATA_PARALLEL_GROUP_WITH_CP = fake_dp_group
     mpu._LAST_RANK_WHEN_USING_PIPELINE = pp_size - 1
 
-    # fused kernel
-    fused_kernels.load(margs)
-
     # random
     CUDA_RNG_STATE_TRACKER = get_cuda_rng_tracker()
     torch.cuda.manual_seed(42)
@@ -291,6 +307,7 @@ def check_for_arg(arg_name, default=None):
     md.previous_tensor_parallel_size = margs.tensor_model_parallel_size
     md.previous_pipeline_parallel_size = margs.pipeline_model_parallel_size
     md.previous_expert_parallel_size = margs.expert_model_parallel_size
+    md.previous_expert_tensor_parallel_size = margs.expert_tensor_parallel_size
     md.previous_decoder_first_pipeline_num_layers = margs.decoder_first_pipeline_num_layers
     md.true_vocab_size = args.true_vocab_size # true (non-padded) vocab size
     md.make_vocab_size_divisible_by = margs.make_vocab_size_divisible_by
@@ -302,16 +319,22 @@ def get_models(count, dtype):
         # for one pp stage
         nonlocal consumed_train_samples
         nonlocal consumed_valid_samples
-        tp_size = margs.tensor_model_parallel_size
         pp_size = margs.pipeline_model_parallel_size
         vp_size = margs.virtual_pipeline_model_parallel_size or 1
 
         models = [[] for _ in range(vp_size)]
         for rank_id in range(count):
-            tp_rank = rank_id % tp_size
-            ep_rank = rank_id // tp_size
+            tp_rank = get_tensor_model_parallel_rank(rank_id, margs)
+            ep_rank = get_expert_model_parallel_rank(rank_id, margs)
+            etp_rank = get_expert_tensor_parallel_rank(rank_id, margs)
             mpu.set_tensor_model_parallel_rank(tp_rank)
             mpu.set_expert_model_parallel_rank(ep_rank)
+            if hasattr(mpu, "set_expert_tensor_parallel_rank"):
+                mpu.set_expert_tensor_parallel_rank(etp_rank)
+            fake_tp_group.set_rank(tp_rank)
+            fake_ep_group.set_rank(ep_rank)
+            fake_etp_group.set_rank(etp_rank)
+            fake_etp_ep_group.set_rank(ep_rank * etp_size + etp_rank)
             if pp_size > 1 and vp_size > 1:
                 model_ = []
                 for vp_rank in range(vp_size):
@@ -353,7 +376,7 @@ def get_models(count, dtype):
     mpu.set_pipeline_model_parallel_rank(0)
     fake_pp_group = _ConverterFakeProcessGroup(rank=0, size=pp_size)
     mpu._PIPELINE_MODEL_PARALLEL_GROUP = fake_pp_group
-    all_models = [get_models(tp_size * ep_size, margs.params_dtype)]
+    all_models = [get_models(mcore_model_parallel_size, margs.params_dtype)]
     models = all_models[0][0] # pp0vpp0
 
     md.consumed_train_samples = consumed_train_samples
@@ -379,15 +402,17 @@ def get_models(count, dtype):
             mpu._PIPELINE_MODEL_PARALLEL_GROUP = fake_pp_group
 
             if pp_rank > 0 and vp_rank == 0:
-                all_models.append(get_models(tp_size * ep_size, margs.params_dtype))
+                all_models.append(get_models(mcore_model_parallel_size, margs.params_dtype))
 
             models = all_models[pp_rank][vp_rank]
             for layer_id in range(len(models[0].decoder.layers)):
                 message = dict()
                 margs.total_layer_num = total_layer_num
 
                 engram_layer_id  = total_layer_num # get_global_layer_id
-                if margs.use_engram and engram_layer_id in margs.engram_layer_ids:
+                if getattr(margs, "use_engram", False) and engram_layer_id in getattr(
+                    margs, "engram_layer_ids", []
+                ):
                     ckpt_plugin.get_engram_ckpt(message, models, engram_layer_id, margs)
 
                 ckpt_plugin.get_attn_ckpt(message, models, layer_id, margs)
@@ -406,9 +431,11 @@ def get_models(count, dtype):
         ckpt_plugin.get_output_layer_ckpt(message, models, margs)
         queue_put("output layer", message)
 
-    message = dict()
-    if margs.mtp_num_layers:
-        for mtp_layer_id in range(margs.mtp_num_layers):
+    mtp_num_layers = getattr(margs, "mtp_num_layers", 0)
+    if getattr(args, "skip_mtp", False):
+        mtp_num_layers = 0
+    if mtp_num_layers:
+        for mtp_layer_id in range(mtp_num_layers):
             message = dict()
             ckpt_plugin.get_mtp_ckpt(message, models, mtp_layer_id, margs)
             queue_put(f"mtp module {mtp_layer_id}", message)

tools/checkpoint/loader_mcore.py

@@ -73,7 +73,6 @@ def queue_put(name, msg):
         '--no-masked-softmax-fusion',
         '--no-bias-gelu-fusion',
         '--no-bias-dropout-fusion',
-        '--no-async-tensor-model-parallel-allreduce',
         '--use-cpu-initialization',
         '--micro-batch-size', '1',
         '--no-load-optim',
@@ -192,9 +191,11 @@ def check_for_arg(arg_name, default=None):
         message = {"weight": hf_model.lm_head.weight.data}
         queue_put("output layer", message)
 
-    message = dict()
-    if margs.mtp_num_layers:
-        for mtp_layer_id in range(margs.mtp_num_layers):
+    mtp_num_layers = getattr(margs, "mtp_num_layers", 0)
+    if getattr(args, "skip_mtp", False):
+        mtp_num_layers = 0
+    if mtp_num_layers:
+        for mtp_layer_id in range(mtp_num_layers):
             message = dict()
             ckpt_plugin.get_hf_mtp_ckpt(message, hf_model, mtp_layer_id, margs)
             queue_put(f"mtp module {mtp_layer_id}", message)