AI Inference klusterin rakentaminen konesalin räkkiin

Kun ollaan tekemisisissä erilaisten laitteistojen parissa, esim jos käytössäsi on RTX PRO -sarjaa ja AMD:n näytönohjaimia, haasteena ei ole vain mallien ajaminen, vaan niiden saaminen toimimaan yhdessä yhtenäisenä ja vikasietoisena rajapintana (API).

Tässä kirjoituksessa käymme läpi, kuinka rakensimme heterogeenisen AI-räkin käyttämällä vLLM:ää moottorina ja HAProxy-ohjelmistoa kuroaksemman umpeen "vihreän" ja "punaisen" tiimin (Nvidia ja AMD) välisen kuilun.

Arkkitehtuuri

Keskeinen ongelma AMD:n ja Nvidian sekoittamisessa on se, ettei samaa LLM:ää voi (tällä hetkellä) jakaa (shard) molempien yli. On ajettava erillisiä klustereita ja ohjattava liikennettä älykkäästi.

Kokoonpanomme koostuu seuraavista:

  1. Nvidia-klusterit: Ajavat vLLM:n CUDA-optimoitua versiota.

  2. AMD-klusterit: Ajavat vLLM:n ROCm-optimoitua versiota.

  3. HAProxy-kerros: Toimii "liikenteenohjaajana", joka tasapainottaa pyyntöjä klustereiden kunnon ja kapasiteetin perusteella.

Vaihe 1: GPU-nodejen valmistelu

Koska vLLM käsittelee Nvidiaa ja AMD:tä eri tavoin, käytämme Docker-teknologiaa ympäristön monimutkaisuuden abstrahoimiseksi. Kaikkien klustereiden edessä käytämme erillisiä palomuurilaitteita. Kaikki GPU nodet toimivat sisäverkossa ja niiden yhteys applikaatio palvelimille hoidetaan kryptatun tunnelin läpi.

Nvidia (CUDA)

Käytämme vLLM:n vakio-imagea. Varmista aina että  on asennettuna uusin Nvidia Container Toolkit.

AMD (ROCm)

AMD vaatii ROCm-kohtaisen buildin. Käytämme virallista rocm/vllm-imagea, joka on optimoitu Instinct- ja RX-sarjoille. 

Vaihe 2: Kuormantasaus HAProxylla

Nyt kun meillä on kaksi itsenäistä vLLM-päätepistettä (portti 8001 Nvidialle ja 8002 AMD:lle), tarvitsemme yhden yhteisen liityntäpisteen. HAProxy on tähän täydellinen, sillä se on uskomattoman nopea ja tukee HTTP-health checkkejä, mikä on tekoälysovelluksissa välttämätöntä

HAProxy-konfiguraatio

Käytämme Least Connection -algoritmia. Miksi? Koska kielimallipyyntöjen prosessointiajat vaihtelevat suuresti. Yksinkertainen Round Robin saattaisi lähettää 4096 tokenin pyynnön nodelle, joka kamppailee jo ennestään massiivisen kehotteen parissa.  

Vaihe 3: Heterogeenisen viiveen hallinta

Ainakin meillä, RTX Pro -korteissa on huomattavasti enemmän VRAM-muistia kuin AMD-korteissa. Tämä mahdollistaa suuremmat eräkoot (batch size). Käytämme HAProxyn weight-parametria lähettääksemme enemmän samanaikaista liikennettä RTX-nodeille, samalla kun pidämme AMD-nodet vastauskykyisinä matalamman latenssin tehtäviä varten.

Räkkiprojektin tärkeimmät opit:

  • Yhtenäinen API: Käyttämällä vLLM:n OpenAI-yhteensopivaa palvelinta molemmissa klustereissa, frontend-sovellus ei edes tiedä vaihtelevansa CUDA:n ja ROCm:n välillä.

  • Vikasietoisuus: Jos avikka AMD-ajuri kaatuu, HAProxyn kuntotarkastus (/health) poistaa kyseisen noden kierrosta automaattisesti 5 sekunnissa, ja Nvidia-klusteri ottaa koko kuorman kantaakseen.

  • Monitorointi: Käytämme HAProxyn tilastosivua visualisoidaksemme reaaliajassa, mikä klusteri on muodostumassa pullonkaulaksi.

Testasimme myös LiteLLM:ää, mutta emme todenneet sitä tarpeisiimme nähden riittävän vakaaksi.

Miksi "rakenna" mieluummin kuin "osta"?

Tämän rakentaminen paikallisesti antoi meille mahdollisuuden pitää data yksityisenä, eikä sitä lähetetä oman infrastruktuurimme ulkopuolelle. Sekoittamalla laitteistoa emme ole myöskään yhden toimittajan toimitusketjun armoilla, jos löydämme hyvän tarjouksen AMD-korteista, hankimme niitä. Jos saamme erän RTX Pro -kortteja, laajennamme Nvidia-klusteria. 

Tärkeintä meille on yksityisyys – data ei koskaan poistu infrastruktuuristamme.  Lisäksi kulut ovat ennakoitavissa koska maksamme vain sähköstä datacentterille.

Antti Koskela 

Youlearn it OY

Blogi