transformers
d49affb3 - Skip caching_allocator_warmup on Neuron (no reuse pool to warm; currently OOMs) (#47029)

Commit
1 day ago
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>
Author
Parents
Loading