Megatron-LM checkpoint를 NeMo checkpoint로 변환하기

Megatron-LM으로 학습한 언어 모델의 체크포인트를 NeMo로 변환하려면 아래와 같이 NeMo에 포함된 컨버터를 사용합니다. NVIDIA에서 설명서를 제공하지만, GPU나 노드가 여러 개 필요한 경우 적용할 수 없어서 아래에서 케이스별로 설명합니다.

기본

<NeMo_ROOT_FOLDER>/examples/nlp/language_modeling/megatron_lm_ckpt_to_nemo.py \
  --checkpoint_folder <path_to_PTL_checkpoints_folder> \
  --checkpoint_name model_optim_rng.pt \
  --nemo_file_path <path_to_output_nemo_file> \
  --model_type <megatron model type> \
  --tensor_model_parallel_size <tensor_model_parallel_size> \
  --pipeline_model_parallel_size <pipeline_model_parallel_size>  \
  --gpus_per_node  <gpus per node> \
  --hparams_file <path_to_hparams_file>

• checkpoint_folder는 iter_xxxxxx까지 지정 필요. latest_checkpointed_iteration.txt를 읽어서 찾아주지 않음
• checkpoint_name은 Megatron-LM의 경우 model_optim_rng.pt임
• nemo_file_path는 megatron_gpt.nemo로 하면 무난
• model_type은 gpt, t5, bert 중 하나. 보통 GPT 변환시 사용
• 참고자료에는 hparam이나 local_rank 언급이 없지만, 실제로는 빠뜨리면 오류 발생. 꼭 있어야 한다
• TP나 PP가 1보다 큰 경우 WORLD_SIZE 시스템 환경 변수를 TP*PP로 지정해 주어야 함.
• hparams.yaml 예시는 NVIDIA NeMo 깃헙을 참고
  • 아래 예시에서 optim은 반드시 있어야 함. 이유는 Megatron-LM 체크포인트를 읽어서 학습을 재개하는 경우 learning rate 스케줄러가 동일한 설정을 따라야 하기 때문

cfg:
  tensor_model_parallel_size: 8
  pipeline_model_parallel_size: 2                                               
  encoder_seq_length: 2048                                                      
  max_position_embeddings: 2048
  num_layers: 48
  hidden_size: 8192
  ffn_hidden_size: 32768
  num_attention_heads: 64
  tokenizer:
    library: huggingface
    type: <tokenizer_path>
    model: null
    vocab_file: null
    merge_file: null
  optim:
    name: fused_adam
    lr: 1e-4
    weight_decay: 0.1
    betas:
      - 0.9
      - 0.95
    sched:
      name: CosineAnnealing
      min_lr: 1e-5            
...

예시

• 아래 세 가지를 준비
  • checkpoint는 /home/user/checkpoints/gpt3/iter_0010000이라고 가정
  • hparam 파일명은 hparams.yaml
  • GPU가 있는 환경에서 실행할 것
• 결과는 model.nemo로 저장

python examples/nlp/language_modeling/megatron_lm_ckpt_to_nemo.py --checkpoint_folder /home/user/checkpoints/gpt3/iter_0010000 --checkpoint_name model_optim_rng.pt --nemo_file_path model.nemo --model_type gpt --tensor_model_parallel_size 1 --pipeline_model_parallel_size 1 --gpus_per_node 1 --hparams_file hparams.yaml --local_rank 0

GPU가 여러 장 필요한 경우

• GPU가 2장 이상 필요한 경우 위 방법으로 실행할 수 없고 torchrun이 필요하다
  • torchrun에 --nproc_per_node를 world size (TP*PP)로 설정하면 충분
• PP가 있으면 에러가 나니까 PP는 머지하고 진행할 것
• 예시
  • OMP_NUM_THREADS=1 PYTHONPATH=~/Megatron-LM:~/work/NeMo torchrun --nproc_per_node=8 examples/nlp/language_modeling/megatron_lm_ckpt_to_nemo.py --checkpoint_folder /home/user/checkpoints/gpt3/pp_merged/iter_0010000 --checkpoint_name model_optim_rng.pt --nemo_file_path megatron_gpt.nemo --model_type gpt --tensor_model_parallel_size 8 --pipeline_model_parallel_size 1 --gpus_per_node 8 --hparams_file hparams_39b.yaml
  • OMP_NUM_THREADS=1은 워닝을 제거하기 위해 필요

노드가 여러 개 필요한 경우

• 2노드 이상이 필요하거나 백그라운드에서 진행하려는 경우 sbatch로 진행할 것

#!/bin/bash
#SBATCH --nodes=2
#SBATCH --gpus-per-node=8
#SBATCH --ntasks-per-node=8
#SBATCH --cpus-per-task=16
 
OMP_NUM_THREADS=1
PYTHONPATH=$HOME/Megatron-LM:$HOME/work/NeMo  # Megatron-LM과 NeMo 소스코드가 있는 위치로 수정할 것
run_cmd="python $PWD/examples/nlp/language_modeling/megatron_lm_ckpt_to_nemo.py --checkpoint_folder /home/user/checkpoints/gpt3/iter_0010000 --checkpoint_name model_optim_rng.pt --nemo_file_path /t1data/users/jay/work/NeMo/megatron_gpt.nemo --model_type gpt --tensor_model_parallel_size 8 --pipeline_model_parallel_size 2 --gpus_per_node 8 --hparams_file $PWD/hparams_39b.yaml"
 
srun -l \
     --no-container-mount-home \
     --container-image=/home/user/enroot/nvcr.io+ea-bignlp+nemofw-training+23.07-py3.sqsh \
     --container-mounts="/home:/home,/tmp/enroot-data/user-$UID:/tmp" \
     --output=%j.%x.log \
     sh -c "$run_cmd"

이렇게 변환한 .nemo 파일을 NeMo framework에서 로드해서 사용하면 됩니다.