Natív Ray-támogatást kapott a PyTorch a Google TPU-in

Egy vezérlősík GPU-khoz és TPU-khoz

Az október 20-i bejelentés szerint a GKE-n dolgozó platform- és gépi tanulási csapatok egyetlen Kubernetes-vezérlősíkon ütemezhetik, skálázhatják és üzemeltethetik az elosztott PyTorch-tréninget, a vLLM-alapú következtetést és a megerősítéses tanulási munkafolyamatokat. A rendszer GPU-kon és TPU-kon is használható.

Az integráció a Google TorchTPU projektjére épül. A TorchTPU a Tensor Processing Unitokhoz készült natív PyTorch-háttér, amely gyors végrehajtásra épülő integrációt és a TPU-kon futó elosztott munkafolyamatokhoz torch.distributed kiterjesztést biztosít. Az Anyscale szerint a Ray korábban elsősorban a JAX-hoz használt, egy programot több adaton futtató SPMD-modellhez igazodott TPU-kon.

Egy Ray-munkás jut egy TPU-chipre

Az új megközelítésben a Ray minden TPU-chiphez külön munkafolyamatot indít. Ez megfelel annak a modellenek, amelyet a PyTorch fejlesztői a GPU-k esetében már megszoktak. Egy Ray-feladat vagy szereplő egyetlen logikai TPU-eszközt kérhet a resources={"TPU": 1} beállítással.

A Ray automatikusan elkülöníti a munkásokhoz rendelt eszközöket a TPU_VISIBLE_CHIPS környezeti változóval. A meglévő PyTorch-kód TPU-ra állításához a forrás szerint a cuda helyett tpu eszköznevet, az elosztott kommunikációhoz pedig az nccl helyett a tpu_dist hátteret kell használni.

A GKE a több gépből álló TPU-szeleteket nagy sebességű Inter-Chip Interconnect kapcsolatokkal összekötött csomópontcsoportokként biztosítja. A Ray SlicePlacementGroup API-ja gondoskodik arról, hogy a munkások ugyanarra a TPU-szeletre, csoportosan kerüljenek.

Tréning és következtetés Ray Trainnel, illetve Ray Serve-dzsel

A Ray Train TorchTrainer eszköze natív TorchTPU-támogatást kapott. Az Anyscale és a Google Cloud szerint így a felhasználók egyetlen gépes és több gépes Cloud TPU-szeleteken is futtathatják meglévő tréningkódjukat a PyTorch-tréninghurkon végzett módosítás nélkül. A TorchTrainer az elosztott erőforrások lefoglalását és a szükséges környezeti változók beállítását is kezeli.

A use_tpu=True beállítással, valamint a kívánt topológiával és gyorsítótípussal a rendszer automatikusan létrehozza a SlicePlacementGroupot. A Ray Train egy munkást indít minden TPU-eszközhöz, beállítja a rang-, méret- és címadatokat, majd inicializálja a tpu_dist folyamatcsoportot. A ray.train.torch.get_device() és a prepare_model() segédfüggvények felismerik a TorchTPU-környezetet, a modellt a TPU-ra helyezik, és DistributedDataParallel burkolást alkalmaznak.

Online következtetésnél a vLLM-TorchTPU a TorchTPU-t vLLM-platformbővítményként regisztrálja. A Ray Serve LLM segítségével OpenAI-kompatibilis vLLM-végpontok telepíthetők egy- és több gépes TPU-szeletekre, ugyanazzal a deklaratív LLMConfig API-val, amely GPU-kon is használható.

A PARK stack TPU-kra is kiterjed

A vállalatok egyre gyakrabban olyan rendszereket építenek, amelyek a tréninget, a nagy áteresztőképességű következtetést és a megerősítéses tanulást egyesítik. Az Anyscale és a Google Cloud ezt a nyílt forráskódú PARK stackkel írja le, amelynek elemei a PyTorch, az AI-modellek, a Ray és a Kubernetes.

A bejelentés szerint a TorchTPU Ray és GKE melletti használata lehetővé teszi, hogy a PyTorch-fejlesztők megtartsák az ismert végrehajtási modellt és Ray API-kat, miközben a platformcsapatok topológiatudatos automatikus skálázást, hibatűrést és többféle gyorsító kezelését kapják. A felek GKE-s útmutatókat tettek elérhetővé több gépes TPU-kon történő modelltréninghez és LLM-kiszolgáláshoz. A közlemény szerint a következő lépések között szerepel a nagy léptékű megerősítéses utótréning támogatásának bővítése, valamint a következő generációs Ironwood, vagyis TPU v7x szeletek használata GKE-n.