Nanotron - 최소한의 3D 병렬 LLM 학습 치트시트
Nanotron (Hugging Face)은 최소한의 대규모 언어 모델 사전 학습 프레임워크로, 3D 병렬화 — 데이터, 텐서, 파이프라인 병렬 — 을 의도적으로 읽을 수 있는 Python으로 구현합니다. 목표는 실제로 이해할 수 있을 정도로 충분히 작고 해킹 가능하면서도 많은 GPU 전체에서 학습을 배포할 수 있을 정도로 충분히 큽니다. 복잡한 최대 규모 프레임워크 없이 분산 사전 학습의 메커니즘을 배우거나 맞춤화하려는 팀에게 강력한 선택입니다.
사전 학습은 계산 집약적입니다. 작게 시작하세요. 스케일링 전에 작은 모델과 몇 개의 GPU에서 구성을 검증하세요.
설치
| 방법 | 명령어 |
|---|
| 소스에서 | git clone https://github.com/huggingface/nanotron && cd nanotron && pip install -e . |
| Flash attention | pip install flash-attn (권장) |
| 요구사항 | PyTorch + CUDA GPU |
| 확인 | 1-2개 GPU에서 예제 구성 실행 |
3D 병렬화
| 차원 | 분할 | 구성 키 |
|---|
| 데이터 병렬 (DP) | 복제본 전체의 배치 | dp |
| 텐서 병렬 (TP) | GPU 전체의 레이어 | tp |
| 파이프라인 병렬 (PP) | GPU 전체의 모델을 단계로 | pp |
곱 dp × tp × pp는 총 GPU 개수와 같아야 합니다.
구성 (YAML)
Nanotron 학습은 모델, 데이터, 병렬화, 최적화기를 설명하는 YAML 구성으로 정의됩니다.
parallelism:
dp: 2
tp: 2
pp: 2
pp_engine: 1f1b # 파이프라인 일정
model:
model_config:
hidden_size: 2048
num_hidden_layers: 24
num_attention_heads: 16
optimizer:
learning_rate_scheduler:
learning_rate: 3.0e-4
optimizer_factory:
adam_beta1: 0.9
adam_beta2: 0.95
tokens:
sequence_length: 2048
micro_batch_size: 4
batch_accumulation_per_replica: 8
학습 실행
# 단일 노드, 다중 GPU
torchrun --nproc_per_node=8 run_train.py --config-file config.yaml
# 스타터 구성 생성
python examples/config_tiny_llama.py
| 키 | 의미 |
|---|
run_train.py | 주요 학습 진입점 |
--config-file | YAML 학습 구성 |
pp_engine | 파이프라인 일정 (예: 1F1B, AFAB) |
micro_batch_size | 단계당 마이크로 배치 |
batch_accumulation_per_replica | 그래디언트 누적 |
데이터
| 옵션 | 주의 |
|---|
| 사전 토큰화된 데이터셋 | 토큰화된 데이터를 가리키도록 구성 |
| HF 데이터셋 | Hugging Face를 통해 스트림/로드 |
| Nanoset | Nanotron의 효율적인 데이터셋 형식 |
| 토크나이저 | 토크나이저 구성을 통해 설정 |
체크포인트 및 복구
| 측면 | 방식 |
|---|
| 저장 빈도 | checkpoints.checkpoint_interval |
| 저장 경로 | checkpoints.checkpoints_path |
| 복구 | --config-file + 저장된 체크포인트 가리키기 |
| 샤드됨 | 체크포인트는 병렬화 인식 |
모니터링
| 메트릭 | 의미 |
|---|
| loss / lm_loss | 학습 손실 |
| tokens/s | 처리량 |
| grad norm | 그래디언트 크기 (안정성) |
| 로깅 | W&B / TensorBoard 통합 |
Nanotron vs 다른 학습 프레임워크
| 측면 | Nanotron | TorchTitan | Megatron-LM |
|---|
| 철학 | 최소한, 해킹 가능 | PyTorch 네이티브 참조 | 기능 풍부, 최대 스케일 |
| 병렬화 | 3D (dp/tp/pp) | 4D (컨텍스트 추가) | 3D+ 광범위 |
| 가독성 | 높음 (배우기) | 높음 | 낮음 (큰 코드베이스) |
| 최적 용도 | 3D 이해/맞춤화 | 네이티브 PyTorch 스케일링 | 프론티어 스케일 학습 |
자료