[model] fix: preserve singleton expert scale dimensions

[yaoyu-33]

Problem

Blockwise-FP8 Hugging Face export with expert parallelism can silently drop a valid singleton block-grid dimension from expert scale tensors. For example, a supported MoE down projection with per-expert scale shape [2, 1] and EP=2 is staged as [2, 2, 1], but the current gather path returns [2, 2]. Grouped export consequently emits [4, 2] instead of the required [4, 2, 1] layout.

The public trigger is an AutoBridge FP8 export of a fused MoE model using blockwise scales and EP greater than one. The malformed scale shape can make the exported checkpoint incompatible with the Hugging Face quantized expert layout or give it incorrect quantization metadata.

Related context: #4804 tracks native quantized export workflows, but does not fix this shape-loss root cause.

Root cause and fix

gather_from_ep_ranks() uses an extra leading dimension to stage tensors collected from each EP rank. Calling unqualified squeeze() removes both that staging dimension and any legitimate singleton dimensions already present in the source tensor.

Use squeeze(0) so only the EP staging axis is removed. No API or supported-format scope changes are included.

Validation

Fail-before/pass-after deterministic CPU reproducer against the actual gather_from_ep_ranks() implementation:

uv run --no-project python /tmp/reproduce_ep_gather_shape.py
before: exit 1 — AssertionError: expected (2, 2, 1), got (2, 2)
after:  exit 0

Focused regression and adjacent expert-mapping tests:

uv run --no-project python -m pytest \
  tests/unit_tests/models/test_param_mapping.py::test_ep_gather_preserves_expert_scale_singleton_dimensions \
  tests/unit_tests/models/test_param_mapping.py::TestFusedExpertMapping \
  tests/unit_tests/models/test_param_mapping.py::TestFusedGatedExpertMapping -q
7 passed, 42 warnings in 1.03s

Additional checks:

git diff --check
uv run pre-commit run --all-files

Both passed.

Scope

This changes only removal of the gather-owned staging axis and adds one focused CPU regression test. It does not change conversion APIs, quantization formats, model registrations, or Megatron-Core.

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

[yaoyu-33]

/ok to test 70da711

[claude]

Light review — LGTM with one optional follow-up.

The fix is correct. gather_from_ep_ranks stacks per-EP-rank tensors with .unsqueeze(0) (staging dim) and then, in the single-expert-per-rank case, must strip only that dim. The old bare .squeeze() removed all singleton dimensions, so an expert scale of shape [E_local, 1] lost its trailing dim and produced a wrongly-shaped export. .squeeze(0) removes only the staging dim, which is what the docstring intends. The new unit test exercises exactly this (scale shaped [2, 1], gathered result [2, 2, 1]) and does not require a GPU.

One non-blocking observation: _offset_gather_from_ep_ranks in src/megatron/bridge/models/ernie_vl/ernie45_vl_bridge.py is a copy of this routine and still uses bare .squeeze() (line 173). It is safe today (only 2D expert weights flow through it) but carries the same latent bug — worth aligning for consistency. Left an inline suggestion.

Suggested test cases

• test_ep_gather_preserves_expert_scale_singleton_dimensions (added in this PR)
• No perf tests impacted.