From b2d2908bed172fd7c44d4d7bf82a0331870f8c82 Mon Sep 17 00:00:00 2001 From: Ando Tomoki Date: Wed, 12 Aug 2026 10:55:00 +0900 Subject: [PATCH 1/2] fix(quant): dequantize MXFP4 on GPU under the distributed conversion backend dequantize_mxfp4_e2m1_packed operated on whatever device the input tensors were already on. AutoBridge.from_hf_pretrained loads HF state dicts on CPU, so under scripts/conversion's GPU backend the entire dequantization (bit unpack, LUT gather, repeat_interleave, multiply) ran on CPU, leaving GPUs idle and making a single CPU process the bottleneck for the whole distributed conversion. Move the packed weight and scale to the current CUDA device first when running under torch.distributed (guarded on is_initialized() so the single-process CPU backend, which never calls init_process_group, is unaffected). Observed ~20-30x slower than expected converting moonshotai/Kimi-K3 (48 GPUs, TP2/PP3/EP8/ETP2); confirmed via matching progress-bar ETA on the full model and directly measured throughput on a cheap toy-model reproduction, both independently converging on ~1 weight/sec instead of completing in minutes. 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 the fix is not 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 issue in practice. Left out of this PR to keep it focused; flagging in case it's worth a follow-up. Signed-off-by: Ando Tomoki --- .../bridge/models/conversion/quantization_utils.py | 13 +++++++++++++ 1 file changed, 13 insertions(+) diff --git a/src/megatron/bridge/models/conversion/quantization_utils.py b/src/megatron/bridge/models/conversion/quantization_utils.py index 379ce8496f..21910302c1 100644 --- a/src/megatron/bridge/models/conversion/quantization_utils.py +++ b/src/megatron/bridge/models/conversion/quantization_utils.py @@ -325,6 +325,19 @@ def dequantize_mxfp4_e2m1_packed( ``scale`` is expected to be one scale per row and per K tile. ``uint8`` E8M0 tensors use exponent bias 127 and are decoded to powers of two. """ + # HF state dicts loaded via AutoBridge.from_hf_pretrained are CPU-resident, + # so without this every op below (bit unpack, LUT gather, repeat_interleave, + # multiply) runs on CPU even under scripts/conversion's GPU backend, leaving + # GPUs idle while a single CPU process becomes the bottleneck for the whole + # distributed conversion (observed ~20-30x slower than expected on a + # multi-GPU MoE conversion). Guarded on is_initialized() so the + # single-process CPU backend, which never calls init_process_group, is + # unaffected. + if torch.cuda.is_available() and torch.distributed.is_initialized(): + device = torch.device("cuda", torch.cuda.current_device()) + weight_packed = weight_packed.to(device) + scale = scale.to(device) + w_u8 = weight_packed.view(torch.uint8) lo = (w_u8 & 0xF).to(torch.int64) hi = (w_u8 >> 4).to(torch.int64) From a42b5be0af37ae9f747e190b61c714110c6c8ed6 Mon Sep 17 00:00:00 2001 From: Ando Tomoki Date: Sat, 15 Aug 2026 01:51:07 +0900 Subject: [PATCH 2/2] fix(quant): make MXFP4 dequant device-explicit and memory-bounded Addresses review feedback on the previous commit: - dequantize_mxfp4_e2m1_packed no longer infers the GPU backend via torch.distributed.is_initialized(). It now takes an explicit `device` parameter (default None preserves the input tensor's device, matching dequantize_int4's convention), and the decision to pass a CUDA device moves to the call site in kimi_k3_bridge.py's _load_one_hf_weight, which is where "are we running under the GPU conversion backend" is an appropriate question to ask. - Rewrote the body to process rows_per_chunk rows at a time (default matches dequantize_mxfp4's), gathering the LUT lookups directly into interleaved slices of a single preallocated output tensor and applying the scale in place, instead of materializing separate full-size fp32 buffers for the stacked LUT gather, the broadcast scale, and the product. This bounds peak transient memory regardless of tensor size instead of scaling it up by roughly an order of magnitude over the packed input, which was flagged as a realistic OOM risk for large expert matrices under low headroom. - Added tests: CPU device preserved by default, explicit CUDA device matches CPU numerically (skipped without a GPU), and chunked output matches unchunked output. Manually re-verified (roundtrip, uint8 E8M0 scale, geometry rejection, CPU-default, explicit CUDA device + non-mutation of inputs, and chunked vs. unchunked equivalence) on a GPU node before pushing. Not addressed here: only the TP/ETP mapping owner rank needs the full dequantized tensor; other ranks in the same TP group currently still do the full transform and discard it. Fixing that means threading ownership into maybe_modify_loaded_hf_weight's call signature, which is shared by ~10 other model bridges - posting on the PR to check the preferred approach before touching that contract. Signed-off-by: Ando Tomoki --- .../models/conversion/quantization_utils.py | 87 ++++++++++++------- .../bridge/models/kimi/kimi_k3_bridge.py | 11 +++ .../models/test_quantization_utils.py | 42 +++++++++ 3 files changed, 108 insertions(+), 32 deletions(-) diff --git a/src/megatron/bridge/models/conversion/quantization_utils.py b/src/megatron/bridge/models/conversion/quantization_utils.py index 21910302c1..f251a12a66 100644 --- a/src/megatron/bridge/models/conversion/quantization_utils.py +++ b/src/megatron/bridge/models/conversion/quantization_utils.py @@ -319,48 +319,71 @@ def dequantize_mxfp4_e2m1_packed( scale: torch.Tensor, *, dtype: torch.dtype = torch.bfloat16, + device: str | torch.device | None = None, + rows_per_chunk: int = 32768 * 1024, ) -> torch.Tensor: """Dequantize MXFP4 E2M1 weights packed two values per byte. ``scale`` is expected to be one scale per row and per K tile. ``uint8`` E8M0 tensors use exponent bias 127 and are decoded to powers of two. + + By default this dequantizes on whatever device ``weight_packed`` is + already on. Pass ``device`` to move the (compact) packed input there + first and dequantize on that device instead, e.g. the current CUDA + device when running under a distributed GPU backend; callers decide + when that's appropriate, this function does not infer it from + ``torch.distributed`` state. + + Processes ``rows_per_chunk`` rows at a time so peak memory stays + bounded regardless of tensor size: unpacking to indices, the LUT + gather, and the scale multiply otherwise each materialize a + full-size fp32 tensor at once (on top of the bf16 output), which for + a single large expert weight can multiply into a large transient + footprint. """ - # HF state dicts loaded via AutoBridge.from_hf_pretrained are CPU-resident, - # so without this every op below (bit unpack, LUT gather, repeat_interleave, - # multiply) runs on CPU even under scripts/conversion's GPU backend, leaving - # GPUs idle while a single CPU process becomes the bottleneck for the whole - # distributed conversion (observed ~20-30x slower than expected on a - # multi-GPU MoE conversion). Guarded on is_initialized() so the - # single-process CPU backend, which never calls init_process_group, is - # unaffected. - if torch.cuda.is_available() and torch.distributed.is_initialized(): - device = torch.device("cuda", torch.cuda.current_device()) - weight_packed = weight_packed.to(device) - scale = scale.to(device) - - w_u8 = weight_packed.view(torch.uint8) - lo = (w_u8 & 0xF).to(torch.int64) - hi = (w_u8 >> 4).to(torch.int64) - - table = torch.tensor(_FP4_E2M1_TABLE_VALUES, dtype=torch.float32, device=weight_packed.device) - logical = torch.stack([table[lo], table[hi]], dim=-1).reshape(weight_packed.shape[0], -1) - - if scale.dtype == torch.uint8: - scale_f32 = torch.ldexp( - torch.ones_like(scale, dtype=torch.float32), - scale.to(torch.int32) - 127, - ) - else: - scale_f32 = scale.to(torch.float32) - if scale_f32.dim() != 2 or scale_f32.shape[0] != logical.shape[0] or logical.shape[1] % scale_f32.shape[1] != 0: + target_device = weight_packed.device if device is None else torch.device(device) + weight_packed = weight_packed.to(target_device) + scale = scale.to(target_device) + + rows_total, packed_cols = weight_packed.shape + logical_cols = packed_cols * 2 + if scale.dim() != 2 or scale.shape[0] != rows_total or logical_cols % scale.shape[1] != 0: raise RuntimeError( f"Unsupported MXFP4 scale geometry: " - f"weight={tuple(weight_packed.shape)} logical={tuple(logical.shape)} scale={tuple(scale.shape)}" + f"weight={tuple(weight_packed.shape)} logical={(rows_total, logical_cols)} scale={tuple(scale.shape)}" ) - block_size = logical.shape[1] // scale_f32.shape[1] - scale_exp = scale_f32.repeat_interleave(block_size, dim=1) + block_size = logical_cols // scale.shape[1] + + table = torch.tensor(_FP4_E2M1_TABLE_VALUES, dtype=torch.float32, device=target_device) + out = torch.empty(rows_total, logical_cols, dtype=dtype, device=target_device) + + for r0 in range(0, rows_total, rows_per_chunk): + r1 = min(r0 + rows_per_chunk, rows_total) + + w_u8 = weight_packed[r0:r1].view(torch.uint8) + idx_lo = (w_u8 & 0xF).to(torch.int64) + idx_hi = (w_u8 >> 4).to(torch.int64) + + logical_chunk = torch.empty(r1 - r0, logical_cols, dtype=torch.float32, device=target_device) + logical_chunk[:, 0::2] = table[idx_lo] + logical_chunk[:, 1::2] = table[idx_hi] + del idx_lo, idx_hi, w_u8 + + scale_slice = scale[r0:r1] + if scale_slice.dtype == torch.uint8: + scale_chunk = torch.ldexp( + torch.ones_like(scale_slice, dtype=torch.float32), + scale_slice.to(torch.int32) - 127, + ) + else: + scale_chunk = scale_slice.to(torch.float32) + scale_exp = scale_chunk.repeat_interleave(block_size, dim=1) + + logical_chunk.mul_(scale_exp) + out[r0:r1] = logical_chunk.to(dtype) + del logical_chunk, scale_chunk, scale_exp - return (logical * scale_exp).to(dtype) + return out def is_mxfp4_e2m1_scale_geometry( diff --git a/src/megatron/bridge/models/kimi/kimi_k3_bridge.py b/src/megatron/bridge/models/kimi/kimi_k3_bridge.py index b66b31ee51..f066284b89 100644 --- a/src/megatron/bridge/models/kimi/kimi_k3_bridge.py +++ b/src/megatron/bridge/models/kimi/kimi_k3_bridge.py @@ -271,9 +271,20 @@ def _load_one_hf_weight(name: str, hf_state_dict: Mapping[str, torch.Tensor]) -> if packed_key and packed_key in hf_state_dict: if scale_key not in hf_state_dict: raise ValueError(f"Missing MXFP4 scale for {packed_key}") + # AutoBridge.from_hf_pretrained loads HF state dicts on CPU. Under + # the distributed GPU conversion backend, dequantize on the + # current CUDA device instead so this doesn't become a CPU-bound + # bottleneck for the whole conversion; the single-process CPU + # backend (which never initializes a process group) is unaffected. + device = ( + torch.device("cuda", torch.cuda.current_device()) + if torch.cuda.is_available() and torch.distributed.is_initialized() + else None + ) return quantization_utils.dequantize_mxfp4_e2m1_packed( hf_state_dict[packed_key], hf_state_dict[scale_key], + device=device, ) return hf_state_dict[name] diff --git a/tests/unit_tests/models/test_quantization_utils.py b/tests/unit_tests/models/test_quantization_utils.py index d947a911c5..5882e3b7f1 100644 --- a/tests/unit_tests/models/test_quantization_utils.py +++ b/tests/unit_tests/models/test_quantization_utils.py @@ -240,6 +240,48 @@ def test_mxfp4_helpers_reject_unsupported_geometry(): quantize_mxfp4_e2m1_like_scale(torch.ones(1, 30), torch.ones(1, 1)) +def _mxfp4_roundtrip_fixture(num_rows: int): + values = torch.tensor( + [0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0, -0.0, -0.5, -1.0, -1.5, -2.0, -3.0, -4.0, -6.0], + dtype=torch.float32, + ).repeat(2) + weight = values.reshape(1, 32).repeat(num_rows, 1).to(torch.bfloat16) + source_scale = torch.arange(1, num_rows + 1, dtype=torch.float32).reshape(num_rows, 1) + return quantize_mxfp4_e2m1_like_scale(weight, source_scale) + (weight,) + + +def test_dequantize_mxfp4_e2m1_packed_preserves_cpu_device_by_default(): + packed, scale, weight = _mxfp4_roundtrip_fixture(num_rows=4) + + result = dequantize_mxfp4_e2m1_packed(packed, scale) + + assert result.device.type == "cpu" + assert torch.equal(result.float(), weight.float()) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") +def test_dequantize_mxfp4_e2m1_packed_moves_to_explicit_cuda_device(): + packed, scale, weight = _mxfp4_roundtrip_fixture(num_rows=4) + + result = dequantize_mxfp4_e2m1_packed(packed, scale, device="cuda") + + assert result.device.type == "cuda" + assert torch.equal(result.float().cpu(), weight.float()) + # Inputs on CPU are untouched; the function moves copies, not in place. + assert packed.device.type == "cpu" + assert scale.device.type == "cpu" + + +def test_dequantize_mxfp4_e2m1_packed_chunking_matches_unchunked(): + packed, scale, weight = _mxfp4_roundtrip_fixture(num_rows=5) + + unchunked = dequantize_mxfp4_e2m1_packed(packed, scale) + chunked = dequantize_mxfp4_e2m1_packed(packed, scale, rows_per_chunk=2) + + assert torch.equal(chunked, unchunked) + assert torch.equal(chunked.float(), weight.float()) + + def test_maybe_dequantize_hf_quantized_weight_dispatches_by_dtype_and_sibling_scale(): fp8_weight = torch.ones(2, 2, dtype=torch.float8_e4m3fn) mxfp4_weight = torch.zeros(1, 16, dtype=torch.int8)