Amazon SageMaker 학습 세션 API로 GPU를 직접 관리하지 않고 자체 학습 루프를 작성하세요

오늘날 대규모 언어 모델(LLM)을 파인튜닝하는 일은 대개 학습 루프 제어와 운영상의 단순성 사이에서 트레이드오프를 강요합니다. 완전 관리형 서비스에서는 기본 모델을 선택하고, 데이터와 하이퍼파라미터를 제공하면 나머지 모든 것은 서비스가 담당하며, 이는 표준화되어 사전에 패키징된 레시피가 문제에 잘 맞을 때 잘 작동합니다. 덜 관리형 서비스는 더 많은 제어 권한을 주지만, 그만큼 기반 인프라와 학습 프레임워크를 더 많이 직접 관리해야 합니다. 때로는 표준 레시피만으로는 부족합니다. 지도 파인튜닝(SFT)을 위한 커스텀 손실 함수나 데이터 커리큘럼, 자신만의 선호도 최적화, 강화 학습(RL) 보상 함수, 또는 지난주에 발표된 정렬(alignment) 알고리즘이 필요할 수 있습니다. 학습 절차를 직접 다듬어야 하는 순간부터는 알고리즘적 자유를 얻기 위해 그 인프라 작업을 떠안게 됩니다.

오늘 Amazon SageMaker AI의 서버리스 모델 커스터마이징은 트레이닝 세션 API를 통해 그 절충안을 제거함으로써 한 단계 낮은 수준에서 추상화 경계를 그어 확장됩니다. 알고리즘을 추상화하는 대신, 컴퓨팅을 추상화합니다. 다음과 같은 작업들 forward_backward(), optim_step(), 그리고 sample() 당신의 기법에 필요한 절차(SFT, 선호 최적화, RL 또는 완전히 새로운 것)를 구성하는 SageMaker API 호출이 됩니다. 이 호출들은 살아 있고 상태를 유지하는(stateful) 세션을 대상으로 하므로, 모델이 여러 단계에 걸쳐 GPU에 로드된 상태로 유지되며, 필요한 만큼의 호출을 이 세션에 대해 수행할 수 있습니다. 이 서비스는 런타임 설정, 분산 실행, 그래디언트 동기화, 장애 복구, 테넌트 격리, 그리고 대규모 GPU 운영 작업을 처리합니다. 학습 루프는 사용자가 작성하고, 인프라는 SageMaker AI가 실행합니다.

서비스와 함께 저희는 다음을 공개합니다: SageMaker 훈련 세션 쿡북, 자체 학습 루프의 청사진으로 사용할 수 있는 바로 실행 가능한 SFT 및 RL 레시피 모음입니다.

이 게시물에서는 프로그래밍 모델과 세 가지 SDK 클라이언트를 소개하고, SageMaker AI의 모델 훈련 및 사용자 지정 옵션들 사이에서 훈련 세션 API가 어떤 위치에 있는지 보여줍니다. 그런 다음 초등학교 수학 문제(GSM8K)에 대한 최소한의 강화학습(RL) 루프를 분해하여 이러한 작업들이 서로 어떻게 맞물리는지 보여주고, 마지막으로 쿡북의 실제 워크로드인 경쟁 프로그래밍 문제에 대한 강화학습 사례를 살펴봅니다. 이 사례에서는 생성된 코드를 단위 테스트와 대조하여 실행함으로써 채점합니다.

SageMaker AI 학습 세션 API란 무엇인가요?

SageMaker AI 학습 세션 API를 사용하면 로컬 머신의 자체 Python 프로세스에서 실행 중인 GPU 기반 세션을 구동할 수 있습니다. 작업을 호출하면 SageMaker AI가 선택한 베이스 모델과 Low-Rank Adaptation(LoRA) 어댑터가 로드된 관리형 가속기에서 이를 실행합니다. 전체 작업을 사전에 기술하는 구성 파일도 없고 클러스터 생성 단계도 없습니다. 코드를 통해 각 단계에서 일어나는 일을 대화형으로 직접 제어할 수 있으므로 빠르게 반복하고 디버깅할 수 있습니다.

책임 분담은 다음과 같습니다:

이 API는 다음의 세 가지 클라이언트를 중심으로 구성되어 있습니다. SageMaker Python SDK, 순서대로 생성됩니다. A ServiceClient 플랫폼 수준 설정(실행 역할, 모델 패키지 그룹, Amazon Simple Storage Service (Amazon S3) 출력 위치, 그리고 선택적으로 AWS Key Management Service (AWS KMS) 키 및 가상 프라이빗 클라우드(VPC) 구성(세션별로 재정의할 수도 있음))을 보유하고 세션을 시작합니다. The ServiceClient 인스턴스화합니다 TrainingClient (학습 및 체크포인팅 작업 예: forward, forward_backward, optim_step, 그리고 save_weights_for_sampler), 이는 다시 a를 인스턴스화합니다 SamplingClient (sample 그리고 compute_logprobs), 롤아웃 및 평가 중 생성과 점수 산정에 사용됩니다). 각 API 호출은 루프 내의 하나의 작업에 매핑됩니다. 활성 세션은 OpenAI 호환 채팅 완성(chat completions) 엔드포인트도 제공하므로, 기존의 OpenAI 클라이언트 코드로 이를 통해 샘플링할 수 있습니다.

그림 1: 학습 세션 API를 위한 SageMaker Python SDK 클라이언트 계층 구조

강화 학습의 경우, 강화 학습루프는 특유의 형태를 가집니다. 다음 그림은 하나의 학습 단계와 각 부분이 실행되는 위치를 보여줍니다. 녹색 단계는 서비스 관리형 GPU에서 API를 호출하고, 보라색 단계인 루프 제어 및 보상 로직은 사용자의 머신에서 실행됩니다.

그림 2: 관리형 GPU에서 수행되는 샘플링, 그래디언트, 체크포인트 작업과 사용자 머신에서 실행되는 루프 제어 및 보상 로직을 보여주는 하나의 학습 단계

각 단계는 프롬프트 배치를 불러오고 (1) GPU에서 현재 모델로부터 샘플링합니다. 그런 다음 여러분의 Python 코드가 (2) 해당 출력들에 점수를 매깁니다: 검증기, 리워드 모델 호출, 테스트 하네스 등 코드로 표현할 수 있는 무엇이든 가능합니다. (3) forward_backward 그리고 optim_step 점수를 GPU에서의 그래디언트 업데이트로 변환하고, (4) save_weights_and_get_sampling_client() 샘플러를 새로 고쳐 다음 단계가 업데이트된 모델에서 샘플링하도록 합니다.

데이터, 루프 제어, 보상은 사용자가 소유하고, SageMaker AI가 GPU 집약적인 샘플링, 그래디언트, 체크포인트 작업을 실행합니다. 동일한 프리미티브로 지도 파인튜닝(SFT) 및 기타 LLM 학습 기법도 지원합니다. SFT의 경우 샘플링과 보상 단계를 제외하고 forward_backward 를 cross_entropy 과 함께 레이블된 예제에 대해 호출하면 됩니다.

이 기능은 누구를 위한 것이며 Amazon SageMaker AI 내에서 어떤 위치를 차지합니까?

Training session API는 빠른 반복 주기와 알고리즘적 자유(예: 커스텀 손실, 새로운 RL 보상 신호, 다단계 파이프라인)가 필요하지만 GPU 인프라를 직접 프로비저닝하거나 운영하고 싶지 않은 경우에 적합한 출발점입니다.

사용량은 처리된 토큰당 청구되므로 루프가 소비하는 컴퓨팅에 대해서만 요금이 부과됩니다. 보상 함수가 여러분의 머신에서 배치 하나를 채점하는 데 10초가 걸린다면, 그 실행 동안에는 GPU 시간이 청구되지 않습니다. 시간당 GPU 청구와 비교하면 두 가지 실질적인 장점이 있습니다. 첫째, 유휴 GPU 비용이 없습니다. 보상 계산, 데이터 로딩, 로깅, 디버깅은 청구 시간 밖에서 이루어지며, GPU 집약적 작업을 호출할 때만 청구됩니다. 둘째, 비용이 월클록 시간이나 클러스터 규모가 아니라 데이터세트와 생성량에 비례하므로 지출 예측이 쉬워지고 실험 간 비교가 용이합니다. 복잡한 보상 모델, 다중 도구 에이전트 롤아웃, 휴먼인더루프 평가 같은 로컬 계산이 많은 워크로드의 경우 세션 전체에 GPU를 예약하는 방식 대비 절감액이 상당할 수 있습니다.

Training session API는 데이터세트뿐 아니라 학습 알고리즘으로 정의되는 작업을 하는 팀을 위해 만들어졌습니다. 루프를 직접 구성하기 때문에 동일한 프리미티브로 레이블된 데이터에 대해 cross_entropy 손실을 적용하는 지도 파인튜닝(SFT), 을 사용하는 강화 학습 ppo, importance_sampling(GRPO, PPO 및 기타 그룹 상대적 방법), 그리고 동일한 구성 요소로 만든 선호도 최적화나 증류 같은 cispo 커스텀 목적 을 지원합니다. 적합한 대상은 다음과 같습니다:RL 방법(GRPO, PPO, 중요도 샘플링 변형), 커스텀 어드밴티지 추정기 또는 새로운 목적을 구현하는 포스트 트레이닝 및 정렬 연구자.

• 비표준 손실, 데이터 커리큘럼, 또는 구성 기반 레시피로 표현할 수 없는 다단계 파이프라인이 필요한 지도 파인튜닝 또는 선호도 최적화 팀.
• 커스텀 롤아웃 로직, 보상 셰이핑, 커리큘럼이 필요한 에이전트 및 멀티턴 RL 개발자.
• 구성 기반 파인튜닝의 한계를 넘었지만 GPU 클러스터의 운영 부담은 떠맡고 싶지 않은 머신 러닝(ML) 엔지니어.
• 기존 학습 루프 코드를 포팅하는 연구자. 작업 기반 프로그래밍 모델이 저수준
• 기존의 학습 루프 코드를 포팅하는 연구자들에게도 유용한데, 연산 기반 프로그래밍 모델이 저수준에 맞춰 작성된 루프에 직접 대응되기 때문입니다. forward_backward, optim_stepAPI에 대해 작성된 루프에 직접 대응하기 때문입니다. sample 지원되는 베이스 모델과 각 모델이 제공되는 AWS 리전은

training session API 문서 에 나와 있습니다..

Training session API가 Amazon SageMaker AI 내에서 차지하는 위치

Training session API는 SageMaker AI의 모델 커스터마이징 옵션 중에서 뚜렷한 위치를 차지합니다. 완전 관리형 Serverless Model Customization, 그리고 Amazon SageMaker Training Jobs 이나 Amazon SageMaker HyperPod.

그림 3: SageMaker AI의 모델 커스터마이징 옵션 중 학습 세션 API가 차지하는 위치

학습 루프 해부: 수학 문제에 대한 GRPO

이 섹션에서는 최소한의 강화 학습 루프를 분해하여 연산들이 어떻게 결합되는지 보여줍니다. 과제는 GSM8K로, 각각 단일 숫자 답을 가진 약 8,500개의 초등학교 수준 수학 단어 문제 벤치마크로, 다음 예시와 같습니다:

Natalia는 4월에 자신의 친구 48명에게 클립을 팔았고, 5월에는 그 절반에 해당하는 클립을 팔았습니다. Natalia는 4월과 5월을 합쳐 모두 몇 개의 클립을 팔았나요? → 72

우리는 OpenAI의 GPT-OSS 20B 모델에 대해 Group Relative Policy Optimization(GRPO)으로 LoRA 어댑터를 학습시킵니다. 각 질문에 대해 모델이 여러 개의 응답을 생성하고, 각 응답의 정확도에 점수가 매겨지며, 그룹 평균을 능가하는 응답이 강화됩니다. 완전히 실행 가능한 예제는 Amazon SageMaker 예제 저장소.

에 있습니다. 세션을 열기 전에 문서의 전제 조건 을 완료하십시오.

세션 열기

세션을 생성한 다음, 여기서 토크나이저와 SamplingClient 를 가져옵니다:

from sagemaker.train.training_session import ServiceClient, SamplingParams, AdamParams

svc = ServiceClient(role_arn=role, region=REGION)

training_client = svc.create_lora_training_client(
    "openai/gpt-oss-20b",                        # a name from svc.list_base_models()
    training_session_name=SESSION_NAME,
    rank=32, alpha=64,                           # LoRA rank and scaling factor
    model_package_group_arn=MODEL_PACKAGE_GROUP_ARN,
    s3_output_path=OUTPUT_S3_PATH,
)
tokenizer = training_client.get_tokenizer()
sampling_client = training_client.create_sampling_client()

보상과 어드밴티지 정의

작업별 로직은 여러분의 머신에서 실행되는 두 개의 짧은 함수입니다:

def compute_reward(response_text, ground_truth):
    predicted = extract_answer(response_text)      # parse the number after "#### "
    return 1.0 if predicted == ground_truth else 0.0

def compute_advantages(rewards):
    rewards = np.array(rewards)
    advantages = rewards - rewards.mean()          # center within the group (the GRPO baseline)
    return advantages / (rewards.std() + 1e-8)     # normalize

루프

각 단계는 Figure 2에 표시된 네 단계를 따릅니다. draw, build_prompt, 그리고 build_training_datum 는 샘플에 포함된 짧은 헬퍼입니다:

for step in range(NUM_STEPS):
    datums, pending = [], []
    # (1) sample K responses per question in one request each (GPU); returns a future right away
    for question, ground_truth in draw(PROMPTS_PER_STEP):          # data loading is local
        prompt = tokenizer.encode(build_prompt(question), add_special_tokens=False)
        pending.append((prompt, ground_truth, sampling_client.sample(
            prompt, num_samples=K,
            sampling_params=SamplingParams(max_tokens=512, temperature=0.7, top_p=0.9))))

    for prompt, ground_truth, future in pending:
        result = future.result()                                    # wait once the batch is in flight
        # (2) score each response, then center the rewards within the group (local)
        rewards = [compute_reward(tokenizer.decode(s.tokens), ground_truth) for s in result.sequences]
        advantages = compute_advantages(rewards)
        datums += [build_training_datum(prompt, s, a) for s, a in zip(result.sequences, advantages)]

    # (3) one gradient step over the batch (GPU); the optimizer step is sequenced behind it
    fwd_bwd_op = training_client.forward_backward(datums, loss_fn="ppo")
    optim_op = training_client.optim_step(AdamParams(learning_rate=1e-5))
    fwd_bwd_op.result(), optim_op.result()

    # (4) refresh the sampler so the next step's rollouts come from the updated policy (GPU)
    sampling_client = training_client.save_weights_and_get_sampling_client()

training_client.stop()

이 루프가 동작하게 만드는 세 가지 세부 사항이 있습니다:

연산은 future를 반환합니다. GPU에 바인딩된 모든 호출은 즉시 APIFuture 를 반환합니다. 요청은 백그라운드 서브미터에 의해 큐에 들어가 전송되며, .result() 가 값을 기다립니다. 루프는 어느 하나를 기다리기 전에 한 단계의 모든 샘플링 요청을 제출하므로, 전체 배치가 동시에 생성됩니다. 학습 연산은 호출한 순서대로 실행되므로, optim_step 는 forward_backward 가 끝나기 전에 제출될 수 있습니다.

데이터 항목은 손실에 필요한 모든 것을 담고 있습니다. build_training_datum 는 프롬프트와 응답 토큰, 에서 sample()가져온 토큰별 로그확률, 어드밴티지, 그리고 프롬프트가 손실에서 제외되도록 하는 weights 마스크를 함께 담고 있습니다.

체크포인트는 모델 패키지입니다. 저장할 때마다 모델 패키지 그룹에 버전이 등록되며, 위 루프의 단계별 샘플러 갱신도 포함됩니다. 학습 후 두 번의 추가 저장으로 배포용 병합 모델과 가중치 및 옵티마이저를 포함한 재개 가능한 상태가 추가됩니다.

training_client.save_weights_for_sampler(merge_weights=True).result()   # merged, deployable model
training_client.save_state().result()                                   # resumable state (weights + optimizer)

재개하려면 최신 체크포인트에서 새 세션을 다음과 같이 시작합니다 create_training_client_from_state_with_optimizer.

실행 중인 세션은 stop() 를 호출하거나 만료될 때까지 동시 세션 할당량에 포함되므로, 오류 경로를 포함하여 작업을 마치면 반드시 중지하거나 with 블록에서 클라이언트를 여십시오. 병합된 체크포인트는 실시간 엔드포인트에 배포됩니다. 체크포인트를 대신 LoRA 어댑터로 제공하려면 쿡북의 배포 노트북.

우리가 샘플을 실행한 결과, 40개 문제에 대해 10회의 학습 스텝만으로 홀드아웃 정확도가 56퍼센트에서 75퍼센트로 향상되었습니다.

실제 워크로드: 경쟁 프로그래밍에 대한 강화 학습(RL)

GSM8K는 메커니즘을 잘 드러내 주지만, 실제 프로덕션 RL 워크로드에는 그 이상이 필요합니다. 더 긴 생성, 멀티턴 상호작용, 외부 채점기, 그리고 GPU를 계속 바쁘게 유지할 만큼의 처리량이 그것입니다. 이 SageMaker 학습 세션 쿡북 동일한 연산 위에 구축된 재사용 가능한 훈련 루프로 엔지니어링한 패키지들입니다. 새로운 RL 작업을 위해 환경을 작성합니다: 모델이 보는 것, 답변이 어떻게 채점되는지, 그리고 또 다른 턴을 받는지 여부입니다. 쿡북은 배칭, 챗 템플릿, 어드밴티지 추정, 겹치는 샘플링과 훈련, 검증, 재개 가능한 체크포인팅, 단계별 메트릭을 처리하고 각 실행의 구성을 기록하여 반복할 수 있도록 합니다. 이 code_rl 레시피는 이를 사용하여 경쟁 프로그래밍 문제로 모델을 학습시킵니다.

과제. 문제들은 DeepCoder 데이터셋(훈련용으로는 PrimeIntellect, TACO, LiveCodeBench; 테스트용으로는 Codeforces와 LiveCodeBench)에서 나옵니다. 모델이 Python 풀이를 작성하면, 이 풀이가 해당 문제의 단위 테스트들과 대비하여 다음에서 실행됩니다. SandboxFusion 컨테이너입니다. 모든 테스트가 통과하면 보상은 1이고 그렇지 않으면 0이며, 응답에 코드 블록이 없으면 0.1의 페널티가 부과됩니다.

멀티턴 도구 사용. 두 번의 턴으로 모델은 도구를 호출할 수 있습니다. check_solution 도구를 호출하고, 테스트 결과를 읽은 후, 최종 성적 전에 답변을 수정합니다. 도구 호출과 결과는 모델의 기본 채팅 형식으로 렌더링되므로, 모델은 추론 시 사용하는 것과 동일한 토큰 형식으로 학습됩니다.

이 레시피가 최소한의 루프에 추가하는 것:

• 중첩된 샘플링과 학습: 롤아웃 워커는 트레이너가 한 스텝을 실행하는 동안에도 계속 생성을 수행합니다. 이전 정책에서 샘플링된 응답 그룹은 최대 k 스텝까지 오래되었을 때 사용되며, 그렇지 않으면 재대기열에 등록되고, cispo 손실이 정책 차이를 보정합니다.
• 스트리밍 그래디언트 계산: forward_backward 그룹이 완료될 때 하위 배치 단위로 제출하며, 단계당 한 번 제출하는 방식이 아닙니다.
• 홀드아웃 검증, 체크포인트 재개, 단계별 메트릭, 선택적으로 다음에 게시됨: SageMaker managed MLflow.

Figure 4: Code RL 설정으로, 왼쪽에 로컬 학습 루프와 샌드박스가 있고 오른쪽에 관리형 GPU가 있습니다

GPT-OSS 20B를 위해 LoRA 어댑터(rank 32)를 50 스텝 동안 학습했으며, 각 스텝은 128개의 문제와 문제당 8개의 응답으로 구성되었고, two turns와 staleness bound k=3을 사용했습니다. 학습 보상은 0.46에서 0.60으로 상승했으며, 200개의 보류된 테스트 문제에 대한 정확도는 40 퍼센트에서 50.5 퍼센트로 상승했습니다.

이를 실행하려면 샌드박스를 시작하고 레시피를 실행하세요:

bash tools/sandbox_up.sh
python -m training_sessions_cookbook.recipes.code_rl.train \
    --config training_sessions_cookbook/recipes/code_rl/configs/coderl_oss20b_service.json \
    --max_steps_off_policy 3 --stale_policy requeue --num_rollout_workers 80 --max_steps 50 \
    --base_model_arn "$BASE_MODEL_ARN" \
    --model_package_group_arn "$MODEL_PACKAGE_GROUP_ARN" \
    --role_arn "$ROLE_ARN" \
    --output_s3_path "$OUTPUT_S3_PATH"

이 쿡북에는 동일한 구조에 대한 더 많은 레시피가 있습니다. 여기에는 채팅 데이터셋(Tulu 3, UltraChat)에 대한 지도 파인튜닝, 전체 규모의 GSM8K에 대한 RL, 바이너리 피드백으로부터의 KTO 선호도 정렬이 포함됩니다. 각 레시피 README에는 명령어, 하이퍼파라미터, 예상 메트릭이 나열되어 있습니다.

결론

SageMaker AI 훈련 세션 API는 추상화 경계를 훈련 연산 수준으로 낮춥니다. 자신의 머신에서 Python으로 훈련 루프를 작성하면, 서비스가 관리형 GPU에서 이를 실행하며 모델 및 LoRA 로딩, 분산 실행, 장애 복구, 체크포인팅을 처리합니다. 완전한 GRPO 루프는 네 가지 연산으로 구성됩니다. sample, forward_backward, optim_step그리고 save_weights_and_get_sampling_client보상과 어드밴티지 로직은 자신의 코드에 담습니다. Code RL 레시피에서 보여주듯, 동일한 연산들이 샌드박스 방식의 채점을 사용하는 멀티턴, 도구 사용 RL로 확장됩니다.

커스텀 보상, 에이전틱 롤아웃, 단계별 디버깅, 새로 발표된 기법 모두 동일한 프로그래밍 모델을 사용합니다. 알고리즘은 여러분의 루프에 있고, SageMaker AI가 GPU를 실행합니다.

시작하려면:

• 다음에서 레시피를 선택합니다: SageMaker training sessions cookbook예를 들어 Code RL.
• Amazon SageMaker examples 리포지토리에서 GSM8K 예제 를 실행하고 자신만의 보상 함수로 교체합니다.
• 전제 조건, 할당량, 리전, 지원 모델 및 API 참조를 확인하려면 training session APIs 문서 를 읽어보세요.

감사의 말

이 게시물에 기여하고 지도해 주신 Arun Nagarajan, Xichao Wang, Shruti Sharma님께 감사드립니다.

저자 소개