MatrixFSDP: pengoptimal matriks tanpa komunikasi di bawah sharding parameter ZeRO-3

Perancangan ulang tata letak sharding yang memungkinkan Muon berjalan lokal pada ZeRO-3 tanpa kolektif matriks pada langkah pengoptimal

Ming Gao; Yanwu Xu; Hao Zhang· 2026· DOI 10.48550/arXiv.2607.05895

Masalah inti

Pengoptimal matriks seperti Muon menarik untuk pelatihan skala besar karena dapat meningkatkan konvergensi dan efisiensi token dibandingkan pengoptimal coordinate-wise. Muon melakukan ini dengan mengortogonalisasi pembaruan matriks yang dihaluskan momentum memakai Newton-Schulz, menghasilkan pembaruan dengan spektrum seimbang yang memerlukan matriks 2D lengkap sebagai masukan. Hal ini memunculkan ketidakcocokan sistem: FSDP/ZeRO-3 menghemat memori dengan membuat pengoptimal melihat shard, bukan matriks utuh. Sistem yang ada karena itu entah merekonstruksi matriks pada setiap langkah pengoptimal, membayar komunikasi seukuran bobot setelah backward, atau membuat pembaruan menjadi lokal dengan penempatan pemilik ZeRO-1 dengan parameter penuh yang residen. Jalur pertama menghadirkan kembali komunikasi yang justru ingin dihindari ZeRO-3; jalur kedua meninggalkan penghematan memori sekala ZeRO-3. MatrixFSDP menempuh jalur ketiga: ia mengubah tempat shard ZeRO-3 berada, bukan pengoptimal yang dihitung.

Pembaruan yang dihasilkan cocok dengan Muon matriks penuh sekaligus mempertahankan memori sekala ZeRO-3. Pada 64 A100, MatrixFSDP menurunkan latensi langkah pengoptimal dibandingkan FSDP2-Muon standar sebesar 4,2x pada satu node dan 54,6x pada delapan node, mencapai percepatan end-to-end hingga 2,15x, dan menjalankan ukuran model di mana penempatan pemilik ZeRO-1 melampaui GPU 80 GB. Angka 54,6x pada delapan node adalah sinyal paling jelas bahwa kolektif matriks pada langkah pengoptimal, bukan pewujudan forward/backward, adalah hambatan dominan pada baseline. Percepatan end-to-end 2,15x menunjukkan bahwa penghapusan kolektif tersebut berterjemah menjadi keuntungan waktu pelatihan bahkan setelah memperhitungkan biaya pewujudan dan reshard parameter yang tersisa.

Pengoptimal matriks seperti Muon menarik untuk pelatihan skala besar karena dapat meningkatkan konvergensi dan efisiensi token dibandingkan pengoptimal coordinate-wise. Muon melakukan ini dengan mengortogonalisasi pembaruan matriks yang dihaluskan momentum memakai Newton-Schulz, menghasilkan pembaruan dengan spektrum seimbang yang memerlukan matriks 2D lengkap sebagai masukan. Hal ini memunculkan ketidakcocokan sistem: FSDP/ZeRO-3 menghemat memori dengan membuat pengoptimal melihat shard, bukan matriks utuh. Sistem yang ada karena itu entah merekonstruksi matriks pada setiap langkah pengoptimal, membayar komunikasi seukuran bobot setelah backward, atau membuat pembaruan menjadi lokal dengan penempatan pemilik ZeRO-1 dengan parameter penuh yang residen. Jalur pertama menghadirkan kembali komunikasi yang justru ingin dihindari ZeRO-3; jalur kedua meninggalkan penghematan memori sekala ZeRO-3. MatrixFSDP menempuh jalur ketiga: ia mengubah tempat shard ZeRO-3 berada, bukan pengoptimal yang dihitung.

Untuk setiap bobot 2D, satu rank data-paralel memiliki seluruh matriks dan rank lain memegang shard kosong; tensor non-matriks dikemas ke pemilik ekor dan tetap memakai AdamW. Reduksi backward biasa kemudian menempatkan masukan Muon lengkap di pemiliknya, sehingga Newton-Schulz berjalan lokal tanpa kolektif matriks pada langkah pengoptimal. Forward dan backward tetap mewujudkan dan melakukan reshard parameter; tantangan runtime adalah membuat tata letak yang tidak merata itu efisien dan benar. MatrixFSDP melakukannya dengan metadata MatrixShard, perencana pemilik yang sadar keseimbangan, kolektif P2P segmen-pemilik yang deterministik, penyematan buffer pemilik, dan reshard checkpoint shard-pemilik.

MatrixFSDP merumuskan ulang masalah Muon-pada-ZeRO-3 sebagai masalah penempatan, bukan masalah komunikasi. Dengan membuat satu rank memiliki setiap matriks 2D secara utuh, ia menyelaraskan kebutuhan masukan pengoptimal dengan pola reduksi yang sudah dilakukan FSDP, sehingga matriks lengkap tiba di pemiliknya sebagai efek samping backward, bukan sebagai kolektif terpisah. Biayanya adalah tata letak shard yang tidak merata: sebagian rank memegang matriks utuh sementara yang lain memegang shard kosong, dan tensor non-matriks harus dikemas ke pemilik ekor agar AdamW tetap benar. Kontribusi rekayasa—metadata MatrixShard, perencana pemilik yang sadar keseimbangan, kolektif P2P segmen-pemilik yang deterministik, penyematan buffer pemilik, dan reshard checkpoint shard-pemilik—adalah yang membuat tata letak tidak merata itu efisien dan benar dalam praktik. Pendekatan ini mempertahankan memori sekala ZeRO-3, tidak seperti penempatan pemilik ZeRO-1, dan menghindari komunikasi pasca-backward seukuran bobot dari skema rekonstruksi matriks. Penurunan latensi langkah pengoptimal yang dilaporkan sebesar 4,2x dan 54,6x serta percepatan end-to-end 2,15x pada 64 A100 menunjukkan bahwa perubahan tata letak adalah jalur yang layak menuju pengoptimal matriks pada skala besar, meskipun perencana yang sadar keseimbangan dan reshard checkpoint tetap menjadi kompleksitas utama yang harus dikelola.

Praktisi dan peneliti ilmu komputer

Membuka konten member…