diff --git a/src/megatron/bridge/models/conversion/quantization_utils.py b/src/megatron/bridge/models/conversion/quantization_utils.py index 379ce8496f..f251a12a66 100644 --- a/src/megatron/bridge/models/conversion/quantization_utils.py +++ b/src/megatron/bridge/models/conversion/quantization_utils.py @@ -319,35 +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. - """ - 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) + 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. + """ + 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) - 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: + 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)