[dev] Fix FSDP TP metadata for LinearCrossEntropyModule

[conver334]

What does this PR do ?

Preserve column tensor-parallel DTensor metadata for LinearCrossEntropyModule when wrapping Megatron modules with FSDP.

Reference: https://github.com/NVIDIA-NeMo/Megatron-Bridge/blob/9c5dcea6e8233be4033d10fcc552880ddf32687a/src/megatron/bridge/models/conversion/param_mapping.py#L1200

Issue tracking

Problem

When converting a Qwen2.5-Math-7B / Qwen3.5-35B-A3B checkpoint through the Megatron-Bridge HF -> Megatron-FSDP roundtrip with TP=2, lm_head.weight was exported with only the local tensor-parallel shard instead of the full vocabulary dimension.
Brief failure output:

  name=lm_head.weight
  export_shape=(76032, 3584)
  original_shape=(152064, 3584)

Root cause

LinearCrossEntropyModule was not registered in the FSDP adapter's column-parallel module registry. As a result, FSDP did not preserve the column tensor-parallel metadata for output_layer.weight.

Before this fix, the parameter metadata looked like:

  module_type=LinearCrossEntropyModule
  shape=(76032, 3584)
  tp_mode=None
  placements=(Shard(dim=0),)
  mesh=DeviceMesh((dp_cp=4), ...)

Fix

Add LinearCrossEntropyModule to _MODULE_TYPE_REGISTRY["column"] in the FSDP adapter.
After this change, output_layer.weight keeps the expected column TP metadata:

  module_type=LinearCrossEntropyModule
  shape=(152064, 3584)
  tp_mode=column
  placements=(Shard(dim=0), Shard(dim=0))
  mesh=DeviceMesh((dp_cp=4, tp=2), ...)

Validation

Megatron-Bridge Megatron-FSDP -> HF roundtrip on Qwen2.5-Math-7B with TP=2, CP=1 is correct.

Contribution process

Pre-checks

• I have added relevant unit tests
• I have added relevant functional tests
• I have added proper typing to my code
• I have added relevant documentation
• I have run the autoformatter.sh on my PR

[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.

[wplf]

LGTM

[Victarry]

Hi @conver334 , could you please also submit a fix PR to the main branch?

[Victarry]

/ok to test 5f8041f

[svcnvidia-nemo-ci]

🔄 Merge queue validation started!

You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/26213051081