Is parallel (multi-process) checkpoint conversion supported?

[liyandong001]

Is your feature request related to a problem? Please describe.
I’m converting a Megatron checkpoint to HuggingFace safetensors using Megatron-Bridge.
To speed up conversion for a large model (e.g., 80B with TP=2, PP=4), I modified the conversion/export script to run in parallel (multi-process) by splitting output files and writing shards concurrently.

After enabling parallel conversion, the generated HuggingFace checkpoint appears to be incorrect:
• Shapes/dtypes look mostly consistent, but some parameters’ values differ from the expected “known-good” HF safetensors (or the model produces incorrect outputs).
• The issue only happens with the parallel version; the original sequential conversion produces correct weights.

I suspect there may be a non-deterministic / non-process-safe step in the conversion pipeline (e.g., parameter collection / ordering, TP/PP gather logic, key-to-filename mapping, or write ordering).


[liyandong001]

I found that the issue is not limited to a single parameter. In my setup (TP=2), three parameters are affected and show consistent block-level permutation after TP gather:
• model.layers..linear_attn.conv1d.weight
• model.layers..linear_attn.in_proj_ba.weight
• model.layers.*.linear_attn.in_proj_qkvz.weight

For these parameters:
• The per-TP shards (tp0, tp1) dumped before gather are internally consistent.
• A manual torch.cat([tp0, tp1], dim=0) exactly matches the exported (incorrect) HF tensor.
• Comparing against a known-correct HF checkpoint shows that the exported tensor is not corrupted, but differs by a deterministic block-level permutation along dim=0 (e.g., 1024-row blocks for conv1d.weight).

This suggests that for these linear_attn-related parameters, Megatron uses an internal layout under tensor parallelism, and an additional layout restoration (inverse permutation) step is required after TP gather when converting to HF format.

I’m currently applying a TP=2 special-case fix to reverse the block permutation after gather as a workaround, but this likely deserves a proper module-specific mapping for linear_attn / Conv1d parameters.

[liyandong001]

Additional findings regarding TP gather and linear_attn / GDN / Conv1d

During further debugging of the parallel conversion path, I found that the current implementation relies on an implicit assumption that does not hold for linear_attn-related parameters.

Affected parameters (confirmed)
The following parameters are incorrect when TP > 1 (confirmed with TP=2):
• model.layers..linear_attn.conv1d.weight
• model.layers..linear_attn.in_proj_qkvz.weight
• model.layers.*.linear_attn.in_proj_ba.weight

Observations
• For these parameters:
• The output produced by generic TP gather (all_gather + cat(dim=0)) exactly matches the incorrect (BAD) weights
• Manually dumping each TP shard before gather and concatenating them reproduces the same BAD result
• This indicates that:
• Megatron stores these weights using a TP-specific internal layout
• Generic TP gather only reconstructs the internal Megatron layout, not the HF semantic layout

Block-level pattern (Conv1d)
• linear_attn.conv1d.weight (shape (8192, 1, 4)) shows a deterministic block permutation:
• Each 1024-row block in BAD matches one block in GOOD, but in a different order
• This is a layout issue, not numerical precision or communication noise

GDNLinear split precondition is violated
• split_gdn_linear_weights() assumes the gathered tensor is laid out as a semantic
[Q, K, V, Z, B, A] concatenation along dim=0
• However, the gathered packed_qkvzba tensor does not satisfy this assumption
• Reusing the split logic without restoring the semantic layout inevitably produces incorrect in_proj_qkvz / in_proj_ba

Root cause
The current implementation effectively assumes:
packed = self._tp_mapping.megatron_to_hf(...)

[yaoyu-33]

Thanks for contribution, replied in the new thread. Can you check @yuzhongw-nvidia 's pr?