diff --git a/src/neuronx_distributed_inference/models/application_base.py b/src/neuronx_distributed_inference/models/application_base.py index 3448e152..91fcb912 100644 --- a/src/neuronx_distributed_inference/models/application_base.py +++ b/src/neuronx_distributed_inference/models/application_base.py @@ -411,6 +411,31 @@ def load_weights(self, compiled_model_path, start_rank_id=None, local_ranks_size cte_model.lora_weight_manager.lora_checkpoint.update_weights_for_lora_cpu(cte_model) lora_cpu_weights = cte_model.lora_weight_manager.lora_checkpoint.shard_cpu_checkpoints(start_rank_id, local_ranks_size, self.neuron_config.tp_degree, cte_model) + # PATCH 2026-08-28 worker-llama70b-measurement: + # For FP8-KV compiled NEFFs, the compiled graph expects kv_mgr.k_scales.N + # and kv_mgr.v_scales.N tensors in the sharded checkpoint (created as + # nn.Parameter by KVCacheManager._init_scale_buffers), but these are + # RUNTIME buffers not present in the model checkpoint. Inject them as + # torch.ones(1) (per-tensor unity) so initialize() finds them. + try: + kv_cfg = getattr(self.neuron_config, "kv_quant_config", None) + if kv_cfg is not None: + num_hidden_layers = getattr(self.config, "num_hidden_layers", None) + if num_hidden_layers is not None and len(weights) > 0: + injected = 0 + for rank_sd in weights: + if not isinstance(rank_sd, dict): + continue + for L in range(num_hidden_layers): + for base in ("kv_mgr.k_scales", "kv_mgr.v_scales"): + key = f"{base}.{L}" + if key not in rank_sd: + rank_sd[key] = torch.ones(1, dtype=torch.bfloat16) + injected += 1 + logger.info(f"[patch-injected] kv_mgr.*_scales unity tensors injected: {injected} across {len(weights)} ranks (num_layer={num_hidden_layers})") + except Exception as _e: + logger.warning(f"[patch-injected] kv_mgr scale injection warning: {_e}") + start_rank_tensor = torch.tensor([start_rank_id], dtype=torch.int32, device="cpu") self.traced_model.nxd_model.initialize(weights, start_rank_tensor) @@ -692,6 +717,21 @@ def set_tensor_capture_step(self, step=0): @classmethod def get_state_dict(cls, model_name_or_path: str, config: InferenceConfig) -> dict: """Gets the state dict for this model.""" + checkpoint_format = os.environ.get("CHECKPOINT_FORMAT", "hf").lower() + normalized_fp8 = checkpoint_format == "normalized_fp8" + if normalized_fp8: + validator = getattr(cls, "_validate_normalized_fp8_checkpoint", None) + if validator is None: + raise ValueError( + "CHECKPOINT_FORMAT=normalized_fp8 is not supported by this model" + ) + if not os.path.isdir(model_name_or_path): + raise ValueError( + "CHECKPOINT_FORMAT=normalized_fp8 requires an indexed " + "SafeTensors directory" + ) + validator(model_name_or_path, config) + if os.path.isdir(model_name_or_path): model_sd = load_state_dict(model_name_or_path) elif os.path.isfile(model_name_or_path): @@ -727,7 +767,8 @@ def get_state_dict(cls, model_name_or_path: str, config: InferenceConfig) -> dic "Recompile the model with save_sharded_checkpoint=True." ) - model_sd = cls.convert_hf_to_neuron_state_dict(model_sd, config) + if not normalized_fp8: + model_sd = cls.convert_hf_to_neuron_state_dict(model_sd, config) if getattr(config, "tie_word_embeddings", False): cls.update_state_dict_for_tied_weights(model_sd) diff --git a/src/neuronx_distributed_inference/models/config.py b/src/neuronx_distributed_inference/models/config.py index 9b58163d..9743838b 100644 --- a/src/neuronx_distributed_inference/models/config.py +++ b/src/neuronx_distributed_inference/models/config.py @@ -299,9 +299,36 @@ def __init__(self, **kwargs) -> None: # KV Quantization self.kv_quant_config = kwargs.pop("kv_quant_config", None) if type(self.kv_quant_config) is dict: - self.kv_quant_config = KVQuantizationConfig( - **self.kv_quant_config - ) + # PATCH 2026-08-28 worker-llama70b-measurement: + # Coerce dict-form enum values (from JSON round-trip) to real enum + # members before instantiating KVQuantizationConfig, so the assert + # in __init__ (direct_cast=True requires PER_TENSOR_SYMMETRIC) sees + # a live QuantizationType instead of a plain dict. + _kv = dict(self.kv_quant_config) + try: + from neuronx_distributed.quantization.quantization_config import QuantizationType as _QT + for _k in ("k_quant_method", "v_quant_method"): + _v = _kv.get(_k) + if isinstance(_v, dict): + _name = _v.get("_name_") or _v.get("name") + _val = _v.get("_value_") or _v.get("value") + if _name is not None: + _kv[_k] = _QT[_name] + elif _val is not None: + _kv[_k] = _QT(_val) + elif isinstance(_v, str): + try: + _kv[_k] = _QT(_v) + except Exception: + _kv[_k] = _QT[_v] + import torch as _torch + _dt = _kv.get("quant_dtype") + if isinstance(_dt, str): + _dt2 = _dt.split(".", 1)[1] if _dt.startswith("torch.") else _dt + _kv["quant_dtype"] = getattr(_torch, _dt2, _dt) + except Exception as _e: + print(f"[nxdi-patch] kv_quant enum coerce warning: {_e}") + self.kv_quant_config = KVQuantizationConfig(**_kv) self.is_chunked_prefill = self.chunked_prefill_config is not None if self.is_chunked_prefill: diff --git a/src/neuronx_distributed_inference/models/gpt_oss/modeling_gpt_oss.py b/src/neuronx_distributed_inference/models/gpt_oss/modeling_gpt_oss.py index f5961c5b..eb4a69fa 100644 --- a/src/neuronx_distributed_inference/models/gpt_oss/modeling_gpt_oss.py +++ b/src/neuronx_distributed_inference/models/gpt_oss/modeling_gpt_oss.py @@ -56,6 +56,9 @@ from neuronx_distributed_inference.modules.kvcache.gpt_oss_kv_cache_manager import GptOssKVCacheManager from neuronx_distributed_inference.modules.moe_v2 import initialize_moe_module from neuronx_distributed_inference.utils.distributed import get_tp_group +from neuronx_distributed_inference.utils.normalized_fp8 import ( + validate_gpt_oss_120b_normalized_fp8_checkpoint, +) from transformers import GptOssForCausalLM @@ -983,6 +986,11 @@ class NeuronGptOssForCausalLM(NeuronBaseForCausalLM): _model_cls = NeuronGptOssModel + @staticmethod + def _validate_normalized_fp8_checkpoint(model_path, config): + """Validate the post-conversion GPT-OSS static-FP8 checkpoint layout.""" + validate_gpt_oss_120b_normalized_fp8_checkpoint(str(model_path), config) + @staticmethod def load_hf_model(model_path, **kwargs): checkpoint_format = os.environ.get("CHECKPOINT_FORMAT", "hf").lower() diff --git a/src/neuronx_distributed_inference/modules/attention/gqa.py b/src/neuronx_distributed_inference/modules/attention/gqa.py index 56383d38..524cf454 100644 --- a/src/neuronx_distributed_inference/modules/attention/gqa.py +++ b/src/neuronx_distributed_inference/modules/attention/gqa.py @@ -102,7 +102,15 @@ def get_shardable_head_counts( def is_per_channel(scale: torch.Tensor) -> bool: """See if the scale is per channel""" - if scale.shape == (1,): + # PATCH 2026-08-28 worker-llama70b-measurement: handle scalar () shape + numel-1 per-tensor scales + if scale is None: + return False + try: + if getattr(scale, "numel", None) is not None and scale.numel() == 1: + return False + except Exception: + pass + if scale.shape == (1,) or scale.shape == () or scale.shape == tuple(): return False return True diff --git a/src/neuronx_distributed_inference/modules/checkpoint.py b/src/neuronx_distributed_inference/modules/checkpoint.py index 8e8b1679..2c656122 100644 --- a/src/neuronx_distributed_inference/modules/checkpoint.py +++ b/src/neuronx_distributed_inference/modules/checkpoint.py @@ -5,6 +5,7 @@ import torch from huggingface_hub import save_torch_state_dict +from safetensors import safe_open from safetensors.torch import load_file _SAFETENSORS_MODEL_INDEX_FILENAME_JSON = "model.safetensors.index.json" @@ -110,6 +111,91 @@ def _load_from_files( return state_dict +def inspect_safetensors_sharded(state_dict_dir: str) -> Dict[str, Dict[str, Any]]: + """Inspect an indexed SafeTensors checkpoint without materializing tensors. + + The indexed file is the authority for tensor ownership. This helper is + intentionally stricter than the legacy loader: every SafeTensors file in + the directory must be indexed, every indexed tensor must exist in its + declared file, and a tensor name may have only one physical owner. The + returned metadata is sufficient for model-specific dtype and shape gates. + """ + if not os.path.isdir(state_dict_dir): + raise ValueError(f"SafeTensors checkpoint directory missing: {state_dict_dir}") + + index_path = os.path.join(state_dict_dir, _SAFETENSORS_MODEL_INDEX_FILENAME_JSON) + if not os.path.isfile(index_path): + raise ValueError(f"Indexed SafeTensors checkpoint required: {index_path}") + + with open(index_path, encoding="utf-8") as f: + index = json.load(f) + weight_map = index.get("weight_map") + if not isinstance(weight_map, dict) or not weight_map: + raise ValueError(f"SafeTensors index has no non-empty weight_map: {index_path}") + if not all( + isinstance(name, str) and isinstance(filename, str) + for name, filename in weight_map.items() + ): + raise ValueError(f"SafeTensors index contains non-string entries: {index_path}") + + filenames = set(weight_map.values()) + unsafe = sorted( + filename + for filename in filenames + if filename != os.path.basename(filename) + or "\\" in filename + or not filename.endswith(".safetensors") + or os.path.isabs(filename) + ) + if unsafe: + raise ValueError(f"SafeTensors index references unsafe files: {unsafe}") + + actual_files = { + filename + for filename in os.listdir(state_dict_dir) + if filename.endswith(".safetensors") + } + if actual_files != filenames: + raise ValueError( + "SafeTensors file inventory disagrees with index: " + f"unindexed={sorted(actual_files - filenames)}, " + f"missing={sorted(filenames - actual_files)}" + ) + + metadata: Dict[str, Dict[str, Any]] = {} + owners: Dict[str, str] = {} + for filename in sorted(filenames): + path = os.path.join(state_dict_dir, filename) + with safe_open(path, framework="pt", device="cpu") as source: + for name in source.keys(): + if name in owners: + raise ValueError( + f"duplicate SafeTensors tensor name {name!r}: " + f"{owners[name]} and {filename}" + ) + if weight_map.get(name) != filename: + raise ValueError( + f"SafeTensors tensor {name!r} is not owned by its indexed file" + ) + tensor = source.get_slice(name) + owners[name] = filename + metadata[name] = { + "filename": filename, + "dtype": tensor.get_dtype(), + "shape": tuple(tensor.get_shape()), + } + + indexed_names = set(weight_map) + actual_names = set(metadata) + if actual_names != indexed_names: + raise ValueError( + "SafeTensors tensor inventory disagrees with index: " + f"unindexed={sorted(actual_names - indexed_names)}, " + f"missing={sorted(indexed_names - actual_names)}" + ) + return metadata + + def load_safetensors_sharded(state_dict_dir: str) -> Dict[str, torch.Tensor]: index_path = os.path.join(state_dict_dir, _SAFETENSORS_MODEL_INDEX_FILENAME_JSON) with open(index_path, "r") as f: diff --git a/src/neuronx_distributed_inference/utils/normalized_fp8.py b/src/neuronx_distributed_inference/utils/normalized_fp8.py new file mode 100644 index 00000000..8ef076bf --- /dev/null +++ b/src/neuronx_distributed_inference/utils/normalized_fp8.py @@ -0,0 +1,143 @@ +"""Validation for explicitly normalized static-FP8 checkpoints.""" + +from __future__ import annotations + +from typing import Any + +from neuronx_distributed_inference.modules.checkpoint import ( + inspect_safetensors_sharded, +) + + +def _expected_names(num_layers: int) -> tuple[set[str], set[str], set[str], set[str]]: + expected = { + "embed_tokens.weight", + "lm_head.weight", + "norm.weight", + "norm.weight_unpadded", + "rank_util.rank", + } + fp8_weights: set[str] = set() + scalar_scales: set[str] = set() + expert_scales: set[str] = set() + + for layer in range(num_layers): + prefix = f"layers.{layer}" + expected.update( + { + f"{prefix}.input_layernorm.weight", + f"{prefix}.input_layernorm.weight_unpadded", + f"{prefix}.post_attention_layernorm.weight", + f"{prefix}.post_attention_layernorm.weight_unpadded", + f"{prefix}.self_attn.learned_sinks.sink", + f"{prefix}.self_attn.rank_util.rank", + f"{prefix}.feed_forward.moe.router.linear_router.weight", + f"{prefix}.feed_forward.moe.router.linear_router.bias", + } + ) + for projection in ("q_proj", "k_proj", "v_proj", "o_proj"): + weight = f"{prefix}.self_attn.{projection}.weight" + scale = weight.removesuffix(".weight") + ".scale" + fp8_weights.add(weight) + scalar_scales.add(scale) + expected.update( + { + weight, + f"{prefix}.self_attn.{projection}.bias", + scale, + } + ) + for projection in ("down_proj", "gate_up_proj"): + base = f"{prefix}.feed_forward.moe.expert_mlps.mlp_op.{projection}" + weight = f"{base}.weight" + scale = f"{base}.scale" + fp8_weights.add(weight) + expert_scales.add(scale) + expected.update({weight, f"{base}.bias", scale}) + + return expected, fp8_weights, scalar_scales, expert_scales + + +def validate_gpt_oss_120b_normalized_fp8_metadata( + metadata: dict[str, dict[str, Any]], config: Any +) -> None: + """Fail closed unless metadata matches the normalized GPT-OSS-120B ABI.""" + neuron_config = config.neuron_config + if not neuron_config.quantized: + raise ValueError("normalized_fp8 requires quantized=true") + if neuron_config.quantization_dtype != "f8e4m3": + raise ValueError("normalized_fp8 requires quantization_dtype=f8e4m3") + if neuron_config.fused_qkv: + raise ValueError("normalized_fp8 requires separate q/k/v projections") + if neuron_config.vocab_parallel: + raise ValueError("normalized_fp8 does not support vocab-parallel embedding") + if config.num_hidden_layers != 36 or config.num_local_experts != 128: + raise ValueError("normalized_fp8 supports the GPT-OSS-120B layout only") + if neuron_config.attention_dp_degree != neuron_config.cp_degree: + raise ValueError("normalized_fp8 requires matching attention DP and CP") + + expected, fp8_weights, scalar_scales, expert_scales = _expected_names( + config.num_hidden_layers + ) + actual = set(metadata) + if actual != expected: + raise ValueError( + "normalized_fp8 tensor namespace mismatch: " + f"missing={sorted(expected - actual)}, " + f"unexpected={sorted(actual - expected)}" + ) + + for name in fp8_weights: + details = metadata[name] + if details["dtype"] != "F8_E4M3": + raise ValueError(f"normalized_fp8 weight is not F8_E4M3: {name}") + if not details["shape"] or any( + dimension <= 0 for dimension in details["shape"] + ): + raise ValueError(f"normalized_fp8 weight has invalid shape: {name}") + + for name in scalar_scales: + details = metadata[name] + if details["dtype"] != "F32" or details["shape"] != (1,): + raise ValueError( + f"normalized_fp8 scalar scale shape/dtype mismatch: {name}" + ) + if details["filename"] == "expert-scales.safetensors": + raise ValueError(f"non-expert scale is in expert scale sidecar: {name}") + + for name in expert_scales: + expected_shape = ( + (128, 1, 1) if name.endswith("down_proj.scale") else (128, 2, 1) + ) + details = metadata[name] + if details["dtype"] != "F32" or details["shape"] != expected_shape: + raise ValueError( + f"normalized_fp8 expert scale shape/dtype mismatch: {name}" + ) + if details["filename"] != "expert-scales.safetensors": + raise ValueError(f"expert scale is not in expert sidecar: {name}") + + for name, details in metadata.items(): + if name.endswith(".scale"): + continue + if details["dtype"] == "F8_E4M3": + if name not in fp8_weights: + raise ValueError(f"unexpected FP8 tensor: {name}") + elif name.endswith(".rank"): + if details["dtype"] != "I32": + raise ValueError(f"rank tensor is not I32: {name}") + elif details["dtype"] not in {"BF16", "F32"}: + raise ValueError(f"unexpected normalized tensor dtype: {name}") + + +def validate_gpt_oss_120b_normalized_fp8_checkpoint( + model_path: str, config: Any +) -> None: + """Validate the indexed checkpoint before loading any tensor payload.""" + metadata = inspect_safetensors_sharded(model_path) + filenames = {details["filename"] for details in metadata.values()} + if "expert-scales.safetensors" in filenames and "model.safetensors" in filenames: + raise ValueError( + "normalized_fp8 requires indexed model shards when the expert scale sidecar is present" + ) + validate_gpt_oss_120b_normalized_fp8_metadata(metadata, config) diff --git a/test/unit/utils/test_normalized_fp8.py b/test/unit/utils/test_normalized_fp8.py new file mode 100644 index 00000000..487b1856 --- /dev/null +++ b/test/unit/utils/test_normalized_fp8.py @@ -0,0 +1,226 @@ +import json +from pathlib import Path +from types import SimpleNamespace + +import pytest +import torch +from safetensors.torch import save_file + +from neuronx_distributed_inference.modules.checkpoint import ( + inspect_safetensors_sharded, +) +from neuronx_distributed_inference.utils.normalized_fp8 import ( + _expected_names, + validate_gpt_oss_120b_normalized_fp8_checkpoint, +) + + +def _config(**neuron_overrides): + neuron_config = { + "quantized": True, + "quantization_dtype": "f8e4m3", + "fused_qkv": False, + "vocab_parallel": False, + "attention_dp_degree": 1, + "cp_degree": 1, + } + neuron_config.update(neuron_overrides) + return SimpleNamespace( + neuron_config=SimpleNamespace(**neuron_config), + num_hidden_layers=36, + num_local_experts=128, + ) + + +def _write_checkpoint( + root: Path, + *, + dtype_overrides=None, + shape_overrides=None, + sidecar_names=None, + rename=None, + model_filename="model-00001-of-00001.safetensors", +): + expected, fp8_weights, scalar_scales, expert_scales = _expected_names(36) + dtype_overrides = dtype_overrides or {} + shape_overrides = shape_overrides or {} + sidecar_names = sidecar_names or expert_scales + rename = rename or {} + + tensors = {} + owners = {} + for original_name in sorted(expected): + name = rename.get(original_name, original_name) + if original_name in fp8_weights: + dtype = dtype_overrides.get(original_name, torch.float8_e4m3fn) + shape = shape_overrides.get(original_name, (1, 1)) + elif original_name in scalar_scales: + dtype = dtype_overrides.get(original_name, torch.float32) + shape = shape_overrides.get(original_name, (1,)) + elif original_name in expert_scales: + dtype = dtype_overrides.get(original_name, torch.float32) + shape = shape_overrides.get( + original_name, + (128, 1, 1) + if original_name.endswith("down_proj.scale") + else (128, 2, 1), + ) + elif original_name.endswith(".rank"): + dtype = dtype_overrides.get(original_name, torch.int32) + shape = shape_overrides.get(original_name, (1,)) + else: + dtype = dtype_overrides.get(original_name, torch.bfloat16) + shape = shape_overrides.get(original_name, (1,)) + tensors[name] = torch.zeros(shape, dtype=dtype) + owners[name] = ( + "expert-scales.safetensors" + if original_name in sidecar_names + else model_filename + ) + + root.mkdir() + model_tensors = { + name: tensor + for name, tensor in tensors.items() + if owners[name] == model_filename + } + expert_tensors = { + name: tensor + for name, tensor in tensors.items() + if owners[name] == "expert-scales.safetensors" + } + save_file(model_tensors, str(root / model_filename)) + save_file(expert_tensors, str(root / "expert-scales.safetensors")) + index = {"weight_map": owners} + (root / "model.safetensors.index.json").write_text( + json.dumps(index), encoding="utf-8" + ) + return root + + +def test_valid_normalized_fp8_checkpoint_passes(tmp_path): + checkpoint = _write_checkpoint(tmp_path / "checkpoint") + + validate_gpt_oss_120b_normalized_fp8_checkpoint(str(checkpoint), _config()) + + +@pytest.mark.parametrize( + "kwargs, message", + [ + ( + { + "rename": { + "layers.0.self_attn.learned_sinks.sink": "layers.0.self_attn.sinks" + } + }, + "namespace mismatch", + ), + ( + {"dtype_overrides": {"layers.0.self_attn.q_proj.weight": torch.bfloat16}}, + "not F8_E4M3", + ), + ( + {"shape_overrides": {"layers.0.self_attn.q_proj.scale": (2,)}}, + "shape/dtype mismatch", + ), + ( + { + "shape_overrides": { + "layers.0.feed_forward.moe.expert_mlps.mlp_op.down_proj.scale": ( + 128, + 2, + 1, + ) + } + }, + "shape/dtype mismatch", + ), + ( + { + "sidecar_names": { + "layers.0.feed_forward.moe.expert_mlps.mlp_op.down_proj.scale" + } + ^ _expected_names(36)[3] + }, + "not in expert sidecar", + ), + ], +) +def test_invalid_normalized_fp8_checkpoint_fails_closed(tmp_path, kwargs, message): + checkpoint = _write_checkpoint(tmp_path / "checkpoint", **kwargs) + + with pytest.raises(ValueError, match=message): + validate_gpt_oss_120b_normalized_fp8_checkpoint(str(checkpoint), _config()) + + +@pytest.mark.parametrize( + "neuron_overrides", + [ + {"quantized": False}, + {"quantization_dtype": "bf16"}, + {"fused_qkv": True}, + {"vocab_parallel": True}, + {"attention_dp_degree": 2}, + ], +) +def test_invalid_normalized_fp8_config_fails_closed(tmp_path, neuron_overrides): + checkpoint = _write_checkpoint(tmp_path / "checkpoint") + + with pytest.raises(ValueError): + validate_gpt_oss_120b_normalized_fp8_checkpoint( + str(checkpoint), _config(**neuron_overrides) + ) + + +def test_cross_file_tensor_collision_fails_before_payload_load(tmp_path): + checkpoint = tmp_path / "collision" + checkpoint.mkdir() + save_file({"duplicate": torch.ones(1)}, str(checkpoint / "one.safetensors")) + save_file( + {"duplicate": torch.ones(1), "other": torch.ones(1)}, + str(checkpoint / "two.safetensors"), + ) + (checkpoint / "model.safetensors.index.json").write_text( + json.dumps( + { + "weight_map": { + "duplicate": "one.safetensors", + "other": "two.safetensors", + } + } + ), + encoding="utf-8", + ) + + with pytest.raises(ValueError, match="duplicate SafeTensors tensor name"): + inspect_safetensors_sharded(str(checkpoint)) + + +def test_canonical_model_file_with_sidecar_fails_closed(tmp_path): + checkpoint = _write_checkpoint( + tmp_path / "checkpoint", model_filename="model.safetensors" + ) + + with pytest.raises(ValueError, match="indexed model shards"): + validate_gpt_oss_120b_normalized_fp8_checkpoint(str(checkpoint), _config()) + + +def test_default_hf_conversion_contract_remains_present(): + application_base = ( + Path(__file__).parents[3] + / "src" + / "neuronx_distributed_inference" + / "models" + / "application_base.py" + ) + source = application_base.read_text(encoding="utf-8") + assert 'os.environ.get("CHECKPOINT_FORMAT", "hf")' in source + assert "if not normalized_fp8:" in source + assert "model_sd = cls.convert_hf_to_neuron_state_dict(model_sd, config)" in source + + gpt_oss_source = ( + application_base.parent / "gpt_oss" / "modeling_gpt_oss.py" + ).read_text(encoding="utf-8") + assert 'if checkpoint_format == "hf":' in gpt_oss_source + assert "_convert_hf_format_state_dict" in gpt_oss_source + assert "_convert_neuron_format_state_dict" in gpt_oss_source