Skip caching_allocator_warmup on Neuron (no reuse pool to warm; currently OOMs) (#47029)
* Skip caching_allocator_warmup on Neuron
caching_allocator_warmup pre-allocates get_total_byte_count bytes per device in
a single torch.empty(). The cuda/xpu branch clamps that with mem_get_info and
mps is skipped (#46239); every other accelerator, including neuron (a
PrivateUse1 backend), falls through to the unclamped allocation.
On Neuron this OOMs the NeuronCore during from_pretrained, before any weight is
placed, and provides no benefit: Neuron keeps no CUDA-style reuse pool for the
warmup to fill (torch_neuronx.memory_stats() reports reserved_bytes == 0 even
while memory is allocated, and pre-warming shows no measured speedup), and
torch.neuron exposes no mem_get_info so the allocation cannot be bounded.
Skip warmup on neuron, mirroring the mps skip.
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
* fix: clean up comments for Neuron device handling
Remove outdated comments regarding Neuron device warmup.
* fix comment
---------
Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>
Co-authored-by: Cyril Vallez <cyril.vallez@gmail.com>