# SPDX-License-Identifier: Apache-2.0 import importlib.util import sys from pathlib import Path from types import ModuleType from unittest.mock import patch import pytest import torch safetensors_torch = pytest.importorskip("safetensors.torch") SFT_PATH = Path(__file__).resolve().parents[2] / "python" / "sft" PACKAGE_NAME = "kt_sft_fp8_loader_under_test" package = ModuleType(PACKAGE_NAME) package.__path__ = [str(SFT_PATH)] sys.modules[PACKAGE_NAME] = package def _load_module(filename: str): name = f"{PACKAGE_NAME}.{Path(filename).stem}" spec = importlib.util.spec_from_file_location(name, SFT_PATH / filename) assert spec is not None and spec.loader is not None module = importlib.util.module_from_spec(spec) sys.modules[spec.name] = module spec.loader.exec_module(module) return module arch = _load_module("arch.py") _load_module("dist_utils.py") weights = _load_module("weights.py") def _moe_config(expert_num=2, intermediate_size=128): return arch.MOEArchConfig( moe_layer_attr="mlp", router_attr="gate", experts_attr="experts", weight_names=("gate_proj", "up_proj", "down_proj"), expert_num=expert_num, intermediate_size=intermediate_size, num_experts_per_tok=1, ) def _write_checkpoint(tmp_path: Path, *, omit_key: str | None = None): tensors = {} prefix = "model.layers.0.mlp.experts" for expert_idx in range(2): for projection, shape in ( ("gate_proj", (128, 128)), ("up_proj", (128, 128)), ("down_proj", (128, 128)), ): weight_key = f"{prefix}.{expert_idx}.{projection}.weight" scale_key = f"{prefix}.{expert_idx}.{projection}.weight_scale_inv" tensors[weight_key] = torch.zeros(shape, dtype=torch.float8_e4m3fn) scale_dtype = torch.bfloat16 if expert_idx == 0 else torch.float32 tensors[scale_key] = torch.ones((1, 1), dtype=scale_dtype) if omit_key is not None: tensors.pop(omit_key) checkpoint = tmp_path / "model.safetensors" safetensors_torch.save_file(tensors, checkpoint) metadata = {"weight_map": {key: checkpoint.name for key in tensors}} return checkpoint, metadata def test_raw_fp8_loader_preserves_per_expert_weights_and_normalizes_scales(tmp_path): checkpoint, metadata = _write_checkpoint(tmp_path) with patch.object( weights.torch, "stack", side_effect=AssertionError("native FP8 loader must not stack tensors"), ): loaded = weights.load_block_fp8_experts_from_checkpoint_files( [str(checkpoint)], metadata, "model.layers", _moe_config(), layer_idx=0, hidden_size=128, ) assert isinstance(loaded, weights.BlockFP8ExpertWeights) assert loaded.block_size == (128, 128) for projection in (loaded.gate_proj, loaded.up_proj, loaded.down_proj): assert len(projection) == 2 assert all(tensor.dtype == torch.float8_e4m3fn for tensor in projection) assert all(tensor.device.type == "cpu" for tensor in projection) assert all(tensor.is_contiguous() for tensor in projection) for scales in (loaded.gate_scale, loaded.up_scale, loaded.down_scale): assert len(scales) == 2 assert all(scale.dtype == torch.float32 for scale in scales) assert all(tuple(scale.shape) == (1, 1) for scale in scales) def test_raw_fp8_loader_can_index_an_unsharded_local_checkpoint(tmp_path): checkpoint, _ = _write_checkpoint(tmp_path) loaded = weights.load_block_fp8_experts_from_checkpoint_files( [str(checkpoint)], None, "model.layers", _moe_config(), layer_idx=0, hidden_size=128, ) assert len(loaded.gate_proj) == 2 def test_raw_fp8_loader_rejects_fused_checkpoint_metadata(tmp_path): checkpoint, metadata = _write_checkpoint(tmp_path) metadata["weight_map"]["model.layers.0.mlp.experts.gate_up_proj"] = checkpoint.name with pytest.raises(ValueError, match="non-fused"): weights.load_block_fp8_experts_from_checkpoint_files( [str(checkpoint)], metadata, "model.layers", _moe_config(), layer_idx=0, hidden_size=128, ) def test_raw_fp8_loader_requires_every_scale(tmp_path): missing = "model.layers.0.mlp.experts.1.down_proj.weight_scale_inv" checkpoint, metadata = _write_checkpoint(tmp_path, omit_key=missing) with pytest.raises(FileNotFoundError, match="missing 1 native FP8 tensor"): weights.load_block_fp8_experts_from_checkpoint_files( [str(checkpoint)], metadata, "model.layers", _moe_config(), layer_idx=0, hidden_size=128, ) def test_raw_fp8_loader_rejects_wrong_block_size(tmp_path): checkpoint, metadata = _write_checkpoint(tmp_path) with pytest.raises(ValueError, match="128, 128"): weights.load_block_fp8_experts_from_checkpoint_files( [str(checkpoint)], metadata, "model.layers", _moe_config(), layer_idx=0, hidden_size=128, block_size=(64, 128), ) def _write_rawint4_checkpoint( tmp_path: Path, *, omit_key: str | None = None, add_zero_point: bool = False, bad_scale: bool = False, ): tensors = {} prefix = "language_model.model.layers.0.mlp.experts" for expert_idx in range(2): for projection in ("gate_proj", "up_proj", "down_proj"): base = f"{prefix}.{expert_idx}.{projection}" tensors[f"{base}.weight_packed"] = torch.arange(32 * 4, dtype=torch.int32).reshape(32, 4) scale_shape = (32, 2) if bad_scale and expert_idx == 1 else (32, 1) scale_dtype = torch.bfloat16 if expert_idx == 0 else torch.float32 tensors[f"{base}.weight_scale"] = torch.ones(scale_shape, dtype=scale_dtype) tensors[f"{base}.weight_shape"] = torch.tensor([32, 32], dtype=torch.int64) if add_zero_point: tensors[f"{prefix}.0.gate_proj.weight_zero_point"] = torch.zeros((32, 1), dtype=torch.int32) if omit_key is not None: tensors.pop(omit_key) checkpoint = tmp_path / "model.safetensors" safetensors_torch.save_file(tensors, checkpoint) metadata = {"weight_map": {key: checkpoint.name for key in tensors}} return checkpoint, metadata def test_rawint4_loader_normalizes_real_compressed_tensors_layout(tmp_path): checkpoint, metadata = _write_rawint4_checkpoint(tmp_path) loaded = weights.load_rawint4_experts_from_checkpoint_files( [str(checkpoint)], metadata, "language_model.model.layers", _moe_config(intermediate_size=32), layer_idx=0, hidden_size=32, ) assert isinstance(loaded, weights.RAWINT4ExpertWeights) assert loaded.group_size == 32 for packed in (loaded.gate_proj, loaded.up_proj, loaded.down_proj): assert packed.dtype == torch.uint8 assert tuple(packed.shape) == (2, 32, 16) assert packed.device.type == "cpu" assert packed.is_contiguous() for scale in (loaded.gate_scale, loaded.up_scale, loaded.down_scale): assert scale.dtype == torch.bfloat16 assert tuple(scale.shape) == (2, 32, 1) assert scale.device.type == "cpu" assert scale.is_contiguous() def test_rawint4_loader_rejects_any_zero_point_tensor(tmp_path): checkpoint, metadata = _write_rawint4_checkpoint(tmp_path, add_zero_point=True) with pytest.raises(ValueError, match="without zero-point tensors"): weights.load_rawint4_experts_from_checkpoint_files( [str(checkpoint)], metadata, "language_model.model.layers", _moe_config(intermediate_size=32), layer_idx=0, hidden_size=32, ) def test_rawint4_loader_requires_exact_scale_shape(tmp_path): checkpoint, metadata = _write_rawint4_checkpoint(tmp_path, bad_scale=True) with pytest.raises(ValueError, match="scale shape mismatch"): weights.load_rawint4_experts_from_checkpoint_files( [str(checkpoint)], metadata, "language_model.model.layers", _moe_config(intermediate_size=32), layer_idx=0, hidden_size=32, ) def test_rawint4_loader_requires_every_shape_tensor(tmp_path): missing = "language_model.model.layers.0.mlp.experts.1.down_proj.weight_shape" checkpoint, metadata = _write_rawint4_checkpoint(tmp_path, omit_key=missing) with pytest.raises(FileNotFoundError, match="missing 1 RAWINT4 tensor"): weights.load_rawint4_experts_from_checkpoint_files( [str(checkpoint)], metadata, "language_model.model.layers", _moe_config(intermediate_size=32), layer_idx=0, hidden_size=32, )