#!/usr/bin/env python3
"""Standalone reproduction for:
GQAQKVColumnParallelLinearSpec._adapt_state_dict raises KeyError on any multi-file
checkpoint when tp_size > num_key_value_heads
Runs on plain CPU. No Trainium hardware, no neuronx-distributed install, and no 65 GB
checkpoint download are required. It only needs `torch`.
How it works
------------
`optimum/neuron/models/training/transformations_utils.py` is loaded directly from source
with its three Neuron-specific imports stubbed out, so the transformation code under test
is the real upstream code, byte for byte. The loader's behaviour is then replayed:
`NeuronModelMixin._load_pretrained_model` reads a sharded safetensors checkpoint one shard
file at a time and calls `adapt_state_dict` on each shard, and `adapt_state_dict` iterates
over every module of the model on every one of those calls.
The shard layout is the real one published for Qwen3-32B (17 files, taken from its
`model.safetensors.index.json`), embedded below so the script needs no network. The model
is dimensionally scaled down (real 64 layers / 64 query heads / 8 key-value heads / tp=32 /
kv_size_multiplier=4, but head_dim 8 instead of 128) so the whole thing fits in a few
hundred MB of RAM.
Three checks are run:
1. Against the unpatched source: the shard-by-shard load raises
KeyError('model.layers.3.self_attn.q_proj.weight').
2. Against the patched source: the shard-by-shard load produces a result that is
bit-identical (torch.equal) to a single-shot load of the complete state dict, for
every one of the 32 tensor-parallel ranks.
3. Against the patched source: `upstanding_sharded_params` is empty afterwards, i.e.
nothing was parked and silently forgotten, and an artificial layout in which q_proj
lands in a different shard file from k/v/o_proj is also handled.
If PATCH.diff sits next to this script and `git` is available, the patched source is
produced by actually applying that diff to a temporary copy, so a green run also proves
the diff applies. Otherwise checks 2 and 3 are skipped with a clear message.
See the bottom of this file for the exact command and the expected output.
"""
from __future__ import annotations
import argparse
import json
import os
import shutil
import subprocess
import sys
import tempfile
import types
import urllib.request
import torch
REL_PATH = "optimum/neuron/models/training/transformations_utils.py"
# ---------------------------------------------------------------------------------------
# Real Qwen3-32B shard layout: 1-based index of the safetensors file holding layer i's
# self_attn weights, from https://huggingface.co/Qwen/Qwen3-32B model.safetensors.index.json
# (17 files total). q_proj, k_proj, v_proj and o_proj of a given layer happen to always
# land in the same file for this checkpoint, which is worth knowing: the bug is NOT about
# a single module's weights straddling a shard boundary. Pass --refresh-index to re-derive
# this list from the Hub instead of trusting the embedded copy.
SHARD_OF_LAYER = [
1, 1, 1, 2, 2, 2, 2, 3, 3, 3, 3, 4, 4, 4, 4, 5, 5, 5, 5, 6, 6, 6, 6,
7, 7, 7, 7, 8, 8, 8, 8, 9, 9, 9, 9, 10, 10, 10, 10, 11, 11, 11, 11,
12, 12, 12, 12, 13, 13, 13, 13, 14, 14, 14, 14, 15, 15, 15, 15, 16,
16, 16, 16, 17,
]
# Qwen3-32B, with head_dim shrunk from 128 to 8 to keep the test small. Everything that the
# GQA sharding math depends on is the real value.
NUM_LAYERS = 64
NUM_HEADS = 64
NUM_KV_HEADS = 8
HEAD_DIM = 8
HIDDEN = NUM_HEADS * HEAD_DIM
TP_SIZE = 32
KV_MULT = max(1, TP_SIZE // NUM_KV_HEADS) # 4
Q_OUT_PER_PARTITION = NUM_HEADS * HEAD_DIM // TP_SIZE # 16
KV_OUT_PER_PARTITION = NUM_KV_HEADS * KV_MULT * HEAD_DIM // TP_SIZE # 8
PROJECTIONS = ("q_proj", "k_proj", "v_proj", "o_proj")
_CURRENT_RANK = [0]
# ---------------------------------------------------------------------------------------
# Stubs standing in for the Neuron-only imports at the top of transformations_utils.py.
# Nothing else in the file is touched.
def get_tensor_model_parallel_rank():
return _CURRENT_RANK[0]
def get_tensor_model_parallel_size():
return TP_SIZE
def is_peft_available():
return False
class _NullLogger:
def __getattr__(self, _name):
return lambda *a, **k: None
class _LoggingStub:
def get_logger(self, *a, **k):
return _NullLogger()
_logging = _LoggingStub()
def load_transformations_utils(path, module_name):
"""Exec the real upstream file with its Neuron imports replaced by the stubs above."""
with open(path) as fh:
source = fh.read()
substitutions = [
(
"from neuronx_distributed.parallel_layers.layers import create_local_weight",
"create_local_weight = None",
),
(
"from neuronx_distributed.parallel_layers.parallel_state import (\n"
" get_tensor_model_parallel_rank,\n"
" get_tensor_model_parallel_size,\n"
")",
"from __main__ import get_tensor_model_parallel_rank, get_tensor_model_parallel_size",
),
("from optimum.utils import logging", "from __main__ import _logging as logging"),
(
"from ...utils.import_utils import is_peft_available",
"from __main__ import is_peft_available",
),
]
for old, new in substitutions:
if old not in source:
raise SystemExit(
f"could not find the expected import block in {path}:\n {old!r}\n"
"The file layout has changed; this reproduction needs updating."
)
source = source.replace(old, new)
for line in source.splitlines():
if line.startswith(("from neuronx_distributed", "import neuronx_distributed",
"from optimum.utils", "from ...utils")):
raise SystemExit(f"import not stubbed: {line}")
module = types.ModuleType(module_name)
module.__dict__["__name__"] = module_name
# Register before exec: the dataclass machinery resolves annotations through
# sys.modules[cls.__module__] while the classes in this file are being created.
sys.modules[module_name] = module
exec(compile(source, path, "exec"), module.__dict__)
return module
# ---------------------------------------------------------------------------------------
def make_spec(tu):
return tu.GQAQKVColumnParallelLinearSpec(
gqa_qkv_projection_name="qkv_proj",
query_projection_name="q_proj",
key_projection_name="k_proj",
value_projection_name="v_proj",
output_projection_name="o_proj",
num_attention_heads=NUM_HEADS,
num_key_value_heads=NUM_KV_HEADS,
kv_size_multiplier=KV_MULT,
q_output_size_per_partition=Q_OUT_PER_PARTITION,
kv_output_size_per_partition=KV_OUT_PER_PARTITION,
fuse_qkv=False,
bias=False,
tp_size=TP_SIZE,
)
def build_full_state_dict():
gen = torch.Generator().manual_seed(0)
state_dict = {}
for layer in range(NUM_LAYERS):
prefix = f"model.layers.{layer}.self_attn"
state_dict[f"{prefix}.q_proj.weight"] = torch.randn(NUM_HEADS * HEAD_DIM, HIDDEN, generator=gen)
state_dict[f"{prefix}.k_proj.weight"] = torch.randn(NUM_KV_HEADS * HEAD_DIM, HIDDEN, generator=gen)
state_dict[f"{prefix}.v_proj.weight"] = torch.randn(NUM_KV_HEADS * HEAD_DIM, HIDDEN, generator=gen)
state_dict[f"{prefix}.o_proj.weight"] = torch.randn(HIDDEN, NUM_HEADS * HEAD_DIM, generator=gen)
return state_dict
def real_shard_layout():
"""[(shard_name, [keys]), ...] in file order, from the real Qwen3-32B index."""
per_file = {}
for layer, shard in enumerate(SHARD_OF_LAYER):
for projection in PROJECTIONS:
key = f"model.layers.{layer}.self_attn.{projection}.weight"
per_file.setdefault(shard, []).append(key)
return [(f"model-{s:05d}-of-00017.safetensors", per_file[s]) for s in sorted(per_file)]
def split_qkv_layout(layout):
"""Adversarial variant: q_proj arrives in a later shard than k/v/o_proj."""
out = []
for name, keys in layout:
rest = [k for k in keys if ".q_proj." not in k]
queries = [k for k in keys if ".q_proj." in k]
out.append((name + ".part-kvo", rest))
out.append((name + ".part-q", queries))
return out
def run_single_shot(tu, full_state_dict):
"""One call with the complete state dict: the reference result."""
spec = make_spec(tu)
state_dict = dict(full_state_dict)
upstanding = {}
for layer in range(NUM_LAYERS):
state_dict = spec._adapt_state_dict(
f"model.layers.{layer}.self_attn", {}, state_dict, upstanding, inplace=True
)
return state_dict, upstanding
def run_sharded(tu, full_state_dict, layout):
"""One call per shard file, over every module, exactly as the real loader does."""
spec = make_spec(tu)
upstanding = {}
result = {}
for _shard_name, keys in layout:
shard_state_dict = {k: full_state_dict[k] for k in keys}
for layer in range(NUM_LAYERS):
shard_state_dict = spec._adapt_state_dict(
f"model.layers.{layer}.self_attn", {}, shard_state_dict, upstanding, inplace=True
)
result.update(shard_state_dict)
return result, upstanding
def tensors_equal(reference, candidate):
if set(reference) != set(candidate):
only_ref = sorted(set(reference) - set(candidate))[:3]
only_cand = sorted(set(candidate) - set(reference))[:3]
return False, f"key mismatch, missing={only_ref} extra={only_cand}"
for key in reference:
if not torch.equal(reference[key], candidate[key]):
return False, f"tensor mismatch at {key}"
return True, ""
# ---------------------------------------------------------------------------------------
def locate_source(explicit):
if explicit:
return explicit
here = os.path.dirname(os.path.abspath(__file__))
# Walk up looking for a checkout of the repository.
probe = here
for _ in range(6):
candidate = os.path.join(probe, REL_PATH)
if os.path.isfile(candidate):
return candidate
probe = os.path.dirname(probe)
try:
import optimum.neuron # noqa: F401
candidate = os.path.join(
os.path.dirname(optimum.neuron.__file__),
"models/training/transformations_utils.py",
)
if os.path.isfile(candidate):
return candidate
except Exception:
pass
raise SystemExit(
f"could not find {REL_PATH}. Run this from a checkout of the repository, or pass "
"--source /path/to/transformations_utils.py"
)
def build_patched_copy(source_path, patch_path):
"""Apply PATCH.diff to a throwaway copy and return its path, or None."""
if not os.path.isfile(patch_path):
return None, f"no patch file at {patch_path}"
if shutil.which("git") is None:
return None, "git not available, cannot apply the patch"
tmpdir = tempfile.mkdtemp(prefix="gqa-repro-")
target = os.path.join(tmpdir, REL_PATH)
os.makedirs(os.path.dirname(target), exist_ok=True)
shutil.copyfile(source_path, target)
proc = subprocess.run(
["git", "apply", "--verbose", patch_path],
cwd=tmpdir, capture_output=True, text=True,
)
if proc.returncode != 0:
return None, f"git apply failed: {proc.stderr.strip()[:400]}"
return target, ""
def refresh_index():
url = "https://huggingface.co/Qwen/Qwen3-32B/resolve/main/model.safetensors.index.json"
with urllib.request.urlopen(url, timeout=60) as response:
weight_map = json.load(response)["weight_map"]
derived = []
for layer in range(NUM_LAYERS):
files = {
weight_map[f"model.layers.{layer}.self_attn.{p}.weight"] for p in PROJECTIONS
}
if len(files) != 1:
raise SystemExit(f"layer {layer} straddles shard files {sorted(files)}")
derived.append(int(sorted(files)[0].split("-")[1]))
if derived != SHARD_OF_LAYER:
print("NOTE: embedded SHARD_OF_LAYER differs from the Hub, using the Hub copy")
SHARD_OF_LAYER[:] = derived
else:
print("index check: embedded shard layout matches the Hub")
def main():
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--source", help="path to transformations_utils.py")
parser.add_argument("--patch", help="path to PATCH.diff")
parser.add_argument("--refresh-index", action="store_true",
help="re-derive the shard layout from the Hub (needs network)")
args = parser.parse_args()
if args.refresh_index:
refresh_index()
source_path = locate_source(args.source)
patch_path = args.patch or os.path.join(os.path.dirname(os.path.abspath(__file__)), "PATCH.diff")
print(f"source under test : {source_path}")
print(f"torch : {torch.__version__}")
print(f"model : {NUM_LAYERS} layers, {NUM_HEADS} q heads, "
f"{NUM_KV_HEADS} kv heads, tp={TP_SIZE}, kv_size_multiplier={KV_MULT}")
layout = real_shard_layout()
print(f"shard layout : real Qwen3-32B, {len(layout)} safetensors files")
print()
full_state_dict = build_full_state_dict()
failures = []
# ---- check 0: the query permutation covers every head exactly once ----------------
tu_plain = load_transformations_utils(source_path, "tu_plain")
spec_cls = tu_plain.GQAQKVColumnParallelLinearSpec
indices = torch.cat([
spec_cls.compute_query_indices_for_rank(TP_SIZE, r, NUM_HEADS, NUM_KV_HEADS, KV_MULT)
for r in range(TP_SIZE)
])
is_permutation = sorted(indices.tolist()) == list(range(NUM_HEADS))
print(f"[0] compute_query_indices_for_rank covers all {NUM_HEADS} heads exactly once "
f"across {TP_SIZE} ranks: {is_permutation}")
if not is_permutation:
failures.append("query index permutation")
# ---- check 1: reproduce the bug ---------------------------------------------------
_CURRENT_RANK[0] = 0
try:
run_sharded(tu_plain, full_state_dict, layout)
print("[1] unpatched sharded load: completed without error <-- BUG NOT REPRODUCED")
bug_reproduced = False
except KeyError as exc:
print(f"[1] unpatched sharded load raises KeyError({exc}) <-- bug reproduced")
bug_reproduced = True
if not bug_reproduced:
failures.append("bug did not reproduce on the unpatched source")
# ---- checks 2 and 3 need the patched source ---------------------------------------
patched_path, reason = build_patched_copy(source_path, patch_path)
if patched_path is None:
print()
print(f"[2] SKIPPED: {reason}")
print("[3] SKIPPED: same reason")
print()
print("RESULT: BUG REPRODUCED, fix not exercised "
"(place PATCH.diff next to this script to exercise it)")
return 0 if bug_reproduced else 1
print(f" patch applied cleanly to a temporary copy: {patch_path}")
tu_fixed = load_transformations_utils(patched_path, "tu_fixed")
all_ranks_ok = True
detail = ""
for rank in range(TP_SIZE):
_CURRENT_RANK[0] = rank
reference, _ = run_single_shot(tu_fixed, full_state_dict)
candidate, upstanding = run_sharded(tu_fixed, full_state_dict, layout)
ok, why = tensors_equal(reference, candidate)
if not ok:
all_ranks_ok, detail = False, f"rank {rank}: {why}"
break
if upstanding:
all_ranks_ok, detail = False, f"rank {rank}: {len(upstanding)} weights left parked"
break
print(f"[2] patched sharded load == single-shot load, bit-identical, "
f"for all {TP_SIZE} ranks: {all_ranks_ok}{(' (' + detail + ')') if detail else ''}")
if not all_ranks_ok:
failures.append("patched result differs from the single-shot reference")
_CURRENT_RANK[0] = 7
reference, _ = run_single_shot(tu_fixed, full_state_dict)
candidate, upstanding = run_sharded(tu_fixed, full_state_dict, split_qkv_layout(layout))
split_ok, why = tensors_equal(reference, candidate)
stash_empty = not upstanding
print(f"[3] patched load with q_proj deliberately placed in a different shard file "
f"from k/v/o_proj: identical={split_ok}{(' (' + why + ')') if why else ''}, "
f"nothing left parked={stash_empty}")
if not (split_ok and stash_empty):
failures.append("straddling-shard case mishandled")
if all_ranks_ok:
name = "model.layers.0.self_attn"
print()
print(" shapes on rank 0 for reference:")
_CURRENT_RANK[0] = 0
reference, _ = run_single_shot(tu_fixed, full_state_dict)
for suffix, expected in (
(f"{name}.qkv_proj.weight_q", (Q_OUT_PER_PARTITION, HIDDEN)),
(f"{name}.qkv_proj.weight_k", (KV_OUT_PER_PARTITION, HIDDEN)),
(f"{name}.qkv_proj.weight_v", (KV_OUT_PER_PARTITION, HIDDEN)),
(f"{name}.o_proj.weight", (HIDDEN, Q_OUT_PER_PARTITION)),
):
got = tuple(reference[suffix].shape)
print(f" {suffix:52s} {str(got):14s} expected {expected}")
print()
if failures:
print("RESULT: FAIL")
for item in failures:
print(f" - {item}")
return 1
print("RESULT: PASS (bug reproduced on the unpatched source, fixed by PATCH.diff)")
return 0
if __name__ == "__main__":
sys.exit(main())
# ---------------------------------------------------------------------------------------
# HOW TO RUN
#
# git clone https://github.com/huggingface/optimum-neuron.git
# cd optimum-neuron
# git checkout 4a80f2f3de15e83e978a6f3c0d43224626d921ca
# cp /path/to/reproduce.py /path/to/PATCH.diff .
# python reproduce.py
#
# Only torch is needed. Runs in well under a minute on CPU.
#
# EXPECTED OUTPUT
#
# source under test : .../optimum/neuron/models/training/transformations_utils.py
# torch : 2.8.0
# model : 64 layers, 64 q heads, 8 kv heads, tp=32, kv_size_multiplier=4
# shard layout : real Qwen3-32B, 17 safetensors files
#
# [0] compute_query_indices_for_rank covers all 64 heads exactly once across 32 ranks: True
# [1] unpatched sharded load raises KeyError('model.layers.3.self_attn.q_proj.weight') <-- bug reproduced
# patch applied cleanly to a temporary copy: ./PATCH.diff
# [2] patched sharded load == single-shot load, bit-identical, for all 32 ranks: True
# [3] patched load with q_proj deliberately placed in a different shard file from k/v/o_proj: identical=True, nothing left parked=True
#
# shapes on rank 0 for reference:
# model.layers.0.self_attn.qkv_proj.weight_q (16, 512) expected (16, 512)
# model.layers.0.self_attn.qkv_proj.weight_k (8, 512) expected (8, 512)
# model.layers.0.self_attn.qkv_proj.weight_v (8, 512) expected (8, 512)
# model.layers.0.self_attn.o_proj.weight (512, 16) expected (512, 16)
#
# RESULT: PASS (bug reproduced on the unpatched source, fixed by PATCH.diff)
#
# Without PATCH.diff present, checks 2 and 3 are skipped and the script reports
# "RESULT: BUG REPRODUCED, fix not exercised".
What happens
Full fine-tuning
Qwen/Qwen3-32Bon Trainium dies insideNeuronModelForCausalLM.from_pretrainedwithThe layer number is whichever layer first falls outside the first shard file, so it
depends on the checkpoint, and it can be
k_projorv_projinstead ofq_projdepending on which branch runs. I hit this on hardware at
tensor_parallel_size=16. Attensor_parallel_size=8the same run loads fine, because Qwen3-32B has 8 key-value headsand the GQA path is not taken at or below that. The attached script also reproduces it at
tensor_parallel_size=32, though I have not run that size on a device.Any model with a checkpoint split over more than one safetensors file should hit this
once TP exceeds the key-value head count. Qwen3-32B ships as 17 files, so it never gets
past the first one.
Reproduction
reproduce.py(attached) needs neither Trainium nor a 65 GB download, onlytorch. Itloads the real
transformations_utils.pywith the three Neuron-only imports stubbed,then replays the loader's shard-by-shard behaviour over the actual published Qwen3-32B
shard layout. The model is scaled down in
head_dimonly; the layer count, head counts,TP size and
kv_size_multiplierare the real ones.The traceback is short because the script calls the spec directly rather than going
through
from_pretrained:Cause
My first guess was that a layer's q/k/v had landed in different shard files. That is not
it. I checked all 64 layers of the published
model.safetensors.index.jsonandq_proj,k_proj,v_projando_projof a given layer are always in the same file.What actually happens is that the spec runs for every layer on every shard file.
_load_pretrained_modelreads the checkpoint one file at a time and calls
adapt_state_dicton each,and
adapt_state_dictwalks
model.named_modules()in full, not just the modules that have weights in theshard it was handed. So while shard file 1 is being processed, layer 3's spec runs too,
and line 1038
pops a key that is still sitting in shard file 2. The same applies to
1013,
1020,
1025,
1048
and 1054,
and to the
o_projlookup at1059,
which would raise the same way.
FusedLinearsSpec._adapt_state_dictalready deals with this, at482-491.
The GQA spec takes the same
upstanding_sharded_paramsargument at995
and never touches it.
The condition that scopes all of this is
llama/modeling_llama.py#L316,self.qkv_linear = (self.num_key_value_heads < tp_size) or (self.num_key_value_heads % tp_size != 0),which is what decides whether the spec gets built. Qwen3 inherits that attention
implementation, so it is affected too.
kv_size_multiplieris not involved. It is derived atconfig.py#L62and gives the right answer (4) for this model at TP=32. Setting it by hand changes
nothing.
Environment
I hit this on
optimum-neuron0.4.3,__sdk_version__2.26.1, Trainium (trn1), bf16,full fine-tune, TP 16. Before filing I diffed against
mainat4a80f2f3de15e83e978a6f3c0d43224626d921ca(__version__ = "0.4.6.dev4"):transformations_utils.pyis byte-identical between the two, 1686 lines each, so nothinghere has been fixed since 0.4.3.
reproduce.pywas run on CPU with torch 2.8.0 andPython 3.12.
Suggested fix
PATCH.diffis attached and applies cleanly to that commit. It gives the GQA spec theshard tolerance
FusedLinearsSpecalready has: skip the module when none of its fourweights are around yet, otherwise park the partial group in
upstanding_sharded_paramsand only transform once all four are in hand. 21 lines added, 1 moved, no reformatting
and no renames. The transformation math is untouched, and
reproduce.pyasserts withtorch.equalthat the patched shard-by-shard result matches a single-shot load of thecomplete state dict for all 32 TP ranks.
One detail worth flagging for review. The early
continuewhen nothing at all is presentis there so the common case never writes to
upstanding_sharded_params.FusedLinearsSpeccurrently assumes that dict holds at most one group at a time, per its length check at
488
and the
ValueErrorat496,
and parking unconditionally would break that assumption. With the early exit the patch
only parks when a module's weights genuinely straddle a shard boundary, which today is a
crash anyway. Relaxing the single-group assumption properly looked like a separate change
so I left it alone.
I also restore the weights into
state_dictand let the original code run, rather thanskipping the module, because a missing weight here is quiet:
_load_pretrained_modelonlywarns about missing keys at
694,
so a group that never completed would train uninitialised attention weights instead of
failing.
_lora_adapt_state_dictat1071
has the same unguarded pops on the LoRA A/B weights and looks like it would fail the same
way on a multi-file adapter checkpoint. I did not test that path and the patch does not
touch it, so treat that as a suspicion rather than part of this report.
reproduce.py