From bc1ca20175f96470be7ad76262e9792cbe44879e Mon Sep 17 00:00:00 2001 From: Elfie Guo Date: Fri, 18 Sep 2026 13:59:15 -0700 Subject: [PATCH 1/5] Add Kimi K3 packed MXFP4 QAT import --- .../create_kimi_k3_mxfp4_fixture.py | 213 +++++++++++ .../test_kimi_k3_mxfp4_checkpoint.py | 97 +++++ .../cpu/test_mx_qat_grouped_experts.py | 91 +++++ .../unit_tests/cpu/test_packed_hf_storage.py | 333 ++++++++++++++++++ .../unit_tests/cpu/test_state_dict_adapter.py | 79 +++++ .../checkpointer/packed_hf_storage.py | 297 ++++++++++++++++ torchtitan/config/transform/__init__.py | 2 + torchtitan/config/transform/quantization.py | 42 +++ torchtitan/models/kimi_k3/config_registry.py | 21 ++ .../models/kimi_k3/state_dict_adapter.py | 103 +++++- torchtitan/quantization/mx_qat/__init__.py | 12 + torchtitan/quantization/mx_qat/experts.py | 57 +++ 12 files changed, 1346 insertions(+), 1 deletion(-) create mode 100755 scripts/checkpoint_conversion/create_kimi_k3_mxfp4_fixture.py create mode 100644 tests/integration_tests/test_kimi_k3_mxfp4_checkpoint.py create mode 100644 tests/unit_tests/cpu/test_mx_qat_grouped_experts.py create mode 100644 tests/unit_tests/cpu/test_packed_hf_storage.py create mode 100644 torchtitan/components/checkpointer/packed_hf_storage.py create mode 100644 torchtitan/quantization/mx_qat/__init__.py create mode 100644 torchtitan/quantization/mx_qat/experts.py diff --git a/scripts/checkpoint_conversion/create_kimi_k3_mxfp4_fixture.py b/scripts/checkpoint_conversion/create_kimi_k3_mxfp4_fixture.py new file mode 100755 index 00000000000..e51ea799973 --- /dev/null +++ b/scripts/checkpoint_conversion/create_kimi_k3_mxfp4_fixture.py @@ -0,0 +1,213 @@ +#!/usr/bin/env python3 +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +"""Create a release-format MXFP4 checkpoint for the Kimi-K3 debug model.""" + +import argparse +import hashlib +import json +import re +import subprocess +from pathlib import Path +from typing import Any + +import torch +import torchao +from safetensors.torch import save_file +from torchao.prototype.mx_formats.mx_tensor import MXTensor + +from torchtitan.models.kimi_k3 import KimiK3StateDictAdapter, model_registry + +_EXPERT_WEIGHT = re.compile( + r"^language_model\.model\.layers\.\d+\.block_sparse_moe\.experts\.\d+\." + r"w[123]\.weight$" +) +_DEFAULT_MAX_SHARD_BYTES = 1 << 30 + + +def _git_revision(path: Path) -> str: + return subprocess.check_output( + ("git", "-C", str(path), "rev-parse", "HEAD"), + text=True, + ).strip() + + +def _tensor_bytes(tensor: torch.Tensor) -> int: + return tensor.numel() * tensor.element_size() + + +def _sha256(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as handle: + for chunk in iter(lambda: handle.read(8 << 20), b""): + digest.update(chunk) + return digest.hexdigest() + + +def convert_hf_state_dict_to_mxfp4( + hf_state_dict: dict[str, torch.Tensor], +) -> tuple[dict[str, torch.Tensor], int]: + """Replace routed-expert weights with released Kimi packed/scale pairs.""" + converted: dict[str, torch.Tensor] = {} + pair_count = 0 + for key, value in hf_state_dict.items(): + value = value.detach().to(device="cpu").contiguous() + if _EXPERT_WEIGHT.fullmatch(key) is None: + converted[key] = value + continue + + mx = MXTensor.to_mx( + value, + elem_dtype=torch.float4_e2m1fn_x2, + block_size=32, + ) + prefix = key.removesuffix(".weight") + converted[f"{prefix}.weight_packed"] = mx.qdata.contiguous() + converted[f"{prefix}.weight_scale"] = ( + mx.scale.view(torch.uint8).contiguous() + ) + pair_count += 1 + return converted, pair_count + + +def _partition_shards( + state_dict: dict[str, torch.Tensor], + max_shard_bytes: int, +) -> list[dict[str, torch.Tensor]]: + shards: list[dict[str, torch.Tensor]] = [] + current: dict[str, torch.Tensor] = {} + current_bytes = 0 + for key in sorted(state_dict): + value = state_dict[key] + value_bytes = _tensor_bytes(value) + if current and current_bytes + value_bytes > max_shard_bytes: + shards.append(current) + current = {} + current_bytes = 0 + current[key] = value + current_bytes += value_bytes + if current: + shards.append(current) + return shards + + +def write_sharded_checkpoint( + state_dict: dict[str, torch.Tensor], + output: Path, + max_shard_bytes: int, +) -> dict[str, Any]: + shards = _partition_shards(state_dict, max_shard_bytes) + weight_map: dict[str, str] = {} + shard_hashes: dict[str, str] = {} + total_size = 0 + for index, shard in enumerate(shards, start=1): + filename = f"model-{index:05d}-of-{len(shards):05d}.safetensors" + path = output / filename + save_file(shard, path) + shard_hashes[filename] = _sha256(path) + for key, value in shard.items(): + weight_map[key] = filename + total_size += _tensor_bytes(value) + index = {"metadata": {"total_size": total_size}, "weight_map": weight_map} + (output / "model.safetensors.index.json").write_text( + json.dumps(index, indent=2, sort_keys=True) + "\n" + ) + return { + "shard_count": len(shards), + "shard_sha256": shard_hashes, + "tensor_count": len(weight_map), + "total_size": total_size, + } + + +def create_fixture( + output: Path, + *, + seed: int, + max_shard_bytes: int = _DEFAULT_MAX_SHARD_BYTES, +) -> dict[str, Any]: + if output.exists() and any(output.iterdir()): + raise ValueError(f"Output directory must be empty: {output}") + output.mkdir(parents=True, exist_ok=True) + + torch.manual_seed(seed) + model_spec = model_registry("debugmodel", seq_len=128) + model = model_spec.model.build() + model.init_states() + model.to(dtype=torch.bfloat16) + adapter = KimiK3StateDictAdapter(model_spec.model, hf_assets_path=None) + hf_state_dict = adapter.to_hf(model.state_dict()) + converted, pair_count = convert_hf_state_dict_to_mxfp4(hf_state_dict) + storage = write_sharded_checkpoint(converted, output, max_shard_bytes) + (output / "config.json").write_text( + json.dumps( + { + "model_type": "kimi_k3", + "text_config": { + "quantization_config": { + "format": "mxfp4-pack-quantized", + "quant_method": "compressed-tensors", + "config_groups": { + "group_0": { + "targets": ["Linear"], + "weights": {"group_size": 32}, + } + }, + "ignore": [ + "re:.*self_attn.*", + "re:.*shared_experts.*", + "re:.*vision_tower.*", + "re:.*mm_projector.*", + "re:.*lm_head.*", + ], + } + }, + }, + indent=2, + sort_keys=True, + ) + + "\n" + ) + + repo_root = Path(__file__).resolve().parents[2] + manifest = { + "block_size": 32, + "expert_pair_count": pair_count, + "model": "kimi_k3_debugmodel", + "seed": seed, + "torch_version": torch.__version__, + "torchao_version": getattr(torchao, "__version__", "unknown"), + "torchtitan_commit": _git_revision(repo_root), + **storage, + } + (output / "fixture-manifest.json").write_text( + json.dumps(manifest, indent=2, sort_keys=True) + "\n" + ) + return manifest + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--output", type=Path, required=True) + parser.add_argument("--seed", type=int, default=20260917) + parser.add_argument( + "--max-shard-bytes", + type=int, + default=_DEFAULT_MAX_SHARD_BYTES, + ) + args = parser.parse_args() + manifest = create_fixture( + args.output, + seed=args.seed, + max_shard_bytes=args.max_shard_bytes, + ) + print(json.dumps(manifest, indent=2, sort_keys=True)) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tests/integration_tests/test_kimi_k3_mxfp4_checkpoint.py b/tests/integration_tests/test_kimi_k3_mxfp4_checkpoint.py new file mode 100644 index 00000000000..4d2e9343454 --- /dev/null +++ b/tests/integration_tests/test_kimi_k3_mxfp4_checkpoint.py @@ -0,0 +1,97 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +import tempfile +import unittest +from pathlib import Path + +import torch +import torch.distributed.checkpoint as dcp +from torchao.prototype.mx_formats.mx_tensor import MXTensor + +from scripts.checkpoint_conversion.create_kimi_k3_mxfp4_fixture import ( + convert_hf_state_dict_to_mxfp4, + write_sharded_checkpoint, +) +from torchtitan.components.checkpointer.packed_hf_storage import ( + PackedPairHuggingFaceStorageReader, + PackedPairSpec, +) + +_HF_WEIGHT = ( + "language_model.model.layers.1.block_sparse_moe.experts.0.w1.weight" +) + + +def _decode_mxfp4(packed, scales, block_size, target_dtype): + return MXTensor( + packed, + scales.view(torch.float8_e8m0fnu), + torch.float4_e2m1fn_x2, + block_size, + target_dtype, + None, + None, + False, + ).dequantize(target_dtype) + + +class KimiK3MXFP4CheckpointIntegrationTest(unittest.TestCase): + def test_release_format_fixture_loads_across_shards(self) -> None: + weight = torch.linspace(-6, 6, 128, dtype=torch.bfloat16).reshape(2, 64) + dense = torch.arange(12, dtype=torch.bfloat16).reshape(3, 4) + expected = MXTensor.to_mx( + weight, + elem_dtype=torch.float4_e2m1fn_x2, + block_size=32, + ).dequantize(torch.bfloat16) + converted, pair_count = convert_hf_state_dict_to_mxfp4( + {_HF_WEIGHT: weight, "dense.weight": dense} + ) + self.assertEqual(pair_count, 1) + self.assertNotIn(_HF_WEIGHT, converted) + + with tempfile.TemporaryDirectory() as directory: + output = Path(directory) + manifest = write_sharded_checkpoint( + converted, + output, + max_shard_bytes=64, + ) + self.assertGreater(manifest["shard_count"], 1) + destination = { + _HF_WEIGHT: torch.empty_like(weight), + "dense.weight": torch.empty_like(dense), + } + reader = PackedPairHuggingFaceStorageReader( + str(output), + PackedPairSpec( + packed_suffix=".weight_packed", + scale_suffix=".weight_scale", + virtual_suffix=".weight", + block_size=32, + packed_values_per_byte=2, + target_dtype=torch.bfloat16, + is_target=lambda key: key == _HF_WEIGHT, + decode=_decode_mxfp4, + ), + ) + dcp.load(destination, storage_reader=reader) + + torch.testing.assert_close( + destination[_HF_WEIGHT], expected, rtol=0, atol=0, equal_nan=True + ) + torch.testing.assert_close(destination["dense.weight"], dense, rtol=0, atol=0) + self.assertFalse( + any( + key.endswith(("weight_packed", "weight_scale")) + for key in destination + ) + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/unit_tests/cpu/test_mx_qat_grouped_experts.py b/tests/unit_tests/cpu/test_mx_qat_grouped_experts.py new file mode 100644 index 00000000000..7e383dba123 --- /dev/null +++ b/tests/unit_tests/cpu/test_mx_qat_grouped_experts.py @@ -0,0 +1,91 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +import unittest +from unittest.mock import patch + +import torch + +from torchtitan.models.common.moe import GroupedExperts +from torchtitan.quantization.mx_qat.experts import ( + _get_mx_qat_grouped_experts_cls, +) +from torchao.prototype.qat import mx_fake_quantized_grouped_mm + + +def _emulated_grouped_mm( + activation: torch.Tensor, + weight: torch.Tensor, + *, + offs: torch.Tensor, +) -> torch.Tensor: + outputs = [] + start = 0 + for expert, stop in enumerate(offs.tolist()): + outputs.append(activation[start:stop] @ weight[expert]) + start = stop + return torch.cat(outputs) + + +class MXQATGroupedExpertsTest(unittest.TestCase): + def test_keeps_bf16_masters_and_applies_both_fake_quantizers(self) -> None: + cls = _get_mx_qat_grouped_experts_cls(GroupedExperts) + module = cls.Config(dim=64, hidden_dim=64, num_experts=2).build() + module.to(dtype=torch.bfloat16) + with torch.no_grad(): + for parameter in module.parameters(): + parameter.normal_(mean=0.0, std=0.1) + + parameter_ids = {id(parameter) for parameter in module.parameters()} + optimizer = torch.optim.AdamW(module.parameters(), lr=1e-4) + optimizer_ids = { + id(parameter) + for group in optimizer.param_groups + for parameter in group["params"] + } + self.assertEqual(parameter_ids, optimizer_ids) + self.assertTrue( + all(parameter.dtype == torch.bfloat16 for parameter in module.parameters()) + ) + self.assertFalse( + any( + key.endswith(("weight_packed", "weight_scale")) + for key in module.state_dict() + ) + ) + + activation = torch.randn( + 4, + 64, + dtype=torch.bfloat16, + requires_grad=True, + ) + offsets = torch.tensor([2, 4], dtype=torch.int32) + original_grouped_mm = torch._grouped_mm + torch._grouped_mm = _emulated_grouped_mm + try: + with patch( + "torchao.prototype.qat.mx_fake_quantized_grouped_mm", + wraps=mx_fake_quantized_grouped_mm, + ) as grouped_mm: + output = module._grouped_mm( + A=activation, + weight_EOI=module.w1_EFD, + offs=offsets, + ) + finally: + torch._grouped_mm = original_grouped_mm + + output.float().sum().backward() + self.assertEqual(grouped_mm.call_count, 1) + self.assertIsNotNone(activation.grad) + self.assertIsNotNone(module.w1_EFD.grad) + self.assertTrue(torch.isfinite(activation.grad).all()) + self.assertTrue(torch.isfinite(module.w1_EFD.grad).all()) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/unit_tests/cpu/test_packed_hf_storage.py b/tests/unit_tests/cpu/test_packed_hf_storage.py new file mode 100644 index 00000000000..1041561bb86 --- /dev/null +++ b/tests/unit_tests/cpu/test_packed_hf_storage.py @@ -0,0 +1,333 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +import tempfile +import unittest +from pathlib import Path + +import torch +import torch.distributed.checkpoint as dcp +from safetensors import safe_open +from safetensors.torch import save_file +from torch.distributed.checkpoint.metadata import MetadataIndex +from torch.distributed.checkpoint.planner import LoadItemType, ReadItem + +from torchtitan.components.checkpointer.packed_hf_storage import ( + PackedPairHuggingFaceStorageReader, + PackedPairSpec, +) + +_PACKED_KEY = "model.layers.0.experts.0.w1.weight_packed" +_SCALE_KEY = "model.layers.0.experts.0.w1.weight_scale" +_VIRTUAL_KEY = "model.layers.0.experts.0.w1.weight" + + +def _is_expert_weight(key: str) -> bool: + return ".experts." in key + + +def _decode_mxfp4( + packed: torch.Tensor, + scales: torch.Tensor, + block_size: int, + target_dtype: torch.dtype, +) -> torch.Tensor: + from torchao.prototype.mx_formats.mx_tensor import MXTensor + + return MXTensor( + packed, + scales.view(torch.float8_e8m0fnu), + torch.float4_e2m1fn_x2, + block_size, + target_dtype, + None, + None, + False, + ).dequantize(target_dtype) + + +def _spec() -> PackedPairSpec: + return PackedPairSpec( + packed_suffix=".weight_packed", + scale_suffix=".weight_scale", + virtual_suffix=".weight", + block_size=32, + packed_values_per_byte=2, + target_dtype=torch.bfloat16, + is_target=_is_expert_weight, + decode=_decode_mxfp4, + ) + + +def _packed_fixture() -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + nibbles = torch.arange(2 * 64, dtype=torch.uint8).reshape(2, 64) % 16 + packed = nibbles[:, 0::2] | (nibbles[:, 1::2] << 4) + scales = torch.tensor([[127, 128], [126, 255]], dtype=torch.uint8) + lookup = 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, + ) + decoded = torch.stack( + (lookup[(packed & 0x0F).long()], lookup[(packed >> 4).long()]), dim=-1 + ).flatten(-2) + expanded_scales = scales.repeat_interleave(32, dim=-1) + expected = torch.ldexp(decoded, expanded_scales.to(torch.int32) - 127) + expected = torch.where(expanded_scales == 255, torch.nan, expected).to( + torch.bfloat16 + ) + return packed, scales, expected + + +class _Planner: + def __init__(self, destination: torch.Tensor) -> None: + self.destination = destination + self.committed = False + + def resolve_tensor(self, _read_item: ReadItem) -> torch.Tensor: + return self.destination + + def commit_tensor(self, _read_item: ReadItem, tensor: torch.Tensor) -> None: + self.committed = tensor.data_ptr() == self.destination.data_ptr() + + +class _RecordingSlice: + def __init__(self, tensor_slice: object, calls: list[tuple[str, object]], key: str): + self.tensor_slice = tensor_slice + self.calls = calls + self.key = key + + def __getitem__(self, slices: object) -> torch.Tensor: + self.calls.append((self.key, slices)) + return self.tensor_slice[slices] # type: ignore[index] + + +class _RecordingFile: + def __init__(self, handle: object) -> None: + self.handle = handle + self.calls: list[tuple[str, object]] = [] + + def get_slice(self, key: str) -> _RecordingSlice: + return _RecordingSlice(self.handle.get_slice(key), self.calls, key) # type: ignore[attr-defined] + + +class PackedPairHuggingFaceStorageReaderMetadataTest(unittest.TestCase): + def _write_checkpoint(self, tensors: dict[str, torch.Tensor]) -> str: + directory = tempfile.TemporaryDirectory() + self.addCleanup(directory.cleanup) + save_file(tensors, Path(directory.name) / "model.safetensors") + return directory.name + + def test_read_metadata_exposes_virtual_weight(self) -> None: + path = self._write_checkpoint( + { + _PACKED_KEY: torch.zeros((2, 32), dtype=torch.uint8), + _SCALE_KEY: torch.full((2, 2), 127, dtype=torch.uint8), + "dense.weight": torch.ones((3, 4), dtype=torch.bfloat16), + } + ) + + metadata = PackedPairHuggingFaceStorageReader(path, _spec()).read_metadata() + + self.assertEqual( + set(metadata.state_dict_metadata), {_VIRTUAL_KEY, "dense.weight"} + ) + virtual = metadata.state_dict_metadata[_VIRTUAL_KEY] + self.assertEqual(virtual.size, torch.Size((2, 64))) + self.assertEqual(virtual.properties.dtype, torch.bfloat16) + self.assertEqual( + {index.fqn for index in metadata.storage_data}, + {_VIRTUAL_KEY, "dense.weight"}, + ) + + def test_read_metadata_rejects_missing_scale(self) -> None: + path = self._write_checkpoint( + {_PACKED_KEY: torch.zeros((2, 32), dtype=torch.uint8)} + ) + + with self.assertRaisesRegex(ValueError, "missing scale tensor"): + PackedPairHuggingFaceStorageReader(path, _spec()).read_metadata() + + def test_read_metadata_rejects_orphan_scale(self) -> None: + path = self._write_checkpoint( + {_SCALE_KEY: torch.full((2, 2), 127, dtype=torch.uint8)} + ) + + with self.assertRaisesRegex(ValueError, "orphan scale tensor"): + PackedPairHuggingFaceStorageReader(path, _spec()).read_metadata() + + def test_read_metadata_rejects_non_uint8_payload(self) -> None: + path = self._write_checkpoint( + { + _PACKED_KEY: torch.zeros((2, 32), dtype=torch.int16), + _SCALE_KEY: torch.full((2, 2), 127, dtype=torch.uint8), + } + ) + + with self.assertRaisesRegex(ValueError, "must use torch.uint8"): + PackedPairHuggingFaceStorageReader(path, _spec()).read_metadata() + + def test_read_metadata_rejects_incompatible_shapes(self) -> None: + path = self._write_checkpoint( + { + _PACKED_KEY: torch.zeros((2, 31), dtype=torch.uint8), + _SCALE_KEY: torch.full((2, 2), 127, dtype=torch.uint8), + } + ) + + with self.assertRaisesRegex(ValueError, "incompatible packed and scale shapes"): + PackedPairHuggingFaceStorageReader(path, _spec()).read_metadata() + + def test_read_metadata_rejects_pair_outside_target_policy(self) -> None: + path = self._write_checkpoint( + { + "model.layers.0.dense.weight_packed": torch.zeros( + (2, 32), dtype=torch.uint8 + ), + "model.layers.0.dense.weight_scale": torch.full( + (2, 2), 127, dtype=torch.uint8 + ), + } + ) + + with self.assertRaisesRegex(ValueError, "outside the packed-weight policy"): + PackedPairHuggingFaceStorageReader(path, _spec()).read_metadata() + + +class PackedPairHuggingFaceStorageReaderReadTest(unittest.TestCase): + def setUp(self) -> None: + self.directory = tempfile.TemporaryDirectory() + packed, scales, self.expected = _packed_fixture() + save_file( + {_PACKED_KEY: packed, _SCALE_KEY: scales}, + Path(self.directory.name) / "model.safetensors", + ) + + def tearDown(self) -> None: + self.directory.cleanup() + + def test_dcp_load_dequantizes_full_tensor(self) -> None: + destination = torch.empty((2, 64), dtype=torch.bfloat16) + + dcp.load( + {_VIRTUAL_KEY: destination}, + storage_reader=PackedPairHuggingFaceStorageReader( + self.directory.name, _spec() + ), + ) + + torch.testing.assert_close( + destination, self.expected, rtol=0, atol=0, equal_nan=True + ) + + def test_unaligned_read_uses_only_intersecting_groups(self) -> None: + reader = PackedPairHuggingFaceStorageReader(self.directory.name, _spec()) + reader.read_metadata() + destination = torch.empty((2, 38), dtype=torch.bfloat16) + planner = _Planner(destination) + request = ReadItem( + type=LoadItemType.TENSOR, + dest_index=MetadataIndex(_VIRTUAL_KEY, [0, 0]), + dest_offsets=torch.Size((0, 0)), + storage_index=MetadataIndex(_VIRTUAL_KEY, [0, 0]), + storage_offsets=torch.Size((0, 7)), + lengths=torch.Size((2, 38)), + ) + + with safe_open( + Path(self.directory.name) / "model.safetensors", framework="pt" + ) as handle: + recording_file = _RecordingFile(handle) + reader._process_read_request(recording_file, request, planner) + + self.assertTrue(planner.committed) + torch.testing.assert_close( + destination, self.expected[:, 7:45], rtol=0, atol=0, equal_nan=True + ) + self.assertEqual( + recording_file.calls, + [ + (_PACKED_KEY, (slice(0, 2), slice(0, 32))), + (_SCALE_KEY, (slice(0, 2), slice(0, 2))), + ], + ) + + def test_aligned_read_dequantizes_second_group(self) -> None: + reader = PackedPairHuggingFaceStorageReader(self.directory.name, _spec()) + reader.read_metadata() + destination = torch.empty((1, 32), dtype=torch.bfloat16) + planner = _Planner(destination) + request = ReadItem( + type=LoadItemType.TENSOR, + dest_index=MetadataIndex(_VIRTUAL_KEY, [0, 0]), + dest_offsets=torch.Size((0, 0)), + storage_index=MetadataIndex(_VIRTUAL_KEY, [0, 0]), + storage_offsets=torch.Size((1, 32)), + lengths=torch.Size((1, 32)), + ) + + with safe_open( + Path(self.directory.name) / "model.safetensors", framework="pt" + ) as handle: + reader._process_read_request(handle, request, planner) + + torch.testing.assert_close( + destination, self.expected[1:2, 32:64], rtol=0, atol=0, equal_nan=True + ) + + def test_e8m0_extreme_bytes_follow_mx_semantics(self) -> None: + packed = torch.full((1, 48), 0x11, dtype=torch.uint8) + scales = torch.tensor([[0, 254, 255]], dtype=torch.uint8) + with tempfile.TemporaryDirectory() as directory: + save_file( + {_PACKED_KEY: packed, _SCALE_KEY: scales}, + Path(directory) / "model.safetensors", + ) + destination = torch.empty((1, 96), dtype=torch.bfloat16) + dcp.load( + {_VIRTUAL_KEY: destination}, + storage_reader=PackedPairHuggingFaceStorageReader(directory, _spec()), + ) + + expected = torch.cat( + ( + torch.full( + (32,), + torch.ldexp(torch.tensor(0.5), torch.tensor(-127)).item(), + ), + torch.full( + (32,), + torch.ldexp(torch.tensor(0.5), torch.tensor(127)).item(), + ), + torch.full((32,), torch.nan), + ) + ).reshape(1, 96) + torch.testing.assert_close( + destination, + expected.to(torch.bfloat16), + rtol=0, + atol=0, + equal_nan=True, + ) + +if __name__ == "__main__": + unittest.main() diff --git a/tests/unit_tests/cpu/test_state_dict_adapter.py b/tests/unit_tests/cpu/test_state_dict_adapter.py index bdfd36f5352..bcedb53307e 100644 --- a/tests/unit_tests/cpu/test_state_dict_adapter.py +++ b/tests/unit_tests/cpu/test_state_dict_adapter.py @@ -4,16 +4,22 @@ # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. +import json import tempfile import unittest +from pathlib import Path import torch import torch.distributed as dist +from torch.distributed.checkpoint import HuggingFaceStorageReader from torch.distributed.device_mesh import init_device_mesh from torch.distributed.tensor import DTensor, Replicate, Shard from torch.testing._internal.distributed.fake_pg import FakeStore from torchtitan.components.checkpointer.base import ModelWrapper +from torchtitan.components.checkpointer.packed_hf_storage import ( + PackedPairHuggingFaceStorageReader, +) from torchtitan.models.deepseek_v3 import deepseekv3_configs from torchtitan.models.deepseek_v3.state_dict_adapter import DeepSeekV3StateDictAdapter from torchtitan.models.deepseek_v4 import model_registry as deepseek_v4_model_registry @@ -21,6 +27,8 @@ from torchtitan.models.deepseek_v4.state_dict_adapter import DeepSeekV4StateDictAdapter from torchtitan.models.gpt_oss import gptoss_configs from torchtitan.models.gpt_oss.state_dict_adapter import GptOssStateDictAdapter +from torchtitan.models.kimi_k3 import model_registry as kimi_k3_model_registry +from torchtitan.models.kimi_k3.state_dict_adapter import KimiK3StateDictAdapter from torchtitan.models.llama3 import llama3_configs from torchtitan.models.llama3.model import Llama3Model from torchtitan.models.llama3.state_dict_adapter import Llama3StateDictAdapter @@ -114,6 +122,77 @@ def test_hf_roundtrip_preserves_tied_embedding_shape(self) -> None: model.load_state_dict(restored, strict=True) +class KimiK3StateDictAdapterTest(unittest.TestCase): + def setUp(self) -> None: + model_spec = kimi_k3_model_registry("debugmodel", seq_len=128) + self.adapter = KimiK3StateDictAdapter( + model_spec.model, + hf_assets_path=None, + ) + + def test_quantized_load_uses_packed_pair_reader(self) -> None: + with tempfile.TemporaryDirectory() as directory: + Path(directory, "config.json").write_text( + json.dumps( + { + "text_config": { + "quantization_config": { + "format": "mxfp4-pack-quantized", + "quant_method": "compressed-tensors", + "config_groups": { + "group_0": { + "targets": ["Linear"], + "weights": {"group_size": 32}, + } + }, + "ignore": ["re:.*shared_experts.*"], + } + } + } + ) + ) + reader = self.adapter.get_hf_storage_reader( + directory, + from_quantized=True, + ) + + self.assertIsInstance(reader, PackedPairHuggingFaceStorageReader) + self.assertEqual(reader.spec.packed_suffix, ".weight_packed") + self.assertEqual(reader.spec.scale_suffix, ".weight_scale") + self.assertEqual(reader.spec.virtual_suffix, ".weight") + self.assertEqual(reader.spec.block_size, 32) + self.assertEqual(reader.spec.target_dtype, torch.bfloat16) + self.assertTrue( + reader.spec.is_target( + "language_model.model.layers.1.block_sparse_moe.experts.3.w1.weight" + ) + ) + self.assertFalse( + reader.spec.is_target( + "language_model.model.layers.1.block_sparse_moe.shared_experts.w1.weight" + ) + ) + self.assertTrue( + reader.spec.is_target( + "language_model.model.layers.1.block_sparse_moe.experts.3.w4.weight" + ) + ) + + def test_quantized_load_rejects_missing_metadata(self) -> None: + with tempfile.TemporaryDirectory() as directory: + with self.assertRaisesRegex(ValueError, "missing quantization_config"): + Path(directory, "config.json").write_text("{}") + self.adapter.get_hf_storage_reader(directory, from_quantized=True) + + def test_unquantized_load_keeps_plain_reader(self) -> None: + reader = self.adapter.get_hf_storage_reader( + "/tmp/kimi-k3-checkpoint", + from_quantized=False, + ) + + self.assertIs(type(reader), HuggingFaceStorageReader) + + class DeepSeekV3StateDictAdapterTest(unittest.TestCase): @classmethod def setUpClass(cls) -> None: diff --git a/torchtitan/components/checkpointer/packed_hf_storage.py b/torchtitan/components/checkpointer/packed_hf_storage.py new file mode 100644 index 00000000000..fba7b63c1b6 --- /dev/null +++ b/torchtitan/components/checkpointer/packed_hf_storage.py @@ -0,0 +1,297 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +"""Hugging Face storage reader for paired packed weights and block scales.""" + +import dataclasses +from collections.abc import Callable +from dataclasses import dataclass +from typing import Any + +import torch +from torch.distributed.checkpoint import HuggingFaceStorageReader +from torch.distributed.checkpoint.metadata import ( + ChunkStorageMetadata, + MetadataIndex, + TensorStorageMetadata, +) +from torch.distributed.checkpoint.planner import LoadPlanner, ReadItem + +__all__ = ["PackedPairHuggingFaceStorageReader", "PackedPairSpec"] + + +@dataclass(frozen=True) +class PackedPairSpec: + """Describe how a packed tensor and its scales expose a virtual weight.""" + + packed_suffix: str + scale_suffix: str + virtual_suffix: str + block_size: int + packed_values_per_byte: int + target_dtype: torch.dtype + is_target: Callable[[str], bool] + decode: Callable[[torch.Tensor, torch.Tensor, int, torch.dtype], torch.Tensor] + + def __post_init__(self) -> None: + if not self.packed_suffix or not self.scale_suffix or not self.virtual_suffix: + raise ValueError("packed, scale, and virtual suffixes must be non-empty") + if len({self.packed_suffix, self.scale_suffix, self.virtual_suffix}) != 3: + raise ValueError("packed, scale, and virtual suffixes must be distinct") + if ( + self.block_size <= 0 + or self.packed_values_per_byte <= 0 + or self.block_size % self.packed_values_per_byte != 0 + ): + raise ValueError( + "block_size must be positive and divisible by packed_values_per_byte" + ) + + +@dataclass(frozen=True) +class _PackedPair: + packed_fqn: str + scale_fqn: str + packed_path: str + scale_path: str + + +class PackedPairHuggingFaceStorageReader(HuggingFaceStorageReader): + """Present selected uint8 packed/scale pairs as unquantized DCP tensors.""" + + def __init__( + self, + path: str, + spec: PackedPairSpec, + thread_count: int = 1, + ) -> None: + super().__init__(path=path, thread_count=thread_count) + self.spec = spec + self._pairs: dict[str, _PackedPair] = {} + + def _replace_suffix(self, fqn: str, source: str, destination: str) -> str: + if not fqn.endswith(source): + raise ValueError(f"{fqn!r} does not end with {source!r}") + return fqn[: -len(source)] + destination + + def _validate_pair( + self, + virtual_fqn: str, + packed_metadata: TensorStorageMetadata, + scale_metadata: TensorStorageMetadata, + ) -> None: + if not self.spec.is_target(virtual_fqn): + raise ValueError( + f"Packed tensor {virtual_fqn!r} is outside the packed-weight policy." + ) + if ( + packed_metadata.properties.dtype != torch.uint8 + or scale_metadata.properties.dtype != torch.uint8 + ): + raise ValueError( + f"Packed tensor pair for {virtual_fqn!r} must use torch.uint8." + ) + packed_shape = packed_metadata.size + scale_shape = scale_metadata.size + if len(packed_shape) != 2 or len(scale_shape) != 2: + raise ValueError( + f"Packed tensor pair for {virtual_fqn!r} must be two-dimensional." + ) + if ( + packed_shape[0] != scale_shape[0] + or packed_shape[1] * self.spec.packed_values_per_byte + != scale_shape[1] * self.spec.block_size + ): + raise ValueError( + f"Packed tensor pair for {virtual_fqn!r} has incompatible packed " + f"and scale shapes: {tuple(packed_shape)} and {tuple(scale_shape)}." + ) + + def _virtual_chunk(self, chunk: ChunkStorageMetadata) -> ChunkStorageMetadata: + offsets = list(chunk.offsets) + sizes = list(chunk.sizes) + offsets[-1] *= self.spec.packed_values_per_byte + sizes[-1] *= self.spec.packed_values_per_byte + return dataclasses.replace( + chunk, + offsets=torch.Size(offsets), + sizes=torch.Size(sizes), + ) + + # pyrefly: ignore [bad-override] + def read_metadata(self) -> Any: + metadata = super().read_metadata() + state_dict_metadata = metadata.state_dict_metadata + storage_paths: dict[str, str] = {} + for index, storage_info in metadata.storage_data.items(): + previous = storage_paths.setdefault(index.fqn, storage_info.relative_path) + if previous != storage_info.relative_path: + raise ValueError( + f"Tensor {index.fqn!r} is split across safetensors files; " + "packed-pair loading requires each source tensor in one file." + ) + packed_fqns = { + fqn for fqn in state_dict_metadata if fqn.endswith(self.spec.packed_suffix) + } + scale_fqns = { + fqn for fqn in state_dict_metadata if fqn.endswith(self.spec.scale_suffix) + } + + pairs: dict[str, _PackedPair] = {} + for packed_fqn in sorted(packed_fqns): + virtual_fqn = self._replace_suffix( + packed_fqn, self.spec.packed_suffix, self.spec.virtual_suffix + ) + scale_fqn = self._replace_suffix( + packed_fqn, self.spec.packed_suffix, self.spec.scale_suffix + ) + if scale_fqn not in scale_fqns: + raise ValueError( + f"Packed tensor {packed_fqn!r} is missing scale tensor " + f"{scale_fqn!r}." + ) + if virtual_fqn in state_dict_metadata: + raise ValueError( + f"Virtual tensor {virtual_fqn!r} conflicts with a stored tensor." + ) + packed_metadata = state_dict_metadata[packed_fqn] + scale_metadata = state_dict_metadata[scale_fqn] + if not isinstance(packed_metadata, TensorStorageMetadata) or not isinstance( + scale_metadata, TensorStorageMetadata + ): + raise TypeError( + f"Packed tensor pair for {virtual_fqn!r} must contain tensors." + ) + self._validate_pair(virtual_fqn, packed_metadata, scale_metadata) + pairs[virtual_fqn] = _PackedPair( + packed_fqn, + scale_fqn, + storage_paths[packed_fqn], + storage_paths[scale_fqn], + ) + + paired_scales = {pair.scale_fqn for pair in pairs.values()} + orphan_scales = sorted(scale_fqns - paired_scales) + if orphan_scales: + raise ValueError(f"Found orphan scale tensor {orphan_scales[0]!r}.") + + virtual_state_dict_metadata = { + fqn: tensor_metadata + for fqn, tensor_metadata in state_dict_metadata.items() + if fqn not in packed_fqns and fqn not in scale_fqns + } + for virtual_fqn, pair in pairs.items(): + packed_metadata = state_dict_metadata[pair.packed_fqn] + virtual_state_dict_metadata[virtual_fqn] = dataclasses.replace( + packed_metadata, + properties=dataclasses.replace( + packed_metadata.properties, dtype=self.spec.target_dtype + ), + size=torch.Size( + ( + packed_metadata.size[0], + packed_metadata.size[1] * self.spec.packed_values_per_byte, + ) + ), + chunks=[self._virtual_chunk(chunk) for chunk in packed_metadata.chunks], + ) + + virtual_storage_data: dict[MetadataIndex, Any] = {} + for index, storage_info in metadata.storage_data.items(): + if index.fqn in scale_fqns: + continue + if index.fqn not in packed_fqns: + virtual_storage_data[index] = storage_info + continue + virtual_fqn = self._replace_suffix( + index.fqn, self.spec.packed_suffix, self.spec.virtual_suffix + ) + virtual_offset = list(index.offset or ()) + if virtual_offset: + virtual_offset[-1] *= self.spec.packed_values_per_byte + virtual_index = dataclasses.replace( + index, + fqn=virtual_fqn, + offset=torch.Size(virtual_offset), + ) + virtual_storage_data[virtual_index] = dataclasses.replace( + storage_info, + shape=torch.Size( + ( + storage_info.shape[0], + storage_info.shape[1] * self.spec.packed_values_per_byte, + ) + ), + dtype=self.spec.target_dtype, + ) + + self._pairs = pairs + return dataclasses.replace( + metadata, + state_dict_metadata=virtual_state_dict_metadata, + storage_data=virtual_storage_data, + ) + + def _process_read_request( + self, + f: Any, + req: ReadItem, + planner: LoadPlanner, + ) -> None: + virtual_fqn = req.storage_index.fqn + pair = self._pairs.get(virtual_fqn) + if pair is None: + super()._process_read_request(f, req, planner) + return + if len(req.storage_offsets) != 2 or len(req.lengths) != 2: + raise ValueError( + f"Packed tensor {virtual_fqn!r} requires a two-dimensional read; " + f"got offsets {tuple(req.storage_offsets)} and lengths " + f"{tuple(req.lengths)}." + ) + + row_start, column_start = req.storage_offsets + row_count, column_count = req.lengths + row_stop = row_start + row_count + column_stop = column_start + column_count + group_start = column_start // self.spec.block_size + group_stop = ( + column_stop + self.spec.block_size - 1 + ) // self.spec.block_size + packed_values_per_group = ( + self.spec.block_size // self.spec.packed_values_per_byte + ) + packed_start = group_start * packed_values_per_group + packed_stop = group_stop * packed_values_per_group + + packed = f.get_slice(pair.packed_fqn)[ + slice(row_start, row_stop), slice(packed_start, packed_stop) + ] + scale_slices = (slice(row_start, row_stop), slice(group_start, group_stop)) + if pair.scale_path == pair.packed_path: + scales = f.get_slice(pair.scale_fqn)[scale_slices] + else: + from safetensors import safe_open # type: ignore[import] + + with safe_open(pair.scale_path, framework="pt") as scale_file: + scales = scale_file.get_slice(pair.scale_fqn)[scale_slices] + + dequantized = self.spec.decode( + packed, + scales, + self.spec.block_size, + self.spec.target_dtype, + ) + crop_start = column_start - group_start * self.spec.block_size + tensor = dequantized[:, crop_start : crop_start + column_count] + target_tensor = planner.resolve_tensor(req).detach() + if target_tensor.size() != tensor.size(): + raise AssertionError( + f"req {req.storage_index} mismatch sizes " + f"{target_tensor.size()} vs {tensor.size()}" + ) + target_tensor.copy_(tensor) + planner.commit_tensor(req, target_tensor) diff --git a/torchtitan/config/transform/__init__.py b/torchtitan/config/transform/__init__.py index c195864ceb4..33f8912fe33 100644 --- a/torchtitan/config/transform/__init__.py +++ b/torchtitan/config/transform/__init__.py @@ -18,6 +18,7 @@ Float8LinearConverter, MXFP8GroupedExpertsConverter, MXFP8LinearConverter, + MXQATGroupedExpertsConverter, NVFP4LinearConverter, QuantizationConverter, ) @@ -37,6 +38,7 @@ "Float8LinearConverter", "MXFP8GroupedExpertsConverter", "MXFP8LinearConverter", + "MXQATGroupedExpertsConverter", "NVFP4LinearConverter", "QuantizationConverter", "validate_converter_compatibility", diff --git a/torchtitan/config/transform/quantization.py b/torchtitan/config/transform/quantization.py index 34ee3b6797d..4ace20060e0 100644 --- a/torchtitan/config/transform/quantization.py +++ b/torchtitan/config/transform/quantization.py @@ -25,6 +25,7 @@ ) from torchtitan.quantization.mxfp8 import _mxfp8_linear_import_error, MXFP8Linear from torchtitan.quantization.mxfp8.experts import _get_mxfp8_grouped_experts_cls +from torchtitan.quantization.mx_qat.experts import _get_mx_qat_grouped_experts_cls from torchtitan.quantization.nvfp4 import NVFP4Linear from torchtitan.quantization.utils import module_filter_fn, swap_token_dispatcher from torchtitan.tools.utils import has_cuda_capability, has_rocm_capability @@ -434,6 +435,47 @@ def convert(self, model_config): return model_config +class MXQATGroupedExpertsConverter(QuantizationConverter): + """Apply emulated MXFP4-weight/MXFP8-activation QAT to grouped experts.""" + + @dataclass(kw_only=True, slots=True) + class Config(QuantizationConverter.Config): + weight_block_size: int = 32 + activation_block_size: int = 32 + + def __init__(self, config: Config): + self.config = config + try: + from torchao.prototype.qat import mx_fake_quantize # noqa: F401 + except ImportError as error: + raise ImportError( + "MX QAT grouped experts require a TorchAO build providing " + "torchao.prototype.qat.mx_fake_quantize." + ) from error + + def convert(self, model_config): + for _fqn, config, parent, attr in model_config.traverse(GroupedExperts.Config): + base_module_cls = type(config)._owner + quantized_cls = _get_mx_qat_grouped_experts_cls(base_module_cls) + config_cls = quantized_cls.Config # type: ignore[attr-defined] + new_config = config_cls( + **{f.name: getattr(config, f.name) for f in fields(config)}, + weight_block_size=self.config.weight_block_size, + activation_block_size=self.config.activation_block_size, + ) + if parent is None: + model_config = new_config + elif isinstance(parent, list): + parent[attr] = new_config + else: + setattr(parent, attr, new_config) + + logger.info( + "Converted GroupedExperts to MXFP4-weight/MXFP8-activation QAT" + ) + return model_config + + class NVFP4LinearConverter(QuantizationConverter): """Replace matching Linear.Config with NVFP4Linear.Config.""" diff --git a/torchtitan/models/kimi_k3/config_registry.py b/torchtitan/models/kimi_k3/config_registry.py index 17ce7b9f211..ebdf58a1462 100644 --- a/torchtitan/models/kimi_k3/config_registry.py +++ b/torchtitan/models/kimi_k3/config_registry.py @@ -6,11 +6,13 @@ from dataclasses import replace +from torchtitan.components.checkpointer import CheckpointManager from torchtitan.components.data import GrainDataLoader, SingleDatasetConfig from torchtitan.components.loss import ChunkedLossWrapper, CrossEntropyLoss from torchtitan.components.optimizer import default_adamw, LRSchedulersContainer from torchtitan.components.tokenizer import MultiModalTokenizer from torchtitan.config import TrainingConfig +from torchtitan.config.transform import MXQATGroupedExpertsConverter from torchtitan.distributed.activation_checkpoint import SelectiveAC from torchtitan.hf_datasets.multimodal.mm_collator import MultiModalCollator from torchtitan.hf_datasets.multimodal.mm_datasets import ( @@ -90,3 +92,22 @@ def kimi_k3_debugmodel( checkpointer=None, activation_checkpoint=SelectiveAC.Config(), ) + + +def kimi_k3_debugmodel_mx_qat( + seq_len: int | None = DEFAULT_DEBUG_MODEL_SEQ_LEN, +) -> Trainer.Config: + """Kimi-K3 debug recipe with MXFP4-weight/MXFP8-activation expert QAT.""" + config = kimi_k3_debugmodel(seq_len=seq_len) + config.model_spec = model_registry( + "debugmodel", + seq_len=seq_len, + converters=[MXQATGroupedExpertsConverter.Config()], + ) + config.checkpointer = CheckpointManager.Config( + interval=5, + initial_load_in_hf=True, + initial_load_in_hf_quantized=True, + last_save_model_only=False, + ) + return config diff --git a/torchtitan/models/kimi_k3/state_dict_adapter.py b/torchtitan/models/kimi_k3/state_dict_adapter.py index 08b76ec2ba7..c3e5b98545e 100644 --- a/torchtitan/models/kimi_k3/state_dict_adapter.py +++ b/torchtitan/models/kimi_k3/state_dict_adapter.py @@ -8,24 +8,102 @@ """HuggingFace checkpoint adapter for unquantized Kimi K3 weights.""" +import json import re +from collections.abc import Callable +from pathlib import Path from typing import Any, TYPE_CHECKING import torch +from torch.distributed.checkpoint import HuggingFaceStorageReader from torch.distributed.tensor import DTensor +from torchtitan.components.checkpointer.packed_hf_storage import ( + PackedPairHuggingFaceStorageReader, + PackedPairSpec, +) from torchtitan.models.utils import MoEStateDictAdapter if TYPE_CHECKING: from .model import KimiK3Model - _UNUSED_HF_LAYER_ZERO_ATTN_RES_KEYS = { "language_model.model.layers.0.self_attention_res_norm.weight", "language_model.model.layers.0.self_attention_res_proj.weight", } +def _released_mxfp4_policy(path: str) -> tuple[int, Callable[[str], bool]]: + config_path = Path(path) / "config.json" + if not config_path.is_file(): + raise ValueError(f"Quantized Kimi checkpoint is missing {config_path}.") + config = json.loads(config_path.read_text()) + text_config = config.get("text_config", config) + quantization = text_config.get("quantization_config") + if not isinstance(quantization, dict): + raise ValueError("Kimi checkpoint is missing quantization_config metadata.") + if ( + quantization.get("format") != "mxfp4-pack-quantized" + or quantization.get("quant_method") != "compressed-tensors" + ): + raise ValueError("Kimi checkpoint does not declare compressed MXFP4 weights.") + + config_groups = quantization.get("config_groups") + if not isinstance(config_groups, dict) or not config_groups: + raise ValueError("Kimi quantization_config has no config groups.") + group_sizes = set() + targets = set() + for group in config_groups.values(): + if not isinstance(group, dict): + raise ValueError("Kimi quantization config group must be an object.") + targets.update(group.get("targets", ())) + weights = group.get("weights") + if not isinstance(weights, dict): + raise ValueError("Kimi quantization config group has no weight config.") + group_sizes.add(weights.get("group_size")) + if targets != {"Linear"} or len(group_sizes) != 1: + raise ValueError( + "Kimi packed import currently requires one Linear MXFP4 group size." + ) + block_size = group_sizes.pop() + if not isinstance(block_size, int): + raise ValueError("Kimi MXFP4 group_size must be an integer.") + + ignore_patterns = [] + for pattern in quantization.get("ignore", ()): + if not isinstance(pattern, str) or not pattern.startswith("re:"): + raise ValueError("Kimi packed import requires regex ignore entries.") + ignore_patterns.append(re.compile(pattern.removeprefix("re:"))) + + def is_target(weight_fqn: str) -> bool: + if not weight_fqn.endswith(".weight"): + return False + module_fqn = weight_fqn.removesuffix(".weight") + return not any(pattern.fullmatch(module_fqn) for pattern in ignore_patterns) + + return block_size, is_target + + +def _decode_mxfp4( + packed: torch.Tensor, + scales: torch.Tensor, + block_size: int, + target_dtype: torch.dtype, +) -> torch.Tensor: + from torchao.prototype.mx_formats.mx_tensor import MXTensor + + return MXTensor( + packed, + scales.view(torch.float8_e8m0fnu), + torch.float4_e2m1fn_x2, + block_size, + target_dtype, + None, + None, + False, + ).dequantize(target_dtype) + + class KimiK3StateDictAdapter(MoEStateDictAdapter): def __init__( self, @@ -136,6 +214,29 @@ def _map_from_hf_layer_key( ) return attention_map.get(abstract_key) + def get_hf_storage_reader( + self, + path: str, + from_quantized: bool = False, + ) -> HuggingFaceStorageReader: + if not from_quantized: + return super().get_hf_storage_reader(path, from_quantized=False) + block_size, is_target = _released_mxfp4_policy(path) + return PackedPairHuggingFaceStorageReader( + path=path, + thread_count=4, + spec=PackedPairSpec( + packed_suffix=".weight_packed", + scale_suffix=".weight_scale", + virtual_suffix=".weight", + block_size=block_size, + packed_values_per_byte=2, + target_dtype=torch.bfloat16, + is_target=is_target, + decode=_decode_mxfp4, + ), + ) + def to_hf(self, state_dict: dict[str, Any]) -> dict[str, Any]: """Convert a TorchTitan state dict to unquantized HuggingFace format.""" state_dict = self._native_fused_linears_to_hf(state_dict) diff --git a/torchtitan/quantization/mx_qat/__init__.py b/torchtitan/quantization/mx_qat/__init__.py new file mode 100644 index 00000000000..487ad815084 --- /dev/null +++ b/torchtitan/quantization/mx_qat/__init__.py @@ -0,0 +1,12 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +"""MX quantization-aware training building blocks.""" + +from .experts import _get_mx_qat_grouped_experts_cls + + +__all__ = ["_get_mx_qat_grouped_experts_cls"] diff --git a/torchtitan/quantization/mx_qat/experts.py b/torchtitan/quantization/mx_qat/experts.py new file mode 100644 index 00000000000..627663b1a82 --- /dev/null +++ b/torchtitan/quantization/mx_qat/experts.py @@ -0,0 +1,57 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +"""MXFP4-weight/MXFP8-activation QAT for grouped experts.""" + +from dataclasses import dataclass + +import torch + + +_mx_qat_experts_cache: dict[type, type] = {} + + +def _get_mx_qat_grouped_experts_cls(parent_cls: type) -> type: + """Return a grouped-expert subclass using stateless TorchAO MX QAT.""" + if parent_cls in _mx_qat_experts_cache: + return _mx_qat_experts_cache[parent_cls] + + parent_config_cls = parent_cls.Config # type: ignore[attr-defined] + + class MXQATGroupedExperts(parent_cls): # type: ignore[valid-type, misc] + @dataclass(kw_only=True, slots=True) + class Config(parent_config_cls): # type: ignore[misc] + weight_block_size: int = 32 + activation_block_size: int = 32 + + def __init__(self, config: Config): + super().__init__(config) + from torchao.prototype.qat import MXFakeQuantizeConfig + + self._weight_fake_quant_config = MXFakeQuantizeConfig( + dtype=torch.float4_e2m1fn_x2, + block_size=config.weight_block_size, + ) + self._activation_fake_quant_config = MXFakeQuantizeConfig( + dtype=torch.float8_e4m3fn, + block_size=config.activation_block_size, + ) + + def _grouped_mm(self, *, A, weight_EOI, offs): + from torchao.prototype.qat import mx_fake_quantized_grouped_mm + + return mx_fake_quantized_grouped_mm( + A.bfloat16(), + weight_EOI.bfloat16(), + offs, + self._activation_fake_quant_config, + self._weight_fake_quant_config, + ) + + MXQATGroupedExperts.__name__ = f"MXQAT{parent_cls.__name__}" + MXQATGroupedExperts.__qualname__ = f"MXQAT{parent_cls.__name__}" + _mx_qat_experts_cache[parent_cls] = MXQATGroupedExperts + return MXQATGroupedExperts From f3678c9fa3febed160b9737788d1641479fe1c6f Mon Sep 17 00:00:00 2001 From: Elfie Guo Date: Mon, 21 Sep 2026 08:33:04 -0700 Subject: [PATCH 2/5] Validate Kimi MXFP4 policy and compose model-independent QAT transforms --- .../create_kimi_k3_mxfp4_fixture.py | 47 ++-- .../create_mxfp4_reference.py | 61 +++++ tests/assets/mxfp4-reference.json | 218 ++++++++++++++++++ .../test_kimi_k3_mxfp4_checkpoint.py | 34 +-- .../cpu/test_mx_qat_grouped_experts.py | 48 ++-- tests/unit_tests/cpu/test_mx_qat_policy.py | 104 +++++++++ tests/unit_tests/cpu/test_mx_qat_transform.py | 180 +++++++++++++++ .../unit_tests/cpu/test_packed_hf_storage.py | 38 +-- .../unit_tests/cpu/test_state_dict_adapter.py | 96 ++++++-- .../checkpointer/packed_hf_storage.py | 14 +- torchtitan/config/transform/README.md | 29 +++ torchtitan/config/transform/__init__.py | 4 +- torchtitan/config/transform/lora.py | 3 +- torchtitan/config/transform/mx_qat.py | 120 ++++++++++ torchtitan/config/transform/quantization.py | 43 ---- torchtitan/models/kimi_k3/config_registry.py | 30 ++- torchtitan/models/kimi_k3/quantization.py | 38 +++ .../models/kimi_k3/state_dict_adapter.py | 207 ++++++++++------- torchtitan/quantization/mx_qat/__init__.py | 1 - torchtitan/quantization/mx_qat/checkpoint.py | 123 ++++++++++ torchtitan/quantization/mx_qat/experts.py | 54 +++-- torchtitan/quantization/mx_qat/linear.py | 55 +++++ 22 files changed, 1258 insertions(+), 289 deletions(-) create mode 100644 scripts/checkpoint_conversion/create_mxfp4_reference.py create mode 100644 tests/assets/mxfp4-reference.json create mode 100644 tests/unit_tests/cpu/test_mx_qat_policy.py create mode 100644 tests/unit_tests/cpu/test_mx_qat_transform.py create mode 100644 torchtitan/config/transform/mx_qat.py create mode 100644 torchtitan/models/kimi_k3/quantization.py create mode 100644 torchtitan/quantization/mx_qat/checkpoint.py create mode 100644 torchtitan/quantization/mx_qat/linear.py diff --git a/scripts/checkpoint_conversion/create_kimi_k3_mxfp4_fixture.py b/scripts/checkpoint_conversion/create_kimi_k3_mxfp4_fixture.py index e51ea799973..c1e6dfa1e01 100755 --- a/scripts/checkpoint_conversion/create_kimi_k3_mxfp4_fixture.py +++ b/scripts/checkpoint_conversion/create_kimi_k3_mxfp4_fixture.py @@ -10,7 +10,6 @@ import argparse import hashlib import json -import re import subprocess from pathlib import Path from typing import Any @@ -19,13 +18,10 @@ import torchao from safetensors.torch import save_file from torchao.prototype.mx_formats.mx_tensor import MXTensor - from torchtitan.models.kimi_k3 import KimiK3StateDictAdapter, model_registry +from torchtitan.models.kimi_k3.quantization import MXFP4_QUANTIZATION_CONFIG +from torchtitan.quantization.mx_qat.checkpoint import MXFP4CheckpointPolicy -_EXPERT_WEIGHT = re.compile( - r"^language_model\.model\.layers\.\d+\.block_sparse_moe\.experts\.\d+\." - r"w[123]\.weight$" -) _DEFAULT_MAX_SHARD_BYTES = 1 << 30 @@ -50,26 +46,28 @@ def _sha256(path: Path) -> str: def convert_hf_state_dict_to_mxfp4( hf_state_dict: dict[str, torch.Tensor], + policy: MXFP4CheckpointPolicy, ) -> tuple[dict[str, torch.Tensor], int]: - """Replace routed-expert weights with released Kimi packed/scale pairs.""" + """Pack the exact resolved policy, including non-expert Linear weights.""" + missing = policy.weight_fqns - hf_state_dict.keys() + if missing: + raise ValueError(f"Fixture is missing selected weights: {sorted(missing)}") converted: dict[str, torch.Tensor] = {} pair_count = 0 for key, value in hf_state_dict.items(): value = value.detach().to(device="cpu").contiguous() - if _EXPERT_WEIGHT.fullmatch(key) is None: + if key not in policy.weight_fqns: converted[key] = value continue mx = MXTensor.to_mx( value, elem_dtype=torch.float4_e2m1fn_x2, - block_size=32, + block_size=policy.block_size, ) prefix = key.removesuffix(".weight") converted[f"{prefix}.weight_packed"] = mx.qdata.contiguous() - converted[f"{prefix}.weight_scale"] = ( - mx.scale.view(torch.uint8).contiguous() - ) + converted[f"{prefix}.weight_scale"] = mx.scale.view(torch.uint8).contiguous() pair_count += 1 return converted, pair_count @@ -141,30 +139,17 @@ def create_fixture( model.to(dtype=torch.bfloat16) adapter = KimiK3StateDictAdapter(model_spec.model, hf_assets_path=None) hf_state_dict = adapter.to_hf(model.state_dict()) - converted, pair_count = convert_hf_state_dict_to_mxfp4(hf_state_dict) + policy = MXFP4CheckpointPolicy.from_config( + MXFP4_QUANTIZATION_CONFIG, adapter.hf_linear_weight_mapping() + ) + converted, pair_count = convert_hf_state_dict_to_mxfp4(hf_state_dict, policy) storage = write_sharded_checkpoint(converted, output, max_shard_bytes) (output / "config.json").write_text( json.dumps( { "model_type": "kimi_k3", "text_config": { - "quantization_config": { - "format": "mxfp4-pack-quantized", - "quant_method": "compressed-tensors", - "config_groups": { - "group_0": { - "targets": ["Linear"], - "weights": {"group_size": 32}, - } - }, - "ignore": [ - "re:.*self_attn.*", - "re:.*shared_experts.*", - "re:.*vision_tower.*", - "re:.*mm_projector.*", - "re:.*lm_head.*", - ], - } + "quantization_config": MXFP4_QUANTIZATION_CONFIG, }, }, indent=2, @@ -176,7 +161,7 @@ def create_fixture( repo_root = Path(__file__).resolve().parents[2] manifest = { "block_size": 32, - "expert_pair_count": pair_count, + "packed_pair_count": pair_count, "model": "kimi_k3_debugmodel", "seed": seed, "torch_version": torch.__version__, diff --git a/scripts/checkpoint_conversion/create_mxfp4_reference.py b/scripts/checkpoint_conversion/create_mxfp4_reference.py new file mode 100644 index 00000000000..4227b7e084d --- /dev/null +++ b/scripts/checkpoint_conversion/create_mxfp4_reference.py @@ -0,0 +1,61 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +"""Regenerate tests/assets/mxfp4-reference.json with compressed-tensors==0.18.0. + +This independent fixture does not import TorchAO or TorchTitan. It exercises +every E2M1 value, multiple E8M0 scales, and rounding/saturation of off-grid values. +""" + +import importlib.metadata +import json +from pathlib import Path + +import torch +from compressed_tensors.compressors.mxfp4 import MXFP4PackedCompressor +from compressed_tensors.quantization import QuantizationArgs, QuantizationScheme + + +def main(): + version = importlib.metadata.version("compressed-tensors") + if version != "0.18.0": + raise RuntimeError(f"Expected compressed-tensors==0.18.0, got {version}") + levels = torch.tensor( + [0, 0.5, 1, 1.5, 2, 3, 4, 6, -0.0, -0.5, -1, -1.5, -2, -3, -4, -6] + ) + scales = torch.tensor([[0.5, 1], [2, 4]], dtype=torch.float32) + values = levels.repeat(8).reshape(2, 64) + values[1] += 0.2 + weight = values * scales.repeat_interleave(32, dim=-1) + scheme = QuantizationScheme( + targets=["Linear"], + weights=QuantizationArgs( + num_bits=4, + type="float", + strategy="group", + group_size=32, + symmetric=True, + dynamic=False, + scale_dtype=torch.uint8, + ), + ) + packed = MXFP4PackedCompressor.compress( + {"weight": weight, "weight_scale": scales}, scheme + ) + decoded = MXFP4PackedCompressor.decompress(packed, scheme)["weight"] + result = { + "serializer": f"compressed-tensors=={version}", + "generator": "scripts/checkpoint_conversion/create_mxfp4_reference.py", + "weight_packed": packed["weight_packed"].tolist(), + "weight_scale": packed["weight_scale"].tolist(), + "dequantized": decoded.tolist(), + } + path = Path(__file__).parents[2] / "tests/assets/mxfp4-reference.json" + path.write_text(json.dumps(result, indent=2) + "\n") + + +if __name__ == "__main__": + main() diff --git a/tests/assets/mxfp4-reference.json b/tests/assets/mxfp4-reference.json new file mode 100644 index 00000000000..b415ba46ff6 --- /dev/null +++ b/tests/assets/mxfp4-reference.json @@ -0,0 +1,218 @@ +{ + "serializer": "compressed-tensors==0.18.0", + "generator": "scripts/checkpoint_conversion/create_mxfp4_reference.py", + "weight_packed": [ + [ + 16, + 50, + 84, + 118, + 144, + 186, + 220, + 254, + 16, + 50, + 84, + 118, + 144, + 186, + 220, + 254, + 16, + 50, + 84, + 118, + 144, + 186, + 220, + 254, + 16, + 50, + 84, + 118, + 144, + 186, + 220, + 254 + ], + [ + 16, + 50, + 84, + 118, + 144, + 186, + 220, + 254, + 16, + 50, + 84, + 118, + 144, + 186, + 220, + 254, + 16, + 50, + 84, + 118, + 144, + 186, + 220, + 254, + 16, + 50, + 84, + 118, + 144, + 186, + 220, + 254 + ] + ], + "weight_scale": [ + [ + 126, + 127 + ], + [ + 128, + 129 + ] + ], + "dequantized": [ + [ + 0.0, + 0.25, + 0.5, + 0.75, + 1.0, + 1.5, + 2.0, + 3.0, + 0.0, + -0.25, + -0.5, + -0.75, + -1.0, + -1.5, + -2.0, + -3.0, + 0.0, + 0.25, + 0.5, + 0.75, + 1.0, + 1.5, + 2.0, + 3.0, + 0.0, + -0.25, + -0.5, + -0.75, + -1.0, + -1.5, + -2.0, + -3.0, + 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, + 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 + ], + [ + 0.0, + 1.0, + 2.0, + 3.0, + 4.0, + 6.0, + 8.0, + 12.0, + 0.0, + -1.0, + -2.0, + -3.0, + -4.0, + -6.0, + -8.0, + -12.0, + 0.0, + 1.0, + 2.0, + 3.0, + 4.0, + 6.0, + 8.0, + 12.0, + 0.0, + -1.0, + -2.0, + -3.0, + -4.0, + -6.0, + -8.0, + -12.0, + 0.0, + 2.0, + 4.0, + 6.0, + 8.0, + 12.0, + 16.0, + 24.0, + 0.0, + -2.0, + -4.0, + -6.0, + -8.0, + -12.0, + -16.0, + -24.0, + 0.0, + 2.0, + 4.0, + 6.0, + 8.0, + 12.0, + 16.0, + 24.0, + 0.0, + -2.0, + -4.0, + -6.0, + -8.0, + -12.0, + -16.0, + -24.0 + ] + ] +} diff --git a/tests/integration_tests/test_kimi_k3_mxfp4_checkpoint.py b/tests/integration_tests/test_kimi_k3_mxfp4_checkpoint.py index 4d2e9343454..ea774b0796e 100644 --- a/tests/integration_tests/test_kimi_k3_mxfp4_checkpoint.py +++ b/tests/integration_tests/test_kimi_k3_mxfp4_checkpoint.py @@ -10,33 +10,21 @@ import torch import torch.distributed.checkpoint as dcp -from torchao.prototype.mx_formats.mx_tensor import MXTensor - from scripts.checkpoint_conversion.create_kimi_k3_mxfp4_fixture import ( convert_hf_state_dict_to_mxfp4, write_sharded_checkpoint, ) +from torchao.prototype.mx_formats.mx_tensor import MXTensor from torchtitan.components.checkpointer.packed_hf_storage import ( PackedPairHuggingFaceStorageReader, PackedPairSpec, ) - -_HF_WEIGHT = ( - "language_model.model.layers.1.block_sparse_moe.experts.0.w1.weight" +from torchtitan.quantization.mx_qat.checkpoint import ( + decode_mxfp4, + MXFP4CheckpointPolicy, ) - -def _decode_mxfp4(packed, scales, block_size, target_dtype): - return MXTensor( - packed, - scales.view(torch.float8_e8m0fnu), - torch.float4_e2m1fn_x2, - block_size, - target_dtype, - None, - None, - False, - ).dequantize(target_dtype) +_HF_WEIGHT = "language_model.model.layers.1.block_sparse_moe.experts.0.w1.weight" class KimiK3MXFP4CheckpointIntegrationTest(unittest.TestCase): @@ -49,7 +37,8 @@ def test_release_format_fixture_loads_across_shards(self) -> None: block_size=32, ).dequantize(torch.bfloat16) converted, pair_count = convert_hf_state_dict_to_mxfp4( - {_HF_WEIGHT: weight, "dense.weight": dense} + {_HF_WEIGHT: weight, "dense.weight": dense}, + MXFP4CheckpointPolicy(frozenset({_HF_WEIGHT})), ) self.assertEqual(pair_count, 1) self.assertNotIn(_HF_WEIGHT, converted) @@ -75,8 +64,8 @@ def test_release_format_fixture_loads_across_shards(self) -> None: block_size=32, packed_values_per_byte=2, target_dtype=torch.bfloat16, - is_target=lambda key: key == _HF_WEIGHT, - decode=_decode_mxfp4, + target_fqns=frozenset({_HF_WEIGHT}), + decode=decode_mxfp4, ), ) dcp.load(destination, storage_reader=reader) @@ -86,10 +75,7 @@ def test_release_format_fixture_loads_across_shards(self) -> None: ) torch.testing.assert_close(destination["dense.weight"], dense, rtol=0, atol=0) self.assertFalse( - any( - key.endswith(("weight_packed", "weight_scale")) - for key in destination - ) + any(key.endswith(("weight_packed", "weight_scale")) for key in destination) ) diff --git a/tests/unit_tests/cpu/test_mx_qat_grouped_experts.py b/tests/unit_tests/cpu/test_mx_qat_grouped_experts.py index 7e383dba123..f4b6d706dc0 100644 --- a/tests/unit_tests/cpu/test_mx_qat_grouped_experts.py +++ b/tests/unit_tests/cpu/test_mx_qat_grouped_experts.py @@ -8,26 +8,9 @@ from unittest.mock import patch import torch - -from torchtitan.models.common.moe import GroupedExperts -from torchtitan.quantization.mx_qat.experts import ( - _get_mx_qat_grouped_experts_cls, -) from torchao.prototype.qat import mx_fake_quantized_grouped_mm - - -def _emulated_grouped_mm( - activation: torch.Tensor, - weight: torch.Tensor, - *, - offs: torch.Tensor, -) -> torch.Tensor: - outputs = [] - start = 0 - for expert, stop in enumerate(offs.tolist()): - outputs.append(activation[start:stop] @ weight[expert]) - start = stop - return torch.cat(outputs) +from torchtitan.models.common.moe import GroupedExperts +from torchtitan.quantization.mx_qat.experts import _get_mx_qat_grouped_experts_cls class MXQATGroupedExpertsTest(unittest.TestCase): @@ -64,20 +47,15 @@ def test_keeps_bf16_masters_and_applies_both_fake_quantizers(self) -> None: requires_grad=True, ) offsets = torch.tensor([2, 4], dtype=torch.int32) - original_grouped_mm = torch._grouped_mm - torch._grouped_mm = _emulated_grouped_mm - try: - with patch( - "torchao.prototype.qat.mx_fake_quantized_grouped_mm", - wraps=mx_fake_quantized_grouped_mm, - ) as grouped_mm: - output = module._grouped_mm( - A=activation, - weight_EOI=module.w1_EFD, - offs=offsets, - ) - finally: - torch._grouped_mm = original_grouped_mm + with patch( + "torchao.prototype.qat.mx_fake_quantized_grouped_mm", + wraps=mx_fake_quantized_grouped_mm, + ) as grouped_mm: + output = module._grouped_mm( + A=activation, + weight_EOI=module.w1_EFD, + offs=offsets, + ) output.float().sum().backward() self.assertEqual(grouped_mm.call_count, 1) @@ -85,6 +63,10 @@ def test_keeps_bf16_masters_and_applies_both_fake_quantizers(self) -> None: self.assertIsNotNone(module.w1_EFD.grad) self.assertTrue(torch.isfinite(activation.grad).all()) self.assertTrue(torch.isfinite(module.w1_EFD.grad).all()) + before = module.w1_EFD.detach().clone() + optimizer.step() + self.assertFalse(torch.equal(before, module.w1_EFD)) + self.assertEqual(parameter_ids, {id(p) for p in module.parameters()}) if __name__ == "__main__": diff --git a/tests/unit_tests/cpu/test_mx_qat_policy.py b/tests/unit_tests/cpu/test_mx_qat_policy.py new file mode 100644 index 00000000000..22e6b782da8 --- /dev/null +++ b/tests/unit_tests/cpu/test_mx_qat_policy.py @@ -0,0 +1,104 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +import json +import unittest +from pathlib import Path + +import torch +from torchtitan.quantization.mx_qat.checkpoint import ( + decode_mxfp4, + MXFP4CheckpointPolicy, +) + + +def _config(): + return { + "format": "mxfp4-pack-quantized", + "quant_method": "compressed-tensors", + "config_groups": { + "group_0": { + "targets": ["Linear"], + "weights": { + "num_bits": 4, + "type": "float", + "strategy": "group", + "group_size": 32, + "dynamic": False, + "symmetric": True, + "scale_dtype": "torch.uint8", + }, + } + }, + "ignore": ["re:.*shared"], + } + + +class MXFP4PolicyTest(unittest.TestCase): + def test_resolves_actual_linear_hierarchy_and_prefix_regex(self): + policy = MXFP4CheckpointPolicy.from_config( + _config(), + [ + "model.experts.0.w1.weight", + "model.latent_proj.weight", + "model.shared.gate.weight", + ], + ) + self.assertEqual( + policy.weight_fqns, + frozenset( + { + "model.experts.0.w1.weight", + "model.latent_proj.weight", + } + ), + ) + self.assertNotIn("model.embed_tokens.weight", policy.weight_fqns) + + def test_rejects_contradictory_numeric_metadata(self): + for field, value in ( + ("num_bits", 8), + ("type", "int"), + ("group_size", 64), + ("scale_dtype", "torch.float32"), + ("symmetric", False), + ("dynamic", True), + ("strategy", "channel"), + ): + with self.subTest(field=field): + config = _config() + config["config_groups"]["group_0"]["weights"][field] = value + with self.assertRaisesRegex(ValueError, "static symmetric float4"): + MXFP4CheckpointPolicy.from_config(config, ["linear.weight"]) + + def test_exact_targets_and_ignore(self): + config = _config() + config["config_groups"]["group_0"]["targets"] = ["model.proj"] + policy = MXFP4CheckpointPolicy.from_config( + config, ["model.proj.weight", "model.other.weight"] + ) + self.assertEqual(policy.weight_fqns, frozenset({"model.proj.weight"})) + + def test_independent_compressed_tensors_reference(self): + path = Path(__file__).parents[2] / "assets" / "mxfp4-reference.json" + fixture = json.loads(path.read_text()) + self.assertEqual(fixture["serializer"], "compressed-tensors==0.18.0") + actual = decode_mxfp4( + torch.tensor(fixture["weight_packed"], dtype=torch.uint8), + torch.tensor(fixture["weight_scale"], dtype=torch.uint8), + 32, + torch.bfloat16, + ) + torch.testing.assert_close( + actual, + torch.tensor(fixture["dequantized"], dtype=torch.bfloat16), + rtol=0, + atol=0, + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/unit_tests/cpu/test_mx_qat_transform.py b/tests/unit_tests/cpu/test_mx_qat_transform.py new file mode 100644 index 00000000000..a8804751e4b --- /dev/null +++ b/tests/unit_tests/cpu/test_mx_qat_transform.py @@ -0,0 +1,180 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +import unittest +from dataclasses import dataclass, replace + +import torch +from torchao.quantization.quantize_.common import KernelPreference +from torchtitan.config.transform import MXQATTransform, transform_model_config_ +from torchtitan.models.common.linear import CastLinear, Linear +from torchtitan.models.common.moe import GroupedExperts +from torchtitan.protocols.module import Module + + +class _Model(Module): + @dataclass(kw_only=True, slots=True) + class Config(Module.Config): + experts: GroupedExperts.Config + projection: Linear.Config + + +def _config(): + return _Model.Config( + experts=GroupedExperts.Config(dim=64, hidden_dim=64, num_experts=2), + projection=Linear.Config(in_features=64, out_features=32), + ) + + +class MXQATTransformTest(unittest.TestCase): + def test_configuration_types_are_resolvable_for_cli(self): + from typing import get_type_hints + + from torchao.prototype.qat import MXFakeQuantizeConfig + + self.assertIs( + get_type_hints(MXQATTransform)["weight_fake_quant_config"], + MXFakeQuantizeConfig, + ) + config = MXQATTransform().transform(_config()) + self.assertIs( + get_type_hints(type(config.experts))["activation_fake_quant_config"], + MXFakeQuantizeConfig, + ) + + def test_idempotence_preserves_parent_config(self): + config = _config() + transform = MXQATTransform() + transform.transform(config) + config_type = type(config.experts) + activation_fn = config.experts.activation_fn + transform.transform(config) + self.assertIs(type(config.experts), config_type) + self.assertIs(config.experts.activation_fn, activation_fn) + self.assertIs(type(config.projection), Linear.Config) + + def test_resolved_selection_and_kernel_preference(self): + config = _config() + transform = MXQATTransform.from_weight_fqns( + config, + { + "experts.w1_EFD", + "experts.w2_EDF", + "experts.w3_EFD", + "projection.weight", + }, + ) + transform.weight_fake_quant_config = replace( + transform.weight_fake_quant_config, kernel_preference=KernelPreference.AUTO + ) + transform.activation_fake_quant_config = replace( + transform.activation_fake_quant_config, + kernel_preference=KernelPreference.AUTO, + ) + transform.transform(config) + self.assertEqual( + config.experts.weight_fake_quant_config.kernel_preference, + KernelPreference.AUTO, + ) + self.assertEqual( + config.experts.activation_fake_quant_config.kernel_preference, + KernelPreference.AUTO, + ) + self.assertEqual( + config.projection.weight_fake_quant_config.kernel_preference, + KernelPreference.AUTO, + ) + + def test_gpt_oss_expert_layout_preserves_biases_and_config(self): + from torchtitan.models.gpt_oss.moe import GptOssGroupedExperts + + config = _config() + config.experts = GptOssGroupedExperts.Config( + dim=64, hidden_dim=64, num_experts=2, swiglu_limit=5.0 + ) + transform = MXQATTransform.from_weight_fqns( + config, {"experts.mlp1_weight_EGD", "experts.mlp2_weight_EDF"} + ) + transform.transform(config) + module = config.experts.build() + self.assertIsInstance(module, GptOssGroupedExperts) + self.assertEqual(module.swiglu_limit, 5.0) + self.assertEqual( + set(module.state_dict()), + {"mlp1_weight_EGD", "mlp1_bias_EG", "mlp2_weight_EDF", "mlp2_bias_ED"}, + ) + + def test_rejects_partial_group_and_unknown_weights(self): + for weights in ({"experts.w1_EFD"}, {"missing.weight"}): + with self.subTest(weights=weights), self.assertRaises(ValueError): + MXQATTransform.from_weight_fqns(_config(), weights) + + def test_rejects_unknown_fqn_before_mutation(self): + for name in ("missing", "experts"): + with self.subTest(name=name): + config = _config() + with self.assertRaisesRegex(ValueError, "did not match"): + MXQATTransform(linear_fqns=(name,)).transform(config) + self.assertIs(type(config.experts), GroupedExperts.Config) + + def test_rejects_duplicate_transform_sequence(self): + with self.assertRaisesRegex(ValueError, "cannot be combined"): + transform_model_config_(_config(), [MXQATTransform(), MXQATTransform()]) + + def test_preserves_custom_linear_contract_by_rejecting_it(self): + config = _config() + config.projection = CastLinear.Config(in_features=64, out_features=32) + with self.assertRaisesRegex(ValueError, "custom forward"): + MXQATTransform(linear_fqns=("projection",)).transform(config) + + def test_lora_wraps_qat_and_preserves_its_forward(self): + from torchao.prototype.qat import mx_fake_quantize + from torchtitan.config.transform import LinearLoRAHandler, LoRATransform + + config = transform_model_config_( + _config(), + [ + LoRATransform(handlers=(LinearLoRAHandler(),), rank=4), + MXQATTransform(linear_fqns=("projection",)), + ], + ) + module = config.projection.build() + with torch.no_grad(): + module.weight.normal_() + module.lora_a.weight.normal_() + module.lora_b.weight.zero_() + value = torch.randn(2, 64, requires_grad=True) + expected = torch.nn.functional.linear( + value, + mx_fake_quantize(module.weight, config.projection.weight_fake_quant_config), + ) + actual = module(value) + torch.testing.assert_close(actual, expected) + actual.sum().backward() + self.assertIsNone(module.weight.grad) + self.assertGreater(torch.count_nonzero(module.lora_b.weight.grad), 0) + self.assertTrue(torch.isfinite(value.grad).all()) + + def test_dense_qat_keeps_parameter_names_and_optimizer_identity(self): + config = _config() + MXQATTransform(grouped_expert_fqns=(), linear_fqns=("projection",)).transform( + config + ) + module = config.projection.build().to(dtype=torch.bfloat16) + with torch.no_grad(): + module.weight.normal_() + parameter = module.weight + optimizer = torch.optim.SGD(module.parameters(), lr=0.1) + module(torch.randn(2, 64, dtype=torch.bfloat16)).float().sum().backward() + before = parameter.detach().clone() + optimizer.step() + self.assertIs(module.weight, parameter) + self.assertEqual(set(module.state_dict()), {"weight"}) + self.assertFalse(torch.equal(before, parameter)) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/unit_tests/cpu/test_packed_hf_storage.py b/tests/unit_tests/cpu/test_packed_hf_storage.py index 1041561bb86..46c84bf1a71 100644 --- a/tests/unit_tests/cpu/test_packed_hf_storage.py +++ b/tests/unit_tests/cpu/test_packed_hf_storage.py @@ -14,41 +14,17 @@ from safetensors.torch import save_file from torch.distributed.checkpoint.metadata import MetadataIndex from torch.distributed.checkpoint.planner import LoadItemType, ReadItem - from torchtitan.components.checkpointer.packed_hf_storage import ( PackedPairHuggingFaceStorageReader, PackedPairSpec, ) +from torchtitan.quantization.mx_qat.checkpoint import decode_mxfp4 _PACKED_KEY = "model.layers.0.experts.0.w1.weight_packed" _SCALE_KEY = "model.layers.0.experts.0.w1.weight_scale" _VIRTUAL_KEY = "model.layers.0.experts.0.w1.weight" -def _is_expert_weight(key: str) -> bool: - return ".experts." in key - - -def _decode_mxfp4( - packed: torch.Tensor, - scales: torch.Tensor, - block_size: int, - target_dtype: torch.dtype, -) -> torch.Tensor: - from torchao.prototype.mx_formats.mx_tensor import MXTensor - - return MXTensor( - packed, - scales.view(torch.float8_e8m0fnu), - torch.float4_e2m1fn_x2, - block_size, - target_dtype, - None, - None, - False, - ).dequantize(target_dtype) - - def _spec() -> PackedPairSpec: return PackedPairSpec( packed_suffix=".weight_packed", @@ -57,8 +33,8 @@ def _spec() -> PackedPairSpec: block_size=32, packed_values_per_byte=2, target_dtype=torch.bfloat16, - is_target=_is_expert_weight, - decode=_decode_mxfp4, + target_fqns=frozenset({_VIRTUAL_KEY}), + decode=decode_mxfp4, ) @@ -159,6 +135,13 @@ def test_read_metadata_exposes_virtual_weight(self) -> None: {_VIRTUAL_KEY, "dense.weight"}, ) + def test_read_metadata_rejects_dense_substitution_for_expected_pair(self) -> None: + path = self._write_checkpoint( + {_VIRTUAL_KEY: torch.ones((2, 64), dtype=torch.bfloat16)} + ) + with self.assertRaisesRegex(ValueError, "requires missing pairs"): + PackedPairHuggingFaceStorageReader(path, _spec()).read_metadata() + def test_read_metadata_rejects_missing_scale(self) -> None: path = self._write_checkpoint( {_PACKED_KEY: torch.zeros((2, 32), dtype=torch.uint8)} @@ -329,5 +312,6 @@ def test_e8m0_extreme_bytes_follow_mx_semantics(self) -> None: equal_nan=True, ) + if __name__ == "__main__": unittest.main() diff --git a/tests/unit_tests/cpu/test_state_dict_adapter.py b/tests/unit_tests/cpu/test_state_dict_adapter.py index bcedb53307e..80e605e65a0 100644 --- a/tests/unit_tests/cpu/test_state_dict_adapter.py +++ b/tests/unit_tests/cpu/test_state_dict_adapter.py @@ -7,6 +7,7 @@ import json import tempfile import unittest +from copy import deepcopy from pathlib import Path import torch @@ -28,6 +29,7 @@ from torchtitan.models.gpt_oss import gptoss_configs from torchtitan.models.gpt_oss.state_dict_adapter import GptOssStateDictAdapter from torchtitan.models.kimi_k3 import model_registry as kimi_k3_model_registry +from torchtitan.models.kimi_k3.quantization import MXFP4_QUANTIZATION_CONFIG from torchtitan.models.kimi_k3.state_dict_adapter import KimiK3StateDictAdapter from torchtitan.models.llama3 import llama3_configs from torchtitan.models.llama3.model import Llama3Model @@ -136,17 +138,7 @@ def test_quantized_load_uses_packed_pair_reader(self) -> None: json.dumps( { "text_config": { - "quantization_config": { - "format": "mxfp4-pack-quantized", - "quant_method": "compressed-tensors", - "config_groups": { - "group_0": { - "targets": ["Linear"], - "weights": {"group_size": 32}, - } - }, - "ignore": ["re:.*shared_experts.*"], - } + "quantization_config": deepcopy(MXFP4_QUANTIZATION_CONFIG), } } ) @@ -163,20 +155,82 @@ def test_quantized_load_uses_packed_pair_reader(self) -> None: self.assertEqual(reader.spec.block_size, 32) self.assertEqual(reader.spec.target_dtype, torch.bfloat16) self.assertTrue( - reader.spec.is_target( - "language_model.model.layers.1.block_sparse_moe.experts.3.w1.weight" - ) + "language_model.model.layers.1.block_sparse_moe.experts.3.w1.weight" + in reader.spec.target_fqns ) self.assertFalse( - reader.spec.is_target( - "language_model.model.layers.1.block_sparse_moe.shared_experts.w1.weight" - ) + "language_model.model.layers.1.block_sparse_moe.shared_experts.w1.weight" + in reader.spec.target_fqns ) - self.assertTrue( - reader.spec.is_target( - "language_model.model.layers.1.block_sparse_moe.experts.3.w4.weight" - ) + self.assertFalse( + "language_model.model.layers.1.block_sparse_moe.experts.3.w4.weight" + in reader.spec.target_fqns + ) + self.assertFalse( + "language_model.model.embed_tokens.weight" in reader.spec.target_fqns + ) + self.assertFalse( + "language_model.model.layers.1.block_sparse_moe.gate.weight" + in reader.spec.target_fqns + ) + self.assertIn( + "language_model.model.layers.1.block_sparse_moe.routed_expert_up_proj.weight", + reader.spec.target_fqns, + ) + self.assertIn( + "language_model.model.layers.1.mlp_res_proj.weight", reader.spec.target_fqns + ) + + def test_qat_selection_matches_import_including_dense_projections(self) -> None: + from torchtitan.config.transform import MXQATTransform + from torchtitan.quantization.mx_qat.checkpoint import MXFP4CheckpointPolicy + + mapping = self.adapter.hf_linear_weight_mapping() + policy = MXFP4CheckpointPolicy.from_config(MXFP4_QUANTIZATION_CONFIG, mapping) + model = self.adapter.kimi_config + transform = MXQATTransform.from_weight_fqns( + model, + {mapping[key] for key in policy.weight_fqns if mapping[key] is not None}, + ) + model = transform.transform(model) + self.adapter._validate_qat_policy(policy) + self.assertTrue(type(model.layers[1].moe.routed_up)._owner._mx_qat) + + def test_qat_rejects_expert_only_selection_for_released_policy(self) -> None: + from torchtitan.config.transform import MXQATTransform + from torchtitan.quantization.mx_qat.checkpoint import MXFP4CheckpointPolicy + + policy = MXFP4CheckpointPolicy.from_config( + MXFP4_QUANTIZATION_CONFIG, self.adapter.hf_linear_weight_mapping() + ) + MXQATTransform().transform(self.adapter.kimi_config) + with self.assertRaisesRegex(ValueError, "selection disagrees"): + self.adapter._validate_qat_policy(policy) + + def test_vision_policy_uses_runtime_layer_names(self) -> None: + from torchtitan.quantization.mx_qat.checkpoint import MXFP4CheckpointPolicy + + mapping = self.adapter.hf_linear_weight_mapping() + name = "vision_tower.encoder.blocks.0.mlp.fc0.weight" + self.assertEqual(mapping[name], "vision_encoder.layers.0.mlp.linear_fc1.weight") + policy = MXFP4CheckpointPolicy.from_config(MXFP4_QUANTIZATION_CONFIG, mapping) + self.assertNotIn(name, policy.weight_fqns) + quantization = deepcopy(MXFP4_QUANTIZATION_CONFIG) + quantization["ignore"] = [] + self.assertIn( + name, MXFP4CheckpointPolicy.from_config(quantization, mapping).weight_fqns + ) + + def test_qat_recipe_is_valid_with_and_without_initial_checkpoint(self) -> None: + from torchtitan.models.kimi_k3.config_registry import kimi_k3_debugmodel_mx_qat + + recipe = kimi_k3_debugmodel_mx_qat(seq_len=16) + self.assertFalse(recipe.checkpointer.initial_load_in_hf_quantized) + recipe = kimi_k3_debugmodel_mx_qat( + seq_len=16, checkpoint_path="/tmp/packed-kimi" ) + self.assertTrue(recipe.checkpointer.initial_load_in_hf_quantized) + self.assertEqual(recipe.checkpointer.initial_load_path, "/tmp/packed-kimi") def test_quantized_load_rejects_missing_metadata(self) -> None: with tempfile.TemporaryDirectory() as directory: diff --git a/torchtitan/components/checkpointer/packed_hf_storage.py b/torchtitan/components/checkpointer/packed_hf_storage.py index fba7b63c1b6..96d437f4cb1 100644 --- a/torchtitan/components/checkpointer/packed_hf_storage.py +++ b/torchtitan/components/checkpointer/packed_hf_storage.py @@ -33,7 +33,7 @@ class PackedPairSpec: block_size: int packed_values_per_byte: int target_dtype: torch.dtype - is_target: Callable[[str], bool] + target_fqns: frozenset[str] decode: Callable[[torch.Tensor, torch.Tensor, int, torch.dtype], torch.Tensor] def __post_init__(self) -> None: @@ -83,7 +83,7 @@ def _validate_pair( packed_metadata: TensorStorageMetadata, scale_metadata: TensorStorageMetadata, ) -> None: - if not self.spec.is_target(virtual_fqn): + if virtual_fqn not in self.spec.target_fqns: raise ValueError( f"Packed tensor {virtual_fqn!r} is outside the packed-weight policy." ) @@ -178,6 +178,12 @@ def read_metadata(self) -> Any: if orphan_scales: raise ValueError(f"Found orphan scale tensor {orphan_scales[0]!r}.") + missing = self.spec.target_fqns - pairs.keys() + if missing: + raise ValueError( + f"Packed-weight policy requires missing pairs: {sorted(missing)[:10]}." + ) + virtual_state_dict_metadata = { fqn: tensor_metadata for fqn, tensor_metadata in state_dict_metadata.items() @@ -258,9 +264,7 @@ def _process_read_request( row_stop = row_start + row_count column_stop = column_start + column_count group_start = column_start // self.spec.block_size - group_stop = ( - column_stop + self.spec.block_size - 1 - ) // self.spec.block_size + group_stop = (column_stop + self.spec.block_size - 1) // self.spec.block_size packed_values_per_group = ( self.spec.block_size // self.spec.packed_values_per_byte ) diff --git a/torchtitan/config/transform/README.md b/torchtitan/config/transform/README.md index f7ecbd9b62e..2bf501f5f34 100644 --- a/torchtitan/config/transform/README.md +++ b/torchtitan/config/transform/README.md @@ -38,6 +38,35 @@ model_config = model_registry("0.6B", attn_backend="varlen") model_config = transform_model_config_(model_config, [LMHeadCastTransform()]) ``` +## MX quantization-aware training + +`MXQATTransform` keeps master parameters and optimizer state in the training +precision. It specializes the grouped-MM hook and, when selected, ordinary +`Linear` projections. Grouped experts fake-quantize weights and activations; +dense projections fake-quantize weights only. + +```python +config = apply_transforms(config, [MXQATTransform()]) +``` + +By default, every grouped-expert config is selected and dense projections are +unchanged. Use exact config FQNs in `grouped_expert_fqns` and `linear_fqns` for +explicit selection. An empty tuple selects none. `from_weight_fqns` translates +adapter-resolved parameter FQNs and rejects unsupported or partially selected +grouped modules. This translation requires parameter and config paths to agree; +models with renamed or repeated configs need adapter-specific translation. + +Pass TorchAO `MXFakeQuantizeConfig` instances through `weight_fake_quant_config` +and `activation_fake_quant_config`. Their existing `kernel_preference` controls +both quantization and grouped execution: `EMULATED` uses dequantized operands, +while `AUTO` uses native MXFP8 grouped forward kernels with a high-precision STE +backward. Both preferences must agree. Native execution requires supported CUDA +hardware and TorchAO kernel dependencies. + +The transform preserves inherited config fields, rejects existing incompatible +execution overrides, and runs before `LoRATransform`. Packed checkpoint import +resolves its own quantization metadata and checks that the QAT selection agrees. + ## What belongs here Use `model_registry` to select the base architecture, attention algorithm, and diff --git a/torchtitan/config/transform/__init__.py b/torchtitan/config/transform/__init__.py index 33f8912fe33..368876e81d5 100644 --- a/torchtitan/config/transform/__init__.py +++ b/torchtitan/config/transform/__init__.py @@ -13,12 +13,12 @@ from .context_parallel import ContextParallelTransform from .converter import ModelConfigConverter, validate_converter_compatibility from .lora import LinearLoRAHandler, LoRATransform +from .mx_qat import MXQATTransform from .quantization import ( Float8GroupedExpertsConverter, Float8LinearConverter, MXFP8GroupedExpertsConverter, MXFP8LinearConverter, - MXQATGroupedExpertsConverter, NVFP4LinearConverter, QuantizationConverter, ) @@ -38,7 +38,7 @@ "Float8LinearConverter", "MXFP8GroupedExpertsConverter", "MXFP8LinearConverter", - "MXQATGroupedExpertsConverter", + "MXQATTransform", "NVFP4LinearConverter", "QuantizationConverter", "validate_converter_compatibility", diff --git a/torchtitan/config/transform/lora.py b/torchtitan/config/transform/lora.py index de4be627cab..1e23f5749fe 100644 --- a/torchtitan/config/transform/lora.py +++ b/torchtitan/config/transform/lora.py @@ -14,7 +14,7 @@ from .base import ModelConfigTransform from .context_parallel import ContextParallelTransform - +from .mx_qat import MXQATTransform logger = logging.getLogger(__name__) @@ -105,6 +105,7 @@ class LoRATransform(ModelConfigTransform): run_after: ClassVar[tuple[type[ModelConfigTransform], ...]] = ( ContextParallelTransform, + MXQATTransform, ) handlers: tuple[_LoRAHandler, ...] diff --git a/torchtitan/config/transform/mx_qat.py b/torchtitan/config/transform/mx_qat.py new file mode 100644 index 00000000000..48f084ebe4c --- /dev/null +++ b/torchtitan/config/transform/mx_qat.py @@ -0,0 +1,120 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +"""Model-independent MX QAT selection and config transformation.""" + +from dataclasses import dataclass, field, replace + +import torch + +from torchtitan.models.common.linear import Linear +from torchtitan.models.common.moe import GroupedExperts +from torchtitan.protocols.module import Module +from torchtitan.quantization.mx_qat.experts import ( + _get_mx_qat_grouped_experts_cls, + activation_config, + MXFakeQuantizeConfig, + weight_config, +) +from torchtitan.quantization.mx_qat.linear import _get_mx_qat_linear_cls + +from .base import convert_config_type, ModelConfigTransform + + +@dataclass(kw_only=True, slots=True) +class MXQATTransform(ModelConfigTransform): + """Preserve BF16 masters and select QAT by exact module FQN. + + None selects every grouped-expert module; an empty tuple selects none. + Dense projections are opt-in and use weight-only fake quantization. + TorchAO configs, including kernel_preference, pass through unchanged. + """ + + grouped_expert_fqns: tuple[str, ...] | None = None + linear_fqns: tuple[str, ...] = () + weight_fake_quant_config: "MXFakeQuantizeConfig" = field( + default_factory=weight_config + ) + activation_fake_quant_config: "MXFakeQuantizeConfig" = field( + default_factory=activation_config + ) + + @classmethod + def from_weight_fqns(cls, model: Module.Config, weight_fqns: set[str], **kwargs): + """Translate adapter-resolved weights without assuming expert names. + + A grouped execution hook quantizes every expert matrix it consumes; + reject partial selection instead of silently quantizing extra weights. + Meta construction inspects registered shapes without allocating weights. + """ + remaining = set(weight_fqns) + groups, linears = [], [] + for fqn, config, _, _ in model.traverse(GroupedExperts.Config): + with torch.device("meta"): + module = config.build() + weights = { + f"{fqn}.{name}".lstrip(".") + for name, parameter in module.named_parameters(recurse=False) + if parameter.ndim == 3 + } + if remaining & weights: + if not weights <= remaining: + raise ValueError( + f"MX QAT requires all matrices in grouped module {fqn}" + ) + groups.append(fqn) + remaining -= weights + for fqn, _, _, _ in model.traverse(Linear.Config): + key = f"{fqn}.weight".lstrip(".") + if key in remaining: + linears.append(fqn) + remaining.remove(key) + if remaining: + raise ValueError(f"MX QAT cannot represent weights: {sorted(remaining)}") + return cls( + grouped_expert_fqns=tuple(groups), linear_fqns=tuple(linears), **kwargs + ) + + def transform(self, model: Module.Config) -> Module.Config: + replacements = [] + missing = set() + for config_type, targets, factory in ( + ( + GroupedExperts.Config, + self.grouped_expert_fqns, + _get_mx_qat_grouped_experts_cls, + ), + (Linear.Config, self.linear_fqns, _get_mx_qat_linear_cls), + ): + matched = set() + for fqn, config, parent, attr in model.traverse(config_type): + if targets is not None and fqn not in targets: + continue + replacement = factory(type(config)._owner) + new_config = convert_config_type(config, replacement) + deltas = {"weight_fake_quant_config": self.weight_fake_quant_config} + if config_type is GroupedExperts.Config: + deltas[ + "activation_fake_quant_config" + ] = self.activation_fake_quant_config + replacements.append((replace(new_config, **deltas), parent, attr)) + matched.add(fqn) + missing.update(set(targets or ()) - matched) + if missing: + raise ValueError(f"MX QAT module FQNs did not match: {sorted(missing)}") + for config, parent, attr in replacements: + if parent is None: + model = config + elif isinstance(parent, list): + parent[attr] = config + else: + setattr(parent, attr, config) + return model + + +# Repeating QAT in one transform sequence is an error; applying it again +# to an existing tree is idempotent. +MXQATTransform.conflicts_with = (MXQATTransform,) diff --git a/torchtitan/config/transform/quantization.py b/torchtitan/config/transform/quantization.py index 4ace20060e0..1e84b67f146 100644 --- a/torchtitan/config/transform/quantization.py +++ b/torchtitan/config/transform/quantization.py @@ -25,14 +25,12 @@ ) from torchtitan.quantization.mxfp8 import _mxfp8_linear_import_error, MXFP8Linear from torchtitan.quantization.mxfp8.experts import _get_mxfp8_grouped_experts_cls -from torchtitan.quantization.mx_qat.experts import _get_mx_qat_grouped_experts_cls from torchtitan.quantization.nvfp4 import NVFP4Linear from torchtitan.quantization.utils import module_filter_fn, swap_token_dispatcher from torchtitan.tools.utils import has_cuda_capability, has_rocm_capability from .converter import ModelConfigConverter - logger = logging.getLogger(__name__) @@ -435,47 +433,6 @@ def convert(self, model_config): return model_config -class MXQATGroupedExpertsConverter(QuantizationConverter): - """Apply emulated MXFP4-weight/MXFP8-activation QAT to grouped experts.""" - - @dataclass(kw_only=True, slots=True) - class Config(QuantizationConverter.Config): - weight_block_size: int = 32 - activation_block_size: int = 32 - - def __init__(self, config: Config): - self.config = config - try: - from torchao.prototype.qat import mx_fake_quantize # noqa: F401 - except ImportError as error: - raise ImportError( - "MX QAT grouped experts require a TorchAO build providing " - "torchao.prototype.qat.mx_fake_quantize." - ) from error - - def convert(self, model_config): - for _fqn, config, parent, attr in model_config.traverse(GroupedExperts.Config): - base_module_cls = type(config)._owner - quantized_cls = _get_mx_qat_grouped_experts_cls(base_module_cls) - config_cls = quantized_cls.Config # type: ignore[attr-defined] - new_config = config_cls( - **{f.name: getattr(config, f.name) for f in fields(config)}, - weight_block_size=self.config.weight_block_size, - activation_block_size=self.config.activation_block_size, - ) - if parent is None: - model_config = new_config - elif isinstance(parent, list): - parent[attr] = new_config - else: - setattr(parent, attr, new_config) - - logger.info( - "Converted GroupedExperts to MXFP4-weight/MXFP8-activation QAT" - ) - return model_config - - class NVFP4LinearConverter(QuantizationConverter): """Replace matching Linear.Config with NVFP4Linear.Config.""" diff --git a/torchtitan/models/kimi_k3/config_registry.py b/torchtitan/models/kimi_k3/config_registry.py index ebdf58a1462..b4bfd9f8a1c 100644 --- a/torchtitan/models/kimi_k3/config_registry.py +++ b/torchtitan/models/kimi_k3/config_registry.py @@ -12,7 +12,7 @@ from torchtitan.components.optimizer import default_adamw, LRSchedulersContainer from torchtitan.components.tokenizer import MultiModalTokenizer from torchtitan.config import TrainingConfig -from torchtitan.config.transform import MXQATGroupedExpertsConverter +from torchtitan.config.transform import apply_transforms, MXQATTransform from torchtitan.distributed.activation_checkpoint import SelectiveAC from torchtitan.hf_datasets.multimodal.mm_collator import MultiModalCollator from torchtitan.hf_datasets.multimodal.mm_datasets import ( @@ -25,9 +25,12 @@ DEFAULT_DEBUG_MODEL_SEQ_LEN, ) from torchtitan.observability.metrics import MetricsProcessor +from torchtitan.quantization.mx_qat.checkpoint import MXFP4CheckpointPolicy from torchtitan.trainer import Trainer from . import KIMI_K3_SPECIAL_TOKENS, model_registry +from .quantization import MXFP4_QUANTIZATION_CONFIG +from .state_dict_adapter import KimiK3StateDictAdapter def _kimi_k3_multimodal_dataloader( @@ -96,18 +99,25 @@ def kimi_k3_debugmodel( def kimi_k3_debugmodel_mx_qat( seq_len: int | None = DEFAULT_DEBUG_MODEL_SEQ_LEN, + *, + checkpoint_path: str | None = None, ) -> Trainer.Config: - """Kimi-K3 debug recipe with MXFP4-weight/MXFP8-activation expert QAT.""" + """Kimi QAT using the released policy and optional packed HF initialization. + + Pass an absolute checkpoint_path to load the packed debug fixture. Without + it, the recipe uses random initialization and remains valid before overrides. + """ config = kimi_k3_debugmodel(seq_len=seq_len) - config.model_spec = model_registry( - "debugmodel", - seq_len=seq_len, - converters=[MXQATGroupedExpertsConverter.Config()], - ) + adapter = KimiK3StateDictAdapter(config.model_spec.model, hf_assets_path=None) + mapping = adapter.hf_linear_weight_mapping() + policy = MXFP4CheckpointPolicy.from_config(MXFP4_QUANTIZATION_CONFIG, mapping) + weights = {mapping[key] for key in policy.weight_fqns if mapping[key] is not None} + transform = MXQATTransform.from_weight_fqns(config.model_spec.model, weights) config.checkpointer = CheckpointManager.Config( interval=5, - initial_load_in_hf=True, - initial_load_in_hf_quantized=True, + initial_load_path=checkpoint_path, + initial_load_in_hf=checkpoint_path is not None, + initial_load_in_hf_quantized=checkpoint_path is not None, last_save_model_only=False, ) - return config + return apply_transforms(config, [transform]) diff --git a/torchtitan/models/kimi_k3/quantization.py b/torchtitan/models/kimi_k3/quantization.py new file mode 100644 index 00000000000..0616927de90 --- /dev/null +++ b/torchtitan/models/kimi_k3/quantization.py @@ -0,0 +1,38 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +"""Released Kimi MXFP4 schema used by the debug recipe and fixture. + +Source: moonshotai/Kimi-K3 config.json, revision c5d1dd4. +Import validates the checkpoint's own metadata; this is not a loader default. +""" + +MXFP4_QUANTIZATION_CONFIG = { + "format": "mxfp4-pack-quantized", + "quant_method": "compressed-tensors", + "config_groups": { + "group_0": { + "targets": ["Linear"], + "weights": { + "num_bits": 4, + "type": "float", + "strategy": "group", + "group_size": 32, + "symmetric": True, + "dynamic": False, + "scale_dtype": "torch.uint8", + }, + }, + }, + "ignore": [ + "re:.*self_attn.*", + "re:.*shared_experts.*", + r"re:.*mlp\.(gate|up|gate_up|down)_proj.*", + "re:.*lm_head.*", + "re:.*vision_tower.*", + "re:.*mm_projector.*", + ], +} diff --git a/torchtitan/models/kimi_k3/state_dict_adapter.py b/torchtitan/models/kimi_k3/state_dict_adapter.py index c3e5b98545e..9dba0700ec7 100644 --- a/torchtitan/models/kimi_k3/state_dict_adapter.py +++ b/torchtitan/models/kimi_k3/state_dict_adapter.py @@ -6,11 +6,10 @@ from __future__ import annotations -"""HuggingFace checkpoint adapter for unquantized Kimi K3 weights.""" +"""HuggingFace checkpoint adapter for dense and packed Kimi K3 weights.""" import json import re -from collections.abc import Callable from pathlib import Path from typing import Any, TYPE_CHECKING @@ -22,7 +21,13 @@ PackedPairHuggingFaceStorageReader, PackedPairSpec, ) +from torchtitan.models.common.linear import Linear, RouterGateLinear +from torchtitan.models.common.moe import GroupedExperts from torchtitan.models.utils import MoEStateDictAdapter +from torchtitan.quantization.mx_qat.checkpoint import ( + decode_mxfp4, + MXFP4CheckpointPolicy, +) if TYPE_CHECKING: from .model import KimiK3Model @@ -33,77 +38,6 @@ } -def _released_mxfp4_policy(path: str) -> tuple[int, Callable[[str], bool]]: - config_path = Path(path) / "config.json" - if not config_path.is_file(): - raise ValueError(f"Quantized Kimi checkpoint is missing {config_path}.") - config = json.loads(config_path.read_text()) - text_config = config.get("text_config", config) - quantization = text_config.get("quantization_config") - if not isinstance(quantization, dict): - raise ValueError("Kimi checkpoint is missing quantization_config metadata.") - if ( - quantization.get("format") != "mxfp4-pack-quantized" - or quantization.get("quant_method") != "compressed-tensors" - ): - raise ValueError("Kimi checkpoint does not declare compressed MXFP4 weights.") - - config_groups = quantization.get("config_groups") - if not isinstance(config_groups, dict) or not config_groups: - raise ValueError("Kimi quantization_config has no config groups.") - group_sizes = set() - targets = set() - for group in config_groups.values(): - if not isinstance(group, dict): - raise ValueError("Kimi quantization config group must be an object.") - targets.update(group.get("targets", ())) - weights = group.get("weights") - if not isinstance(weights, dict): - raise ValueError("Kimi quantization config group has no weight config.") - group_sizes.add(weights.get("group_size")) - if targets != {"Linear"} or len(group_sizes) != 1: - raise ValueError( - "Kimi packed import currently requires one Linear MXFP4 group size." - ) - block_size = group_sizes.pop() - if not isinstance(block_size, int): - raise ValueError("Kimi MXFP4 group_size must be an integer.") - - ignore_patterns = [] - for pattern in quantization.get("ignore", ()): - if not isinstance(pattern, str) or not pattern.startswith("re:"): - raise ValueError("Kimi packed import requires regex ignore entries.") - ignore_patterns.append(re.compile(pattern.removeprefix("re:"))) - - def is_target(weight_fqn: str) -> bool: - if not weight_fqn.endswith(".weight"): - return False - module_fqn = weight_fqn.removesuffix(".weight") - return not any(pattern.fullmatch(module_fqn) for pattern in ignore_patterns) - - return block_size, is_target - - -def _decode_mxfp4( - packed: torch.Tensor, - scales: torch.Tensor, - block_size: int, - target_dtype: torch.dtype, -) -> torch.Tensor: - from torchao.prototype.mx_formats.mx_tensor import MXTensor - - return MXTensor( - packed, - scales.view(torch.float8_e8m0fnu), - torch.float4_e2m1fn_x2, - block_size, - target_dtype, - None, - None, - False, - ).dequantize(target_dtype) - - class KimiK3StateDictAdapter(MoEStateDictAdapter): def __init__( self, @@ -214,6 +148,124 @@ def _map_from_hf_layer_key( ) return attention_map.get(abstract_key) + def hf_linear_weight_mapping(self) -> dict[str, str | None]: + """Map the supported HF Linear hierarchy to Titan parameter FQNs. + + Reuse the adapter's architecture mapping. RouterGateLinear represents + HF KimiMoEGate's raw parameter, not an HF Linear. Grouped expert tensors + represent one HF Linear per expert; no quantization-name regex is used. + """ + linear_weights = { + f"{fqn}.weight" + for fqn, config, _, _ in self.kimi_config.traverse(Linear.Config) + if not isinstance(config, RouterGateLinear.Config) + } + # MoonViT builds repeated layers from one block config and renames MLP + # fields. Inspect one meta block to recover its actual parameter names. + vision = self.kimi_config.vision_encoder + if vision is not None: + with torch.device("meta"): + block = vision.block.build() + linear_weights.update( + f"vision_encoder.layers.{layer}.{name}.weight" + for layer in range(vision.num_layers) + for name, module in block.named_modules() + if isinstance(module, Linear) + ) + grouped = { + fqn: config + for fqn, config, _, _ in self.kimi_config.traverse(GroupedExperts.Config) + } + result: dict[str, str | None] = {} + for mapping in (self.from_hf_map, self.mla_from_hf_map, self.kda_from_hf_map): + for hf_template, titan_template in mapping.items(): + layers = ( + range( + vision.num_layers + if hf_template.startswith("vision_tower.") + and vision is not None + else len(self.kimi_config.layers) + ) + if "{}" in titan_template + else (None,) + ) + for layer in layers: + titan_key = titan_template.format(layer) + if titan_key in linear_weights: + result[hf_template.format(layer)] = titan_key + elif hf_template.count("{}") == 2: + module = titan_key.rsplit(".", 1)[0] + if module in grouped: + for expert in range(grouped[module].num_experts): + result[hf_template.format(layer, expert)] = titan_key + # HF has an unused layer-zero residual projection absent from Titan. + if self.kimi_config.layers[0].attention_res_proj is None: + result[ + "language_model.model.layers.0.self_attention_res_proj.weight" + ] = None + # The HF vision projection is fused; Titan stores its three slices. + # Current Kimi recipes require vision to remain unquantized. + if self.kimi_config.vision_encoder is not None: + for layer in range(self.kimi_config.vision_encoder.num_layers): + result[ + f"vision_tower.encoder.blocks.{layer}.wqkv.weight" + ] = f"vision_encoder.layers.{layer}.attn.wqkv.weight" + return result + + def mxfp4_policy(self, path: str) -> MXFP4CheckpointPolicy: + config_path = Path(path) / "config.json" + if not config_path.is_file(): + raise ValueError(f"Quantized Kimi checkpoint is missing {config_path}.") + config = json.loads(config_path.read_text()) + quantization = config.get("text_config", config).get("quantization_config") + if not isinstance(quantization, dict): + raise ValueError("Kimi checkpoint is missing quantization_config metadata.") + return MXFP4CheckpointPolicy.from_config( + quantization, self.hf_linear_weight_mapping() + ) + + @staticmethod + def _validate_qat_weight_config(config, policy: MXFP4CheckpointPolicy) -> None: + if ( + config.dtype != torch.float4_e2m1fn_x2 + or config.block_size != policy.block_size + ): + raise ValueError("MX QAT weight format disagrees with checkpoint policy") + + def _validate_qat_policy(self, policy: MXFP4CheckpointPolicy) -> None: + """Reject a QAT recipe whose selected parameters differ from import.""" + mapping = self.hf_linear_weight_mapping() + selected = set() + has_qat = False + for fqn, config, _, _ in self.kimi_config.traverse(Linear.Config): + if getattr(type(config)._owner, "_mx_qat", False): + has_qat = True + self._validate_qat_weight_config( + config.weight_fake_quant_config, policy + ) + selected.add(f"{fqn}.weight") + for fqn, config, _, _ in self.kimi_config.traverse(GroupedExperts.Config): + if getattr(type(config)._owner, "_mx_qat", False): + has_qat = True + self._validate_qat_weight_config( + config.weight_fake_quant_config, policy + ) + selected.update( + key + for key in mapping.values() + if key and key.rsplit(".", 1)[0] == fqn + ) + if not has_qat: + return # Packed import into a BF16 model remains supported. + expected = { + mapping[key] for key in policy.weight_fqns if mapping[key] is not None + } + if selected != expected: + raise ValueError( + "MX QAT selection disagrees with checkpoint policy: " + f"missing={sorted(expected - selected)}, unexpected={sorted(selected - expected)}" + ) + def get_hf_storage_reader( self, path: str, @@ -221,7 +273,8 @@ def get_hf_storage_reader( ) -> HuggingFaceStorageReader: if not from_quantized: return super().get_hf_storage_reader(path, from_quantized=False) - block_size, is_target = _released_mxfp4_policy(path) + policy = self.mxfp4_policy(path) + self._validate_qat_policy(policy) return PackedPairHuggingFaceStorageReader( path=path, thread_count=4, @@ -229,11 +282,11 @@ def get_hf_storage_reader( packed_suffix=".weight_packed", scale_suffix=".weight_scale", virtual_suffix=".weight", - block_size=block_size, + block_size=policy.block_size, packed_values_per_byte=2, target_dtype=torch.bfloat16, - is_target=is_target, - decode=_decode_mxfp4, + target_fqns=policy.weight_fqns, + decode=decode_mxfp4, ), ) diff --git a/torchtitan/quantization/mx_qat/__init__.py b/torchtitan/quantization/mx_qat/__init__.py index 487ad815084..9f6e8003223 100644 --- a/torchtitan/quantization/mx_qat/__init__.py +++ b/torchtitan/quantization/mx_qat/__init__.py @@ -8,5 +8,4 @@ from .experts import _get_mx_qat_grouped_experts_cls - __all__ = ["_get_mx_qat_grouped_experts_cls"] diff --git a/torchtitan/quantization/mx_qat/checkpoint.py b/torchtitan/quantization/mx_qat/checkpoint.py new file mode 100644 index 00000000000..7e57304bdc5 --- /dev/null +++ b/torchtitan/quantization/mx_qat/checkpoint.py @@ -0,0 +1,123 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +"""Resolve compressed-tensors MXFP4 policy against a model's HF Linear weights.""" + +import re +from collections.abc import Iterable, Mapping +from dataclasses import dataclass +from typing import Any + +import torch + + +@dataclass(frozen=True) +class MXFP4CheckpointPolicy: + weight_fqns: frozenset[str] + block_size: int = 32 + + @classmethod + def from_config( + cls, quantization: Mapping[str, Any], linear_weights: Iterable[str] + ) -> "MXFP4CheckpointPolicy": + """Resolve an exact set; the adapter supplies actual HF Linear weights. + + Only the released static, symmetric 1x32 E2M1/E8M0 format is supported. + Names alone cannot distinguish Linear weights from embeddings or norms. + """ + if ( + quantization.get("format") != "mxfp4-pack-quantized" + or quantization.get("quant_method") != "compressed-tensors" + ): + raise ValueError("Checkpoint must declare compressed-tensors MXFP4 weights") + groups = quantization.get("config_groups") + if not isinstance(groups, dict) or not groups: + raise ValueError("MXFP4 quantization_config has no config groups") + targets = [] + expected = { + "num_bits": 4, + "type": "float", + "strategy": "group", + "group_size": 32, + "symmetric": True, + "dynamic": False, + "scale_dtype": "torch.uint8", + } + for group in groups.values(): + if not isinstance(group, dict): + raise ValueError("MXFP4 config group must be an object") + if group.get("format", "mxfp4-pack-quantized") != "mxfp4-pack-quantized": + raise ValueError("All groups must use mxfp4-pack-quantized") + weights = group.get("weights") + if not isinstance(weights, dict) or any( + weights.get(key) != value for key, value in expected.items() + ): + raise ValueError( + "MXFP4 weights require static symmetric float4, group_size=32 and uint8 E8M0 scales" + ) + if any( + weights.get(key) is not None + for key in ("actorder", "block_structure", "zp_dtype") + ): + raise ValueError( + "MXFP4 activation ordering, block structures and zero points are unsupported" + ) + if any( + group.get(key) is not None + for key in ("input_activations", "output_activations") + ): + raise ValueError( + "Checkpoint activation quantization is unsupported; configure runtime QAT separately" + ) + group_targets = group.get("targets") + if not isinstance(group_targets, list) or not group_targets: + raise ValueError("MXFP4 group targets must be a nonempty list") + targets.extend(group_targets) + ignore = quantization.get("ignore", []) + if not isinstance(ignore, list) or not all( + isinstance(p, str) for p in targets + ignore + ): + raise ValueError("MXFP4 targets and ignore entries must be strings") + + def matches(pattern: str, module: str) -> bool: + if pattern.startswith("re:"): + return re.match(pattern[3:], module) is not None + return pattern == "Linear" or pattern == module + + return cls( + frozenset( + name + for name in linear_weights + if any( + matches(pattern, name.removesuffix(".weight")) + for pattern in targets + ) + and not any( + matches(pattern, name.removesuffix(".weight")) for pattern in ignore + ) + ) + ) + + +def decode_mxfp4( + packed: torch.Tensor, + scales: torch.Tensor, + block_size: int, + target_dtype: torch.dtype, +) -> torch.Tensor: + """Adapt the HF byte representation to TorchAO's numerical codec.""" + from torchao.prototype.mx_formats.mx_tensor import MXTensor + + return MXTensor( + qdata=packed, + scale=scales.view(torch.float8_e8m0fnu), + elem_dtype=torch.float4_e2m1fn_x2, + block_size=block_size, + orig_dtype=target_dtype, + kernel_preference=None, + act_quant_kwargs=None, + is_swizzled_scales=False, + ).dequantize(target_dtype) diff --git a/torchtitan/quantization/mx_qat/experts.py b/torchtitan/quantization/mx_qat/experts.py index 627663b1a82..e6139ceea93 100644 --- a/torchtitan/quantization/mx_qat/experts.py +++ b/torchtitan/quantization/mx_qat/experts.py @@ -6,41 +6,67 @@ """MXFP4-weight/MXFP8-activation QAT for grouped experts.""" -from dataclasses import dataclass +from dataclasses import dataclass, field import torch +try: + from torchao.prototype.qat import MXFakeQuantizeConfig +except ImportError: + # Keep non-QAT recipes importable when optional TorchAO is not installed. + from typing import Any as MXFakeQuantizeConfig + + +def weight_config(): + from torchao.prototype.qat import MXFakeQuantizeConfig + + return MXFakeQuantizeConfig(dtype=torch.float4_e2m1fn_x2) + + +def activation_config(): + from torchao.prototype.qat import MXFakeQuantizeConfig + + return MXFakeQuantizeConfig(dtype=torch.float8_e4m3fn) + _mx_qat_experts_cache: dict[type, type] = {} def _get_mx_qat_grouped_experts_cls(parent_cls: type) -> type: """Return a grouped-expert subclass using stateless TorchAO MX QAT.""" + if getattr(parent_cls, "_mx_qat", False): + return parent_cls if parent_cls in _mx_qat_experts_cache: return _mx_qat_experts_cache[parent_cls] + from torchtitan.models.common.moe import GroupedExperts + + if parent_cls._grouped_mm is not GroupedExperts._grouped_mm: + raise ValueError( + f"MX QAT cannot replace an existing grouped-MM override on {parent_cls.__name__}" + ) parent_config_cls = parent_cls.Config # type: ignore[attr-defined] class MXQATGroupedExperts(parent_cls): # type: ignore[valid-type, misc] + _mx_qat = True + @dataclass(kw_only=True, slots=True) class Config(parent_config_cls): # type: ignore[misc] - weight_block_size: int = 32 - activation_block_size: int = 32 + weight_fake_quant_config: "MXFakeQuantizeConfig" = field( + default_factory=weight_config + ) + activation_fake_quant_config: "MXFakeQuantizeConfig" = field( + default_factory=activation_config + ) def __init__(self, config: Config): super().__init__(config) - from torchao.prototype.qat import MXFakeQuantizeConfig - - self._weight_fake_quant_config = MXFakeQuantizeConfig( - dtype=torch.float4_e2m1fn_x2, - block_size=config.weight_block_size, - ) - self._activation_fake_quant_config = MXFakeQuantizeConfig( - dtype=torch.float8_e4m3fn, - block_size=config.activation_block_size, - ) + self._weight_fake_quant_config = config.weight_fake_quant_config + self._activation_fake_quant_config = config.activation_fake_quant_config - def _grouped_mm(self, *, A, weight_EOI, offs): + def _grouped_mm( + self, *, A: torch.Tensor, weight_EOI: torch.Tensor, offs: torch.Tensor + ) -> torch.Tensor: from torchao.prototype.qat import mx_fake_quantized_grouped_mm return mx_fake_quantized_grouped_mm( diff --git a/torchtitan/quantization/mx_qat/linear.py b/torchtitan/quantization/mx_qat/linear.py new file mode 100644 index 00000000000..5afffd06292 --- /dev/null +++ b/torchtitan/quantization/mx_qat/linear.py @@ -0,0 +1,55 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +"""Weight-only MX QAT for ordinary Linear projections.""" + +from dataclasses import dataclass, field + +import torch +import torch.nn.functional as F + +from .experts import MXFakeQuantizeConfig, weight_config + +_linear_cache: dict[type, type] = {} + + +def _get_mx_qat_linear_cls(parent_cls: type) -> type: + if getattr(parent_cls, "_mx_qat", False): + return parent_cls + if parent_cls in _linear_cache: + return _linear_cache[parent_cls] + if parent_cls.forward is not torch.nn.Linear.forward: + raise ValueError( + f"MX QAT cannot replace custom forward of {parent_cls.__name__}" + ) + parent_config_cls = parent_cls.Config + + class MXQATLinear(parent_cls): + _mx_qat = True + + @dataclass(kw_only=True, slots=True) + class Config(parent_config_cls): + weight_fake_quant_config: "MXFakeQuantizeConfig" = field( + default_factory=weight_config + ) + + def __init__(self, config: Config): + super().__init__(config) + self._weight_fake_quant_config = config.weight_fake_quant_config + + def forward(self, input: torch.Tensor) -> torch.Tensor: + from torchao.prototype.qat import mx_fake_quantize + + return F.linear( + input, + mx_fake_quantize(self.weight, self._weight_fake_quant_config), + self.bias, + ) + + MXQATLinear.__name__ = f"MXQAT{parent_cls.__name__}" + MXQATLinear.__qualname__ = MXQATLinear.__name__ + _linear_cache[parent_cls] = MXQATLinear + return MXQATLinear From a8fd33d452f8d26f6c901c1f515d5e593fbe9852 Mon Sep 17 00:00:00 2001 From: Elfie Guo Date: Mon, 21 Sep 2026 22:33:02 -0700 Subject: [PATCH 3/5] Keep MX QAT recipes on the standard model-running path --- tests/unit_tests/cpu/test_mx_qat_recipe.py | 82 +++++++++++++++++++ tests/unit_tests/cpu/test_mx_qat_transform.py | 15 ++++ torchtitan/config/transform/mx_qat.py | 9 ++ torchtitan/models/kimi_k3/README.md | 66 +++++++++++++++ torchtitan/models/kimi_k3/config_registry.py | 9 ++ 5 files changed, 181 insertions(+) create mode 100644 tests/unit_tests/cpu/test_mx_qat_recipe.py diff --git a/tests/unit_tests/cpu/test_mx_qat_recipe.py b/tests/unit_tests/cpu/test_mx_qat_recipe.py new file mode 100644 index 00000000000..c98bc8d9e7b --- /dev/null +++ b/tests/unit_tests/cpu/test_mx_qat_recipe.py @@ -0,0 +1,82 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +import sys +import types +import unittest +from unittest.mock import patch + +import torch +from torchao.prototype.qat import MXFakeQuantizeConfig +from torchao.quantization.quantize_.common import KernelPreference +from torchtitan.config import ConfigManager +from torchtitan.models.common.moe import GroupedExperts +from torchtitan.models.kimi_k3.config_registry import kimi_k3_debugmodel_mx_qat + + +class MXQATRecipeTest(unittest.TestCase): + def test_standard_model_entrypoint_and_training_override(self): + config = ConfigManager().parse_args( + [ + "--module", + "kimi_k3", + "--config", + "kimi_k3_debugmodel_mx_qat", + "--training.steps", + "1", + ] + ) + self.assertEqual(config.model_spec.name, "kimi_k3") + self.assertEqual(config.training.steps, 1) + self.assertFalse(config.checkpointer.initial_load_in_hf_quantized) + self.assertTrue( + type(config.model_spec.model.layers[1].moe.routed_up)._owner._mx_qat + ) + + def test_custom_recipe_uses_existing_torchao_configs(self): + def qat(): + return kimi_k3_debugmodel_mx_qat( + seq_len=16, + checkpoint_path="/tmp/packed-kimi", + weight_fake_quant_config=MXFakeQuantizeConfig( + dtype=torch.float4_e2m1fn_x2, + kernel_preference=KernelPreference.AUTO, + ), + activation_fake_quant_config=MXFakeQuantizeConfig( + dtype=torch.float8_e4m3fn, + kernel_preference=KernelPreference.AUTO, + ), + ) + + module = types.ModuleType("my_kimi_runs") + module.qat = qat + with patch.dict(sys.modules, {module.__name__: module}): + config = ConfigManager().parse_args( + ["--module", module.__name__, "--config", "qat"] + ) + self.assertEqual(config.checkpointer.initial_load_path, "/tmp/packed-kimi") + self.assertTrue(config.checkpointer.initial_load_in_hf) + self.assertTrue(config.checkpointer.initial_load_in_hf_quantized) + experts = list(config.model_spec.model.traverse(GroupedExperts.Config)) + self.assertTrue(experts) + for _, expert, _, _ in experts: + self.assertEqual( + expert.weight_fake_quant_config.kernel_preference, KernelPreference.AUTO + ) + self.assertEqual( + expert.activation_fake_quant_config.kernel_preference, + KernelPreference.AUTO, + ) + self.assertEqual( + config.model_spec.model.layers[ + 1 + ].moe.routed_up.weight_fake_quant_config.kernel_preference, + KernelPreference.AUTO, + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/unit_tests/cpu/test_mx_qat_transform.py b/tests/unit_tests/cpu/test_mx_qat_transform.py index a8804751e4b..b4ee0aaf4c9 100644 --- a/tests/unit_tests/cpu/test_mx_qat_transform.py +++ b/tests/unit_tests/cpu/test_mx_qat_transform.py @@ -107,6 +107,21 @@ def test_gpt_oss_expert_layout_preserves_biases_and_config(self): {"mlp1_weight_EGD", "mlp1_bias_EG", "mlp2_weight_EDF", "mlp2_bias_ED"}, ) + def test_backend_mismatch_fails_before_model_config_changes(self): + config = _config() + transform = MXQATTransform() + transform.weight_fake_quant_config = replace( + transform.weight_fake_quant_config, kernel_preference=KernelPreference.AUTO + ) + with self.assertRaisesRegex(ValueError, "matching.*kernel_preference"): + transform.transform(config) + self.assertIs(type(config.experts), GroupedExperts.Config) + # Weight-only dense QAT has no activation backend to match. + transform.grouped_expert_fqns = () + transform.linear_fqns = ("projection",) + transform.transform(config) + self.assertTrue(type(config.projection)._owner._mx_qat) + def test_rejects_partial_group_and_unknown_weights(self): for weights in ({"experts.w1_EFD"}, {"missing.weight"}): with self.subTest(weights=weights), self.assertRaises(ValueError): diff --git a/torchtitan/config/transform/mx_qat.py b/torchtitan/config/transform/mx_qat.py index 48f084ebe4c..d07f8126a96 100644 --- a/torchtitan/config/transform/mx_qat.py +++ b/torchtitan/config/transform/mx_qat.py @@ -93,6 +93,15 @@ def transform(self, model: Module.Config) -> Module.Config: for fqn, config, parent, attr in model.traverse(config_type): if targets is not None and fqn not in targets: continue + if ( + config_type is GroupedExperts.Config + and self.activation_fake_quant_config.kernel_preference + != self.weight_fake_quant_config.kernel_preference + ): + raise ValueError( + "MX QAT grouped experts require matching activation and weight " + "kernel_preference. Set both TorchAO configs to the same preference." + ) replacement = factory(type(config)._owner) new_config = convert_config_type(config, replacement) deltas = {"weight_fake_quant_config": self.weight_fake_quant_config} diff --git a/torchtitan/models/kimi_k3/README.md b/torchtitan/models/kimi_k3/README.md index e3c802cfd52..f2662dcee2e 100644 --- a/torchtitan/models/kimi_k3/README.md +++ b/torchtitan/models/kimi_k3/README.md @@ -58,3 +58,69 @@ Test scripts: - `scripts/checkpoint_conversion/numerical_tests_kimi_k3.py` -- Hugging Face vs. TorchTitan comparison - `tests/unit_tests/gpu/test_kimi_k3.py` -- KDA and FSDP2 correctness + +## Running MX QAT + +MX QAT uses the same launcher, trainer, data configuration, and optimizer as the +ordinary Kimi debug recipe. Install a compatible TorchAO version that provides +`MXFakeQuantizeConfig` and `mx_fake_quantized_grouped_mm`, then select the QAT +recipe: + +```bash +NGPU=1 MODULE=kimi_k3 CONFIG=kimi_k3_debugmodel_mx_qat ./run_train.sh +``` + +The debug recipe starts from random initialization and uses `EMULATED` by +default. It selects the expert and dense projections required by the released +MXFP4 policy, keeps BF16 master parameters, and applies MXFP8 activation fake +quantization to grouped experts. No manual expert replacement or parameter-name +list is needed. Prepare the normal Kimi tokenizer and dataset dependencies as +for `kimi_k3_debugmodel`; QAT does not replace that data setup. + +For packed checkpoint initialization, write a normal Python run configuration. +Use a checkpoint matching the debug architecture, such as the fixture generated +by `scripts/checkpoint_conversion/create_kimi_k3_mxfp4_fixture.py`; the full +released checkpoint does not match the debug model. + +```python +# my_kimi_runs.py +from torchtitan.models.kimi_k3.config_registry import kimi_k3_debugmodel_mx_qat + + +def qat(): + config = kimi_k3_debugmodel_mx_qat( + checkpoint_path="/absolute/path/to/packed-debug-checkpoint", + ) + config.training.steps = 100 + # Set tokenizer, data, and parallelism here, just as for a normal run. + return config +``` + +```bash +NGPU=1 MODULE=my_kimi_runs CONFIG=qat ./run_train.sh +``` + +`checkpoint_path` sets the existing HF and quantized-load options together. The +loader validates the checkpoint policy against the selected QAT modules. + +For native grouped execution, pass the existing TorchAO configurations to the +same recipe function; both `kernel_preference` values must agree: + +```python +import torch +from torchao.prototype.qat import MXFakeQuantizeConfig +from torchao.quantization.quantize_.common import KernelPreference + +config = kimi_k3_debugmodel_mx_qat( + weight_fake_quant_config=MXFakeQuantizeConfig( + dtype=torch.float4_e2m1fn_x2, kernel_preference=KernelPreference.AUTO, + ), + activation_fake_quant_config=MXFakeQuantizeConfig( + dtype=torch.float8_e4m3fn, kernel_preference=KernelPreference.AUTO, + ), +) +``` + +`AUTO` requires supported SM100 CUDA kernels and currently at most 32 local +experts. `EMULATED` uses dequantized operands for GEMM. Backend configuration +changes belong in the recipe; the model launch command stays the same. diff --git a/torchtitan/models/kimi_k3/config_registry.py b/torchtitan/models/kimi_k3/config_registry.py index b4bfd9f8a1c..107f9eb9462 100644 --- a/torchtitan/models/kimi_k3/config_registry.py +++ b/torchtitan/models/kimi_k3/config_registry.py @@ -26,6 +26,7 @@ ) from torchtitan.observability.metrics import MetricsProcessor from torchtitan.quantization.mx_qat.checkpoint import MXFP4CheckpointPolicy +from torchtitan.quantization.mx_qat.experts import MXFakeQuantizeConfig from torchtitan.trainer import Trainer from . import KIMI_K3_SPECIAL_TOKENS, model_registry @@ -101,11 +102,15 @@ def kimi_k3_debugmodel_mx_qat( seq_len: int | None = DEFAULT_DEBUG_MODEL_SEQ_LEN, *, checkpoint_path: str | None = None, + weight_fake_quant_config: MXFakeQuantizeConfig | None = None, + activation_fake_quant_config: MXFakeQuantizeConfig | None = None, ) -> Trainer.Config: """Kimi QAT using the released policy and optional packed HF initialization. Pass an absolute checkpoint_path to load the packed debug fixture. Without it, the recipe uses random initialization and remains valid before overrides. + Optional TorchAO configs control fake quantization and kernel_preference; + model-specific parameter selection stays inside the recipe. """ config = kimi_k3_debugmodel(seq_len=seq_len) adapter = KimiK3StateDictAdapter(config.model_spec.model, hf_assets_path=None) @@ -113,6 +118,10 @@ def kimi_k3_debugmodel_mx_qat( policy = MXFP4CheckpointPolicy.from_config(MXFP4_QUANTIZATION_CONFIG, mapping) weights = {mapping[key] for key in policy.weight_fqns if mapping[key] is not None} transform = MXQATTransform.from_weight_fqns(config.model_spec.model, weights) + if weight_fake_quant_config is not None: + transform.weight_fake_quant_config = weight_fake_quant_config + if activation_fake_quant_config is not None: + transform.activation_fake_quant_config = activation_fake_quant_config config.checkpointer = CheckpointManager.Config( interval=5, initial_load_path=checkpoint_path, From 15b8a65975cac8ab39855ff3d554ae000a5061cc Mon Sep 17 00:00:00 2001 From: Elfie Guo Date: Tue, 22 Sep 2026 00:20:23 -0700 Subject: [PATCH 4/5] Normalize released Kimi checkpoint shapes before distributed loading --- .../create_kimi_k3_mxfp4_fixture.py | 9 +- .../validate_kimi_k3_mxfp4_checkpoint.py | 85 ++++++++ .../test_kimi_k3_mxfp4_checkpoint.py | 36 +++- .../unit_tests/cpu/test_kimi_k3_checkpoint.py | 184 ++++++++++++++++++ tests/unit_tests/cpu/test_mx_qat_recipe.py | 11 +- tests/unit_tests/cpu/test_mx_qat_transform.py | 35 ++++ .../unit_tests/cpu/test_packed_hf_storage.py | 108 ++++++++-- .../unit_tests/cpu/test_state_dict_adapter.py | 10 +- .../{packed_hf_storage.py => hf_storage.py} | 99 +++++++++- torchtitan/models/kimi_k3/README.md | 26 +++ torchtitan/models/kimi_k3/config_registry.py | 4 +- .../models/kimi_k3/state_dict_adapter.py | 64 ++++-- torchtitan/quantization/mx_qat/linear.py | 16 +- 13 files changed, 623 insertions(+), 64 deletions(-) create mode 100644 scripts/checkpoint_conversion/validate_kimi_k3_mxfp4_checkpoint.py create mode 100644 tests/unit_tests/cpu/test_kimi_k3_checkpoint.py rename torchtitan/components/checkpointer/{packed_hf_storage.py => hf_storage.py} (72%) diff --git a/scripts/checkpoint_conversion/create_kimi_k3_mxfp4_fixture.py b/scripts/checkpoint_conversion/create_kimi_k3_mxfp4_fixture.py index c1e6dfa1e01..7a8ce75777b 100755 --- a/scripts/checkpoint_conversion/create_kimi_k3_mxfp4_fixture.py +++ b/scripts/checkpoint_conversion/create_kimi_k3_mxfp4_fixture.py @@ -18,8 +18,9 @@ import torchao from safetensors.torch import save_file from torchao.prototype.mx_formats.mx_tensor import MXTensor -from torchtitan.models.kimi_k3 import KimiK3StateDictAdapter, model_registry +from torchtitan.models.kimi_k3 import model_registry from torchtitan.models.kimi_k3.quantization import MXFP4_QUANTIZATION_CONFIG +from torchtitan.models.kimi_k3.state_dict_adapter import KimiK3StateDictAdapter from torchtitan.quantization.mx_qat.checkpoint import MXFP4CheckpointPolicy _DEFAULT_MAX_SHARD_BYTES = 1 << 30 @@ -133,11 +134,11 @@ def create_fixture( output.mkdir(parents=True, exist_ok=True) torch.manual_seed(seed) - model_spec = model_registry("debugmodel", seq_len=128) - model = model_spec.model.build() + model_config = model_registry("debugmodel", seq_len=128) + model = model_config.build() model.init_states() model.to(dtype=torch.bfloat16) - adapter = KimiK3StateDictAdapter(model_spec.model, hf_assets_path=None) + adapter = KimiK3StateDictAdapter(model_config, hf_assets_path=None) hf_state_dict = adapter.to_hf(model.state_dict()) policy = MXFP4CheckpointPolicy.from_config( MXFP4_QUANTIZATION_CONFIG, adapter.hf_linear_weight_mapping() diff --git a/scripts/checkpoint_conversion/validate_kimi_k3_mxfp4_checkpoint.py b/scripts/checkpoint_conversion/validate_kimi_k3_mxfp4_checkpoint.py new file mode 100644 index 00000000000..5509caae36c --- /dev/null +++ b/scripts/checkpoint_conversion/validate_kimi_k3_mxfp4_checkpoint.py @@ -0,0 +1,85 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +"""Preflight every Kimi HF tensor shape without allocating model weights. + +Read safetensors headers through the same reader used by DCP. Packed tensors +are checked as logical weights, and known padding tails are read and validated. +This checks loading compatibility, not numerical accuracy or training behavior. +""" + +import argparse +import json +from pathlib import Path + +import torch +from torch.distributed.checkpoint.metadata import Metadata, TensorStorageMetadata +from torchtitan.models.kimi_k3 import kimi_k3_configs, model_registry +from torchtitan.models.kimi_k3.state_dict_adapter import KimiK3StateDictAdapter + + +def validate_shapes(metadata: Metadata, expected: dict[str, torch.Size]) -> None: + """Require every model tensor and reject unexplained checkpoint tensors.""" + actual = metadata.state_dict_metadata + # The adapter explicitly ignores these legacy HF buffers on import. + ignored = {name for name in actual if name.endswith("rotary_emb.inv_freq")} + missing = sorted(expected.keys() - actual.keys()) + unexpected = sorted(actual.keys() - expected.keys() - ignored) + mismatches = [] + for name in sorted(expected.keys() & actual.keys()): + tensor = actual[name] + if not isinstance(tensor, TensorStorageMetadata): + mismatches.append(f"{name}: expected a tensor") + elif tensor.size != expected[name]: + mismatches.append( + f"{name}: checkpoint {tuple(tensor.size)}, model {tuple(expected[name])}" + ) + if missing or unexpected or mismatches: + raise ValueError( + "Checkpoint shape preflight failed: " + f"missing={missing[:20]} (total {len(missing)}), " + f"unexpected={unexpected[:20]} (total {len(unexpected)}), " + f"mismatches={mismatches[:20]} (total {len(mismatches)})" + ) + + +def validate_checkpoint( + checkpoint: Path, *, model_flavor: str = "Kimi-K3", from_quantized: bool = True +) -> dict[str, int | str]: + config = model_registry(model_flavor, seq_len=128) + adapter = KimiK3StateDictAdapter(config, hf_assets_path=None) + with torch.device("meta"): + model = config.build() + expected = { + name: tensor.shape + for name, tensor in adapter.to_hf(model.state_dict()).items() + } + reader = adapter.get_hf_storage_reader(str(checkpoint), from_quantized) + metadata = reader.read_metadata() + validate_shapes(metadata, expected) + return {"model_flavor": model_flavor, "validated_tensors": len(expected)} + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--checkpoint", type=Path, required=True) + parser.add_argument("--model-flavor", choices=kimi_k3_configs, default="Kimi-K3") + parser.add_argument("--unquantized", action="store_true") + args = parser.parse_args() + print( + json.dumps( + validate_checkpoint( + args.checkpoint, + model_flavor=args.model_flavor, + from_quantized=not args.unquantized, + ), + indent=2, + ) + ) + + +if __name__ == "__main__": + main() diff --git a/tests/integration_tests/test_kimi_k3_mxfp4_checkpoint.py b/tests/integration_tests/test_kimi_k3_mxfp4_checkpoint.py index ea774b0796e..65a9e3bfc41 100644 --- a/tests/integration_tests/test_kimi_k3_mxfp4_checkpoint.py +++ b/tests/integration_tests/test_kimi_k3_mxfp4_checkpoint.py @@ -14,9 +14,13 @@ convert_hf_state_dict_to_mxfp4, write_sharded_checkpoint, ) +from scripts.checkpoint_conversion.validate_kimi_k3_mxfp4_checkpoint import ( + validate_shapes, +) from torchao.prototype.mx_formats.mx_tensor import MXTensor -from torchtitan.components.checkpointer.packed_hf_storage import ( - PackedPairHuggingFaceStorageReader, +from torchtitan.components.checkpointer.hf_storage import ( + HuggingFaceStorageReaderWithViews, + LogicalPrefixSpec, PackedPairSpec, ) from torchtitan.quantization.mx_qat.checkpoint import ( @@ -31,13 +35,20 @@ class KimiK3MXFP4CheckpointIntegrationTest(unittest.TestCase): def test_release_format_fixture_loads_across_shards(self) -> None: weight = torch.linspace(-6, 6, 128, dtype=torch.bfloat16).reshape(2, 64) dense = torch.arange(12, dtype=torch.bfloat16).reshape(3, 4) + a_log = torch.arange(96, dtype=torch.float32) + dt_bias = torch.arange(128, dtype=torch.bfloat16) expected = MXTensor.to_mx( weight, elem_dtype=torch.float4_e2m1fn_x2, block_size=32, ).dequantize(torch.bfloat16) converted, pair_count = convert_hf_state_dict_to_mxfp4( - {_HF_WEIGHT: weight, "dense.weight": dense}, + { + _HF_WEIGHT: weight, + "dense.weight": dense, + "A_log": torch.cat((a_log, torch.zeros(32))), + "dt_bias": dt_bias, + }, MXFP4CheckpointPolicy(frozenset({_HF_WEIGHT})), ) self.assertEqual(pair_count, 1) @@ -54,8 +65,10 @@ def test_release_format_fixture_loads_across_shards(self) -> None: destination = { _HF_WEIGHT: torch.empty_like(weight), "dense.weight": torch.empty_like(dense), + "A_log": torch.empty_like(a_log), + "dt_bias": torch.empty_like(dt_bias), } - reader = PackedPairHuggingFaceStorageReader( + reader = HuggingFaceStorageReaderWithViews( str(output), PackedPairSpec( packed_suffix=".weight_packed", @@ -67,13 +80,28 @@ def test_release_format_fixture_loads_across_shards(self) -> None: target_fqns=frozenset({_HF_WEIGHT}), decode=decode_mxfp4, ), + logical_prefixes={"A_log": LogicalPrefixSpec(96, 128)}, ) + metadata = reader.read_metadata() + shapes = {key: value.shape for key, value in destination.items()} + validate_shapes(metadata, shapes) + with self.assertRaisesRegex(ValueError, "dt_bias.*checkpoint.*model"): + validate_shapes(metadata, {**shapes, "dt_bias": torch.Size((129,))}) + with self.assertRaisesRegex(ValueError, "missing=.*absent"): + validate_shapes(metadata, {**shapes, "absent": torch.Size((1,))}) + with self.assertRaisesRegex(ValueError, "unexpected=.*dt_bias"): + validate_shapes( + metadata, + {key: shape for key, shape in shapes.items() if key != "dt_bias"}, + ) dcp.load(destination, storage_reader=reader) torch.testing.assert_close( destination[_HF_WEIGHT], expected, rtol=0, atol=0, equal_nan=True ) torch.testing.assert_close(destination["dense.weight"], dense, rtol=0, atol=0) + torch.testing.assert_close(destination["A_log"], a_log, rtol=0, atol=0) + torch.testing.assert_close(destination["dt_bias"], dt_bias, rtol=0, atol=0) self.assertFalse( any(key.endswith(("weight_packed", "weight_scale")) for key in destination) ) diff --git a/tests/unit_tests/cpu/test_kimi_k3_checkpoint.py b/tests/unit_tests/cpu/test_kimi_k3_checkpoint.py new file mode 100644 index 00000000000..5adcddc4158 --- /dev/null +++ b/tests/unit_tests/cpu/test_kimi_k3_checkpoint.py @@ -0,0 +1,184 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +import tempfile +import unittest +from pathlib import Path + +import torch +import torch.distributed as dist +import torch.distributed.checkpoint as dcp +import torch.multiprocessing as mp +from safetensors.torch import save_file +from scripts.checkpoint_conversion.validate_kimi_k3_mxfp4_checkpoint import ( + validate_checkpoint, +) +from torch.distributed.device_mesh import init_device_mesh +from torch.distributed.tensor import DTensor, Shard +from torchtitan.components.checkpointer.hf_storage import ( + HuggingFaceStorageReaderWithViews, + LogicalPrefixSpec, +) +from torchtitan.models.kimi_k3 import model_registry as kimi_k3_model_registry +from torchtitan.models.kimi_k3.state_dict_adapter import KimiK3StateDictAdapter + + +def _run_kimi_uneven_dt_bias_roundtrip(rank: int, rendezvous: str) -> None: + dist.init_process_group( + backend="gloo", + init_method=f"file://{rendezvous}", + rank=rank, + world_size=2, + ) + try: + config = kimi_k3_model_registry("debugmodel", seq_len=128) + delta_config = config.layers[1].delta_attention + assert delta_config is not None + delta_config.num_heads = 1 + adapter = KimiK3StateDictAdapter(config, hf_assets_path=None) + + mesh = init_device_mesh("cpu", (2,), mesh_dim_names=("dp_shard",)) + full_dt_bias = torch.arange(128, dtype=torch.bfloat16).reshape(1, 128) + local_dt_bias = full_dt_bias if rank == 0 else full_dt_bias[:0] + sharded_dt_bias = DTensor.from_local( + local_dt_bias, + mesh, + (Shard(0),), + run_check=False, + shape=torch.Size((1, 128)), + stride=(128, 1), + ) + + hf_state_dict = adapter.to_hf( + { + "layers.1.delta_attention.dt_bias": sharded_dt_bias, + "layers.1.attention_res_norm.weight": torch.ones(1), + "layers.1.attention_res_proj.weight": torch.ones(1), + } + ) + hf_dt_bias = hf_state_dict["language_model.model.layers.1.self_attn.dt_bias"] + assert isinstance(hf_dt_bias, DTensor) + assert hf_dt_bias.shape == torch.Size((128,)) + assert hf_dt_bias.placements == (Shard(0),) + + # Load a new bias and a padded A_log through the real two-rank planner. + checkpoint = Path(rendezvous).parent + loaded_bias = full_dt_bias + 100 + if rank == 0: + save_file( + { + "language_model.model.layers.1.self_attn.dt_bias": loaded_bias.flatten(), + "A_log": torch.cat((torch.arange(96).float(), torch.zeros(32))), + }, + checkpoint / "model.safetensors", + ) + dist.barrier() + a_log = DTensor.from_local( + torch.empty(48), + mesh, + (Shard(0),), + run_check=False, + shape=torch.Size((96,)), + stride=(1,), + ) + dcp.load( + { + "language_model.model.layers.1.self_attn.dt_bias": hf_dt_bias, + "A_log": a_log, + }, + storage_reader=HuggingFaceStorageReaderWithViews( + str(checkpoint), logical_prefixes={"A_log": LogicalPrefixSpec(96, 128)} + ), + ) + torch.testing.assert_close( + a_log.to_local(), + torch.arange(rank * 48, (rank + 1) * 48).float(), + rtol=0, + atol=0, + ) + + restored = adapter.from_hf(hf_state_dict)["layers.1.delta_attention.dt_bias"] + assert isinstance(restored, DTensor) + assert restored.shape == torch.Size((1, 128)) + assert restored.placements == (Shard(0),) + assert restored.device_mesh == mesh + assert restored.to_local().shape == local_dt_bias.shape + torch.testing.assert_close( + restored.to_local(), + loaded_bias if rank == 0 else loaded_bias[:0], + rtol=0, + atol=0, + ) + torch.testing.assert_close(restored.full_tensor(), loaded_bias, rtol=0, atol=0) + # Other placement layouts are rejected before any collective reshape. + invalid = DTensor.from_local( + full_dt_bias[:, :64], + mesh, + (Shard(1),), + run_check=False, + shape=torch.Size((1, 128)), + stride=(128, 1), + ) + try: + adapter._reshape_dt_bias(invalid, (-1,)) + except ValueError as error: + assert "supports only" in str(error) + else: + raise AssertionError("Expected unsupported placement rejection") + finally: + dist.destroy_process_group() + + +class KimiK3CheckpointTest(unittest.TestCase): + def test_preflight_checks_nonpacked_shapes_against_model(self): + key = "language_model.model.layers.1.self_attn.dt_bias" + with tempfile.TemporaryDirectory() as directory: + save_file({key: torch.zeros(1)}, Path(directory) / "model.safetensors") + with self.assertRaisesRegex( + ValueError, "mismatches=.*dt_bias.*checkpoint.*model" + ): + validate_checkpoint( + Path(directory), model_flavor="debugmodel", from_quantized=False + ) + + @unittest.skipUnless(dist.is_gloo_available(), "Requires Gloo") + def test_dt_bias_with_fewer_heads_than_ranks(self): + with tempfile.TemporaryDirectory() as directory: + mp.spawn( + _run_kimi_uneven_dt_bias_roundtrip, + args=(f"{directory}/rendezvous",), + nprocs=2, + join=True, + ) + + def test_unquantized_reader_normalizes_release_padding(self): + config = kimi_k3_model_registry("debugmodel", seq_len=128) + config.layers[1].delta_attention.num_heads = 96 + adapter = KimiK3StateDictAdapter(config, hf_assets_path=None) + key = "language_model.model.layers.1.self_attn.A_log" + prefix = torch.arange(96, dtype=torch.float32) + with tempfile.TemporaryDirectory() as directory: + save_file( + {key: torch.cat((prefix, torch.zeros(32)))}, + Path(directory) / "model.safetensors", + ) + reader = adapter.get_hf_storage_reader(directory, from_quantized=False) + self.assertIsInstance(reader, HuggingFaceStorageReaderWithViews) + self.assertIsNone(reader.spec) + destination = {key: torch.empty_like(prefix)} + dcp.load(destination, storage_reader=reader) + restored = adapter.from_hf(destination) + torch.testing.assert_close( + restored["layers.1.delta_attention.A_log"], prefix, rtol=0, atol=0 + ) + # Export the model's canonical vector, never the release's padded shape. + restored["layers.1.attention_res_norm.weight"] = torch.ones(1) + restored["layers.1.attention_res_proj.weight"] = torch.ones(1) + self.assertEqual(adapter.to_hf(restored)[key].shape, (96,)) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/unit_tests/cpu/test_mx_qat_recipe.py b/tests/unit_tests/cpu/test_mx_qat_recipe.py index c98bc8d9e7b..43a2dd401f4 100644 --- a/tests/unit_tests/cpu/test_mx_qat_recipe.py +++ b/tests/unit_tests/cpu/test_mx_qat_recipe.py @@ -15,6 +15,7 @@ from torchtitan.config import ConfigManager from torchtitan.models.common.moe import GroupedExperts from torchtitan.models.kimi_k3.config_registry import kimi_k3_debugmodel_mx_qat +from torchtitan.models.kimi_k3.model import KimiK3Model class MXQATRecipeTest(unittest.TestCase): @@ -29,12 +30,10 @@ def test_standard_model_entrypoint_and_training_override(self): "1", ] ) - self.assertEqual(config.model_spec.name, "kimi_k3") + self.assertIsInstance(config.model, KimiK3Model.Config) self.assertEqual(config.training.steps, 1) self.assertFalse(config.checkpointer.initial_load_in_hf_quantized) - self.assertTrue( - type(config.model_spec.model.layers[1].moe.routed_up)._owner._mx_qat - ) + self.assertTrue(type(config.model.layers[1].moe.routed_up)._owner._mx_qat) def test_custom_recipe_uses_existing_torchao_configs(self): def qat(): @@ -60,7 +59,7 @@ def qat(): self.assertEqual(config.checkpointer.initial_load_path, "/tmp/packed-kimi") self.assertTrue(config.checkpointer.initial_load_in_hf) self.assertTrue(config.checkpointer.initial_load_in_hf_quantized) - experts = list(config.model_spec.model.traverse(GroupedExperts.Config)) + experts = list(config.model.traverse(GroupedExperts.Config)) self.assertTrue(experts) for _, expert, _, _ in experts: self.assertEqual( @@ -71,7 +70,7 @@ def qat(): KernelPreference.AUTO, ) self.assertEqual( - config.model_spec.model.layers[ + config.model.layers[ 1 ].moe.routed_up.weight_fake_quant_config.kernel_preference, KernelPreference.AUTO, diff --git a/tests/unit_tests/cpu/test_mx_qat_transform.py b/tests/unit_tests/cpu/test_mx_qat_transform.py index b4ee0aaf4c9..4a106d6ffa2 100644 --- a/tests/unit_tests/cpu/test_mx_qat_transform.py +++ b/tests/unit_tests/cpu/test_mx_qat_transform.py @@ -30,6 +30,41 @@ def _config(): class MXQATTransformTest(unittest.TestCase): + def test_stacked_linear_preserves_forward_shape_and_gradients(self): + from torchao.prototype.qat import mx_fake_quantize + + config = _config() + config.projection = Linear.Config( + in_features=64, out_features=32, num_linears=2, bias=True + ) + transform = MXQATTransform(grouped_expert_fqns=(), linear_fqns=("projection",)) + transform.transform(config) + module = config.projection.build() + x = torch.randn(3, 64, requires_grad=True) + expected_weight = module.weight.detach().clone().requires_grad_() + expected_bias = module.bias.detach().clone().requires_grad_() + expected_x = x.detach().clone().requires_grad_() + expected = torch.nn.functional.linear( + expected_x, + mx_fake_quantize( + expected_weight, transform.weight_fake_quant_config + ).flatten(0, 1), + expected_bias.flatten(), + ).unflatten(-1, (2, 32)) + actual = module(x) + self.assertEqual(actual.shape, (3, 2, 32)) + self.assertIs(type(module).forward, Linear.forward) + torch.testing.assert_close(actual, expected) + grad = torch.randn_like(actual) + actual.backward(grad) + expected.backward(grad) + for actual_grad, expected_grad in ( + (module.weight.grad, expected_weight.grad), + (module.bias.grad, expected_bias.grad), + (x.grad, expected_x.grad), + ): + torch.testing.assert_close(actual_grad, expected_grad) + def test_configuration_types_are_resolvable_for_cli(self): from typing import get_type_hints diff --git a/tests/unit_tests/cpu/test_packed_hf_storage.py b/tests/unit_tests/cpu/test_packed_hf_storage.py index 46c84bf1a71..83dc5ea50c5 100644 --- a/tests/unit_tests/cpu/test_packed_hf_storage.py +++ b/tests/unit_tests/cpu/test_packed_hf_storage.py @@ -6,6 +6,7 @@ import tempfile import unittest +from dataclasses import replace from pathlib import Path import torch @@ -14,8 +15,9 @@ from safetensors.torch import save_file from torch.distributed.checkpoint.metadata import MetadataIndex from torch.distributed.checkpoint.planner import LoadItemType, ReadItem -from torchtitan.components.checkpointer.packed_hf_storage import ( - PackedPairHuggingFaceStorageReader, +from torchtitan.components.checkpointer.hf_storage import ( + HuggingFaceStorageReaderWithViews, + LogicalPrefixSpec, PackedPairSpec, ) from torchtitan.quantization.mx_qat.checkpoint import decode_mxfp4 @@ -106,7 +108,7 @@ def get_slice(self, key: str) -> _RecordingSlice: return _RecordingSlice(self.handle.get_slice(key), self.calls, key) # type: ignore[attr-defined] -class PackedPairHuggingFaceStorageReaderMetadataTest(unittest.TestCase): +class HuggingFaceStorageReaderWithViewsMetadataTest(unittest.TestCase): def _write_checkpoint(self, tensors: dict[str, torch.Tensor]) -> str: directory = tempfile.TemporaryDirectory() self.addCleanup(directory.cleanup) @@ -122,7 +124,7 @@ def test_read_metadata_exposes_virtual_weight(self) -> None: } ) - metadata = PackedPairHuggingFaceStorageReader(path, _spec()).read_metadata() + metadata = HuggingFaceStorageReaderWithViews(path, _spec()).read_metadata() self.assertEqual( set(metadata.state_dict_metadata), {_VIRTUAL_KEY, "dense.weight"} @@ -140,7 +142,7 @@ def test_read_metadata_rejects_dense_substitution_for_expected_pair(self) -> Non {_VIRTUAL_KEY: torch.ones((2, 64), dtype=torch.bfloat16)} ) with self.assertRaisesRegex(ValueError, "requires missing pairs"): - PackedPairHuggingFaceStorageReader(path, _spec()).read_metadata() + HuggingFaceStorageReaderWithViews(path, _spec()).read_metadata() def test_read_metadata_rejects_missing_scale(self) -> None: path = self._write_checkpoint( @@ -148,7 +150,7 @@ def test_read_metadata_rejects_missing_scale(self) -> None: ) with self.assertRaisesRegex(ValueError, "missing scale tensor"): - PackedPairHuggingFaceStorageReader(path, _spec()).read_metadata() + HuggingFaceStorageReaderWithViews(path, _spec()).read_metadata() def test_read_metadata_rejects_orphan_scale(self) -> None: path = self._write_checkpoint( @@ -156,7 +158,7 @@ def test_read_metadata_rejects_orphan_scale(self) -> None: ) with self.assertRaisesRegex(ValueError, "orphan scale tensor"): - PackedPairHuggingFaceStorageReader(path, _spec()).read_metadata() + HuggingFaceStorageReaderWithViews(path, _spec()).read_metadata() def test_read_metadata_rejects_non_uint8_payload(self) -> None: path = self._write_checkpoint( @@ -167,7 +169,7 @@ def test_read_metadata_rejects_non_uint8_payload(self) -> None: ) with self.assertRaisesRegex(ValueError, "must use torch.uint8"): - PackedPairHuggingFaceStorageReader(path, _spec()).read_metadata() + HuggingFaceStorageReaderWithViews(path, _spec()).read_metadata() def test_read_metadata_rejects_incompatible_shapes(self) -> None: path = self._write_checkpoint( @@ -178,7 +180,7 @@ def test_read_metadata_rejects_incompatible_shapes(self) -> None: ) with self.assertRaisesRegex(ValueError, "incompatible packed and scale shapes"): - PackedPairHuggingFaceStorageReader(path, _spec()).read_metadata() + HuggingFaceStorageReaderWithViews(path, _spec()).read_metadata() def test_read_metadata_rejects_pair_outside_target_policy(self) -> None: path = self._write_checkpoint( @@ -193,10 +195,10 @@ def test_read_metadata_rejects_pair_outside_target_policy(self) -> None: ) with self.assertRaisesRegex(ValueError, "outside the packed-weight policy"): - PackedPairHuggingFaceStorageReader(path, _spec()).read_metadata() + HuggingFaceStorageReaderWithViews(path, _spec()).read_metadata() -class PackedPairHuggingFaceStorageReaderReadTest(unittest.TestCase): +class HuggingFaceStorageReaderWithViewsReadTest(unittest.TestCase): def setUp(self) -> None: self.directory = tempfile.TemporaryDirectory() packed, scales, self.expected = _packed_fixture() @@ -213,7 +215,7 @@ def test_dcp_load_dequantizes_full_tensor(self) -> None: dcp.load( {_VIRTUAL_KEY: destination}, - storage_reader=PackedPairHuggingFaceStorageReader( + storage_reader=HuggingFaceStorageReaderWithViews( self.directory.name, _spec() ), ) @@ -223,7 +225,7 @@ def test_dcp_load_dequantizes_full_tensor(self) -> None: ) def test_unaligned_read_uses_only_intersecting_groups(self) -> None: - reader = PackedPairHuggingFaceStorageReader(self.directory.name, _spec()) + reader = HuggingFaceStorageReaderWithViews(self.directory.name, _spec()) reader.read_metadata() destination = torch.empty((2, 38), dtype=torch.bfloat16) planner = _Planner(destination) @@ -255,7 +257,7 @@ def test_unaligned_read_uses_only_intersecting_groups(self) -> None: ) def test_aligned_read_dequantizes_second_group(self) -> None: - reader = PackedPairHuggingFaceStorageReader(self.directory.name, _spec()) + reader = HuggingFaceStorageReaderWithViews(self.directory.name, _spec()) reader.read_metadata() destination = torch.empty((1, 32), dtype=torch.bfloat16) planner = _Planner(destination) @@ -288,7 +290,7 @@ def test_e8m0_extreme_bytes_follow_mx_semantics(self) -> None: destination = torch.empty((1, 96), dtype=torch.bfloat16) dcp.load( {_VIRTUAL_KEY: destination}, - storage_reader=PackedPairHuggingFaceStorageReader(directory, _spec()), + storage_reader=HuggingFaceStorageReaderWithViews(directory, _spec()), ) expected = torch.cat( @@ -313,5 +315,81 @@ def test_e8m0_extreme_bytes_follow_mx_semantics(self) -> None: ) +class LogicalPrefixTest(unittest.TestCase): + def setUp(self): + self.directory = tempfile.TemporaryDirectory() + self.addCleanup(self.directory.cleanup) + self.path = Path(self.directory.name) / "model.safetensors" + self.prefix = torch.arange(96, dtype=torch.float32) + + def reader(self): + return HuggingFaceStorageReaderWithViews( + self.directory.name, + logical_prefixes={"A_log": LogicalPrefixSpec(96, 128)}, + ) + + def test_canonical_and_padded_vectors_load_through_dcp(self): + for padding in (0, 32): + with self.subTest(padding=padding): + save_file( + {"A_log": torch.cat((self.prefix, torch.zeros(padding)))}, self.path + ) + reader = self.reader() + metadata = reader.read_metadata() + tensor = metadata.state_dict_metadata["A_log"] + self.assertEqual(tensor.size, (96,)) + self.assertEqual(tensor.chunks[0].sizes, (96,)) + self.assertEqual( + next(iter(metadata.storage_data.values())).shape, (96,) + ) + destination = {"A_log": torch.empty(96)} + dcp.load(destination, storage_reader=reader) + torch.testing.assert_close( + destination["A_log"], self.prefix, rtol=0, atol=0 + ) + + def test_rejects_nonzero_or_nan_tail(self): + for invalid in (1e-30, float("nan"), float("inf")): + with self.subTest(invalid=invalid): + tail = torch.zeros(32) + tail[-1] = invalid + save_file({"A_log": torch.cat((self.prefix, tail))}, self.path) + with self.assertRaisesRegex(ValueError, "padding must be exactly zero"): + self.reader().read_metadata() + + def test_rejects_unknown_physical_shape_or_missing_tensor(self): + for shape in ((95,), (97,), (129,), (1, 128)): + with self.subTest(shape=shape): + save_file({"A_log": torch.zeros(shape)}, self.path) + with self.assertRaisesRegex(ValueError, "expected physical shape"): + self.reader().read_metadata() + save_file({"other": torch.zeros(128)}, self.path) + with self.assertRaisesRegex(ValueError, "requires tensor"): + self.reader().read_metadata() + + def test_read_at_logical_boundary_never_reads_padding(self): + save_file({"A_log": torch.cat((self.prefix, torch.zeros(32)))}, self.path) + reader = self.reader() + reader.read_metadata() + destination = torch.empty(7) + planner = _Planner(destination) + request = ReadItem( + type=LoadItemType.TENSOR, + dest_index=MetadataIndex("A_log", [0]), + dest_offsets=torch.Size((0,)), + storage_index=MetadataIndex("A_log", [0]), + storage_offsets=torch.Size((89,)), + lengths=torch.Size((7,)), + ) + with safe_open(self.path, framework="pt") as handle: + recording = _RecordingFile(handle) + reader._process_read_request(recording, request, planner) + self.assertEqual(recording.calls, [("A_log", (slice(89, 96),))]) + request = replace(request, lengths=torch.Size((8,))) + with self.assertRaisesRegex(ValueError, "exceeds logical prefix"): + reader._process_read_request(recording, request, planner) + torch.testing.assert_close(destination, self.prefix[89:], rtol=0, atol=0) + + if __name__ == "__main__": unittest.main() diff --git a/tests/unit_tests/cpu/test_state_dict_adapter.py b/tests/unit_tests/cpu/test_state_dict_adapter.py index 80e605e65a0..21cbbfe7caa 100644 --- a/tests/unit_tests/cpu/test_state_dict_adapter.py +++ b/tests/unit_tests/cpu/test_state_dict_adapter.py @@ -18,8 +18,8 @@ from torch.testing._internal.distributed.fake_pg import FakeStore from torchtitan.components.checkpointer.base import ModelWrapper -from torchtitan.components.checkpointer.packed_hf_storage import ( - PackedPairHuggingFaceStorageReader, +from torchtitan.components.checkpointer.hf_storage import ( + HuggingFaceStorageReaderWithViews, ) from torchtitan.models.deepseek_v3 import deepseekv3_configs from torchtitan.models.deepseek_v3.state_dict_adapter import DeepSeekV3StateDictAdapter @@ -126,9 +126,9 @@ def test_hf_roundtrip_preserves_tied_embedding_shape(self) -> None: class KimiK3StateDictAdapterTest(unittest.TestCase): def setUp(self) -> None: - model_spec = kimi_k3_model_registry("debugmodel", seq_len=128) + model_config = kimi_k3_model_registry("debugmodel", seq_len=128) self.adapter = KimiK3StateDictAdapter( - model_spec.model, + model_config, hf_assets_path=None, ) @@ -148,7 +148,7 @@ def test_quantized_load_uses_packed_pair_reader(self) -> None: from_quantized=True, ) - self.assertIsInstance(reader, PackedPairHuggingFaceStorageReader) + self.assertIsInstance(reader, HuggingFaceStorageReaderWithViews) self.assertEqual(reader.spec.packed_suffix, ".weight_packed") self.assertEqual(reader.spec.scale_suffix, ".weight_scale") self.assertEqual(reader.spec.virtual_suffix, ".weight") diff --git a/torchtitan/components/checkpointer/packed_hf_storage.py b/torchtitan/components/checkpointer/hf_storage.py similarity index 72% rename from torchtitan/components/checkpointer/packed_hf_storage.py rename to torchtitan/components/checkpointer/hf_storage.py index 96d437f4cb1..22758456c5b 100644 --- a/torchtitan/components/checkpointer/packed_hf_storage.py +++ b/torchtitan/components/checkpointer/hf_storage.py @@ -4,7 +4,7 @@ # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. -"""Hugging Face storage reader for paired packed weights and block scales.""" +"""Expose packed weights and checked tensor prefixes as logical DCP tensors.""" import dataclasses from collections.abc import Callable @@ -20,7 +20,23 @@ ) from torch.distributed.checkpoint.planner import LoadPlanner, ReadItem -__all__ = ["PackedPairHuggingFaceStorageReader", "PackedPairSpec"] +__all__ = ["HuggingFaceStorageReaderWithViews", "LogicalPrefixSpec", "PackedPairSpec"] + + +@dataclass(frozen=True) +class LogicalPrefixSpec: + """Accept a canonical vector or a known zero-padded physical vector. + + Only the explicitly declared padded length may be truncated. Model adapters + own these lengths and the FQNs to which the rule applies. + """ + + logical_length: int + padded_length: int + + def __post_init__(self) -> None: + if not 0 < self.logical_length < self.padded_length: + raise ValueError("Expected 0 < logical_length < padded_length") @dataclass(frozen=True) @@ -59,19 +75,77 @@ class _PackedPair: scale_path: str -class PackedPairHuggingFaceStorageReader(HuggingFaceStorageReader): - """Present selected uint8 packed/scale pairs as unquantized DCP tensors.""" +class HuggingFaceStorageReaderWithViews(HuggingFaceStorageReader): + """Apply optional packed decoding and checked prefix views before planning.""" def __init__( self, path: str, - spec: PackedPairSpec, + spec: PackedPairSpec | None = None, thread_count: int = 1, + *, + logical_prefixes: dict[str, LogicalPrefixSpec] | None = None, ) -> None: super().__init__(path=path, thread_count=thread_count) self.spec = spec + self.logical_prefixes = dict(logical_prefixes or {}) self._pairs: dict[str, _PackedPair] = {} + def _apply_logical_prefixes(self, metadata: Any) -> Any: + if not self.logical_prefixes: + return metadata + from safetensors import safe_open + + tensors = dict(metadata.state_dict_metadata) + storage = dict(metadata.storage_data) + prefix_storage: dict[str, list[tuple[Any, Any]]] = { + fqn: [] for fqn in self.logical_prefixes + } + for index, info in storage.items(): + if index.fqn in prefix_storage: + prefix_storage[index.fqn].append((index, info)) + for fqn, spec in self.logical_prefixes.items(): + tensor = tensors.get(fqn) + if not isinstance(tensor, TensorStorageMetadata): + raise ValueError(f"Logical prefix requires tensor {fqn!r}") + if tuple(tensor.size) not in ( + (spec.logical_length,), + (spec.padded_length,), + ): + raise ValueError( + f"{fqn}: expected physical shape [{spec.logical_length}] or " + f"[{spec.padded_length}], got {tuple(tensor.size)}" + ) + if not tensor.properties.dtype.is_floating_point: + raise ValueError(f"{fqn}: logical prefix requires floating-point data") + if tuple(tensor.size) == (spec.logical_length,): + continue + # Prefix normalization handles a whole source tensor, not an + # independently sharded HF export. DCP destination shards are fine. + entries = prefix_storage[fqn] + if ( + len(entries) != 1 + or len(tensor.chunks) != 1 + or tuple(tensor.chunks[0].offsets) != (0,) + or tuple(tensor.chunks[0].sizes) != (spec.padded_length,) + ): + raise ValueError(f"{fqn}: padded source must be one complete vector") + index, info = entries[0] + with safe_open(info.relative_path, framework="pt") as handle: + tail = handle.get_slice(fqn)[spec.logical_length : spec.padded_length] + if not bool(torch.all(tail == 0)): + raise ValueError(f"{fqn}: discarded padding must be exactly zero") + logical_shape = torch.Size((spec.logical_length,)) + tensors[fqn] = dataclasses.replace( + tensor, + size=logical_shape, + chunks=[dataclasses.replace(tensor.chunks[0], sizes=logical_shape)], + ) + storage[index] = dataclasses.replace(info, shape=logical_shape) + return dataclasses.replace( + metadata, state_dict_metadata=tensors, storage_data=storage + ) + def _replace_suffix(self, fqn: str, source: str, destination: str) -> str: if not fqn.endswith(source): raise ValueError(f"{fqn!r} does not end with {source!r}") @@ -83,6 +157,7 @@ def _validate_pair( packed_metadata: TensorStorageMetadata, scale_metadata: TensorStorageMetadata, ) -> None: + assert self.spec is not None if virtual_fqn not in self.spec.target_fqns: raise ValueError( f"Packed tensor {virtual_fqn!r} is outside the packed-weight policy." @@ -111,6 +186,7 @@ def _validate_pair( ) def _virtual_chunk(self, chunk: ChunkStorageMetadata) -> ChunkStorageMetadata: + assert self.spec is not None offsets = list(chunk.offsets) sizes = list(chunk.sizes) offsets[-1] *= self.spec.packed_values_per_byte @@ -124,6 +200,9 @@ def _virtual_chunk(self, chunk: ChunkStorageMetadata) -> ChunkStorageMetadata: # pyrefly: ignore [bad-override] def read_metadata(self) -> Any: metadata = super().read_metadata() + metadata = self._apply_logical_prefixes(metadata) + if self.spec is None: + return metadata state_dict_metadata = metadata.state_dict_metadata storage_paths: dict[str, str] = {} for index, storage_info in metadata.storage_data.items(): @@ -250,8 +329,18 @@ def _process_read_request( virtual_fqn = req.storage_index.fqn pair = self._pairs.get(virtual_fqn) if pair is None: + prefix = self.logical_prefixes.get(virtual_fqn) + if prefix is not None and ( + len(req.storage_offsets) != 1 + or len(req.lengths) != 1 + or req.storage_offsets[0] < 0 + or req.lengths[0] < 0 + or req.storage_offsets[0] + req.lengths[0] > prefix.logical_length + ): + raise ValueError(f"{virtual_fqn}: read exceeds logical prefix") super()._process_read_request(f, req, planner) return + assert self.spec is not None if len(req.storage_offsets) != 2 or len(req.lengths) != 2: raise ValueError( f"Packed tensor {virtual_fqn!r} requires a two-dimensional read; " diff --git a/torchtitan/models/kimi_k3/README.md b/torchtitan/models/kimi_k3/README.md index f2662dcee2e..eecc0c2cd55 100644 --- a/torchtitan/models/kimi_k3/README.md +++ b/torchtitan/models/kimi_k3/README.md @@ -124,3 +124,29 @@ config = kimi_k3_debugmodel_mx_qat( `AUTO` requires supported SM100 CUDA kernels and currently at most 32 local experts. `EMULATED` uses dequantized operands for GEMM. Backend configuration changes belong in the recipe; the model launch command stays the same. + +## Released checkpoint preflight + +Before allocating a training job, check the safetensors headers against the +model's HF tensor shapes, including ordinary non-quantized tensors: + +```bash +python -m scripts.checkpoint_conversion.validate_kimi_k3_mxfp4_checkpoint \ + --checkpoint /absolute/path/to/Kimi-K3 --model-flavor Kimi-K3 +``` + +The preflight constructs the model on the meta device and uses the same storage +reader as training. It validates packed/scale pairs and every expected logical +tensor shape. It reads only headers and the small padding tails, not full model +weights. Use `--model-flavor debugmodel` for a matching synthetic fixture, or +`--unquantized` for ordinary HF weights. + +The released 96-head KDA checkpoint stores `A_log` as a 128-element vector. +Loading accepts either the canonical 96-element vector or that exact padded +shape, and requires all 32 discarded values to be zero. The reader exposes 96 +elements before DCP planning; model parameters and HF export remain canonical. +This rule also applies to unquantized HF imports. Other unexpected shapes or +nonzero padding are errors. See the [release inspection](https://huggingface.co/moonshotai/Kimi-K3/discussions/150). + +Shape preflight does not establish numerical equivalence or successful +released-model training. Those require separate evaluation runs. diff --git a/torchtitan/models/kimi_k3/config_registry.py b/torchtitan/models/kimi_k3/config_registry.py index 107f9eb9462..8bdc4d2c12b 100644 --- a/torchtitan/models/kimi_k3/config_registry.py +++ b/torchtitan/models/kimi_k3/config_registry.py @@ -113,11 +113,11 @@ def kimi_k3_debugmodel_mx_qat( model-specific parameter selection stays inside the recipe. """ config = kimi_k3_debugmodel(seq_len=seq_len) - adapter = KimiK3StateDictAdapter(config.model_spec.model, hf_assets_path=None) + adapter = KimiK3StateDictAdapter(config.model, hf_assets_path=None) mapping = adapter.hf_linear_weight_mapping() policy = MXFP4CheckpointPolicy.from_config(MXFP4_QUANTIZATION_CONFIG, mapping) weights = {mapping[key] for key in policy.weight_fqns if mapping[key] is not None} - transform = MXQATTransform.from_weight_fqns(config.model_spec.model, weights) + transform = MXQATTransform.from_weight_fqns(config.model, weights) if weight_fake_quant_config is not None: transform.weight_fake_quant_config = weight_fake_quant_config if activation_fake_quant_config is not None: diff --git a/torchtitan/models/kimi_k3/state_dict_adapter.py b/torchtitan/models/kimi_k3/state_dict_adapter.py index 9dba0700ec7..387737588e8 100644 --- a/torchtitan/models/kimi_k3/state_dict_adapter.py +++ b/torchtitan/models/kimi_k3/state_dict_adapter.py @@ -15,15 +15,17 @@ import torch from torch.distributed.checkpoint import HuggingFaceStorageReader -from torch.distributed.tensor import DTensor +from torch.distributed.tensor import DTensor, Replicate, Shard -from torchtitan.components.checkpointer.packed_hf_storage import ( - PackedPairHuggingFaceStorageReader, +from torchtitan.components.checkpointer.hf_storage import ( + HuggingFaceStorageReaderWithViews, + LogicalPrefixSpec, PackedPairSpec, ) from torchtitan.models.common.linear import Linear, RouterGateLinear from torchtitan.models.common.moe import GroupedExperts from torchtitan.models.utils import MoEStateDictAdapter +from torchtitan.protocols.state_dict_adapter import dtensor_safe from torchtitan.quantization.mx_qat.checkpoint import ( decode_mxfp4, MXFP4CheckpointPolicy, @@ -266,19 +268,48 @@ def _validate_qat_policy(self, policy: MXFP4CheckpointPolicy) -> None: f"missing={sorted(expected - selected)}, unexpected={sorted(selected - expected)}" ) + def _reshape_dt_bias( + self, value: torch.Tensor, shape: tuple[int, ...] + ) -> torch.Tensor: + """Preserve leading-axis FSDP shards, including ranks with no heads.""" + if isinstance(value, DTensor) and any( + not isinstance(p, Replicate) and not (type(p) is Shard and p.dim == 0) + for p in value.placements + ): + raise ValueError("KDA dt_bias reshape supports only Replicate and Shard(0)") + return self._reshape_dt_bias_replicated(value, shape) + + @dtensor_safe + def _reshape_dt_bias_replicated( + self, value: torch.Tensor, shape: tuple[int, ...] + ) -> torch.Tensor: + # Reuse the adapter's gather/restore helper only for this small bias. + return value.reshape(shape) + def get_hf_storage_reader( self, path: str, from_quantized: bool = False, ) -> HuggingFaceStorageReader: - if not from_quantized: + # The released 96-head KDA stores A_log in a 128-element vector. + # Other architectures must not inherit this checkpoint-specific rule. + prefixes = { + f"language_model.model.layers.{index}.self_attn.A_log": LogicalPrefixSpec( + logical_length=layer.delta_attention.num_heads, + padded_length=128, + ) + for index, layer in enumerate(self.kimi_config.layers) + if layer.delta_attention is not None + and layer.delta_attention.num_heads == 96 + and layer.delta_attention.head_dim == 128 + } + if not from_quantized and not prefixes: return super().get_hf_storage_reader(path, from_quantized=False) - policy = self.mxfp4_policy(path) - self._validate_qat_policy(policy) - return PackedPairHuggingFaceStorageReader( - path=path, - thread_count=4, - spec=PackedPairSpec( + spec = None + if from_quantized: + policy = self.mxfp4_policy(path) + self._validate_qat_policy(policy) + spec = PackedPairSpec( packed_suffix=".weight_packed", scale_suffix=".weight_scale", virtual_suffix=".weight", @@ -287,7 +318,11 @@ def get_hf_storage_reader( target_dtype=torch.bfloat16, target_fqns=policy.weight_fqns, decode=decode_mxfp4, - ), + ) + return HuggingFaceStorageReaderWithViews( + path=path, + spec=spec, + logical_prefixes=prefixes, ) def to_hf(self, state_dict: dict[str, Any]) -> dict[str, Any]: @@ -367,7 +402,7 @@ def to_hf(self, state_dict: dict[str, Any]) -> dict[str, Any]: unmapped.append(key) continue if abstract_key == "layers.{}.delta_attention.dt_bias": - value = value.reshape(-1) + value = self._reshape_dt_bias(value, (-1,)) hf_state_dict[hf_abstract_key.format(layer_num)] = value continue @@ -515,9 +550,8 @@ def from_hf(self, hf_state_dict: dict[str, Any]) -> dict[str, Any]: ].delta_attention if delta_config is None: raise ValueError(f"HF key '{key}' targets a non-KDA layer.") - value = value.reshape( - delta_config.num_heads, - delta_config.head_dim, + value = self._reshape_dt_bias( + value, (delta_config.num_heads, delta_config.head_dim) ) state_dict[new_abstract_key.format(layer_num)] = value continue diff --git a/torchtitan/quantization/mx_qat/linear.py b/torchtitan/quantization/mx_qat/linear.py index 5afffd06292..f1957c98115 100644 --- a/torchtitan/quantization/mx_qat/linear.py +++ b/torchtitan/quantization/mx_qat/linear.py @@ -9,7 +9,8 @@ from dataclasses import dataclass, field import torch -import torch.nn.functional as F + +from torchtitan.models.common.linear import Linear from .experts import MXFakeQuantizeConfig, weight_config @@ -21,7 +22,7 @@ def _get_mx_qat_linear_cls(parent_cls: type) -> type: return parent_cls if parent_cls in _linear_cache: return _linear_cache[parent_cls] - if parent_cls.forward is not torch.nn.Linear.forward: + if parent_cls.forward is not Linear.forward: raise ValueError( f"MX QAT cannot replace custom forward of {parent_cls.__name__}" ) @@ -40,14 +41,13 @@ def __init__(self, config: Config): super().__init__(config) self._weight_fake_quant_config = config.weight_fake_quant_config - def forward(self, input: torch.Tensor) -> torch.Tensor: + def _flatten_weight_and_bias( + self, + ) -> tuple[torch.Tensor, torch.Tensor | None]: from torchao.prototype.qat import mx_fake_quantize - return F.linear( - input, - mx_fake_quantize(self.weight, self._weight_fake_quant_config), - self.bias, - ) + weight, bias = super()._flatten_weight_and_bias() + return mx_fake_quantize(weight, self._weight_fake_quant_config), bias MXQATLinear.__name__ = f"MXQAT{parent_cls.__name__}" MXQATLinear.__qualname__ = MXQATLinear.__name__ From a61b4172c733810066e74becd6d70da1d8d34bc6 Mon Sep 17 00:00:00 2001 From: Elfie Guo Date: Wed, 23 Sep 2026 22:26:18 -0700 Subject: [PATCH 5/5] Derive Kimi MX QAT selection from actual checkpoint packed pairs --- .../test_kimi_k3_mxfp4_checkpoint.py | 13 ++++- tests/unit_tests/cpu/mx_qat_test_utils.py | 35 ++++++++++++ tests/unit_tests/cpu/test_mx_qat_policy.py | 35 ++++++++++++ tests/unit_tests/cpu/test_mx_qat_recipe.py | 28 ++++++++-- .../unit_tests/cpu/test_state_dict_adapter.py | 55 +++++++++++++------ torchtitan/models/kimi_k3/README.md | 18 ++++-- torchtitan/models/kimi_k3/config_registry.py | 13 +++-- .../models/kimi_k3/state_dict_adapter.py | 34 ++++++++++-- torchtitan/quantization/mx_qat/checkpoint.py | 50 ++++++++++++++++- 9 files changed, 242 insertions(+), 39 deletions(-) create mode 100644 tests/unit_tests/cpu/mx_qat_test_utils.py diff --git a/tests/integration_tests/test_kimi_k3_mxfp4_checkpoint.py b/tests/integration_tests/test_kimi_k3_mxfp4_checkpoint.py index 65a9e3bfc41..d4ffdd6502b 100644 --- a/tests/integration_tests/test_kimi_k3_mxfp4_checkpoint.py +++ b/tests/integration_tests/test_kimi_k3_mxfp4_checkpoint.py @@ -4,6 +4,7 @@ # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. +import json import tempfile import unittest from pathlib import Path @@ -23,6 +24,7 @@ LogicalPrefixSpec, PackedPairSpec, ) +from torchtitan.models.kimi_k3.quantization import MXFP4_QUANTIZATION_CONFIG from torchtitan.quantization.mx_qat.checkpoint import ( decode_mxfp4, MXFP4CheckpointPolicy, @@ -62,6 +64,15 @@ def test_release_format_fixture_loads_across_shards(self) -> None: max_shard_bytes=64, ) self.assertGreater(manifest["shard_count"], 1) + # Both linears are eligible, but only the actual indexed pair is + # packed. This is the release's mixed MXFP4/BF16 storage contract. + index = json.loads((output / "model.safetensors.index.json").read_text()) + policy = MXFP4CheckpointPolicy.from_manifest( + MXFP4_QUANTIZATION_CONFIG, + [_HF_WEIGHT, "dense.weight"], + index["weight_map"], + ) + self.assertEqual(policy.weight_fqns, frozenset({_HF_WEIGHT})) destination = { _HF_WEIGHT: torch.empty_like(weight), "dense.weight": torch.empty_like(dense), @@ -77,7 +88,7 @@ def test_release_format_fixture_loads_across_shards(self) -> None: block_size=32, packed_values_per_byte=2, target_dtype=torch.bfloat16, - target_fqns=frozenset({_HF_WEIGHT}), + target_fqns=policy.weight_fqns, decode=decode_mxfp4, ), logical_prefixes={"A_log": LogicalPrefixSpec(96, 128)}, diff --git a/tests/unit_tests/cpu/mx_qat_test_utils.py b/tests/unit_tests/cpu/mx_qat_test_utils.py new file mode 100644 index 00000000000..99dd8645748 --- /dev/null +++ b/tests/unit_tests/cpu/mx_qat_test_utils.py @@ -0,0 +1,35 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +"""Small checkpoint metadata fixtures; tensor payloads are deliberately absent.""" + +import json +from pathlib import Path + +from torchtitan.models.common.moe import GroupedExperts +from torchtitan.models.kimi_k3.quantization import MXFP4_QUANTIZATION_CONFIG + + +def write_mixed_checkpoint_metadata(path: Path, adapter) -> dict[str, str]: + """Mirror release storage: packed expert matrices, ordinary dense weights.""" + groups = { + fqn for fqn, _, _, _ in adapter.kimi_config.traverse(GroupedExperts.Config) + } + weight_map = {} + for hf_name, target in adapter.hf_linear_weight_mapping().items(): + if target is not None and target.rsplit(".", 1)[0] in groups: + prefix = hf_name.removesuffix(".weight") + weight_map[prefix + ".weight_packed"] = "weights.safetensors" + weight_map[prefix + ".weight_scale"] = "scales.safetensors" + else: + weight_map[hf_name] = "dense.safetensors" + (path / "config.json").write_text( + json.dumps({"text_config": {"quantization_config": MXFP4_QUANTIZATION_CONFIG}}) + ) + (path / "model.safetensors.index.json").write_text( + json.dumps({"weight_map": weight_map}) + ) + return weight_map diff --git a/tests/unit_tests/cpu/test_mx_qat_policy.py b/tests/unit_tests/cpu/test_mx_qat_policy.py index 22e6b782da8..f2492d9c353 100644 --- a/tests/unit_tests/cpu/test_mx_qat_policy.py +++ b/tests/unit_tests/cpu/test_mx_qat_policy.py @@ -38,6 +38,41 @@ def _config(): class MXFP4PolicyTest(unittest.TestCase): + def test_manifest_selects_only_actual_pairs_from_eligible_linears(self): + weights = ["model.expert.weight", "model.residual.weight", "model.fused.weight"] + policy = MXFP4CheckpointPolicy.from_manifest( + _config(), + weights, + { + "model.expert.weight_packed": "weights.safetensors", + "model.expert.weight_scale": "scales.safetensors", + "model.residual.weight": "dense.safetensors", + "model.fused.weight": "dense.safetensors", + }, + ) + self.assertEqual(policy.weight_fqns, frozenset({"model.expert.weight"})) + + def test_manifest_rejects_invalid_pairs(self): + pair = {"model.expert.weight_packed": "a", "model.expert.weight_scale": "b"} + cases = [ + ({"model.expert.weight_packed": "a"}, "missing scales"), + ({"model.expert.weight_scale": "a"}, "orphan scales"), + ({**pair, "model.expert.weight": "c"}, "both packed and ordinary"), + ( + {"model.shared.weight_packed": "a", "model.shared.weight_scale": "b"}, + "outside", + ), + ({"unknown.weight_packed": "a", "unknown.weight_scale": "b"}, "outside"), + ({"model.expert.weight_packed": None}, "weight_map"), + ] + for manifest, message in cases: + with self.subTest(manifest=manifest), self.assertRaisesRegex( + ValueError, message + ): + MXFP4CheckpointPolicy.from_manifest( + _config(), ["model.expert.weight", "model.shared.weight"], manifest + ) + def test_resolves_actual_linear_hierarchy_and_prefix_regex(self): policy = MXFP4CheckpointPolicy.from_config( _config(), diff --git a/tests/unit_tests/cpu/test_mx_qat_recipe.py b/tests/unit_tests/cpu/test_mx_qat_recipe.py index 43a2dd401f4..ef56e0630ff 100644 --- a/tests/unit_tests/cpu/test_mx_qat_recipe.py +++ b/tests/unit_tests/cpu/test_mx_qat_recipe.py @@ -5,8 +5,10 @@ # LICENSE file in the root directory of this source tree. import sys +import tempfile import types import unittest +from pathlib import Path from unittest.mock import patch import torch @@ -14,8 +16,12 @@ from torchao.quantization.quantize_.common import KernelPreference from torchtitan.config import ConfigManager from torchtitan.models.common.moe import GroupedExperts +from torchtitan.models.kimi_k3 import model_registry from torchtitan.models.kimi_k3.config_registry import kimi_k3_debugmodel_mx_qat from torchtitan.models.kimi_k3.model import KimiK3Model +from torchtitan.models.kimi_k3.state_dict_adapter import KimiK3StateDictAdapter + +from tests.unit_tests.cpu.mx_qat_test_utils import write_mixed_checkpoint_metadata class MXQATRecipeTest(unittest.TestCase): @@ -36,10 +42,18 @@ def test_standard_model_entrypoint_and_training_override(self): self.assertTrue(type(config.model.layers[1].moe.routed_up)._owner._mx_qat) def test_custom_recipe_uses_existing_torchao_configs(self): + directory = tempfile.TemporaryDirectory() + self.addCleanup(directory.cleanup) + checkpoint_path = directory.name + write_mixed_checkpoint_metadata( + Path(checkpoint_path), + KimiK3StateDictAdapter(model_registry("debugmodel"), None), + ) + def qat(): return kimi_k3_debugmodel_mx_qat( seq_len=16, - checkpoint_path="/tmp/packed-kimi", + checkpoint_path=checkpoint_path, weight_fake_quant_config=MXFakeQuantizeConfig( dtype=torch.float4_e2m1fn_x2, kernel_preference=KernelPreference.AUTO, @@ -56,7 +70,7 @@ def qat(): config = ConfigManager().parse_args( ["--module", module.__name__, "--config", "qat"] ) - self.assertEqual(config.checkpointer.initial_load_path, "/tmp/packed-kimi") + self.assertEqual(config.checkpointer.initial_load_path, checkpoint_path) self.assertTrue(config.checkpointer.initial_load_in_hf) self.assertTrue(config.checkpointer.initial_load_in_hf_quantized) experts = list(config.model.traverse(GroupedExperts.Config)) @@ -69,11 +83,13 @@ def qat(): expert.activation_fake_quant_config.kernel_preference, KernelPreference.AUTO, ) + self.assertFalse( + getattr(type(config.model.layers[1].moe.routed_up)._owner, "_mx_qat", False) + ) + adapter = KimiK3StateDictAdapter(config.model, None) + reader = adapter.get_hf_storage_reader(checkpoint_path, from_quantized=True) self.assertEqual( - config.model.layers[ - 1 - ].moe.routed_up.weight_fake_quant_config.kernel_preference, - KernelPreference.AUTO, + reader.spec.target_fqns, adapter.mxfp4_policy(checkpoint_path).weight_fqns ) diff --git a/tests/unit_tests/cpu/test_state_dict_adapter.py b/tests/unit_tests/cpu/test_state_dict_adapter.py index 21cbbfe7caa..7926321a6d4 100644 --- a/tests/unit_tests/cpu/test_state_dict_adapter.py +++ b/tests/unit_tests/cpu/test_state_dict_adapter.py @@ -16,7 +16,6 @@ from torch.distributed.device_mesh import init_device_mesh from torch.distributed.tensor import DTensor, Replicate, Shard from torch.testing._internal.distributed.fake_pg import FakeStore - from torchtitan.components.checkpointer.base import ModelWrapper from torchtitan.components.checkpointer.hf_storage import ( HuggingFaceStorageReaderWithViews, @@ -39,6 +38,8 @@ from torchtitan.models.qwen3.state_dict_adapter import Qwen3StateDictAdapter from torchtitan.protocols.state_dict_adapter import StateDictAdapter +from tests.unit_tests.cpu.mx_qat_test_utils import write_mixed_checkpoint_metadata + class NativeFusedLinearStateDictAdapterTest(unittest.TestCase): def test_stacked_helpers_support_nonleading_projection_dim(self) -> None: @@ -132,17 +133,12 @@ def setUp(self) -> None: hf_assets_path=None, ) + def _write_checkpoint_metadata(self, path): + return write_mixed_checkpoint_metadata(Path(path), self.adapter) + def test_quantized_load_uses_packed_pair_reader(self) -> None: with tempfile.TemporaryDirectory() as directory: - Path(directory, "config.json").write_text( - json.dumps( - { - "text_config": { - "quantization_config": deepcopy(MXFP4_QUANTIZATION_CONFIG), - } - } - ) - ) + self._write_checkpoint_metadata(directory) reader = self.adapter.get_hf_storage_reader( directory, from_quantized=True, @@ -173,11 +169,11 @@ def test_quantized_load_uses_packed_pair_reader(self) -> None: "language_model.model.layers.1.block_sparse_moe.gate.weight" in reader.spec.target_fqns ) - self.assertIn( + self.assertNotIn( "language_model.model.layers.1.block_sparse_moe.routed_expert_up_proj.weight", reader.spec.target_fqns, ) - self.assertIn( + self.assertNotIn( "language_model.model.layers.1.mlp_res_proj.weight", reader.spec.target_fqns ) @@ -196,7 +192,9 @@ def test_qat_selection_matches_import_including_dense_projections(self) -> None: self.adapter._validate_qat_policy(policy) self.assertTrue(type(model.layers[1].moe.routed_up)._owner._mx_qat) - def test_qat_rejects_expert_only_selection_for_released_policy(self) -> None: + def test_qat_rejects_expert_only_selection_when_dense_weights_are_packed( + self, + ) -> None: from torchtitan.config.transform import MXQATTransform from torchtitan.quantization.mx_qat.checkpoint import MXFP4CheckpointPolicy @@ -226,11 +224,32 @@ def test_qat_recipe_is_valid_with_and_without_initial_checkpoint(self) -> None: recipe = kimi_k3_debugmodel_mx_qat(seq_len=16) self.assertFalse(recipe.checkpointer.initial_load_in_hf_quantized) - recipe = kimi_k3_debugmodel_mx_qat( - seq_len=16, checkpoint_path="/tmp/packed-kimi" - ) - self.assertTrue(recipe.checkpointer.initial_load_in_hf_quantized) - self.assertEqual(recipe.checkpointer.initial_load_path, "/tmp/packed-kimi") + with tempfile.TemporaryDirectory() as directory: + self._write_checkpoint_metadata(directory) + recipe = kimi_k3_debugmodel_mx_qat(seq_len=16, checkpoint_path=directory) + self.assertTrue(recipe.checkpointer.initial_load_in_hf_quantized) + self.assertEqual(recipe.checkpointer.initial_load_path, directory) + self.assertFalse( + getattr( + type(recipe.model.layers[1].moe.routed_up)._owner, "_mx_qat", False + ) + ) + + def test_manifest_qat_rejects_mixed_experts_in_one_parameter(self): + with tempfile.TemporaryDirectory() as directory: + manifest = self._write_checkpoint_metadata(directory) + prefix = "language_model.model.layers.1.block_sparse_moe.experts.0.w1" + del manifest[prefix + ".weight_packed"] + del manifest[prefix + ".weight_scale"] + manifest[prefix + ".weight"] = "dense.safetensors" + Path(directory, "model.safetensors.index.json").write_text( + json.dumps({"weight_map": manifest}) + ) + policy = self.adapter.mxfp4_policy(directory) + with self.assertRaisesRegex( + ValueError, "cannot mix packed and BF16 experts" + ): + self.adapter.qat_weight_fqns(policy) def test_quantized_load_rejects_missing_metadata(self) -> None: with tempfile.TemporaryDirectory() as directory: diff --git a/torchtitan/models/kimi_k3/README.md b/torchtitan/models/kimi_k3/README.md index eecc0c2cd55..ea29468be61 100644 --- a/torchtitan/models/kimi_k3/README.md +++ b/torchtitan/models/kimi_k3/README.md @@ -71,8 +71,8 @@ NGPU=1 MODULE=kimi_k3 CONFIG=kimi_k3_debugmodel_mx_qat ./run_train.sh ``` The debug recipe starts from random initialization and uses `EMULATED` by -default. It selects the expert and dense projections required by the released -MXFP4 policy, keeps BF16 master parameters, and applies MXFP8 activation fake +default. Without a checkpoint it selects all config-eligible expert and dense +projections, keeps BF16 master parameters, and applies MXFP8 activation fake quantization to grouped experts. No manual expert replacement or parameter-name list is needed. Prepare the normal Kimi tokenizer and dataset dependencies as for `kimi_k3_debugmodel`; QAT does not replace that data setup. @@ -100,8 +100,18 @@ def qat(): NGPU=1 MODULE=my_kimi_runs CONFIG=qat ./run_train.sh ``` -`checkpoint_path` sets the existing HF and quantized-load options together. The -loader validates the checkpoint policy against the selected QAT modules. +`checkpoint_path` sets the existing HF and quantized-load options together and +reads `config.json` plus `model.safetensors.index.json` before selecting QAT +modules. The config defines eligible weights; only actual packed/scale pairs in +the index select QAT. Eligible weights stored in BF16 remain ordinary weights. +The released checkpoint packs routed expert matrices while residual and fused +projection weights remain BF16. The synthetic debug fixture deliberately packs +all eligible weights, so its selected set can differ from the release. + +The loader validates the same manifest-derived policy against the selected QAT +modules and actual tensor headers. Missing pairs, orphan scales, and pairs +outside the config policy remain errors. A grouped parameter mixing packed and +BF16 experts can be imported, but cannot be selected for grouped QAT. For native grouped execution, pass the existing TorchAO configurations to the same recipe function; both `kernel_preference` values must agree: diff --git a/torchtitan/models/kimi_k3/config_registry.py b/torchtitan/models/kimi_k3/config_registry.py index 8bdc4d2c12b..294cb1b26ad 100644 --- a/torchtitan/models/kimi_k3/config_registry.py +++ b/torchtitan/models/kimi_k3/config_registry.py @@ -105,18 +105,23 @@ def kimi_k3_debugmodel_mx_qat( weight_fake_quant_config: MXFakeQuantizeConfig | None = None, activation_fake_quant_config: MXFakeQuantizeConfig | None = None, ) -> Trainer.Config: - """Kimi QAT using the released policy and optional packed HF initialization. + """Kimi QAT using checkpoint storage policy when initializing from HF. Pass an absolute checkpoint_path to load the packed debug fixture. Without - it, the recipe uses random initialization and remains valid before overrides. + it, random initialization applies QAT to all config-eligible weights. + With a checkpoint, only actual manifest packed pairs select QAT weights. Optional TorchAO configs control fake quantization and kernel_preference; model-specific parameter selection stays inside the recipe. """ config = kimi_k3_debugmodel(seq_len=seq_len) adapter = KimiK3StateDictAdapter(config.model, hf_assets_path=None) mapping = adapter.hf_linear_weight_mapping() - policy = MXFP4CheckpointPolicy.from_config(MXFP4_QUANTIZATION_CONFIG, mapping) - weights = {mapping[key] for key in policy.weight_fqns if mapping[key] is not None} + policy = ( + adapter.mxfp4_policy(checkpoint_path) + if checkpoint_path is not None + else MXFP4CheckpointPolicy.from_config(MXFP4_QUANTIZATION_CONFIG, mapping) + ) + weights = adapter.qat_weight_fqns(policy) transform = MXQATTransform.from_weight_fqns(config.model, weights) if weight_fake_quant_config is not None: transform.weight_fake_quant_config = weight_fake_quant_config diff --git a/torchtitan/models/kimi_k3/state_dict_adapter.py b/torchtitan/models/kimi_k3/state_dict_adapter.py index 387737588e8..c0b95fd0885 100644 --- a/torchtitan/models/kimi_k3/state_dict_adapter.py +++ b/torchtitan/models/kimi_k3/state_dict_adapter.py @@ -222,10 +222,36 @@ def mxfp4_policy(self, path: str) -> MXFP4CheckpointPolicy: quantization = config.get("text_config", config).get("quantization_config") if not isinstance(quantization, dict): raise ValueError("Kimi checkpoint is missing quantization_config metadata.") - return MXFP4CheckpointPolicy.from_config( - quantization, self.hf_linear_weight_mapping() + index_path = Path(path) / "model.safetensors.index.json" + if not index_path.is_file(): + raise ValueError(f"Quantized Kimi checkpoint is missing {index_path}.") + index = json.loads(index_path.read_text()) + return MXFP4CheckpointPolicy.from_manifest( + quantization, self.hf_linear_weight_mapping(), index.get("weight_map") ) + def qat_weight_fqns(self, policy: MXFP4CheckpointPolicy) -> set[str]: + """Translate a manifest policy without silently quantizing BF16 experts. + + Several HF experts share one Titan parameter. QAT can only select that + parameter when every corresponding HF weight is packed. + """ + mapping = self.hf_linear_weight_mapping() + selected = { + mapping[key] for key in policy.weight_fqns if mapping[key] is not None + } + partial = { + target + for key, target in mapping.items() + if target in selected and key not in policy.weight_fqns + } + if partial: + raise ValueError( + "MX QAT cannot mix packed and BF16 experts within a shared parameter: " + f"{sorted(partial)}" + ) + return selected + @staticmethod def _validate_qat_weight_config(config, policy: MXFP4CheckpointPolicy) -> None: if ( @@ -259,9 +285,7 @@ def _validate_qat_policy(self, policy: MXFP4CheckpointPolicy) -> None: ) if not has_qat: return # Packed import into a BF16 model remains supported. - expected = { - mapping[key] for key in policy.weight_fqns if mapping[key] is not None - } + expected = self.qat_weight_fqns(policy) if selected != expected: raise ValueError( "MX QAT selection disagrees with checkpoint policy: " diff --git a/torchtitan/quantization/mx_qat/checkpoint.py b/torchtitan/quantization/mx_qat/checkpoint.py index 7e57304bdc5..ff863b3c3be 100644 --- a/torchtitan/quantization/mx_qat/checkpoint.py +++ b/torchtitan/quantization/mx_qat/checkpoint.py @@ -23,10 +23,11 @@ class MXFP4CheckpointPolicy: def from_config( cls, quantization: Mapping[str, Any], linear_weights: Iterable[str] ) -> "MXFP4CheckpointPolicy": - """Resolve an exact set; the adapter supplies actual HF Linear weights. + """Resolve eligible weights, not the checkpoint's actual packed set. Only the released static, symmetric 1x32 E2M1/E8M0 format is supported. Names alone cannot distinguish Linear weights from embeddings or norms. + Use from_manifest for checkpoint import; eligible weights may be BF16. """ if ( quantization.get("format") != "mxfp4-pack-quantized" @@ -101,6 +102,53 @@ def matches(pattern: str, module: str) -> bool: ) ) + @classmethod + def from_manifest( + cls, + quantization: Mapping[str, Any], + linear_weights: Iterable[str], + weight_map: Mapping[str, str], + ) -> "MXFP4CheckpointPolicy": + """Validate actual index pairs against eligibility and select only those. + + A config target is permission to quantize, not proof of packed storage. + The storage reader subsequently validates the indexed pairs' physical + presence, dtype, and shape. Ordinary weights remain ordinary tensors. + """ + eligible = cls.from_config(quantization, linear_weights) + if not isinstance(weight_map, Mapping) or any( + not isinstance(name, str) or not isinstance(shard, str) or not shard + for name, shard in weight_map.items() + ): + raise ValueError("Checkpoint index requires a tensor-to-shard weight_map") + packed = { + name.removesuffix(".weight_packed") + ".weight" + for name in weight_map + if name.endswith(".weight_packed") + } + scaled = { + name.removesuffix(".weight_scale") + ".weight" + for name in weight_map + if name.endswith(".weight_scale") + } + if packed != scaled: + raise ValueError( + "Checkpoint index has unpaired MXFP4 tensors: " + f"missing scales={sorted(packed - scaled)[:10]}, " + f"orphan scales={sorted(scaled - packed)[:10]}" + ) + outside = packed - eligible.weight_fqns + if outside: + raise ValueError( + f"Checkpoint packed weights are outside the config policy: {sorted(outside)[:10]}" + ) + conflicts = packed & weight_map.keys() + if conflicts: + raise ValueError( + f"Checkpoint index has both packed and ordinary weights: {sorted(conflicts)[:10]}" + ) + return cls(frozenset(packed), block_size=eligible.block_size) + def decode_mxfp4( packed: torch.Tensor,