Nanotron

Nanotron

Nanotron은 수백 개의 GPU에 걸친 대규모 분산 학습을 위한 프레임워크예요. Transformers 모델을 Nanotron 형식으로 변환해 사전학습하는 방법을 살펴볼게요.

출처: 문서

본문

Nanotron은 텐서 병렬화(tensor), 파이프라인 병렬화(parallel), 데이터 병렬화(data)를 지원하는 분산 학습 프레임워크예요 (3D 병렬화). 수백 개의 GPU에 걸친 대규모 학습 작업을 위해 설계됐어요.

convert_hf_to_nanotron.py 스크립트로 어떤 Transformers 모델이든 사전학습용 최적화 Nanotron transformer 모델 구현으로 변환할 수 있어요.

torchrun --nproc_per_node=1 examples/llama/convert_hf_to_nanotron.py \
    --checkpoint_path=meta-llama/Llama-2-7b-hf \
    --save_path=./llama-7b-nanotron

Transformers 통합 과정

  1. from_pretrained() 함수로 Llama 같은 지원되는 Transformers 모델을 불러와요. 이때 체크포인트 디렉터리의 config.json 파일을 읽고 LlamaConfig를 만든답니다.
  2. Nanotron은 LlamaConfig를 자신의 설정 포맷에 매핑하고 Nanotron 모델을 만들어요.
  3. Transformers 가중치를 Nanotron 가중치로 변환해요. 가중치 매핑(weight mapping)이 Nanotron 파라미터 이름을 Transformers 파라미터 이름에 어떻게 매핑할지 안내하는데, QKV projection과 gate/up projection을 합치는(fusing) 변환도 이때 처리돼요.

또한 Nanotron은 전처리와 생성 중 텍스트를 토큰 id로 바꿀 때 AutoTokenizer에 의존해요.

리소스 (Resources)

더 알아보기 (Learn more)