Fix DeepSeek-V3 checkpoint export compatibility

[lihongyang1990]

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

[CLAassistant]

All committers have signed the CLA.

[lxd-cumt]

LGTM

[LiJunscs]

The checkpoint convert method is legacy. It does not handle the tensor_model_parallel_size and expert_tensor_parallel_size. Current codes load tp * ep parts of checkpoint. But mcore only max(tp, ep * etp) parts of checkpoint when ckpt_format is torch. See NVIDIA/Megatron-LM#4200.

[lihongyang1990]

The checkpoint convert method is legacy. It does not handle the tensor_model_parallel_size and expert_tensor_parallel_size. Current codes load tp * ep parts of checkpoint. But mcore only max(tp, ep * etp) parts of checkpoint when ckpt_format is torch. See NVIDIA/Megatron-LM#4200.

Thanks for pointing this out.

I have updated the converter to handle the MCore legacy torch checkpoint shard layout. The mcore loader/saver no longer assumes tp * ep
checkpoint parts; it now uses max(tp, ep * etp) for ckpt_format=torch, and the rank mapping distinguishes tensor parallel rank, expert
parallel rank, and expert tensor parallel rank.

The DeepSeek-V3/Mixtral checkpoint plugins were also updated accordingly:

• dense/attention/shared-expert tensors are gathered/split by tensor_model_parallel_size
• routed expert tensors are gathered/split by expert_tensor_parallel_size
• expert_model_parallel_size is used only for expert ownership

The reason I did not catch this initially is that my test config had:

tp = 1
ep = 4
etp = 1

So both formulas produced the same number of checkpoint parts:

tp * ep = 4
max(tp, ep * etp) = max(1, 4 * 1) = 4

Because of that, the old logic worked for my local case but would fail for layouts such as tp=4, ep=2, etp=1, where:

tp * ep = 8
max(tp, ep * etp) = 4

The PR has been updated with this fix.

[LiJunscs]

The checkpoint convert method is legacy. It does not handle the tensor_model_parallel_size and expert_tensor_parallel_size. Current codes load tp * ep parts of checkpoint. But mcore only max(tp, ep * etp) parts of checkpoint when ckpt_format is torch. See NVIDIA/Megatron-LM#4200.

Thanks for pointing this out.

I have updated the converter to handle the MCore legacy torch checkpoint shard layout. The mcore loader/saver no longer assumes tp * ep checkpoint parts; it now uses max(tp, ep * etp) for ckpt_format=torch, and the rank mapping distinguishes tensor parallel rank, expert parallel rank, and expert tensor parallel rank.

The DeepSeek-V3/Mixtral checkpoint plugins were also updated accordingly:

• dense/attention/shared-expert tensors are gathered/split by tensor_model_parallel_size
• routed expert tensors are gathered/split by expert_tensor_parallel_size
• expert_model_parallel_size is used only for expert ownership

The reason I did not catch this initially is that my test config had:

tp = 1
ep = 4
etp = 1

So both formulas produced the same number of checkpoint parts:

tp * ep = 4
max(tp, ep * etp) = max(1, 4 * 1) = 4

Because of that, the old logic worked for my local case but would fail for layouts such as tp=4, ep=2, etp=1, where:

tp * ep = 8
max(tp, ep * etp) = 4

The PR has been updated with this fix.

[LiJunscs]

Sorry, I mistakenly closed this issue.
And I create a new issue about saving the ckpt when tp > ep * etp.
#1204

[lihongyang1990]

Sorry, I mistakenly closed this issue. And I create a new issue about saving the ckpt when tp > ep * etp. #1204

Thanks for the clarification.

I have removed the broader TP/EP/ETP checkpoint converter changes from this PR and kept this PR focused on the DeepSeek-V3/Moonlight HF export
compatibility fixes.

This PR now only keeps the conversion-side fixes needed for the DeepSeek-V3/Moonlight HF export path:

• add --skip-mtp so HF export can skip MTP modules when the target HF implementation only contains the main LM layers
• guard Engram-specific checkpoint fields so non-Engram checkpoints do not process Engram paths
• align DeepSeek-V3 HF config export with the current Megatron argument name for RMSNorm epsilon
• keep HF config vocab size consistent with --true-vocab-size
• remove unsupported legacy Megatron conversion args / fused kernel loading for the current Megatron version
• use the current vocab_size_with_padding helper name

The generic TP/EP/ETP checkpoint saving issue is now left to #1204, since that is a training checkpoint save-path fix and should be handled
separately from this HF export compatibility PR.

[LiJunscs]

Sorry, I mistakenly closed this issue. And I create a new issue about saving the ckpt when tp > ep * etp. #1204

Thanks for the clarification.

I have removed the broader TP/EP/ETP checkpoint converter changes from this PR and kept this PR focused on the DeepSeek-V3/Moonlight HF export compatibility fixes.

This PR now only keeps the conversion-side fixes needed for the DeepSeek-V3/Moonlight HF export path:

• add --skip-mtp so HF export can skip MTP modules when the target HF implementation only contains the main LM layers
• guard Engram-specific checkpoint fields so non-Engram checkpoints do not process Engram paths
• align DeepSeek-V3 HF config export with the current Megatron argument name for RMSNorm epsilon
• keep HF config vocab size consistent with --true-vocab-size
• remove unsupported legacy Megatron conversion args / fused kernel loading for the current Megatron version
• use the current vocab_size_with_padding helper name

The generic TP/EP/ETP checkpoint saving issue is now left to #1204, since that is a training checkpoint save-path fix and should be handled separately from this HF export compatibility PR.

I think you misunderstand my meaning. The #1204 is only fix the bug of saving checkpoint, not the converting chechkpoint between mcore and hf. You previous commits are necessary. The main branch of now saves checkpoint only when edp_rank = 0, it does not save ckpt when dp = 0 and edp != 0(tp > ep * etp), some ckpts are missing.
My pr makes ckpts saved correctly when tp > ep * etp, your pr makes ckpt convert correctly.
Our prs work together to fix the whole bugs. I am sorry for the trouble I've caused you.

[lihongyang1990]

Sorry, I mistakenly closed this issue. And I create a new issue about saving the ckpt when tp > ep * etp. #1204

Thanks for the clarification.
I have removed the broader TP/EP/ETP checkpoint converter changes from this PR and kept this PR focused on the DeepSeek-V3/Moonlight HF export compatibility fixes.
This PR now only keeps the conversion-side fixes needed for the DeepSeek-V3/Moonlight HF export path:

• add --skip-mtp so HF export can skip MTP modules when the target HF implementation only contains the main LM layers
• guard Engram-specific checkpoint fields so non-Engram checkpoints do not process Engram paths
• align DeepSeek-V3 HF config export with the current Megatron argument name for RMSNorm epsilon
• keep HF config vocab size consistent with --true-vocab-size
• remove unsupported legacy Megatron conversion args / fused kernel loading for the current Megatron version
• use the current vocab_size_with_padding helper name

The generic TP/EP/ETP checkpoint saving issue is now left to #1204, since that is a training checkpoint save-path fix and should be handled separately from this HF export compatibility PR.

I think you misunderstand my meaning. The #1204 is only fix the bug of saving checkpoint, not the converting chechkpoint between mcore and hf. You previous commits are necessary. The main branch of now saves checkpoint only when edp_rank = 0, it does not save ckpt when dp = 0 and edp != 0(tp > ep * etp), some ckpts are missing. My pr makes ckpts saved correctly when tp > ep * etp, your pr makes ckpt convert correctly. Our prs work together to fix the whole bugs. I am sorry for the trouble I've caused you.

Thanks for the clarification, no worries.

I understand now: #1204 fixes checkpoint saving, while this PR still needs the MCore <-> HF TP/EP/ETP conversion fixes. I have restored those
converter changes and updated the PR.