chore(sdk): upgrade to Neuron SDK 2.31 - #1113
Open
tengomucho wants to merge 22 commits into
Open
Conversation
|
The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update. |
Align the neuronx extra with the versions shipped in the SDK 2.31 DLC: - neuronx-cc 2.21.33363.0 -> 2.26.6360.0 - torch-neuronx 2.8.0.2.10.16998 -> 2.9.0.2.15.32035 - torch 2.8.0 -> 2.9.1, torchvision 0.23 -> 0.24 - neuronx_distributed 0.15.22404 -> 0.19.28492 - libneuronxla 2.2.12677.0 -> 2.2.17544.0 - numpy upper bound 1.26.4 -> 2.4.6 The 2.26 compiler only ships cp311/cp312/cp313 wheels, so drop Python 3.10 and advertise 3.12. Bump torchcodec to 0.8.1 as well: 0.7.0 is ABI-pinned to torch 2.8 and aborts the process with std::bad_alloc under torch 2.9.1. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
The SDK version is part of the test model hub repository names, so it must be updated along with the dependencies to avoid reusing artifacts compiled with the previous compiler. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
torch_neuronx keeps a process-global HLO instance counter that gets baked into the compiled module hash, so exporting the same model twice in one process always produces different hashes and the Hub cache lookup never hits. Run each export in its own subprocess.
_sample computed top_k=1 for greedy, but built the CPU logits warper from the raw generation_config, so greedy decoding actually ran topk(logits, 50) and took element 0. On an exact bfloat16 tie that breaks towards whatever the topk sort returns instead of the lowest token id, which is what transformers does. token_selector.select() and _assisted_decoding already used a plain argmax, so the same model could disagree with itself between speculative and regular greedy decoding. Select from the full logits when do_sample is False, and keep the fused warper on the sampling path only. This is host side: the graph inputs are unchanged and sampling_params stays (batch_size, 3) whatever the sampling parameters are, so no recompilation is triggered. qwen2-4x1024 and qwen2-1x8192 now reproduce the CPU output exactly, and granite-4x1024 and qwen3-1x8192 turned out to already match, so their entries were stale. The greedy expectations go from 9 passed/4 xfailed to 11 passed/2 xfailed. qwen3-4x1024 and qwen3-tp1-4x1024 still differ, and no modeling change can help: CPU resolves that token with a 0.0317 logit gap, a quarter of a bfloat16 ULP at that magnitude (0.125). Both now generate what the CPU model generates in bfloat16. The gemma3 long-sequence xfail is kept as is, but its margin is 3.44 ULP, so unlike the others it is not a rounding tie and still needs an explanation. Also run test_speculation_same_model in a subprocess: loading a compiled model onto NeuronCores sets the Neuron runtime world_size for the process lifetime, so the preceding tests in the file make it fail with "Could not load the model" when the whole file is run. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Neuron picks CPU's second-best candidate at the third generated token (17.10 vs 16.93 logits, a near-tie), producing coherent but different text for smolvlm-16-images-cross-chunk. Xfail on the known observed string, same pattern as the decoder greedy tests.
Moves import paths (FlexibleArgumentParser, WorkerBase, set_random_seed), ModelConfig.task -> runner_type, and the CLI's --task -> --runner/--convert split. Also fixes behavioral breaks the import-only survey missed: WorkerBase.__init__ now requires local_rank/rank/distributed_init_method directly, load_model must run inside set_current_vllm_config (CustomOp instantiation asserts on it), SchedulerConfig lost max_model_len (moved to ModelConfig), and 0.16 defaults to async scheduling, which needs an execute_model/ sample_tokens split our worker doesn't implement -- disabled it. pyproject.toml's vllm pin stays uncommitted for now to avoid paying the full-suite re-export cost more than once.
Bump CI's install_neuronx_runtime action and the vLLM Dockerfile to the SDK 2.31 DLC runtime versions (tools/runtime-lib/collectives, plus dkms in the Dockerfile only -- the CI runner gets its driver from the AMI). Verified via a real image build: - Ubuntu 22.04 ships Python 3.10; optimum-neuron needs >=3.11 since SDK 2.31 ships no cp310 wheels. Provision Python 3.12 via uv (not 3.11: neuronx-cc pins numpy<2 there, clashing with vllm's numpy>=2). - torch_xla's compiled extension needs libpython on the linker path, which uv's standalone interpreter ships but doesn't register. - neuronx-cc's compiler binary now needs libarchive13, not in the base image. Also add .dockerignore: with no venv/.git exclusions, docker build was sending the whole repo (7GB+ locally) as build context, which likely also slows down CI's vllm-docker-tests job since setup_venv creates a venv in the repo root before the image build runs. Also ignore tests/PostSPMDPassesExecutionDuration.txt, a new compiler artifact under SDK 2.31.
vllm==0.11.0 requires torch==2.8.0, unresolvable against the SDK 2.31 torch==2.9.1 pin, so every vLLM CI job was failing at install. 0.16.0 (and the 0.14-0.16 range) pins torch==2.9.1/torchvision==0.24.1, matching what's already required -- no torch change needed. This is the only pyproject.toml edit in this migration, landed last since it invalidates the decoder test-model re-export cache (~1h, once).
The large-d NKI flash attention kernel stored the softmax numerator in bfloat16 before the PV matmul, quantizing every attention weight to 8 mantissa bits. That error accumulates over attended positions, which is why it only ever hit a long prompt: under SDK 2.31 it was enough to change a sampled token at ~5k tokens, and the long sequence test was marked as a known divergence from the CPU float32 reference. Keep p_local, its reduction output and p_local_transposed in fp32 so full precision reaches the PV matmul. Measured on gemma3-270m at batch 1 / sequence 8192, first divergence over 50 greedy tokens and median prefill latency: bfloat16 numerator step 2 (3.44 ULP) 142 ms fp32 numerator only step 2 (bit identical) 161 ms flash kernel disabled step 20 (1.15 ULP) 317 ms fp32 through matmul no divergence 222 ms Keeping only the matmul operand in bfloat16 reproduces the original output bit for bit, which localizes the loss to that operand rather than to the reduction or the softmax denominator. Prefill is ~1.6x slower than before and still ~1.4x faster than the compiler-native path. SBUF grows by ~512 KB, independent of head_dim. Only models reaching the large-d kernel are affected: head_dim > 128 with a prefill of at least 4096 (LNC1) or 2048 (LNC2). The generation now matches the reference, so the long sequence test no longer needs its known-divergence escape hatch.
check_and_update_config assigned a closure to model_config.verify_with_parallel_config. Closures cannot be pickled, so when vLLM spawns the EngineCore process it fails with "Can't pickle local object". Replace it with a module-level function so the patched config survives pickling.
vLLM's V1 engine forks an EngineCore process by default (VLLM_WORKER_MULTIPROC_METHOD=fork). Forking after torch and the Neuron runtime have initialized their native thread pools leaves the child with a dead neuron::ThreadPool, so weight loading deadlocks in neuron::parallel_load (pthread_barrier_wait) or aborts with "Invalid thread pool!". Force spawn in the platform plugin's register() so every vLLM-on-Neuron usage starts EngineCore from a clean interpreter. setdefault respects an explicit user override. This mirrors what optimum-cli neuron serve already does. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
convbert, hubert, wav2vec2 and yolos segfault/abort inside torch_neuronx HLO generation with Neuron SDK 2.31 (torch-neuronx 2.9 / torch-xla 2.9). The crash is in the compiler's tracer, not in optimum-neuron code, so it cannot be caught and kills the whole pytest process. Skip these model types before export until the SDK is fixed. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
The SDK 2.31 DLC is based on Ubuntu 24.04, and the 22.04 runners no longer match the environment the packages are built for. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
The SDK 2.31 neuronx-cc only ships cp311/cp312/cp313 wheels, and Ubuntu 24.04 runners ship system Python 3.12. Pin the setup_venv and sanity-check venvs to 3.12. Match the CPU torch wheel with the SDK's torch 2.9.1 (previously 2.8). The CPU wheel sanity-check is cached against these package versions. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
neuronx-cc's weight-layout-optimization step compiles NKI kernels via torch_neuronx.xla_impl.trace.hlo_compile, which runs the compiler with subprocess.run(command) and no cwd= (trace.py). The walrus backend then materializes content-addressed neuronxcc.private_nkl.* kernel dirs in the process CWD, dirtying the caller's working directory on every cache-miss export. Wrap builder.trace() in a context manager that chdirs to a throwaway temp dir during compilation so those droppings are discarded. Verified on inf2: exporting katuni4ka/tiny-random-phi3 (fp32, cold NKI cache) without the fix leaks a 16-hex kernel dir into the CWD; with the fix the same fresh compile leaves the CWD clean. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
The Ubuntu 24.04 inf2 runner lacks two shared libs SDK 2.31 needs: libpython3.12.so.1.0 (torch_xla._XLAC import) and libarchive.so.13 (neuronxcc walrus_driver). Verified green on run 32349690521. Co-Authored-By: Claude <noreply@anthropic.com>
`T5DecoderWrapper.reorder_cache` was called with the module's own
`ParameterList` holding the KV cache, and assigned the gathered tensors
back into it. During the trace this unregisters those parameters,
replacing them with intermediate XLA tensors.
`torch_neuronx` swaps the parameters for XLA placeholders while tracing
and restores them afterwards by walking `named_parameters()`, which is
also how it restores the input/output alias keys. Since the cache
parameters were gone by then, the aliases kept their placeholder keys,
and `parallel_model_trace` failed to send them back to the parent
process through its multiprocessing queue:
RuntimeError: _share_filename_: only available on CPU
Build and return a new list instead. The caller already used the return
value and never relied on the in-place assignment, so the traced graph
is unchanged.
This only showed up when exporting with both `num_beams > 1` (the only
path calling `reorder_cache`) and `tensor_parallel_size > 1` (the only
path using a multiprocessing queue), i.e. `test_encoder_decoder_tp2`.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
The list of model types crashing the Neuron SDK 2.31 tracer lived in `tests/exporters/test_transformers.py`, but the same crash affects other test suites. Move it next to the other test helpers and expose a `skip_if_sdk_231_trace_crash()` guard so every suite can reuse it instead of duplicating the list. No behaviour change: the exporters tests skip exactly the same models. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
The conv-based models that crash the Neuron SDK 2.31 tracer were only skipped in the exporters tests, so the transformers inference and pipeline suites still exported them and died with SIGSEGV, taking the whole pytest process down: - `pytest tests/pipelines` -> convbert, in `F.unfold` - `pytest -m "not slow" tests/inference/transformers/test_modeling.py` -> yolos, in `F.interpolate` - `pytest -m slow tests/inference/transformers/test_modeling.py` -> convbert, in `F.unfold` Guard the shared `NeuronModelTestMixin._setup()`, which every modeling test goes through, and the `inf_encoder_model` fixture used by the pipeline tests. The latter is now parametrized on the architecture instead of the model id so the skip can be keyed on it. Both suites pass on main, i.e. with SDK 2.30, so only the tracer version changed. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
The `vllm_docker_launcher` context manager stopped and removed the
container after its `yield`, without a `try`/`finally`. When a test
raised inside the `with` block, the exception was thrown back in at the
yield and the whole teardown was skipped, leaving the container running
and holding its Neuron cores.
The next test using the same device then died during startup:
NRT:nrt_allocate_neuron_cores Logical Neuron Core(s) not available
- Requested:lnc8-lnc9 Available:0 Logical Core size:1
(cores busy, ret=-16)
RuntimeError: The PyTorch Neuron Runtime could not be initialized.
So a single test failure cascaded into the next one. Wrap the yield so
the container is always stopped and removed.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
`test_vllm_docker_service_sampling_parameters` asserts that a sampled answer differs from the greedy one. The output distribution of Llama-3.2-1B on this prompt is peaked enough that a single draw at temperature 1.0 / top_p 0.9 sometimes reproduces the greedy answer, so the assertion is flaky. Bumping the temperature from 0.8 to 1.0 in 7d9d1ac made it rarer but did not remove it. Draw up to five samples and stop as soon as one differs, which is what the test actually means to check: that the sampling parameters are taken into account. A single draw matching greedy is not a failure. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
tengomucho
marked this pull request as ready for review
August 24, 2026 07:55
JingyaHuang
reviewed
Aug 24, 2026
| # Install uv at a specific version on a given path | ||
| RUN curl -LsSf https://astral.sh/uv/0.9.27/install.sh | XDG_BIN_HOME=/usr/local/bin sh | ||
|
|
||
| # optimum-neuron requires Python >= 3.11 (SDK 2.31's compiler ships no cp310 wheels), |
Collaborator
There was a problem hiding this comment.
Why don't we upgrade to ubuntu24?
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Draft PR to move optimum-neuron to Neuron SDK 2.31.0 and vLLM 0.16.0.
Note that part of the work in this PR has been created with the help of AI, that was used also to keep the CI working.
Changes
7bb4756f) and align CI Neuron runtime pins (d2dbef4d); bump package version to 0.4.7.dev0 withsdk_version = 2.31.0(744e2e3f); move vLLM to 0.16.0 (f0105ecb).70066cb7); force spawn for EngineCore to avoid an NRT fork deadlock (3e0bac3f) and make the ModelConfig parallel-config patch picklable for spawn (a91b7982).ff1fb2ff).146d3721).fdee2751); add known-divergence escape for the VLM cross-chunk case (94d11a9e).Known issue (needs investigation)
The SDK 2.31 tracer segfaults/aborts on several conv-based encoder models (convbert, hubert, wav2vec2, yolos). This is an upstream compiler regression, not branch code — affected exporter tests are skipped for now in
26508ed5.This might be related to and fixed by #1119, but I think it is better to work on that after this branch has been merged.