Fix tracing crashes for yolos, wav2vec2, hubert and convbert (aws-neuron-sdk#1265)

[jimburtoft]

Fixes the tracing crashes tracked in aws-neuron/aws-neuron-sdk#1265, which are what currently pin this project below Neuron SDK 2.27.

What was happening

yolos, wav2vec2, hubert and convbert don't raise during export — they kill the Python process, so there's nothing to catch.

All four are the same bug, and it isn't the individual ops. An aten op with no XLA lowering falls back to CPU through at::native::cpu_fallback, which calls XLANativeFunctions::_to_cpu() on the operands. That transfer only works if the operand has already been materialized into a PjRt buffer:

• operand derived from a model input → a buffer exists → fine
• operand derived from a parameter or buffer → torch-xla still holds an un-materialized placeholder whose buffer is nullptr → Check failed: pjrt_data->buffer != nullptr → SIGABRT or SIGSEGV

So a model only breaks when it applies an unlowered op to a weight rather than to an activation:

Model	Unlowered op	Operand
wav2vec2	aten::_weight_norm_interface	weight_norm parametrization on pos_conv_embed.conv
hubert	aten::_weight_norm_interface	same
yolos	aten::upsample_bicubic2d	position_embeddings
convbert	aten::im2col	conv-weight-derived tensor

The check: the same ops at the same shapes trace fine when fed a model input, and crash when fed a parameter or a registered buffer. Only provenance changes.

Two things worth correcting from the issue as filed:

• The two reported failure classes are one bug. The "segfault" and the RuntimeError: PjRt buffer is null are the same null-buffer dereference surfacing differently depending on timing. wav2vec2 was filed under segfault but actually aborts.
• cvt no longer reproduces on 2.31. It passes untouched, so it needs no patch and could be un-skipped in a separate PR.

What this PR does

Adds optimum/exporters/neuron/neuron_trace_patches.py, applied from export_neuronx, which rewrites each call site using ops that do have XLA lowerings so no CPU fallback is attempted:

• weight_norm — fold the parametrization so w = g * v / ||v|| is evaluated once eagerly instead of on every forward. Applied to every model rather than just these two: any architecture using weight normalization (speech encoders, vocoders, some GANs) hits this identically, and folding is a no-op otherwise.
• yolos — the bicubic interpolation depends only on a parameter and on img_size, which is fixed for a traced model, so it's constant with respect to the graph. Snapshot the embeddings before tracing, then precompute and cache. The snapshot matters: reading the live tensor mid-trace synchronizes an XLA tensor whose computation is still in flight and fails with Check failed: handle->HasValue().
• convbert — reimplement its narrow unfold case with pad/slice/stack/reshape. It raises on argument combinations it doesn't cover rather than silently computing something else.

All are monkey-patches applied at runtime, so no vendored transformers changes.

Verification

tests/exporters/test_transformers.py on SDK 2.31 / inf2.8xlarge, one pytest process per model so a crash is attributable rather than taking down the run:

	before	after
yolos	rc=139 (SIGSEGV)	pass
wav2vec2	rc=134 (SIGABRT)	pass
convbert	rc=139 (SIGSEGV)	pass
hubert	rc=134 (SIGABRT)	pass
cvt	pass	pass

Together: 18 passed, 36 skipped (previously 2 failed / 16 passed, and the run aborted partway when invoked as one process).

Numerical equivalence, against CPU, on the tiny-random fixtures and the real checkpoints:

Checkpoint	before	after	cos_sim	patched-vs-unpatched on CPU
facebook/wav2vec2-base-960h	SIGABRT	pass	1.0	0.0
facebook/hubert-base-ls960	SIGABRT	pass	1.000005	0.0
YituTech/conv-bert-base	SIGSEGV	pass	1.0	0.0
hustvl/yolos-tiny	SIGSEGV	pass	1.0 / 1.0 / 1.000004	0.0

The last column is the point: the patches are numerically inert. Residual cos_sim deviation is ordinary FP32 Neuron-vs-CPU noise. Also verified with dynamic_batch_size=True, since test_compare_to_transformers_dyn_bs was one of the failing tests.

New tests in tests/exporters/test_neuron_trace_patches.py (13 cases, CPU-only, no Neuron needed, ~13s): unfold_via_slices exactness vs nn.functional.unfold, one case per clause of its argument guard, weight-norm folding equivalence, and output equivalence for all four models.

requires-python

Also relaxes requires-python to <3.13. The previous <3.12 bound reflected the Python 3.10 venvs on older Neuron DLAMIs, but the SDK 2.31 DLAMI ships Python 3.12.3 — the package wouldn't install there at all. Everything above was run on 3.12.3 in the stock aws_neuronx_venv_pytorch_2_9 venv. Happy to split this into its own PR if you'd rather keep it separate.

Still worth fixing upstream

These are workarounds; the real issue is in torch-xla. _to_cpu() could materialize the placeholder instead of asserting on it — parameters are known constants at trace time. Failing that, replacing the CHECK with a thrown exception would at least make it catchable from Python instead of killing the process. I've noted this on the Neuron issue.

Environment

SDK 2.31 (DLAMI 20260708), inf2.8xlarge, torch-neuronx 2.9.0.2.15.32035, neuronx-cc 2.26.6360.0, torch 2.9.1, transformers 4.57.6, Python 3.12.3.