Dokumen Catatan rilis Neuron 2.33 dari AWS mencantumkan 19 Oktober sebagai tanggal beta publik TorchNeuron. Pengumuman PyTorch, yang juga bertanggal 19 Oktober, menyatakan peluncuran beta publik dan kode sumber terbuka mulai tersedia pada 20 Oktober. Paket torch-neuronx mendaftarkan perangkat neuron ke PyTorch untuk Trainium2 dan Trainium3, sehingga operasi model dan tensor yang didukung dapat menggunakan API PyTorch yang sudah dikenal. Bagi insinyur yang mempertimbangkan chip ini, perubahan yang berguna adalah jalur yang lebih singkat dari program PyTorch yang sudah ada menuju uji coba pada perangkat. Apakah program itu berjalan dengan benar dan efisien tetap bergantung pada operasi, bentuk, dan lingkungan perangkat lunaknya.

Jalur PyTorch sebelumnya untuk Trainium menggunakan pustaka torch-neuronx berbasis XLA dan kode perangkat torch_xla. AWS menyatakan backend native menggunakan ekstensi PrivateUse1 milik PyTorch dan mengganti xm.xla_device() dengan penempatan pada perangkat neuron. Ini merupakan perubahan antarmuka yang berarti bagi proyek yang sudah menggunakan Neuron, meskipun nama paketnya familier. Rilis ini berlaku untuk pelatihan dan inferensi pada instans Trn2 dan Trn3; matriks dukungannya tidak mencantumkan Trn1.

Apa yang berubah dalam program PyTorch

Menurut catatan rilis AWS, sebuah model dipindahkan dengan .to("neuron") dan tensor masukan dibuat pada device="neuron". Operasi torch.nn, autograd, dan pengoptimal yang didukung kemudian dijalankan melalui backend. Kode yang ditulis untuk API PyTorch yang tidak bergantung pada perangkat, yaitu torch.accelerator, dapat mendeteksi Neuron tanpa penulisan ulang khusus perangkat, menurut ikhtisar TorchNeuron.

Ada dua jalur eksekusi. Mode eager menjalankan operasi PyTorch saat dipanggil Python, dengan mempertahankan alur kontrol, keluaran cetak, dan breakpoint biasa. Setiap operasi dikompilasi saat pertama digunakan, sedangkan pemanggilan berikutnya yang cocok dapat menggunakan cache. Insinyur juga dapat mengompilasi model atau bagian tertentu dengan torch.compile(model, backend="neuron", dynamic=False). Backend native mengompilasi graf forward dan backward yang didukung; kode di luar bagian yang ditangkap tetap berjalan dalam mode eager. Pengumuman PyTorch juga menjelaskan pengaturan torch.compiler.set_default_backend("neuron") sekali sebelum kompilasi, sehingga pemanggilan kompilasi berikutnya dapat menghilangkan argumen backend.

Rilis ini mencakup backend neuron untuk torch.distributed, dengan operasi kolektif dan titik ke titik yang terdokumentasi. AWS mencantumkan DistributedDataParallel, FSDP2, DTensor, DeviceMesh, serta paralelisme pipeline sebagai API yang didukung. Jalur-jalur itu berguna untuk ditelaah dalam proyek pelatihan terdistribusi, tetapi catatan rilisnya tidak membuktikan bahwa pekerjaan multinode yang sudah ada akan berjalan tanpa perubahan atau mencapai sasaran throughput tertentu.

Evaluasi dimulai dengan lingkungan yang sesuai

Dokumen panduan instalasi cepat AWS mensyaratkan instans Trn2 atau Trn3 serta Neuron Base DLAMI atau Ubuntu AMI yang didukung. Matriks rilisnya menetapkan Linux x86_64 pada Ubuntu 24.04 atau 26.04. Dukungan Ubuntu 22.04 berakhir untuk rilis ini. Python 3.12, 3.13, dan 3.14 tercantum, dengan torch-neuronx 2.13.3.0 untuk PyTorch 2.13.0 atau torch-neuronx 2.14.3.0 untuk PyTorch 2.14.0. Keluarga paket harus sesuai dengan versi PyTorch.

Panduan cepat itu memasang torch-neuronx==2.14.* bersama neuronx-cc==2.* dan nki dari indeks paket Neuron dan PyTorch. Setelah itu, panduan memeriksa ketersediaan perangkat dan membandingkan operasi matriks kecil pada Neuron dengan hasil CPU. Langkah tersebut memverifikasi jalur perangkat yang terpasang dalam contoh AWS. Itu tidak memvalidasi model pembaca. Evaluasi lanjutan yang masuk akal adalah menjalankan jalur forward dan backward model yang sebenarnya, memeriksa fallback CPU yang tercatat dengan torch_neuronx.get_fallback_ops(), membandingkan keluaran dengan lingkungan saat ini, serta mengukur kompilasi dan eksekusi kondisi tunak secara terpisah. Langkah-langkah terakhir itu merupakan pertanyaan evaluasi yang diajukan BIG CHANGE, bukan pengujian yang kami lakukan.

Bagian pekerjaan yang masih dapat berubah karena beta

Dalam mode eager, operasi yang tidak didukung dapat beralih ke CPU dan mengeluarkan peringatan. Pada torch.compile, operasi yang ditangkap dalam graf tetapi tidak dapat diturunkan ke bentuk eksekusi oleh backend Neuron akan menggagalkan kompilasi secara default. AWS mengarahkan pengembang ke tabel kompatibilitas operator untuk melihat cakupan per operasi. Karena itu, eksekusi eager yang selesai tidak dengan sendirinya membuktikan bahwa seluruh eksekusi berlangsung pada perangkat atau kompatibel dengan kompilasi.

Rilis tersebut juga menyatakan torch.compile memerlukan dynamic=False; bentuk masukan baru dapat memerlukan kompilasi tambahan. AWS menyarankan padding atau pengelompokan ke bucket untuk mengurangi variasi bentuk. Eksekusi pertama mencakup kompilasi jika tidak ada artefak yang cocok dalam cache. Rincian ini penting saat mengukur waktu beban kerja: latensi kompilasi dan pola bentuk masukan dapat mengubah hasil evaluasi.

Batas eksplisit lainnya mempersempit klaim kompatibilitas. Beta ini tidak mendukung torch.sparse, torch.func, torch.nested, atau torch.distributions. Hanya beberapa transformasi torch.fft yang berjalan pada Neuron, dan float64 dijalankan sebagai float32 dengan peringatan. Integrasi TorchAO dan dukungan tipe data komposit MXFP8 tidak tersedia dalam rilis ini. Tim yang menggunakan antarmuka tersebut perlu memeriksa jalur kodenya sendiri sebelum menganggap penempatan pada perangkat sebagai rencana migrasi.

Perubahan besarnya

TorchNeuron menempatkan Trainium di balik antarmuka perangkat dan kompilasi native PyTorch dalam beta publik. Hal itu membuat uji coba pertama lebih langsung bagi program PyTorch yang didukung serta memberi pengguna Neuron saat ini jalur berbeda dari kode berbasis XLA mereka. Bukti penentu adopsi akan berasal dari setiap beban kerja: apakah operatornya tetap berada pada perangkat yang dimaksud, apakah hasil numeriknya tetap dapat diterima, dan apakah kinerja serta biaya yang terukur sepadan dengan instans dan pekerjaan migrasi. BIG CHANGE menelaah dokumentasi dan pengumuman untuk analisis ini; kami tidak menjalankan TorchNeuron atau memverifikasi klaim kinerja AWS secara independen.

Sumber dan bacaan lanjutan