Pemindahan model Google Gemma-4 31B ke AWS Inferentia2 inf2.24xlarge menghasilkan kecocokan token-per-token yang sempurna terhadap referensi CPU—namun setiap kalimat yang dihasilkan adalah omong kosong. Kesenjangan antara “cocok” dan “berfungsi” kini menjadi peringatan bagi siapa pun yang mencoba memasukkan LLM masif ke dalam chip inferensi khusus Amazon.

Mengapa kecocokan token-per-token saja tidak cukup

Pengembang membandingkan setiap token output dari perangkat Inferentia dengan token yang dihasilkan oleh model saat dijalankan di CPU. Aliran datanya identik, sehingga perangkat keras tampak telah mereproduksi implementasi referensi secara persis. Kenyataannya, kedua aliran tersebut memasukkan prompt yang salah format ke dalam model yang kehilangan chat template-nya dan diberikan penanda giliran (turn markers) yang salah. Template yang hilang membuat model terjebak dalam loop tak terbatas, menghasilkan teks yang tidak masuk akal. Perangkat keras telah menjalankan tugasnya—ia mereproduksi bug yang ada pada kode referensi.

Pelajaran sederhananya adalah: SEQ_MATCH (kesamaan token sekuensial) tidak sama dengan kebenaran (correctness). Jika implementasi referensi rusak, replika perangkat keras yang setia akan mewarisi kegagalan yang sama. Validasi harus melampaui paritas tingkat token; diperlukan pemeriksaan fungsional end-to-end dengan input yang diformat dengan benar.

Buffer yang menyamar sebagai parameter

Selama fase pemuatan, model loader melewatkan komponen bernama layer_scalar. Kode tersebut mendaftarkan objek ini sebagai buffer dan bukan sebagai parameter dalam definisi model PyTorch. Buffer adalah tensor statis yang tidak diperbarui selama pelatihan, dan banyak loader mengabaikannya saat mengonversi ke format yang kompatibel dengan Neuron. Melewatkannya menyebabkan faktor penskalaan untuk beberapa lapisan tetap pada nilai default, sehingga mendistorsi perhitungan matematika di seluruh jaringan. Tidak ada error yang muncul; model berhasil dikompilasi, dan pipeline inferensi berjalan, tetapi hasil numeriknya tidak akurat.

Bagi siapa pun yang memindahkan model besar ke Inferentia, auditlah setiap tensor non-parameter. Meskipun sebuah tensor tidak dimaksudkan untuk dipelajari, ia mungkin tetap penting untuk komputasi forward-pass yang benar. Memverifikasi penyertaan buffer secara manual dapat mencegah kesalahan skala senyap (silent scale errors) yang sulit didiagnosis.

Volatilitas spot-instance dan kompilasi selama 39 menit

Menjalankan model dengan 31 miliar parameter pada spot instance terlihat murah, tetapi penghematannya disertai dengan peristiwa pengambilan kembali (reclaim) yang tidak terduga. Waktu kompilasi pengembang—sekitar 39 menit untuk menerjemahkan model ke dalam kode yang kompatibel dengan Neuron—lenyap saat AWS mengambil kembali instance tersebut. Untuk bertahan dari gangguan, mereka membangun jaring pengaman tiga lapis:

  • ModelBuilder menjaga penggunaan memori tetap dalam batas host 384 GB, menghindari crash yang akan memaksa restart.
  • Mirroring S3 segera untuk file bobot mentah maupun “neffs” (file eksekusi Neuron) yang telah dikompilasi memungkinkan instance baru melanjutkan tepat di tempat instance sebelumnya berhenti.
  • Sebuah multi-region poller memindai wilayah AWS untuk kapasitas spot yang tersedia dan meluncurkan instance baru segera setelah tersedia.

Langkah-langkah ini mengubah kompilasi titik-tunggal yang rapuh menjadi pipeline tangguh yang mampu bertahan dalam fluktuasi pasar spot.

Jebakan sharding dengan tata letak attention yang bercampur

Gemma-4 31B menggunakan dua konfigurasi attention. Beberapa lapisan menggunakan empat key-value (KV) heads, yang lain menggunakan jumlah yang berbeda. Membagi model secara merata ke dalam delapan rank paralel akan gagal jika jumlah KV head suatu lapisan tidak dapat dibagi habis. Mencoba melakukan sharding pada lapisan dengan 4 head ke dalam delapan rank akan memaksa setiap rank untuk menangani setengah head—sebuah ketidakmungkinan matematis yang memicu ketidakcocokan bentuk (shape mismatches) dan error saat runtime.

Solusinya adalah dengan mereplikasi lapisan yang di-shard secara global (lapisan dengan jumlah head yang kompatibel) ke semua rank dan hanya melakukan sharding pada lapisan “sliding” yang jumlah head-nya memungkinkan pembagian merata. Strategi hibrida ini menjaga efisiensi tensor-parallel sambil menghindari pembagian KV head yang ilegal, menghilangkan error tensor-parallelization yang menghambat upaya sebelumnya.

Kesimpulan

Memindahkan LLM raksasa ke Inferentia lebih dari sekadar latihan kompilasi-dan-jalankan. Hal ini menuntut pengujian fungsional yang ketat melampaui sekadar kesamaan token, verifikasi teliti bahwa setiap tensor—baik parameter maupun buffer—ditangani dengan benar, dan strategi deployment yang mengantisipasi pengambilan kembali spot-instance. Terakhir, sharding harus menghormati geometri attention internal model; jika tidak, paralelisme yang menjanjikan kecepatan justru akan menjadi sumber kegagalan senyap.