Skapa kostnadseffektiv ML-utbildningsinfrastruktur

Våra ML-utbildningskostnader var ohållbara

Allt detta började på grund av en utmaning jag stod inför på mitt företag. Vår produkt drivs av AI-teknik inom GAN-området, och vårt team består av ML-forskare och ingenjörer. Under vår resa för att etablera kärnteknologin började vi köra många ML-träningsexperiment av flera forskare parallellt. Snart började vi se en enorm ökning i våra molnkostnader. Det var inte hållbart – vi var tvungna att göra något, och det snabbt. Men innan jag kommer till den kostnadseffektiva ML-träningslösningen, låt oss förstå varför utbildningskostnaderna var så höga.

Vi började använda ett mycket populärt GAN-nätverk som heter stylegan.

Tabellen ovan visar hur lång tid det tar att träna detta nätverk beroende på antalet GPU:er och önskad utdataupplösning. Låt oss anta att vi har åtta grafikprocessorer och vill ha en utgångsupplösning på 512×512; vi behöver en EC2 av typen “p3.16xlarge” som kostar $24,48 per timme, så vi betalar $6 120 för detta experiment. Men det finns mer innan vi kan göra det här experimentet. Forskare måste upprepa flera cykler med att köra kortare experiment, utvärdera resultaten, ändra parametrarna och börja om från början.

Så nu kan utbildningskostnaden för endast en forskare vara allt från $8–12k per månad. Multiplicera detta med N forskare, och vår månatliga förbränningshastighet ligger utanför diagrammet.

Vi var tvungna att göra något

Att bränna dessa summor varje månad är inte hållbart för en liten startup, så vi var tvungna att hitta en lösning som dramatiskt skulle minska kostnaderna – men också förbättra utvecklarhastigheten och skala snabbt.

Här är en översikt över vår lösning:

Forskare: kommer att utlösa träningsjobb via ett Python-skript (skriptet kommer att vara deklarativa instruktioner för att bygga ett ML-experiment).

Utbildningsjobb: kommer att schemaläggas på AWS ovanpå Spot-instansen och kommer att hanteras helt av vår infrastruktur.

Spårbarhet: under träning kommer mätvärden som GPU-statistik/framsteg att skickas till forskaren via Slack, och modellkontrollpunkter kommer automatiskt att laddas upp för att ses via Tensorboard.

Att utveckla infrastrukturen

Låt oss först se över den minsta enheten i infrastrukturen, docker-bilden.

Bilden är konstruerad av tre steg som upprepar varje träningspass och har ett Python-gränssnitt för abstraktion. För en integration kommer Algo-forskaren att lägga till ett anrop till någon träningskod i “Träna-funktionen”; sedan, när denna docker-bild kompileras, hämtar den träningsdata från en S3-hink och sparar den på den lokala maskinen → Anropa en träningsfunktion → Spara resultaten tillbaka till S3.

Denna logik ovan är faktiskt en klass som anropas när dockaren startar. Allt användaren behöver göra är att åsidosätta tågfunktionen. För det gav vi en enkel abstraktion:

from resarch.some_algo import train_my_alg
from algo.training.session import TrainingSession


class Session(TrainingSession):
    def __init__(self):
        super().__init__(path_to_training_res_folder="/...")

    def train(self):
        super().train()
        train_my_alg(restore=self.training_env_var.resume_needed)

Att ärva från TrainingSession innebär att alla tunga lyft görs för användaren. Importera samtalet till träningsfunktionen (rad 1). Lägg till sökvägen där kontrollpunkterna sparas (rad 7). Denna väg kommer att backas upp av infrastrukturen till s3 under träning. Åsidosätt “träna”-funktionen och anrop någon algoträningskod (rad 9–11).

Startar a mer kostnadseffektivt ML-utbildningsjobb

För att starta ett utbildningsjobb tillhandahöll vi ett enkelt deklarativt skript via Python SDK:

from algo.training.helpers import run
from algo.training.env_var import TrainingEnvVar, DataSourceEnvVar

env_vars = TrainingEnvVar(...)

run(env_vars=env_vars)

TrainingEnvVar – Deklarativa instruktioner för experimentet. kör – Kommer att avfyra SNS-ämne som kommer att starta ett flöde för att köra ett träningsjobb på AWS.

Utlöser ett experimentjobb

SNS-meddelande med all träningsmetadata skickad (3). Detta är samma meddelande som används av infra om vi behöver återuppta jobbet på en annan plats. Meddelandet konsumeras av SQS för att bevara staten och lambda som avfyrar en punktförfrågan. Spot-förfrågningar är asynkrona, vilket innebär att uppfyllandet kan ta tid. När en spot-instans är igång skickas en CloudWatch-händelse. Spotuppfyllelsehändelsen utlöser en Lambda (4) som är ansvarig för att dra ett meddelande från SQS(5) med alla instruktioner för träningsjobbet.

Svara på avbrott i kostnadseffektiva ML-utbildningsjobb

Innan AWS-spotinstansen kommer att tas från oss får vi ett CloudWatch-meddelande. För det här fallet lade vi till en Lambda-trigger som ansluter till instansen och kör en återställningsfunktion inuti docker-bilden (1) som startar ovanstående flöde igen från toppen.

Starta kostnadseffektiv ML-utbildning

Lambda (6) utlöses av en CloudWatch-händelse:

{
  "source": ["aws.ec2"],
  "detail-type": ["EC2 Spot Instance Request Fulfillment"]
}

Den ansluter sedan till den nya punktinstansen för att starta ett träningsjobb från den sista punkten där det slutade eller starta ett nytt jobb om SNS (3)-meddelandet skickades av forskaren.

Efter sex månader i produktion var resultaten dramatiska

Ovanstående mätvärden visar utvecklingsfasen när vi tillbringade två veckor med att bygga ovanstående kostnadseffektiva ML-träningsinfrastruktur, följt av utvecklingsanvändningen av vårt team.

Låt oss zooma in på en forskare med vår plattform. I juli och augusti använde de inte infra och körde K små experiment som kostade ~$650. I september körde de samma K-experiment++ men vi halverade kostnaden. I oktober mer än fördubblade de sina experiment och kostnaden var bara runt 600 dollar.

Idag använder alla Bria-forskare vår interna infra samtidigt som de drar fördel av dramatiskt minskade kostnader och en avsevärt förbättrad forskningshastighet.

Den här historien dök ursprungligen upp på Medium.com. Copyright 2021

DataDecisionMakers

Välkommen till VentureBeat-communityt!

DataDecisionMakers är där experter, inklusive tekniska personer som arbetar med data, kan dela datarelaterade insikter och innovation.

Om du vill läsa om banbrytande idéer och aktuell information, bästa praxis och framtiden för data- och datateknik, gå med oss ​​på DataDecisionMakers.

Du kan till och med överväga att bidra med en egen artikel!

Läs mer från DataDecisionMakers