JAX telah muncul sebagai alat yang ampuh bagi para inovator AI, memungkinkan pelatihan model kompleks di berbagai aplikasi. Performa dan skalabilitasnya sangat bermanfaat bagi perusahaan seperti Lightricks, yang sedang merintis kemajuan dalam ekonomi kreator dengan model pembangkitan video berkinerja tinggi. Tim Lightricks LTX-Video, yang dipimpin oleh Yoav HaCohen dan Yaki Bitterman, menghadapi tantangan penskalaan yang signifikan dengan kerangka kerja mereka sebelumnya, yang mereka atasi dengan bermigrasi ke JAX.

Lightricks ingin meningkatkan pembuatan konten dengan Generative AI , tetapi menemui kendala saat mengembangkan model dengan 13 miliar parameter. Kerangka kerja tim sebelumnya tidak dapat menyediakan jumlah langkah pelatihan dan skalabilitas yang dibutuhkan. Perpindahan strategis ke JAX menghasilkan peningkatan langsung, yang secara efektif meningkatkan skala dan memungkinkan peningkatan 40% dalam jumlah langkah pelatihan per hari. Perubahan ini tidak hanya mempercepat pengembangan model, tetapi juga menggandakan produktivitas tim.

Kesuksesan JAX berkat ekosistem komprehensif yang mencakup perangkat-perangkat seperti Flax, Optax, dan Orbax, yang semuanya bekerja sama dengan mulus. Lightricks memanfaatkan perangkat-perangkat ini untuk membangun tumpukan baru yang tangguh demi R&D yang efisien. Lightricks juga memanfaatkan inti TPU Google untuk pelatihan, memanfaatkan sepenuhnya desain sistem terintegrasi untuk mencapai kinerja optimal.

Pengalaman Lightricks dengan JAX telah membuktikan dengan kuat kemampuan kerangka kerja ini dalam membantu tim mengatasi tantangan teknis dan mencapai peningkatan kinerja yang signifikan. Tim JAX di Google tetap berkomitmen untuk mendukung komunitas komputasi ilmiah dan mendorong pengguna untuk berbagi pengalaman dan berkontribusi pada perkembangan kerangka kerja ini.