Skip caching_allocator_warmup on Neuron (no reuse pool to warm; currently OOMs)

[dacorvo]

What does this PR do?

caching_allocator_warmup pre-allocates get_total_byte_count bytes per device in a single
torch.empty(...). The cuda/xpu branch bounds that allocation against device capacity via
mem_get_info; mps is skipped (#46239). Every other accelerator device — including neuron
(AWS Trainium/Inferentia, a PrivateUse1 backend) — falls through to the unclamped allocation, which
OOMs the NeuronCore during from_pretrained, before any weight is loaded:

RuntimeError: NRT allocation failed for size 233658313344 (OOM)

Concrete failures (TP load, bf16, Trainium2 trn2.48xlarge, ~24 GB per NeuronCore):

Model	TP	Allocation attempted	Per-core budget
openai/gpt-oss-120b	16	233.66 GB (≈116.83B params × bf16)	24 GB
Qwen/Qwen3-VL-32B-Instruct	8	66.71 GB	24 GB
CohereLabs/North-Mini-Code-1.0	8	59.32 GB	24 GB
openai/gpt-oss-20b	4	41.83 GB	24 GB

HBM is ~0 GB at the point of failure — this is the warmup pre-allocation, before any weight is placed.
With the skip, loading proceeds normally: the spurious warmup allocation is gone and the model loads up to
its real per-device footprint. In a full support sweep, models that previously failed here in warmup now
proceed to actual load/compile — several pass (e.g. Qwen/Qwen2.5-32B-Instruct tp8), and any that still
run out of memory do so on their genuine footprint, not on a bogus full-model pre-allocation.

Why skip on Neuron? (measured on Trainium2, Neuron SDK)

Two independent reasons, both verified empirically:

1. The warmup provides no benefit. Neuron keeps no CUDA-style reuse pool for a warm-up to fill:
   torch_neuronx.memory_stats() reports reserved_bytes = 0 even while 2 GB is allocated; alloc/free are
   1:1 with the NRT runtime; and pre-warming a 4 GB pool then allocating 256×16 MB is the same speed as cold
   (2.9 ms vs 3.0 ms, ~1.0×). The warmup's premise — amortize many slow device mallocs into one — does not
   hold; Neuron allocation is a cheap direct runtime call.

2. The allocation can't be bounded, so it OOMs. torch.neuron exposes no mem_get_info, so the cuda/xpu
   clamp (min(byte_count, total_device_memory - 1.2 GiB)) can't run. Combined with get_total_byte_count
   over-estimating under tp_plan="auto" (params the auto plan doesn't cover — MoE experts, vision towers —
   stay at full unsharded size), the unclamped torch.empty requests ~the whole model on one core.

Skipping therefore costs nothing and fixes the OOM — the same trade-off as the MPS skip (#46239).

Notes / alternatives

• Minimal, precedent-following fix (one new elif, mirroring MPS).
• A cleaner generalization would skip warmup on any accelerator whose module lacks mem_get_info (covers
  Neuron and future backends). Happy to take that shape if preferred.
• get_total_byte_count over-estimating under tp_plan="auto" is a latent over-reservation on cuda/xpu too
  (masked there by the clamp); out of scope here.

Before submitting

• Read the contributor guideline.
• Verified on AWS Trainium2: the OOM repro + allocator/timing measurements above, and a full support
  sweep with the skip applied — the warmup allocations are gone, models load up to their real footprint,
  and TP>1 models that previously passed still pass (no regression). transformers CI has no Neuron runner,
  as with the MPS skip.

Who can review?

@Cyrilvallez @SunMarc (reviewed the related warmup PRs #36380 / #46145; cf. the MPS skip #46239)

[github-actions]

---

CI recap

Dashboard: View test results in Grafana
Latest run: 28611015907
Result: success | Grafana metrics are not available yet.

[HuggingFaceDocBuilderDev]

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.

[dacorvo]

Pushed without my consent by an overzealous agent, sorry about that.