[ckpt] fix: support vocab padding with Megatron-FSDP TP

[yaoyu-33]

Summary

• compute the padded global vocabulary size correctly when the Megatron target is an MFSDP DTensor
• add TP=2 coverage to the existing MFSDP Hugging Face round-trip functional test

Root cause

DTensor.shape already describes the global tensor across TP ranks. The vocabulary-padding path multiplied that dimension by tp_size again, so a TP=2 import padded a 151,936-token embedding to 303,872 rows and scattered 151,936 rows to a local 75,968-row buffer.

This fixes the regression introduced by #1473 and tracked internally as NVBug 6565261.

Validation

• Before the fix, the Qwen3-0.6B MFSDP TP=2 round trip reproduced expected (75968, 1024), got (151936, 1024).
• After the fix, the same real-checkpoint round trip completed successfully and wrote all 226 tensors.
• uv run --no-sync python -m pytest tests/functional_tests/test_groups/converter/test_hf_fsdp_conversion.py -k FSDP_TP2 -v -s — 1 passed
• uv run pre-commit run --all-files — passed

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[claude]

LGTM.

The fix is correct. expect_dim0_size now accounts for DTensor.shape being the global shape across TP ranks in the Megatron-FSDP path, while a regular Megatron parameter stores only its local TP shard and still needs multiplication by tp_size. This matches the existing DTensor output_shape handling that uses orig_param.shape below it, and the global-size chunk along dim 0 produces per-rank shards matching each local shape. The new FSDP_TP2 parametrization gives the TP>1 path real coverage now that the Megatron-LM dependency is available.

No logic concerns, no typos, and no perf configs are touched.

Suggested test cases:

• test_hf_fsdp_roundtrip 1-1-1-FSDP_base
• test_hf_fsdp_roundtrip 2-1-2-FSDP_TP2

No perf tests impacted.