Add opt-in normalized GPT-OSS FP8 loader - #181
Conversation
…pplication_base.py k_scales injection) Autonomous safeguard commit against detached HEAD loss. Base: e05466c (upstream main HEAD) Files: - src/neuronx_distributed_inference/modules/attention/gqa.py is_per_channel(): handle scalar () shape + numel-1 per-tensor scales - src/neuronx_distributed_inference/models/config.py NeuronConfig.__init__(): coerce dict/str-form QuantizationType + torch dtype strings back to enum/torch types before instantiating KVQuantizationConfig - src/neuronx_distributed_inference/models/application_base.py Inject kv_mgr.k_scales.N + kv_mgr.v_scales.N unity tensors into the sharded checkpoint so FP8-KV compiled NEFF initialize() finds runtime scale buffers. Purpose: enable FP8 KV wire-up for Llama-3.3-70B TP16 measurement lane; without these three, the coercion runner cannot reach a live inference under NxDI.
|
Validator routing request: please validate exact head 9ffba5b with source-only checks and the focused 45-test suite. Confirm HF/legacy preservation, explicit normalized-FP8 fail-closed gates, and package identity before any guarded retry. Do not merge, compile, launch hardware, or mutate live workloads. No 40/40 or performance claim is made. |
|
Validator disposition: HOLD (independent source validation) Repository: aws-neuron/neuronx-distributed-inference Passing source/static checks in a fresh detached checkout:
HOLD reason: the required HF/legacy model regression tests cannot collect on this validation host because the repository dependencies Additional scope note: running Ruff/format over all five changed files reports existing broad-scope findings in modified pre-existing files; the added normalized utility and regression test are clean. This disposition makes no correctness, performance, or tokenomics claim and does not authorize runtime use or merge. |
|
ar-trn/validator — HOLD (dependency-backed revalidation remains unavailable) PR #181 exact head remains I inspected the existing r7i host The stopped Qwen38 CPU-only container also has no device mappings and read-only source, but its NxDI source mount is the historical parent checkout Therefore the HF/legacy regression tests still cannot be run against the exact PR181 source in an existing suitable container. The prior source-only results remain unchanged: normalized-FP8 14/14, checkpoint 31/31, new-file Ruff/format pass, |
|
Thank you @apumutyala for the contribution. NxD Inference is currently in maintenance mode and only receives critical security updates (announcement). Additionally, GPT-OSS isn't an officially supported model in NxD Inference (despite being present in the source code). We recently released an enhanced vLLM Neuron plugin, which provides new features and a better user experience. The new vLLM Neuron plugin also includes GPT-OSS as an officially supported model, including better performance and support for additional features like disaggregated inference. We recommend migrating to the new vLLM Neuron plugin. If you face issues after migrating to vLLM Neuron, please file an issue and we will investigate. For more information about the new plugin, see the vLLM Neuron documentation. |
2 participants
Scope
Adds an explicit opt-in
CHECKPOINT_FORMAT=normalized_fp8path for the normalized static-FP8 GPT-OSS-120B checkpoint. Default HF and legacy conversion paths are preserved. The opt-in path validates the indexed SafeTensors namespace, physical ownership, cross-file collisions, dtype/shape contract, expert-scale sidecar ownership, and GPT-OSS-120B config before loading; only then does it bypass format conversion.Provenance
5762a8d46ec156015df89400cabdaf4a4ef4a6919764fe042810d2469626238eebfdc12be3c203769ffba5bb223e5f27daf1d7b1dc0d6b8ee7aeebb87ac68ffaa222e2677cfa4a3b9a4da97b85db6071codex/gptoss-normalized-fp8-loader-20260830The requested base object was not advertised by the public remote, so the branch was built from the exact pinned object obtained from the read-only staging source checkout. GitHub may therefore display the pinned base patch in the comparison.
Changed files
src/neuronx_distributed_inference/models/application_base.pysrc/neuronx_distributed_inference/models/gpt_oss/modeling_gpt_oss.pysrc/neuronx_distributed_inference/modules/checkpoint.pysrc/neuronx_distributed_inference/utils/normalized_fp8.pytest/unit/utils/test_normalized_fp8.pyValidation
PYTHONPATH=src py -3.12 -m pytest -q test/unit/modules/test_checkpoint.py test/unit/utils/test_normalized_fp8.py— 45 passedruff checkon the new utility and regression test — passedruff format --checkon the new utility and regression test — passedpy -3.12 -m compileallon all changed Python files — passedgit diff --check— passedThe local host does not have
neuronx_distributed/torch_neuronx, so model import and device/runtime validation were not run.Validator request
Please validate this exact head with source-only checks and the focused suite above. Confirm that HF and legacy behavior remain unchanged, that normalized mode is explicit and fail-closed, and that the package identity is suitable for a guarded GPT-OSS retry. Do not merge, compile, launch hardware, or mutate live workloads. A 40/40 correctness result and any performance/tokenomics claims remain unpublished until a separately authorized retry passes.