Amazon SageMaker AI, 학습 루프를 직접 구현할 수 있는 training session API 추가
Write your own training loop — without managing GPUs — with Amazon SageMaker training session APIs
AWS가 Amazon SageMaker AI의 서버리스 모델 커스터마이징에 training session API를 추가했다. forward_backward(), optim_step(), sample()을 API 호출로 조합해 SFT나 강화학습 학습 루프를 로컬 Python으로 작성할 수 있다. GPU 할당, 분산 실행, 장애 복구는 SageMaker AI가 담당한다.
training session API는 알고리즘이 아니라 컴퓨팅 리소스를 추상화하는 것이 특징이다. 라이브 상태 유지형 세션에서는 모델이 GPU에 올라간 채로 사용자가 필요한 만큼 호출을 반복할 수 있다. 베이스 모델과 LoRA 어댑터가 로드되며, 런타임 설정, 분산 실행, 그래디언트 동기화, 장애 복구, 테넌트 격리는 서비스 측에서 처리한다. 클라이언트는 ServiceClient, TrainingClient, SamplingClient 세 종류이며, 활성 세션은 OpenAI 호환 채팅 완성 엔드포인트도 제공한다.
과금은 토큰 처리량 기준이다. 보상 계산, 데이터 로딩, 디버깅이 로컬에서 진행되는 동안에는 GPU 시간이 과금되지 않아 유휴 시간 비용이 발생하지 않는다. 복잡한 보상 모델이나 다중 도구 에이전트형 롤아웃처럼 로컬 연산이 무거운 작업에서 차이가 크다고 알려져 있다. 위치상으로는 설정만으로 수행하는 Serverless Model Customization보다 제어 수준이 높고, SageMaker Training Jobs나 HyperPod보다 운영 부담이 작다.
기사는 GSM8K를 사용해 OpenAI의 GPT-OSS 20B용 LoRA 어댑터를 GRPO로 학습하는 최소 루프를 설명한다. 샘플링, 로컬 보상 채점, 그래디언트 업데이트, 샘플러 갱신의 4단계로 구성되며, 각 GPU 호출은 future를 반환하므로 병렬 처리할 수 있다. 체크포인트는 모델 패키지로 등록된다. 이 샘플에서는 10스텝, 40문제로 홀드아웃 정확도가 56%에서 75%로 올랐다.
실운영에 가까운 예로 cookbook의 code_rl 레시피를 소개한다. DeepCoder 데이터셋의 경쟁 프로그래밍 문제를 사용해, 생성된 코드를 SandboxFusion 컨테이너에서 단위 테스트로 채점한다. 2턴에 걸쳐 check_solution 도구를 호출할 수 있고, 샘플링과 학습을 겹쳐 수행하는 비동기 실행도 지원한다. GPT-OSS 20B를 50스텝 학습한 결과 학습 보상은 0.46에서 0.60으로, 200문제 테스트 정확도는 40%에서 50.5%가 됐다. cookbook에는 이 밖에도 Tulu 3, UltraChat 기반 SFT, KTO를 이용한 선호 정렬 레시피가 있다.
출처: AWS Machine Learning Blog · aws.amazon.com