SDK V3 학습: ModelTrainer로 학습 Job 제출하기¶
Scope
V3에서 학습 Job을 제출하는 방법만 다룹니다. V2와의 전체 차이와 마이그레이션 함정은 SDK V3 개요, 배포는 SDK V3 배포, LoRA 설계와 하이퍼파라미터 같은 학습 내용 자체는 파인튜닝에 있습니다.
ModelTrainer로 학습 Job 제출¶
왼쪽 V2는 framework_version="1.12.0", py_version="py38"을 주면 SDK가 내부 lookup 표로 DLC 이미지를 찾아냅니다. 오른쪽 V3는 training_image를 직접 넘깁니다. SDK가 프레임워크를 알 필요가 없어졌습니다. 나머지 인자는 SourceCode, Compute, InputData 세 config 객체로 흩어집니다.
tracks/*/02_train_sft_sagemaker.ipynb에서 추린 형태입니다. 하이퍼파라미터 dict와 role, sagemaker_session은 V2와 같은 자리에 남고, 나머지가 전부 config 객체로 이동합니다.
import boto3
from sagemaker.core.helper.session_helper import Session
from sagemaker.core.image_uris import retrieve
from sagemaker.train.model_trainer import ModelTrainer
from sagemaker.core.training.configs import (
SourceCode, Compute, InputData, StoppingCondition,
)
sess = Session(boto3.Session(region_name="us-east-1"))
# 이 프로젝트는 .env의 DLC_IMAGE_URI(완전 URI)를 그대로 씁니다. retrieve는 그 env가
# 없을 때의 fallback입니다: common/dlc.resolve_training_image()의 우선순위.
image_uri = retrieve(framework="pytorch", region="us-east-1",
version="2.8.0", py_version="py312",
image_scope="training", instance_type="ml.g6.2xlarge")
trainer = ModelTrainer(
training_image=image_uri,
source_code=SourceCode(source_dir="scripts", entry_script="train.py",
requirements="requirements.txt"),
compute=Compute(instance_type="ml.g6.2xlarge", instance_count=1),
hyperparameters={"model_id": "google/gemma-4-E4B-it", "epochs": 2,
"use_qlora": True, "merge_adapter": True},
environment={"HF_TOKEN": "..."},
role=role, sagemaker_session=sess,
base_job_name="gemma-extraction-train",
# 생략하면 SDK가 1시간을 넣습니다: 아래 함정 절 참고
stopping_condition=StoppingCondition(max_runtime_in_seconds=4 * 3600),
)
trainer.train(input_data_config=[InputData(channel_name="train",
data_source=train_s3)],
wait=False, logs=False)
print(trainer._latest_training_job.training_job_name)
import sagemaker
from sagemaker.huggingface import HuggingFace # V3에 없습니다
from sagemaker.inputs import TrainingInput
estimator = HuggingFace(
entry_point="train.py", source_dir="scripts",
instance_type="ml.g6.2xlarge", instance_count=1,
transformers_version="4.36", # 이미지를 버전 인자로 골랐습니다
pytorch_version="2.1", py_version="py310",
hyperparameters={"model_id": "...", "epochs": 2},
environment={"HF_TOKEN": "..."},
role=role, sagemaker_session=sagemaker.Session(),
base_job_name="gemma-extraction-train",
max_run=4 * 3600, # V3의 StoppingCondition
)
estimator.fit({"train": TrainingInput(train_s3)}, wait=False, logs=False)
print(estimator.latest_training_job.name)
ModelTrainer에는 hyperparameters가 --key value CLI 인자로 직렬화돼 train.py에 들어갑니다. --use_qlora True 형태이므로 argparse에서 action="store_true"를 쓰면 parsing에 실패합니다. 이 프로젝트의 str2bool 처리 이유는 파인튜닝에 있습니다.
ModelTrainer 하나로 합쳐진 estimator들¶
왼쪽의 estimator 7종(PyTorch, TensorFlow, HuggingFace, XGBoost, SKLearn, MXNet 등)은 ModelTrainer 하나로 통합됩니다. 오른쪽 아래의 specialized trainer는 별도 도구입니다.
둘의 역할이 다릅니다.
ModelTrainer |
특화 trainer | |
|---|---|---|
| 하는 말 | "내 학습 코드와 컨테이너가 있으니 Amazon SageMaker AI 인프라에서 돌려라" | "foundation model을 이 기법으로 파인튜닝하고 싶고, 인프라는 신경 쓰고 싶지 않다" |
| 성격 | 범용 compute orchestrator | 정해진 모델, 기법, 파라미터만 받는 고수준 워크플로 |
| 내가 주는 것 | 이미지, 스크립트, 하이퍼파라미터 | 모델과 데이터 |
이 프로젝트는 ModelTrainer 쪽입니다. TRL SFTTrainer와 PEFT를 직접 조합하고 최신 Gemma를 바로 쓰기 위해 train.py를 들고 가기 때문입니다(파인튜닝에 그 선택 근거가 있습니다).
특화 trainer는 sagemaker.train에서 바로 import됩니다(3.16.0 확인).
이 프로젝트는 아래 기능을 쓰지 않습니다
특화 trainer, 평가, AI Registry, Batch queue는 symbol과 signature만 확인했으며 실제 학습은 실행하지 않았습니다. 이 프로젝트에서 실행해 확인한 경로는 ModelTrainer + custom train.py입니다.
아래는 "V3에 이런 것이 생겼다"는 지도이니, 쓰실 때는 SDK 저장소의 현행 시그니처를 확인하세요.
평가가 SDK 안으로 들어왔습니다¶
V2에서 파인튜닝 결과를 표준 benchmark로 평가하려면 데이터셋 준비, 평가 loop, metric 계산, 결과 저장을 직접 구현해야 했습니다. V3는 세 종류의 evaluator를 제공합니다.
| evaluator | 무엇을 하나 |
|---|---|
BenchMarkEvaluator |
표준 벤치마크 11종을 기본 제공 |
LLMAsJudgeEvaluator |
Bedrock Evaluations 기반. foundation model을 judge로 골라 품질, 안전성 채점 |
CustomScorerEvaluator |
자체 채점 로직을 끼워 넣기 |
셋 다 결과를 MLflow에 자동 기록합니다. 이 프로젝트는 대신 코스별 메트릭을 직접 계산합니다(common/eval_utils.py). 추출은 arg_f1, 분류는 macro-F1처럼 태스크에 맞춘 지표가 필요해서입니다.
AI Registry: 데이터셋과 evaluator에 버전을 붙입니다¶
S3 경로를 주고받으며 "모두 같은 버전을 쓰고 있겠지" 하고 믿는 대신, 데이터셋과 evaluator를 버전이 붙은 hub content로 등록합니다.
AI Registry Evaluator는 실행 주체가 아니라 저장과 메타데이터 엔티티입니다. SageMaker Hub 안의 버전 레코드이고, 그 안에 평가 로직 참조(reward 프롬프트 문자열 또는 Lambda ARN)를 담습니다. 그래서 reward 프롬프트를 한 번 등록해 두고, LLMAsJudgeEvaluator나 RLAIFTrainer를 설정할 때 그 ARN으로 가리키는 식으로 씁니다.
설치본에서는 sagemaker.ai_registry subpackage가 이 영역을 담당합니다.
AWS Batch 큐에 학습 Job을 넣기¶
Job이 많아 scheduling이 필요할 때, ModelTrainer를 SageMaker AI에 바로 제출하지 않고 AWS Batch queue에 넣을 수 있습니다. priority queueing, fair-share scheduling, retry는 Batch가 담당합니다.
ModelTrainer는 똑같이 만들고, .train()을 부르는 대신 큐에 넘깁니다.
from sagemaker.train.aws_batch.training_queue import TrainingQueue
# TrainingQueue: submit / map / get_job / list_jobs / list_jobs_by_share
학습에서 배포로 넘어가는 지점¶
ModelBuilder(model=trainer): trainer 객체를 그대로 넘깁니다. S3 경로를 손으로 옮겨 적을 필요가 없고, 이 자리가 학습과 서빙을 잇는 이음매입니다.
이후 단계는 SDK V3 배포에서 다룹니다.


