Amazon oznámil TorchNeuron, open-source nativní backend PyTorch pro svou rodinu čipů AI Trainium, v kroku, který odstraňuje jednu z největších praktických překážek při přijímání vlastního křemíku pro úlohy strojového učení. Podle oficiálního příspěvku na blogu PyTorch zveřejněného před konferencí PyTorch v Severní Americe 20. října mohou nyní vývojáři spouštět stávající kód modelu PyTorch na Trainium s jednořádkovou změnou, spíše než přepisovat své pracovní postupy kolem proprietárních nástrojů.

Oznámení přitáhlo okamžitou pozornost, protože se zaměřuje na to, co analytici dlouho popisovali jako nejtrvalejší výhodu Nvidie: nikoli na hrubý výkon čipu, ale na softwarový ekosystém CUDA, který váže vývojáře ke svým GPU. Reuters zarámoval vydání jako Amazon, že jeho čipy AI budou fungovat s nejoblíbenějším vývojovým rámcem v AI, zatímco TechCrunch označil TorchNeuron za dosud největší šanci společnosti na Nvidii. Chcete-li získat širší pohled na to, jak se trh s akcelerátory posouvá, naše probíhající informace o hardwaru AI sleduje tento vývoj, jak se vyvíjí.

Co TorchNeuron ve skutečnosti mění

Hlavním příslibem TorchNeuron je jednoduchost. V příkladu zveřejněném na blogu PyTorch běží existující tréninková smyčka na Trainium po ničem jiném, než přesunutí modelu na nové zařízení pomocí `.to('neuron')`. Kód napsaný proti `torch.accelerator`, doporučenému PyTorch API pro rozpoznávání zařízení, běží beze změn, podle příspěvku. Pro zkompilované úlohy nastaví vývojáři výchozí backend jednou pomocí `torch.compiler.set_default_backend("neuron")` a poté vyvolají `model.compile()` bez dalších změn.

TorchNeuron pod kapotou registruje Trainium prostřednictvím PrivateUse1, oficiálního rozšiřujícího bodu PyTorch pro vlastní hardwarové backendy, což znamená, že se čip chová jako standardní zařízení PyTorch spíše než jako doplněk přišroubovaný k frameworku. AWS popisuje integraci jako zahrnující několik režimů provádění:

  • Eager provádění pro výzkum a ladění, s experimentálním „adaptivním dychtivým režimem“, který snižuje režii spuštění
  • torch.compile pro produkční výkon, zachycování modelů do grafů, které kompilátor Neuron optimalizuje, s nativní podporou FlexAttention
  • Rozhraní Neuron Kernel Interface (NKI) pro výzkumníky, kteří potřebují vlastní jádra na nízké úrovni

Backend funguje jak na Trainium2, tak na novějších Trainium3 čipech, podle oznámení.

Přímý útok na softwarový příkop společnosti Nvidia

Schopnost Nvidie v oblasti AI computingu nikdy nespočívala pouze na rychlejším hardwaru. CUDA, dvě desetiletí stará softwarová platforma společnosti, je místem, kde se píše, ladí a nasazuje většina výzkumných kódů AI a každý čip vyzyvatele – od řady Instinct od AMD po TPU od Googlu a Trainium od Amazonu – se musel potýkat s náklady na to, aby vývojáři opustili. Pokrytí startu TorchNeuron v GeekWire, Benzinga a GuruFocus jednotně zarámovalo oznámení jako útok na tento příkop.

Amazon se nesnaží vytáhnout vývojáře z PyTorch, aby je získal. Místo toho přináší Trainium do PyTorch, takže obrovská uživatelská základna frameworku může čip přijmout, aniž by změnila návyky. Příspěvek na blogu PyTorch obsahoval podporu od společnosti Meta, která zachovává rámec: „Nativní podpora pro AWS Trainium rozšiřuje tuto možnost a umožňuje výzkumníkům a inženýrům experimentovat a opakovat pomocí dychtivého provádění a distribuovaných API, které již znají,“ uvedla v příspěvku Jana van Greunen, ředitelka PyTorch Engineering ve společnosti Meta.

Toto uspořádání je pro cloudové podnikání AWS důležité. Instance Trainium byly prodávány jako levnější alternativa k GPU Nvidia pro školení a vyvozování a čipy již přilákaly zákazníky s AI. Tím, že se čipy cítí jako nativní pro PyTorch, přímo řeší námitku, která zpomalila širší přijetí: úsilí o portování.

Technický rozsah, od notebooků po školení v clusterovém měřítku

Schopnosti uvedené na blogu PyTorch naznačují, že Amazon míří daleko nad rámec amatérských experimentů. TorchNeuron s přesností BF16, FP32 a FP16 pokrývá sady operátorů, které lze vidět napříč textovými, multimodálními a vizuálními modely, stejně jako generování obrázků a videa, modely stavového prostoru a světové modely; Podpora TorchAO a FP8 je v soukromé beta verzi před obecnou dostupností v příštím vydání.

Distribuované školení, disciplína, která mění kód jednoho GPU na běhy na úrovni clusteru, je podporováno prostřednictvím standardních PyTorch API — DDP, FSDP2, SimpleFSDP, paralelismus tenzorů a potrubí a sharding ve stylu ZeRO, spolu s DTensor a DeviceMesh pro vícerozměrné skupiny procesů. Streamy a události mapují sémantiku PyTorch na hardwarové fronty Trainium, uvedl příspěvek.

Pro závěr, zásobníky Neuron společnosti Amazon staví na stejném základě: vLLM Neuron slouží velkým jazykovým modelům, zatímco vLLM Omni Neuron slouží k difúzním modelům pro generování obrazu a videa. Rozhraní Neuron Kernel Interface umožňuje výzkumníkům psát vlastní operátory pro dychtivé i kompilované provádění a AWS je opensourcingový NKI-lib, sada jader optimalizovaných pro Trainium pokrývající blokové výpočty FlashAttention a MoE pro mainstreamové LLM a multimodální scénáře.

Součástí je také výkonnostní nářadí. Neuron Explorer rozšiřuje standardní `torch.profiler` o sledování provádění z rámce na hardware, propojuje každou událost zpět se zdrojovou linií a může exportovat metriky ve formátu Parquet, které mohou být zpracovány agenty pro automatizovanou analýzu úzkých míst.

Dostupnost a co přijde dál

Podle blogového příspěvku je TorchNeuron připraven ke spuštění open source na konferenci PyTorch v Severní Americe 20. října, přičemž kompatibilita se rozšiřuje na Hugging Face Transformers v5. AWS popisuje Trainium jako portfolio čipů navržených od základů pro cloudový výcvik a vyvozování AI.

Strategickým kontextem je trh, na kterém je každý velký poskytovatel cloudu současně největším kupcem akcelerátorů Nvidia a potenciálním konkurentem. Microsoft, Google a Meta všechny vyvíjejí vlastní křemík a čínské firmy včetně Huawei budují své vlastní softwarové balíky, aby unikli gravitaci CUDA. Tím, že se Amazon sešel s vývojáři uvnitř PyTorch, místo aby je požádal o migraci, zvolil do této soutěže nejméně rušivou cestu – a nejpravděpodobnější, pokud vývojářské zkušenosti vydrží, přesunout skutečné tréninkové zátěže z hardwaru Nvidia.

Zda TorchNeuron skutečně naruší výhodu CUDA, bude záviset na výkonu v měřítku a dlouhém ocasu reálných modelů, otázky, na které startovací materiály neodpovídají. Ale bariéra, která bránila většině vývojářů PyTorch dokonce vyzkoušet Trainium – přepsání – byla nyní snížena na jediný řádek kódu.

---

Stay Ahead of AI

Získejte nejnovější zprávy, analýzy a průlomové informace o umělé inteligenci – vše na jednom místě.

Přečtěte si další zprávy o AI →