Skip to content

[bugfix] Fix Qwen4Exp PLE HF export and TP checkpoint metadata - #192

Merged
hjh0119 merged 1 commit into
modelscope:mainfrom
Zhikaiiii:fix/qwen4-ple-export-tp-resharding
Sep 15, 2026
Merged

hjh0119 merged 1 commit into
modelscope:mainfrom
Zhikaiiii:fix/qwen4-ple-export-tp-resharding

Conversation

@Zhikaiiii

@Zhikaiiii Zhikaiiii commented Sep 14, 2026 •

Copy link
Copy Markdown
Contributor

Summary

  • Fix missing trained Qwen4Exp PLE ngram weights in HF export when cpu_offload=False: export the current embedding TP partition and assemble the HF shards across TP ranks using the existing table export path.
  • Fix PLE checkpoint metadata for TP resharding (e.g. training with TP8 and exporting with TP2): recursively call child sharded_state_dict methods so VocabParallelEmbedding supplies its vocabulary-axis TP metadata instead of being marked replicated. Keep hash buffers replicated.
  • Add tests/test_ple_checkpoint.py covering updated-weight export with simulated TP1/2/8 ranks (with and without FP8 scale), and TP8/TP2 checkpoint layouts.

Validation

  • Pre-commit checks on both changed files.
  • Eight isolated cases passed using the source methods and real PyTorch/Megatron metadata helpers, with simulated TP ranks and collectives.

@hjh0119

hjh0119 commented Sep 15, 2026

Copy link
Copy Markdown
Collaborator

thanks!

@hjh0119
hjh0119 merged commit 2370d15 into modelscope:main Sep 15, 2026
1 check passed
hjh0119 added a commit to hjh0119/mcore-bridge that referenced this pull request Sep 15, 2026
…): resolve ple.py -- keep modelscope#192 unified naming + chunked all_reduce with CPU accumulation
hjh0119 added a commit to hjh0119/mcore-bridge that referenced this pull request Sep 15, 2026
- ple.py: split export_table_to_hf into a generator (iter_export_table_to_hf)
  that yields one shard at a time; the bulk API remains as a thin wrapper
- qwen4_exp.py: stage the shard iterator in _set_layer_ple and interleave the
  cross-PP broadcast per shard (_iter_ple_table_export); consumed at _convert's
  yield point so no rank accumulates the 100GB-scale table in host memory
- gpt_bridge.py: drain the staged iterator at the per-layer yield point and
  make the per-layer yield lazy

Verified: 13 unit tests (modelscope#192's test_ple_checkpoint included), medium e2e
TP2xPP2xEP2xETP2 (0 missing keys), and a load->export roundtrip that is
bit-exact on all 16 PLE shards and 12 sampled other weights
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants