
SageMaker umožňuje písať tréningové slučky v Pythone so spravovanými GPU
Amazon SageMaker AI podľa článku AWS ponúka API tréningových relácií, cez ktoré možno zostaviť vlastnú tréningovú slučku v Pythone. Operácie ako forward_backward, optim_step a sample pracujú so živou reláciou, ktorá medzi krokmi uchováva model na GPU. Služba zabezpečuje nastavenie prostredia, distribuované vykonávanie, synchronizáciu gradientov, obnovu po zlyhaní a izoláciu nájomníkov.
Algoritmus v lokálnom procese
Reláciu možno riadiť z lokálneho procesu Pythonu, pričom zvolený základný model a adaptér LoRA sú načítané na spravovaných akcelerátoroch. Programovací model nevyžaduje vytvorenie klastra ani konfiguračný súbor opisujúci celý tréning vopred.
SDK používa tri postupne vytvárané klienty. ServiceClient uchováva nastavenia platformy a spúšťa reláciu, TrainingClient poskytuje tréning a ukladanie kontrolných bodov a SamplingClient ponúka sample a compute_logprobs. Aktívna relácia sprístupňuje aj rozhranie chat completions kompatibilné s OpenAI.
V ukážkovej RL slučke GPU generujú odpovede a lokálny kód ich hodnotí. Operácie forward_backward a optim_step aktualizujú model; save_weights_and_get_sampling_client následne obnoví sampler. Pri supervidovanom dolaďovaní sa generovanie a odmeny vynechajú a na označené príklady sa použije strata cross_entropy.
Poradie operácií a účtovanie
Volania viazané na GPU okamžite vracajú objekty APIFuture; požiadavky sa zaraďujú a odosielajú na pozadí. Metóda result() čaká na dokončenie. Generovanie možno zadávať súbežne, tréningové operácie však prebiehajú v poradí zadania. Tréningový záznam v príklade obsahuje tokeny zadania a odpovede, logaritmické pravdepodobnosti tokenov, výhodu a masku vylučujúcu zadanie zo straty.
Podľa AWS sa účtuje za tokeny spracované pri operáciách na GPU. Počas lokálneho výpočtu odmien, načítavania dát, logovania či ladenia sa čas GPU neúčtuje. API podporujú SFT, straty ppo, importance_sampling a cispo aj vlastné ciele vrátane preferenčnej optimalizácie a destilácie. AWS ich umiestňuje medzi Serverless Model Customization a prevádzkovo náročnejšie Training Jobs či HyperPod.
Kontrolné body a recepty
Každé uloženie kontrolného bodu registruje verziu balíka modelu. Príklad ukladá zlúčený model na nasadenie aj stav s váhami a optimalizátorom na pokračovanie cez create_training_client_from_state_with_optimizer. Zlúčené modely možno nasadiť na endpointy v reálnom čase. Relácia sa započítava do kvóty súbežných relácií až do zastavenia alebo vypršania.
Sprievodná kuchárka rieši dávky, šablóny chatu, odhad výhod, prekrývanie generovania s tréningom, validáciu, obnovu, metriky krokov a záznam konfigurácií. Obsahuje chatové SFT s Tulu 3 a UltraChat, GSM8K RL a preferenčné zarovnanie KTO z binárnej spätnej väzby.
Výsledky uvádza AWS
V ukážke GSM8K sa adaptér LoRA modelu GPT-OSS 20B trénuje pomocou GRPO: posilňujú sa odpovede s úspešnosťou nad priemerom skupiny. AWS uvádza, že 10 krokov na 40 úlohách zvýšilo presnosť na oddelených testovacích dátach z 56 % na 75 %.
Recept code_rl využíva úlohy DeepCoder a lokálny kontajner SandboxFusion na testovanie riešení v Pythone. Za splnenie všetkých testov udeľuje 1, inak 0, a pri chýbajúcom bloku kódu odpočíta 0,1. Dva ťahy umožňujú zavolať check_solution, prečítať výsledky a opraviť odpoveď.
Recept prekrýva generovanie s tréningom, obmedzuje zastaranosť odpovedí, používa cispo a počíta gradienty v menších dávkach. AWS pri LoRA hodnosti 32 uvádza 50 krokov, 128 úloh, osem odpovedí na úlohu a limit zastaranosti tri kroky. Tréningová odmena vzrástla z 0,46 na 0,60 a presnosť na 200 oddelených úlohách zo 40 % na 50,5 %. Metriky možno voliteľne publikovať do spravovaného MLflow. Ide o výsledky konkrétnych behov uvedené AWS, nie o všeobecnú záruku zlepšenia. Podklady neuvádzajú zoznam podporovaných modelov, regióny, číselné kvóty ani ceny za token.