Kunngjør PyTorch/XLA 2.4: En b... Notat

Kunngjør PyTorch/XLA 2.4: En bedre Pallas og utvikleropplevelse, pluss 'eager mode'

PyTorch/XLA 2.4 tilbyr betydelige forbedringer for dyp læring på TPU-er og GPU-er. Denne utgivelsen introduserer forbedringer til Pallas, et custom kernel-språk, som nå støtter både TPU-er og GPU-er, og forbedrer ytelsen med Python-kode. Nye API-kall, som torch_xla.sync(), forenkler integrering i eksisterende PyTorch-arbeidsflyter. En eksperimentell ivrig modus tillater umiddelbar eksekusjon av operasjoner på målhardware, selv om TPU-er krever et "mark_step"-kall for emulering. Forbedringer til Pallas inkluderer Flash Attention og Paged Attention-støtte, og innebygde Megablocks' blokk-sparse kernels for gruppe-matrise-multiplikasjon. I tillegg introduserer en ny TPU-kommandolinje-grensesnitt, tpu-info, som forenkler feilsøking ved å vise bruk og enhetsinformasjon, likt Nvidia's nvidia-smi-verktøy. Til tross for disse endringene, forblir eksisterende kode kompatibel, noe som gjør oppgraderingen smidig for utviklere. Disse oppdateringene har til hensikt å forbedre brukervennligheten og ytelsen i PyTorch/XLA.