Fix DeepSeek-V3 checkpoint export compatibility

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/loader_mcore.py

@@ -46,7 +46,6 @@ 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
@@ -78,7 +77,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 +123,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:
@@ -259,9 +264,6 @@ def check_for_arg(arg_name, default=None):
     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)
@@ -387,7 +389,9 @@ def get_models(count, dtype):
                 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 +410,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_transformers.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)

tools/checkpoint/qwen2_5_vl/hf2mcore_qwen2.5_vl_convertor.sh

@@ -198,7 +198,6 @@ cmd="torchrun ${DISTRIBUTED_ARGS} hf2mcore_qwen2.5_vl.py \
     --num-attention-heads ${NUM_ATTN_HEADS} \
     --max-position-embeddings ${MAX_POSITION_EMBEDDINGS} \
     --seq-length 1 \
-    --no-async-tensor-model-parallel-allreduce \
     --tokenizer-type Qwen2VLTokenizer \
     --extra-vocab-size ${EXTRA_VOCAB_SIZE} \
     --no-bias-swiglu-fusion \

tools/checkpoint/qwen3_vl/hf2mcore_qwen3_vl_convertor.sh

@@ -233,7 +233,6 @@ cmd="torchrun ${DISTRIBUTED_ARGS} hf2mcore_qwen3_vl.py \
     --num-attention-heads ${NUM_ATTN_HEADS} \
     --max-position-embeddings ${MAX_POSITION_EMBEDDINGS} \
     --seq-length 1 \
-    --no-async-tensor-model-parallel-allreduce \
     --tokenizer-type Qwen2VLTokenizer \
     --extra-vocab-size ${EXTRA_VOCAB_SIZE} \
     --no-bias-swiglu-fusion \

tools/checkpoint/saver_mcore.py

@@ -54,8 +54,7 @@ def save_checkpoint(queue, args):
         from megatron.training.arguments import parse_args, validate_args
         from megatron.training.checkpointing import save_checkpoint, get_checkpoint_name
         from megatron.training.global_vars import set_global_variables, get_args
-        from megatron.training.tokenizer.tokenizer import _vocab_size_with_padding
-        from megatron.legacy import fused_kernels
+        from megatron.training.tokenizer.tokenizer import vocab_size_with_padding
         from megatron.core import mpu
         from megatron.core.tensor_parallel.random import (
                 get_cuda_rng_tracker, _DATA_PARALLEL_RNG_TRACKER_NAME,
@@ -169,7 +168,6 @@ def check_message(msg):
         '--no-masked-softmax-fusion',
         '--no-bias-gelu-fusion',
         '--no-bias-dropout-fusion',
-        '--no-async-tensor-model-parallel-allreduce',
         '--use-cpu-initialization',
         '--transformer-impl', 'transformer_engine',
         '--micro-batch-size', '1',
@@ -223,7 +221,7 @@ def check_message(msg):
         args_to_keep = ['tensor_model_parallel_size', 'pipeline_model_parallel_size', 'expert_model_parallel_size', 'world_size', 'params_dtype',
                         'num_layers_per_virtual_pipeline_stage', 'virtual_pipeline_model_parallel_size',
                         'masked_softmax_fusion', 'bias_gelu_fusion', 'bias_dropout_fusion',
-                        'sequence_parallel', 'async_tensor_model_parallel_allreduce',
+                        'sequence_parallel',
                         'no_load_optim', 'no_load_rng', 'no_save_optim', 'no_save_rng',
                         'vocab_file', 'tokenizer_model',
                         'save_interval', 'save', 'load', 'use_mcore_models', 'num_experts',
@@ -292,11 +290,14 @@ def check_message(msg):
     margs.model_type = model_plugin.model_type
 
     if md.true_vocab_size is not None:
-        margs.padded_vocab_size = _vocab_size_with_padding(md.true_vocab_size, margs)
+        margs.padded_vocab_size = vocab_size_with_padding(md.true_vocab_size, margs)
     else:
         # margs.padded_vocab_size will be set in ckpt_plugin.set_embedding_ckpt func
         margs.padded_vocab_size = None
 
+    if getattr(args, "skip_mtp", False):
+        margs.mtp_num_layers = 0
+
     """
     use megatron args build object and init env
     """
@@ -349,10 +350,6 @@ def check_message(msg):
     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)

tools/checkpoint/saver_transformers.py

@@ -55,7 +55,7 @@ def save_checkpoint(queue, args):
 
     try:
         from megatron.training.arguments import parse_args, validate_args
-        from megatron.training.tokenizer.tokenizer import _vocab_size_with_padding
+        from megatron.training.tokenizer.tokenizer import vocab_size_with_padding
     except ModuleNotFoundError:
         print("Unable to import Megatron, please specify the path to Megatron using --megatron-path. Exiting.")
         queue.put("exit")
@@ -158,7 +158,6 @@ def check_message(msg):
         '--no-masked-softmax-fusion',
         '--no-bias-gelu-fusion',
         '--no-bias-dropout-fusion',
-        '--no-async-tensor-model-parallel-allreduce',
         '--use-cpu-initialization',
         '--transformer-impl', 'transformer_engine',
         '--micro-batch-size', '1',
@@ -205,7 +204,7 @@ def check_message(msg):
         args_to_keep = ['tensor_model_parallel_size', 'pipeline_model_parallel_size', 'expert_model_parallel_size', 'world_size', 'params_dtype',
                         'num_layers_per_virtual_pipeline_stage', 'virtual_pipeline_model_parallel_size',
                         'masked_softmax_fusion', 'bias_gelu_fusion', 'bias_dropout_fusion',
-                        'sequence_parallel', 'async_tensor_model_parallel_allreduce',
+                        'sequence_parallel',
                         'no_load_optim', 'no_load_rng', 'no_save_optim', 'no_save_rng',
                         'vocab_file', 'tokenizer_model',
                         'save_interval', 'save', 'load', 'use_mcore_models', 'num_experts',
@@ -256,11 +255,15 @@ def check_message(msg):
     margs.transformer_impl = "transformer_engine"
 
     if md.true_vocab_size is not None:
-        margs.padded_vocab_size = _vocab_size_with_padding(md.true_vocab_size, margs)
+        margs.padded_vocab_size = vocab_size_with_padding(md.true_vocab_size, margs)
+        margs.vocab_size = md.true_vocab_size
     else:
         # margs.padded_vocab_size will be set in ckpt_plugin.set_embedding_ckpt func
         margs.padded_vocab_size = None
 
+    if getattr(args, "skip_mtp", False):
+        margs.mtp_num_layers = 0
+
     """
     use megatron args build object and init env
     """
@@ -283,7 +286,9 @@ def check_message(msg):
         msg = queue_get(f"transformer layer {layer_id}")
 
         margs.total_layer_num = layer_id
-        if margs.use_engram and layer_id in margs.engram_layer_ids:
+        if getattr(margs, "use_engram", False) and layer_id in getattr(
+            margs, "engram_layer_ids", []
+        ):
             ckpt_plugin.set_hf_engram_ckpt(msg, hf_model, layer_id, md, margs)
         ckpt_plugin.set_hf_attn_ckpt(msg, hf_model, layer_id, md, margs)
         ckpt_plugin.set_hf_mlp_ckpt(msg, hf_model, layer_id, md, margs)