
SageMaker pozwala sterować treningiem z Pythona na zarządzanych GPU
Interfejsy sesji treningowych Amazon SageMaker AI pozwalają według AWS tworzyć własne pętle treningowe w Pythonie. Operacje takie jak forward_backward, optim_step i sample działają na aktywnej sesji zachowującej stan i model załadowany na GPU między krokami. Usługa obsługuje środowisko wykonawcze, wykonanie rozproszone, synchronizację gradientów, odtwarzanie po awarii oraz izolację najemców.
Lokalny kod steruje sesją
Sesję można prowadzić z lokalnego procesu Pythona, gdy wybrany model bazowy i adapter LoRA znajdują się na zarządzanych akceleratorach. Nie trzeba tworzyć klastra ani przygotowywać pliku konfiguracji opisującego całe zadanie z góry.
SDK wykorzystuje kolejno trzy klienty. ServiceClient przechowuje ustawienia platformy i uruchamia sesję, TrainingClient udostępnia trening oraz zapis punktów kontrolnych, a SamplingClient zapewnia sample i compute_logprobs. Aktywne sesje mają także endpoint chat completions zgodny z OpenAI.
W pokazanej pętli RL zarządzane GPU generują odpowiedzi, a lokalny Python je ocenia. forward_backward i optim_step aktualizują model; następnie save_weights_and_get_sampling_client odświeża sampler. Przy nadzorowanym dostrajaniu pomija się generowanie i obliczanie nagród, stosując stratę cross_entropy do oznaczonych przykładów.
Kolejność operacji i rozliczenia
Wywołania związane z GPU natychmiast zwracają obiekty APIFuture. Żądania trafiają do kolejki i są wysyłane w tle, a result() czeka na zakończenie. Generowanie może odbywać się równolegle, lecz operacje treningowe zachowują kolejność zgłoszenia. Przykładowy rekord zawiera tokeny polecenia i odpowiedzi, logarytmy prawdopodobieństw tokenów, przewagę oraz maskę wyłączającą tokeny polecenia z obliczania straty.
AWS podaje, że opłaty zależą od tokenów przetworzonych podczas operacji na GPU. Czas GPU nie jest naliczany podczas lokalnego obliczania nagród, ładowania danych, logowania i debugowania. Dostępne są SFT, straty RL ppo, importance_sampling i cispo oraz własne cele, w tym optymalizacja preferencji i destylacja. AWS pozycjonuje sesje jako zapewniające większą kontrolę algorytmu niż Serverless Model Customization i mniej obowiązków operacyjnych niż Training Jobs lub HyperPod.
Punkty kontrolne i przepisy
Każdy zapis punktu kontrolnego rejestruje wersję pakietu modelu. Przykład zapisuje model scalony do wdrożenia oraz wagi i stan optymalizatora do wznowienia przez create_training_client_from_state_with_optimizer. Scalone punkty kontrolne można wdrażać na endpointach czasu rzeczywistego. Sesje zajmują limit współbieżności aż do zatrzymania lub wygaśnięcia.
Towarzyszący zbiór przepisów obsługuje partie danych, szablony czatu, estymację przewagi, nakładanie generowania i treningu, walidację, wznawianie, metryki kroków i zapis konfiguracji. Obejmuje chatowe SFT z Tulu 3 i UltraChat, pełny trening RL na GSM8K oraz dopasowanie preferencji KTO z binarnych ocen.
Wyniki zgłoszone przez AWS
Przykład GSM8K trenuje adapter LoRA dla GPT-OSS 20B metodą GRPO, wzmacniając odpowiedzi o poprawności powyżej średniej grupy. AWS podaje, że dziesięć kroków na 40 zadaniach zwiększyło trafność na wydzielonych danych testowych z 56% do 75%.
Przepis code_rl korzysta z zadań programistycznych DeepCoder i ocenia wygenerowane rozwiązania w Pythonie testami w lokalnym kontenerze SandboxFusion. Przyznaje 1 za zaliczenie wszystkich testów i 0 w innym przypadku, z karą 0,1 za brak bloku kodu. Dwie tury pozwalają wywołać check_solution, sprawdzić wyniki i poprawić odpowiedź.
Przepis nakłada generowanie na trening, ogranicza nieaktualność odpowiedzi, stosuje cispo do korekty różnic polityk i liczy gradienty w podpartiach. AWS opisuje adapter LoRA rangi 32, 50 kroków, 128 zadań, osiem odpowiedzi na zadanie, dwie tury i limit opóźnienia trzech kroków. Nagroda treningowa wzrosła z 0,46 do 0,60, a trafność na 200 wydzielonych zadaniach z 40% do 50,5%. Metryki można opcjonalnie publikować w zarządzanym MLflow SageMaker.
To wyniki konkretnych przebiegów raportowane przez AWS, bez gwarancji dla innych zastosowań. Materiał nie podaje listy obsługiwanych modeli, regionów, liczbowych limitów ani cen za token.