fix(quant): dequantize MXFP4 on GPU under the distributed conversion backend

[mikan-atomoki]

What does this PR do ?

Dequantize MXFP4-packed weights on GPU instead of CPU when dequantize_mxfp4_e2m1_packed runs under the distributed GPU conversion backend, fixing a ~20-30x slowdown.

Changelog

• src/megatron/bridge/models/conversion/quantization_utils.py: dequantize_mxfp4_e2m1_packed now moves weight_packed/scale to the current CUDA device before dequantizing, when running under torch.distributed (guarded on torch.distributed.is_initialized()). Previously it operated on whatever device the input tensors were already on; since AutoBridge.from_hf_pretrained loads HF state dicts on CPU, the whole dequantization (bit unpack, LUT gather, repeat_interleave, multiply) ran on CPU even under scripts/conversion's GPU backend, leaving GPUs idle and making a single CPU process the bottleneck for the entire distributed conversion. The single-process CPU backend (which never calls init_process_group) is unaffected.

GitHub Actions CI

N/A (external contributor without CI trigger permissions; happy to have a maintainer trigger CI).

Before your PR is "Ready for review"

Pre checks:

• Make sure you read and followed Contributor guidelines
• Did you write any new necessary tests? — No. The regression is about wall-clock throughput under a real multi-GPU distributed process group, not something the existing unit tests for this function (which check numerical correctness of dequantization, unaffected by this change) exercise. Verified manually instead (see Additional Information).
• Did you add or update any necessary documentation? — No. This is a performance fix with no behavior, API, or config change.
• Does the PR affect components that are optional to install? (Ex: Numba, Pynini, Apex etc) — No, only core torch/torch.distributed.
  • Reviewer: Does the PR have correct import guards for all optional libraries? — N/A, see above.

Additional Information

Found while converting moonshotai/Kimi-K3 (48 GPUs, TP2/PP3/EP8/ETP2). Observed ~20-30x slower than expected: the progress-bar ETA on the full model and directly measured throughput on a cheap toy-model reproduction both independently converged on ~1 weight/sec instead of completing in minutes. Root-caused with a single-process probe showing GPU utilization near 0% and CPU pegged during the dequantize step. Fixed and re-verified end to end at the same 48-GPU scale.

The same function is also reached from DeepSeek V4's bridge via maybe_dequantize_hf_quantized_weight, so this isn't Kimi-K3-specific.

Note for reviewers: dequantize_int4 (used by Kimi K2.5 VL) takes an explicit device parameter with the same intent, but its only caller (kimi_k25_vl_bridge.py) passes device=hf_state_dict[packed_key].device — i.e. the tensor's own (CPU) device — so it likely has the same CPU-bound behavior in practice. Left out of this PR to keep it focused on the reproduced/verified case; flagging in case it's worth a follow-up.

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

The GPU-dequantization direction makes sense, but I think the current implementation needs a bounded-memory/backend-aware design before we merge it.

torch.distributed.is_initialized() is not a GPU conversion backend contract. This helper runs before hf_to_megatron() decides the TP/ETP source/owner, so every rank that has the task first moves and dequantizes the full HF tensor. In TP/ETP mappings, non-source ranks then discard that full result while rank 0 splits/scatters it. It can also change the device semantics of this otherwise tensor-local helper in any future distributed caller that happens to have CUDA visible.

The transient memory is much larger than the packed input: both nibble indices are int64, both LUT gathers and logical are FP32, repeat_interleave materializes a full FP32 scale tensor, the multiply creates another full FP32 tensor, and the BF16 result coexists during the final cast. The rough peak is about 20-22.5 bytes per logical element, before the Megatron model shard and any mapping/scatter copies. For a Kimi K3 expert matrix of roughly 3584 x 7168, that is around 0.5 GiB of transient GPU memory per participating rank for one weight; gate/up retention and low-headroom model shards make OOM realistic.

Could we instead:

1. Pass an explicit execution/target device from the GPU backend (preserving the input device by default), rather than infer the backend from process-group state.
2. Dequantize in bounded row chunks, following dequantize_mxfp4(): move only each packed/scale slice to CUDA, write the BF16 slice to the intended destination, and release intermediates.
3. Arrange for only the mapping source/owner rank to perform the full source transform. If the final destination is CPU, copy each completed chunk back; if this is the GPU backend, keep the result on GPU. Moving the complete result back only at the end would still create an avoidable peak and PCIe burst.
4. Add focused coverage for CPU device preservation, CUDA numerical/device behavior, and the chunked path.

Direct use of the GPU conversion backend is the right fast path for these large models; it just needs to make GPU execution explicit and memory-bounded.

[mikan-atomoki]

Thanks for the detailed review. Pushed a commit addressing the device-explicitness and memory-bounding points:

• dequantize_mxfp4_e2m1_packed now takes an explicit device param (default None preserves the input's device); torch.distributed.is_initialized() is gone from the function itself and the decision moved to the call site in kimi_k3_bridge.py.
• Rewrote the body to match dequantize_mxfp4()'s chunked pattern: preallocate the output once, gather LUT lookups directly into interleaved output slices, apply the scale in place, process rows_per_chunk rows at a time. This bounds peak transient memory instead of scaling it up with several coexisting full-size fp32 buffers.
• Added CPU-default, explicit-CUDA-device, and chunking-equivalence tests, and manually re-verified all of the existing + new cases (roundtrip, uint8 E8M0 scale, geometry rejection, CPU default, explicit CUDA device + non-mutation of inputs, chunked vs. unchunked) on a GPU node before pushing.

On point 3 (only the TP/ETP mapping owner rank needs the full dequantized tensor) - I didn't want to guess at the right shape for this before checking with you, since it's not contained to this function. As I understand the call path, maybe_modify_loaded_hf_weight (and this function through it) runs in model_bridge.py's loop before task.mapping.hf_to_megatron() decides the TP/ETP source rank, and maybe_modify_loaded_hf_weight is overridden by ~10 model bridges with the same 2-arg signature. The most direct fix I can see is threading ownership info (e.g. an is_owner: bool = True-style default-safe kwarg) through that shared signature so non-owner ranks can skip the expensive transform - which is mechanical for the bridges that don't care, but does touch the shared contract.

Is that the direction you'd want, or is there a narrower way to gate this you'd prefer? Happy to send it as a separate PR once we agree on the shape, rather than bundling a framework-wide signature change into this one.