Michael Barlow

코드로 구현하는 모델 학습 (Model Training as Code)

요약(TL;DR): 모델 학습은 많은 전문화된 단계와 팀을 필요로 할 만큼 복잡해졌으며, 이들 사이의 수동적인 조정은 확장성(scale)을 갖기 어렵습니다. Aleph Alpha에서는 전체 학습 파이프라인(pipeline)을 코드로 구현하여 모델 학습을 협업 소프트웨어 프로젝트로 전환하는 모델 팩토리(model factory)인 Savanna를 구축했습니다. Savanna에서 엔드 투 엔드(end-to-end) 학습 실행은 격리(hermetic)되어 있으며, 클릭 한 번으로 실행 가능합니다. 이 포스트에서는 Savanna와 이것이 필요한 이유, 그리고 이를 효과적으로 만드는 엔지니어링 문화에 대해 설명합니다.

서론 (Introduction)

모델 학습은 빠르게 변화하고 있습니다. 새로운 단계들이 파이프라인에 계속 추가되고 기존 단계들은 더욱 복잡해지고 있으며, 이는 세 가지 주요 이유로 인해 모델 학습을 엔지니어링 과제로 만듭니다. 첫째, 복잡성이 증가할수록 오류의 가능성도 커집니다. 데이터, 코드 또는 설정(configuration)에서의 버그나 불일치는 전체 학습 실행을 실패하게 하거나 발산(diverge)하게 만들 수 있습니다. 둘째, 실패 비용이 계속 상승하고 있습니다. 모델은 더 커지고, GPU 가격은 오르며, 실행당 처리되는 데이터는 점점 더 많아집니다. 수천 시간의 GPU 시간을 소모하고 있을 때, "실수"는 매우 값비싼 단어입니다.

세 번째이자 가장 어려운 과제는 조직적인 문제입니다. 복잡성은 이미 한 개인의 역량을 넘어섰기에, 저희와 같은 연구소들은 크고 전문화된 팀들을 구축해 왔습니다. 문제는 그 팀들을 조정하는 것이 됩니다. 어떻게 하면 구성원들이 자신의 전문 분야에서 최신 연구를 자율적으로 탐색하면서도, 파이프라인을 망가뜨리거나 서로의 작업을 방해하지 않고 자신의 변경 사항을 프로덕션 파이프라인(production pipeline)에 통합할 수 있을까요? 그리고 개별 단계의 개선 사항이 최종적으로 더 나은 모델로 이어지도록 어떻게 보장할 수 있을까요?

전통적인 수동 모델 학습 프로세스는 이러한 질문들에 대해 좋은 답을 내놓지 못합니다.

수동 모델 학습의 숨겨진 비용 (The hidden cost of manual model training)

전통적인 수동 프로세스를 자세히 살펴보겠습니다. 거시적인 관점에서 모델 학습은 단순해 보입니다. 인터넷의 정보를 흡수하기 위한 사전 학습 (pre-training)이 이루어지고, 이어서 지시 이행 (instruction-following)을 배우기 위한 사후 학습 (post-training)이 진행됩니다. 하지만 실제로 첫 시도에 바로 좋은 모델을 학습시킬 수는 없습니다. 적절한 데이터 혼합 (data mix), 아키텍처 (architecture), 그리고 학습 레시피 (training recipe)에 도달하는 과정은 평가 (evaluation)에 의해 가이드되는 반복적이고 연산 집약적인 (compute-bound) 프로세스이며, 아래의 각 화살표는 상대적인 GPU 비용에 따라 가중치가 부여됩니다:

이러한 각 구성 요소는 여러 개의 전담 팀이 필요할 정도로 충분히 복잡합니다. 예를 들어, 현대적인 사후 학습 (post-training)은 지도 미세 조정 (supervised fine-tuning, SFT) 단계와 그 뒤를 잇는 강화 학습 (reinforcement learning, RL) 단계로 구성됩니다. SFT와 RL은 서로 다른 기술 스택과 도구를 필요로 하지만, 모델을 학습시키기 위해서는 반드시 통합되어야 합니다.

수동 실험실 환경에서 단일 모델이 파이프라인을 통과하는 여정이 어떠할지 가정해 봅시다:

데이터 팀이 새로운 데이터 혼합을 완료하고 Slack을 통해 데이터베이스 경로를 사전 학습 (pre-training) 팀에 전달하면, 사전 학습 팀은 몇 주간의 실행을 시작합니다. 2주가 지났을 때, 저장 용량 할당량 (storage quota)이 가득 차면서 학습 실행이 중단됩니다. 파일 시스템이 수동으로 관리되기 때문에, 파일 이름에 do_not_delete가 포함된 30TB 규모의 데이터셋을 삭제해도 안전한지 아무도 확신하지 못하며, 사전 학습 팀이 이 문제를 해결하는 동안 GPU는 유휴 상태 (idle)로 방치됩니다. 마침내 재실행할 때, 그들은 플래그 (flag) 설정을 잊지 않았기를 바라며 기억과 Slack 대화 내용을 바탕으로 원래의 설정을 재구성합니다. 이것이 첫 번째 숨겨진 비용입니다: 모든 수동 단계는 인간의 실수가 발생할 수 있는 기회입니다.

사전 학습 (Pre-training)이 마침내 완료되면, 사전 학습 팀은 체크포인트 (Checkpoint)를 SFT 팀에 전달합니다. 좋은 레시피 (Recipe)를 찾기 위해, SFT 팀은 서로 다른 설정과 데이터 혼합 (Data mixes)을 사용하여 병렬 학습 스윕 (Sweep)을 수동으로 시작합니다. 체크포인트가 들어올 때마다 팀은 각 체크포인트에 대해 평가 스크립트 (Evaluation script)를 실행하고, 그 결과와 분석 내용을 Slack에서 공유합니다. 어떤 레시피는 유망해 보이고, 어떤 것은 그렇지 않으며, 이들은 괜찮은 레시피를 좁혀 나갈 때까지 몇 주 동안 이 과정을 반복합니다. 자신들도 모르는 사이에, 팀은 몇 달 전 이전 사전 학습된 체크포인트를 위해 이미 완료했던 일부 실험들을 반복했습니다.

이것이 두 번째 숨겨진 비용입니다: 팀이 학습한 내용을 잊어버린다는 것입니다. 하이퍼파라미터 (Hyperparameter)의 현재 값 뒤에 숨겨진 추론 과정에 대한 영구적인 기록도 없고, 데이터 혼합과 이를 구성하는 데이터셋 사이의 공식적인 연결 고리도 없으며, 모델을 생성한 학습 레시피와 모델을 연결하는 명확한 귀속 (Attribution)도 없습니다. 수동적인 실험실 환경에서 이러한 계보 (Lineage)는 Slack, 파일 시스템, 실험 관리자 (Experiment manager), 그리고 다양한 위키 (Wiki) 페이지에 흩어져 있으며, 시간이 지남에 따라 쉽게 유실됩니다.

SFT 팀은 자신들의 체크포인트를 RL 팀에 전달하고, RL 팀은 이를 베이스로 하여 학습 실행 (Training run)을 시작합니다. 최종 모델의 성능이 기대에 미치지 못합니다. RL 레시피가 지난달의 SFT 체크포인트에 과적합 (Overfit)된 것일까요, 아니면 SFT 체크포인트 자체가 문제일까요? 2주간의 디버깅 (Debugging) 끝에, RL 팀은 후자가 원인임을 확인합니다. 어느 팀도 상대방의 단계를 실행할 수 없기 때문에, 각 팀은 최종 모델보다는 파이프라인 (Pipeline)의 각자 맡은 부분만을 최적화했습니다. 그리고 통합 (Integration)이 수동적인 인수인계로 이루어지기 때문에 드물게 발생하며, 이로 인해 매번 한 달 치의 차이 (Divergence)를 조정해야 하는 상황이 발생합니다. 이것이 세 번째 숨겨진 비용입니다: 수동적이고 드문 인수인계가 팀의 소유권 (Ownership)을 파편화합니다.

수동적인 모델 학습이 확장 가능하지 않다는 점은 명확하며, 자동화만이 유일하게 결여된 요소도 아닙니다. 이러한 숨겨진 비용들은 각각 동일한 근본적인 문제에서 기인합니다. 즉, 파이프라인 (Pipeline)이 공유되고 지속 가능한 산출물 (Artefact)이 아닌 팀원들의 머릿속에만 존재한다는 점입니다. 규모를 확장하려면 파이프라인 자체가 팀이 협업할 수 있는 다루기 쉬운 형태여야 합니다.

소개: 코드로 구현하는 모델 학습 (Model Training as Code)

'Savanna'라는 코드명을 가진 우리의 모델 팩토리 (Model Factory)는 전체 모델 학습 파이프라인과 프로세스를 명령형 코드 (Imperative code)로 구현합니다. 우리는 이것을 코드로 구현하는 모델 학습 (Model Training as Code, MTaC)이라고 부릅니다. 다음은 Savanna의 간단한 사후 학습 (Post-training) 파이프라인을 의사 코드 (Pseudocode)로 나타낸 모습입니다:

async post_train(config: PostTrainConfig) -> PostTrainEvaluation:
sft_checkpoint = await sft(config.sft)
sft_eval = spawn evaluate(config.eval, sft_checkpoint)
rl_checkpoint = await rl(config.rl, sft_checkpoint)
rl_eval = spawn evaluate(config.eval, rl_checkpoint)
return PostTrainEvaluation(await sft_eval, await rl_eval)

파이프라인을 코드로 격상시킴으로써, 여러분은 세 가지 이점인 결합성 (Composability), 합의 (Consensus), 그리고 출처 (Provenance)를 얻을 수 있습니다.

• 결합성 (Composability): 결합성은 수동 단계를 타입이 지정된 입력과 출력을 가진 함수로 표현함으로써 얻어집니다. 이를 통해 추상화를 구축하고 클릭 한 번으로 실행되는 엔드 투 엔드 (end-to-end) 파이프라인으로 결합할 수 있습니다. 파이프라인을 수정하는 것은 함수를 편집하는 것만큼 간단해지며, 중간 체크포인트를 평가하는 것과 같은 반복적인 작업은 for 루프로 자동화할 수 있고, 다른 파라미터화 (parametrisation)를 사용하여 파이프라인의 하위 집합이나 축소된 버전을 실행할 수 있으므로 테스트도 직관적입니다.

• 합의 (Consensus): 합의는 버전 관리 (version control)에서 비롯됩니다. 메인 브랜치 (main branch)는 모델을 학습시키는 방법에 대한 팀의 집단적 최선의 이해를 나타냅니다. 코드는 전체 학습 레시피 (training recipe)를 포함하고 있으므로, 학습 실행을 시작할 때 재구성해야 할 설정이나 잊어버려서는 안 될 플래그 (flag)가 없습니다.

• 출처 (Provenance): 출처는 코드 주석과 커밋 히스토리 (commit history)에서 나옵니다. 이는 메인 (main)으로 이어지는 학습 내용과 결정 사항들을 인코딩합니다. 과거의 학습 실행들은 이를 생성한 코드가 체크아웃(check out)하여 재실행할 수 있는 커밋에 고정되어 있으므로 재현 가능성을 유지합니다.

어떤 팀이든 스스로 전체 파이프라인을 실행할 수 있게 되면, 다른 팀의 단계를 실행할 수 있으며 자신의 부분만이 아니라 모델 전체를 대상으로 반복 (iterate)할 수 있습니다. 연구소(Labs)들은 규모를 확장할 때 일반적으로 모델 학습을 시간 순서에 따라 분해하며, 각 팀이 파이프라인의 특정 단계를 담당합니다. MTaC는 팀이 다국어 능력과 같은 모델의 동작을 엔드 투 엔드로 소유하는 역량 기반 분해 (capability-based decomposition)의 가능성을 열어줍니다.

적게 통합하고, 자주 통합하라

파이프라인이 코드로 구현되면 표준적인 코드 협업 모범 사례가 적용되며, MTaC를 최대한 활용하려면 이를 따라야 합니다. 이 중 가장 중요한 것은 트렁크 기반 개발 (trunk-based development)으로, 변경 사항을 가능한 한 빨리, 그리고 작은 단위로 main에 반영하는 방식입니다. 이를 통해 팀은 가장 빠른 기회에 서로의 작업물을 바탕으로 구축할 수 있고, 접근 방식이 틀렸을 경우 빠르게 실패(fail fast)할 수 있습니다. 만약 대신 수명이 긴 브랜치 (long-lived branches)에 변경 사항을 축적한다면, 이전과 동일한 통합 부채 (integration debt)를 지게 됩니다.

Savanna 사용하기

Savanna는 GitHub 내에서 동작하며, 이곳의 CI (지속적 통합)가 모델 학습의 진입점 (entrypoint) 역할을 합니다. 브랜치에 푸시(push)하거나 GitHub UI에서 수동으로 실행을 시작하여 이 CI를 트리거할 수 있습니다. 우리의 최적 모델을 학습시키는 것은 main 브랜치에서 CI를 트리거하는 것만큼 간단합니다. 또한, 우리는 풀 리퀘스트 (pull request) 시 CI를 사용하여 소규모 엔드 투 엔드 (end-to-end) 학습 실행으로 파이프라인을 빠르게 검증합니다. 이 과정은 5분 이내에 완료되므로, Savanna에 기여하는 것이 결코 느리게 느껴지지 않으며 변경 사항에 대한 확신을 줍니다. 파이프라인의 의미론적 회귀 (semantic regressions)를 포착하기 위해, 우리는 매일 밤 더 큰 규모의 엔드 투 엔드 학습 테스트 실행을 수행하며, 이를 통해 결과 모델이 평가 스위트 (evaluation suite)에서 측정 가능한 개선을 달성하는지 확인(asserting)함으로써 학습 로직을 검증합니다.

MTaC는 git blame을 통해 의사결정 계보 (decision lineage)를 제공하며, Savanna는 이를 아티팩트 계보 (artefact lineage)로 확장합니다. 실행 (runs)은 허메틱 (hermetic, 격리된)하며, 모든 비코드 아티팩트 (예: 데이터, 모델, 토크나이저 (tokenisers))는 불변 (immutable)하며 레지스트리 (registry)에 버전 관리되어 저장되므로, 매번 동일한 것을 얻고 있음을 확신할 수 있습니다. 실행을 시작하면 참조된 아티팩트, 학습 로그, 메트릭 (metrics), 평가 결과가 모두 해당 실행 및 결과 체크포인트 (checkpoint)와 연결되므로, 출력을 입력에 쉽게 귀속시킬 수 있습니다. 특정 데이터셋으로 어떤 모델이 학습되었는지 확인하려면 Slack 검색이 아닌 아티팩트 계보 그래프를 사용하면 됩니다.

Savanna는 대규모 실행을 쉽게 만들어 주지만, 이러한 실행의 형태를 발전시키려면 데이터, 하이퍼파라미터 (hyperparameters), 파이프라인 단계, 환경, 샤딩 토폴로지 (sharding topologies) 또는 기타 모든 측면을 더 작은 규모에서 실험적으로 수정할 수 있는 능력이 필요합니다. Savanna에서 실험은 브랜치에 변경 사항을 푸시하고 거기서 CI를 실행하는 것만큼 간단합니다. 새로운 데이터셋을 시도하고 싶다면, 절제 학습 설정 (ablation training config)에 추가하고 CI를 실행하면 됩니다. 결과 평가가 개선을 보여준다면, 이 변경 사항을 main에 병합하여 다음 대규모 실행의 일부로 만들 수 있습니다.

하이퍼파라미터 스윕 (Hyperparameter sweeps) 또한 쉽습니다. 파이프라인이 단순한 함수이기 때문에, 서로 다른 매개변수화 (parameterisations)를 사용하여 프로그래밍 방식으로 여러 번 호출할 수 있습니다. 예를 들어, 아래는 앞서 보여준 사후 학습 (post-training) 파이프라인에서 어떤 SFT 및 RL 학습률 (learning rates)이 가장 잘 작동하는지 찾기 위한 실험입니다:

async post_training_learning_rate_experiment(config: PostTrainConfig) -> str:
post_train_runs = []
for sft_learning_rate in (1e-4, 1e-5):
for rl_learning_rate in (1e-4, 1e-5):
run_config = config
.with_sft_learning_rate(sft_learning_rate)
.with_rl_learning_rate(rl_learning_rate)
post_train_runs.append(spawn post_train(run_config))
post_train_evaluations = gather post_train_runs
experiment_report_url = create_report(post_train_evaluations)
return experiment_report_url

그러면 Savanna가 스윕을 오케스트레이션 (orchestrate)하고 최종 결과를 수집하여 분석할 수 있는 보고서로 만들어 줍니다. 다음은 Savanna의 워크플로 엔진 (workflow engine) UI에서 시각화된 이 실험의 모습입니다:

파란색 노드는 "실행 중 (running)"이며 주황색 노드는 "캐시 대기 중 (awaiting cache)"입니다. post_train을 네 번 호출하지만, SFT 단계는 두 번만 실행됩니다. Savanna의 워크플로 엔진은 단계들이 동일한 입력을 가질 때를 인식하며, 해당 실행 단계가 완료되면 캐시에서 출력을 읽어옵니다. 이를 통해 중복 계산의 조합 폭발 (combinatorial explosion)을 걱정하지 않고도 파이프라인의 여러 측면을 수정하는 스윕을 실행할 수 있습니다.

사전 학습 (pre-training) 실행이나 오래 걸리는 RL 작업의 경우, 체크포인트 (checkpoints)를 점진적으로 방출하고 평가하여 학습 과정 전반에 걸쳐 모델이 어떻게 변화하는지 더 잘 이해할 수 있도록 합니다. 학습이 끝나면 Savanna는 벤치마크 제품군 (benchmark suite)을 통해 모델을 자동으로 평가하고, 모든 결과와 지표 (metrics)를 실행 보고서에 수집합니다. 또한 Savanna는 모델 리더보드 (leaderboard)를 자동으로 업데이트하며, 우리는 이를 하이델베르크 본사의 대형 스크린에 표시하여 새로운 최고 모델을 학습할 때 모두가 그 흥분을 공유할 수 있도록 합니다.

Savanna의 작동 방식

실행(run)을 트리거하면, Savanna는 Kubernetes에서 실행되는 당사가 선택한 워크플로 엔진(workflow engine)인 Flyte 상에서 작업을 시작합니다. Flyte는 임의의 태스크(arbitrary tasks)를 실행하고, 내구성이 있는 실행(durable execution), 시퀀싱(sequencing), 병렬성(parallelism), 재시도(retries), 캐싱(caching)을 제공하며, 실행 과정을 시각화할 수 있는 UI를 제공하는 편리한 방법을 지원합니다.