Megatron-LM으로 언어 모델 훈련하기

개요

Megatron-LM은 NVIDIA의 Applied Deep Learning Research 팀이 개발한 고성능 프레임워크로, GPU에서 대형 트랜스포머 모델을 효율적으로 사전 훈련하도록 설계되었습니다. Hugging Face Trainer API나 accelerate 라이브러리보다 설정이 복잡하지만, Megatron-LM은 특화된 데이터 처리와 하드웨어 수준 최적화를 통해 큰 속도 향상을 제공합니다.

Megatron-LM의 기술 최적화

Megatron-LM은 최적화된 데이터 로딩과 커널 융합이라는 두 가지 주요 메커니즘을 통해 표준 PyTorch 루프에 비해 뛰어난 훈련 효율성을 달성합니다.

효율적인 DataLoader

Megatron-LM은 훈련 시작 전에 데이터를 토크나이즈하고 셔플하는 DataLoader를 사용합니다. 인덱스를 한 번만 계산하고 훈련 파라미터에 따라 epoch 수를 결정합니다. 이 접근 방식은 새로운 epoch을 위해 전체 데이터셋을 다시 순회할 필요를 없애며, 보다 부드러운 학습 곡선과 훈련 시간 감소를 가져옵니다.

융합된 CUDA 커널

메모리 오버헤드를 최소화하기 위해 Megatron-LM은 융합된 CUDA 커널을 사용합니다. 표준 PyTorch 연산에서는 데이터가 메모리에서 가져와 계산되고 매 연산마다 다시 저장됩니다. 커널 융합은 유사한 연산을 하나의 하드웨어 연산으로 결합하여 중간 결과를 GPU 레지스터에 저장하게 하여 메인 메모리로 복사하는 비용을 줄입니다.

또한 Megatron-LM은 NVIDIA Apex 라이브러리의 융합된 AdamW 구현을 사용하며, 이는 표준 PyTorch 구현보다 빠른 성능을 제공합니다.

구현 워크플로우

Megatron-LM에서 모델(예: GPT-2)을 훈련하는 과정은 네 단계 파이프라인으로 구성됩니다: 환경 설정, 데이터 전처리, 훈련, 모델 변환.

환경 설정

권장 설정은 NGC의 NVIDIA PyTorch Container를 사용하는 것으로, 모든 필수 종속성이 포함되어 있습니다. 수동 설치 시 최신 버전의 PyTorch, CUDA, NCCL, NVIDIA Apex, 그리고 nltk 라이브러리가 필요합니다. 사용자는 또한 Megatron-LM 디렉터리 내에 토크나이저의 vocab.json 및 merges.txt 파일을 제공해야 합니다.

데이터 전처리

데이터는 느슨한 JSON 형식(한 줄에 하나의 샘플)으로 변환된 후 처리됩니다. Megatron-LM은 tools/preprocess_data.py를 사용해 데이터를 토크나이즈, 셔플 및 바이너리 형식으로 변환합니다. 이 과정은 훈련 단계에 필요한 .idx 및 .bin 파일을 생성합니다. dataset-impl은 'lazy', 'cached', 또는 'mmap'으로 설정할 수 있습니다.

훈련 실행

훈련은 torch.distributed.launch를 통해 시작됩니다. 110M 파라미터 모델(예: CodeParrot-small)의 경우 8 GPU에서 다음 구성으로 약 12시간이 소요됩니다:

  • 아키텍처: 12 레이어, 768 hidden size, 12 attention heads.
  • 시퀀스 길이: 1024.
  • 옵티마이저: 코사인 학습률 감소를 적용한 AdamW.
  • 배치 크기: 마이크로 배치 크기 12, 전체 배치 크기 192.

모델 병렬 전략

단일 GPU에 맞지 않을 정도로 큰 모델의 경우 Megatron-LM은 두 가지 모델 병렬 방식을 지원합니다:

  1. 텐서 병렬: tensor-model-parallel-size 파라미터를 통해 단일 트랜스포머 모듈의 실행을 여러 GPU에 분산합니다.
  2. 파이프라인 병렬: pipeline-model-parallel-size 파라미터를 통해 트랜스포머 모듈을 동일한 크기의 단계로 나누어 GPU에 배치합니다.

Hugging Face Transformers와의 통합

Megatron-LM으로 훈련된 모델을 평가하거나 프로덕션에 사용하려면 transformers 라이브러리가 지원하는 형식으로 변환해야 합니다. 이는 model_optim_rng.pt 체크포인트 파일을 convert_megatron_gpt2_checkpoint.py 스크립트를 사용해 pytorch_model.bin 파일로 변환함으로써 이루어집니다.

변환이 완료되면 모델을 AutoModelForCausalLM으로 로드할 수 있습니다. 모델 병렬을 사용해 훈련된 매우 큰 모델의 경우 accelerate 라이브러리의 device_map="auto" 인자를 사용해 가용 GPU와 CPU RAM에 가중치를 자동으로 배분할 수 있습니다.

프레임워크 선택 가이드

Megatron-LM은 높은 최적화 덕분에 대형 모델 사전 훈련이나 장기 파인튜닝에 가장 적합합니다. 그러나 전처리 및 변환 단계에서 오버헤드가 발생합니다. 중간 규모 모델의 짧은 파인튜닝에는 장치에 구애받지 않고 유연성이 높은 Hugging Face Trainer API와 accelerate 라이브러리를 권장합니다.

Sources