Add opt-in normalized GPT-OSS FP8 loader

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)
 

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:

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()

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
 

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:

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)