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

Equal to #4888 in dev branch

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.

[github-actions]

This PR has been automatically converted to draft because all PRs must start as drafts.

When you are ready for review, click Ready for Review to begin the review process. This will:

1. Add the oncall reviewer (optional reviewer)
2. Add required review teams based on your changes

See the contribution guide for more details.