TorchTitan 업데이트: 모델 지원 확대, 대규모 학습 성능 개선

1. 모델별 코드를 따로 관리하는 부담 축소

TorchTitan은 생성형 인공지능의 사전학습과 사전학습 뒤 목적에 맞게 조정하는 후학습을 위한 플랫폼이다. 이번에는 Qwen3.5와 DeepSeek-V4 등 최신 모델을 지원하며, 모델 구성 요소와 병렬화 방식, 하드웨어별 GPU 연산을 수행하는 커널을 바꿔도 모델별 코드를 별도로 유지하지 않도록 구성 체계를 제공한다.

병렬화는 학습 작업을 여러 GPU에 나누는 방식이다. 개발팀은 같은 모델을 바탕으로 작업을 나누는 방법이나 연산 부품을 교체할 수 있다. 이는 새 모델을 붙일 때마다 학습 기반 전체를 다시 만드는 부담을 줄이는 방향이다.

2. TitanRL로 학습과 추론의 모델 구현 통합

TitanRL은 사전학습·후학습에 쓰는 모델을 vLLM 내부의 추론에도 그대로 사용한다. 모델 구조를 바꿀 때 학습용과 추론용 구현을 함께 수정해야 하는 중복 작업을 줄인다. 학습에 쓸 응답을 생성하는 작업은 완료를 기다리지 않고 진행하는 비동기 방식으로 진행되며, 학습과 생성 각각을 여러 병렬화 방식으로 확장할 수 있다.

팀은 입력별로 일부 계산만 활성화하는 전문가 모듈로 계산을 나누는 모델과 모든 입력에 전체 매개변수를 쓰는 밀집 모델에서 모델이 한 번에 참고하는 입력 범위인 컨텍스트·코드 실행 등을 포함한 작업으로 학습 데이터 처리의 한 반복 단위인 수백 스텝의 안정적인 학습을 시험했다고 밝혔다. 공개 사례 중 하나는 Qwen3.5-35B-A3B 기본 모델을 GB300 40대, 64K 컨텍스트로 학습·검증한 소프트웨어 문제 해결 실험이다.

팀은 배치 불변 모드(batch-invariant mode)를 사용하는 지원 온폴리시 구성, 즉 현재 정책으로 생성한 데이터를 학습에 쓰는 구성에서는 학습기와 생성기의 확률에 로그를 취한 값인 로그 확률이 이진수 자리까지 같은 수준으로 일치한다고 밝혔다. 팀은 이를 수학과 터미널 에이전트 강화학습의 수치 불일치를 살펴보는 연구에 활용했다.

3. 대규모 처리량은 실행 조건과 함께 제시

회사 발표 기준, 사내 GB300 클러스터에서 DeepSeek V3 671B에 다음 토큰 예측을 보조하는 MTP1 모듈을 붙이고 8비트 부동소수점 수 표현 형식인 MXFP8로 수를 낮은 정밀도 형식으로 표현하는 양자화를 적용한 머신러닝 성능을 재는 공개 평가 기준인 MLPerf 학습 벤치마크는 GPU 256대 구성에서 GPU당 초당 5,324.21토큰을 기록했다. 계산을 여러 단계로 나누는 파이프라인 단계는 2, 데이터를 나눠 여러 GPU에서 처리하는 정도인 데이터 병렬도는 128, 전문가 모듈 계산을 나눠 처리하는 정도인 전문가 병렬도는 64이며 한 장치에서 한 번에 처리하는 데이터 묶음인 로컬 배치와 배치를 나눈 작은 데이터 묶음인 마이크로배치 수는 각각 120이었다.

Muse Spark 1.2 OSS의 회사 발표 결과에서는 96K 컨텍스트·GPU 256대 구성의 GPU당 처리량이 학습된 라우팅에서 1,500.61, 전문가 간 부하를 강제로 균등하게 나누는 라우팅에서 1,839.52토큰/초였다. 192K·512대 구성에서는 각각 1,335.86과 1,588.22토큰/초였다. 서로 다른 컨텍스트·GPU·병렬 구성의 결과는 같은 조건의 성능 순위로 비교할 수 없다.

4. 통신 절감의 열쇠는 모델을 나누는 경계

FlexShard는 모델 상태를 GPU 사이에 나누는 경계를 더 세밀하게 정한다. 회사 발표 기준, DeepSeek V3를 H100 256대에서 실행할 때 여러 GPU에 걸친 학습 상태 갱신 기능인 분산 Muon 옵티마이저에 Owned·BlockShard 배치를 적용해 갱신 전 All Gather 통신을 없애면서 기존 완전 분할 데이터 병렬 학습 방식인 FSDP만 쓴 동등 구현보다 모델 매개변수를 갱신하는 알고리즘인 옵티마이저 구간 시간을 최대 79% 줄였다.

이익과 비용은 데이터를 나누는 방식에 달려 있다. 기존 방식은 분할 경계와 양자화 블록이 어긋나면 더 큰 형식으로 데이터를 전송하거나, 압축 전송에 필요한 보정값을 추가로 주고받아야 했다. 새 방식은 경계를 양자화 블록에 맞추는 대신 추가 통신을 없앴고, 회사 발표 기준 Llama3 405B·H100 256대에서 전체 학습 스텝 시간을 최대 20% 줄였다.

5. 체크포인트 저장 대기까지 단축

새 체크포인트 라이브러리는 첫 저장 전에 데이터를 설명하는 정보인 메타데이터를 공유하고 이후 저장에서는 여러 프로세스가 함께 주고받는 집단 통신을 생략한다. 학습 상태를 저장할 때마다 여러 프로세스가 발을 맞추던 부담을 줄이는 방식이다. 회사 발표 기준, H100 512대에서 아래와 같이 기존 방식보다 저장 대기와 저장·불러오기 시간이 짧아졌다.

표: 기존 방식과 새 체크포인트 방식의 소요 시간

조건: H100 GPU 512대, DSv3 671B, 체크포인트 간격 50(회사 발표 기준)

단위: 트레이너 대기 시간은 ms, 저장·불러오기 시간은 s

P99는 측정값을 작은 순서로 놓았을 때 99% 지점의 값이다. 저장 때문에 학습이 멈추는 시간이 줄면 체크포인트를 더 자주 남길 여지가 생긴다. 적용은 TorchTitan 저장소를 복제하고 안내 문서를 따라 시작한다.

이번 업데이트의 가치는 모델 구현의 중복과 분산 학습의 대기 시간을 함께 줄여, 새 모델 실험을 대규모 강화학습으로 잇는 개발 경로를 다듬은 데 있다.