colab에서 transformers 라이브러리로 LLM 학습시 checkpoint 사용하기

코랩 colab에서 transformers 라이브러리로 LLM 학습시 checkpoint 사용하는 방법을 정리해봤습니다.

2024. 2. 26  최초작성

2024. 2. 27

 

 

LLM(Large Language Models)에 대한  파인 튜닝, LoRA, QLoRA을 최근 코랩에서 진행하다보니 그동안 신경안쓰던 체크포인트를 사용할 일이 생겼습니다.

 

코랩에서 학습을 진행하다보면 런타임 연결이 끊어져서 학습이 중단되는 경우가 종종 발생하기 때문입니다. 

 

모델을 학습하는 중간 결과물인 체크포인트를 저장하면 학습이 중단되었을 경우 이어서 학습을 진행할 수 있더군요. 

transformers 라이브러리를 사용할시에는 TrainingArguments에서 체크포인트에 대한 설정을 다음처럼 할 수 있습니다.

transformers.TrainingArguments(

 

        .............................................................

 

        # 다음 4줄을 추가합니다.

        output_dir="outputs",  # 체크포인트가 저장될 경로입니다.  

        save_strategy="steps", # steps로 설정해야 합니다. 

        save_total_limit=1, # 체크포인트 폴더를 유지할 개수입니다. checkpoint-xx 이름을 갖습니다.

        save_steps=10,        # 10 스텝마다 체크포인트를 저장하게 됩니다. 

 

    ),

학습을 이어서 하는 경우가 아니라면 가급적이면  outputs 폴더를 지우고 학습을 진행하는게 좋을 듯합니다. 셀에서 다음처럼 입력하여 삭제할 수 있습니다. 또는 파이썬 코드에서 할 수도 있습니다. 

 

!rm -rf outputs

처음 학습시에는 다음처럼 합니다.

trainer.train()

학습이 멈춘 경우 셀에 다음 명령을 입력하여 체크포인트가 생겼는지 확인합니다. 

!ls outputs/

outputs/ 경로에 checkpoint-xx 이름의 폴더가 보입니다. xx는 학습을 진행한 스텝수입니다. 

예를 들어 학습 시작시 다음처럼 하면 100스텝에서 저장된 체크포인트를 사용하여 학습을 진행합니다. 

trainer.train(resume_from_checkpoint="outputs/checkpoint-100")

실행시켜보면 101 스텝부터 학습이 진행되는 것을 볼 수 있습니다.  이전 학습을 강제로 중단한 경우 메모리가 제대로 해제안될 수 도 있으므로 메뉴에서 "런타임>세션 다시 시작"을 선택한후 진행해야 합니다.

체크포인트를 저장할 저장소를 코랩의 저장공간이 아닌 구글 드라이브로 할 수도 있습니다. 

하지만  Colab Notebooks/outputs 폴더에 체크포인트 폴더가 바로 보이지 않는 문제가 있습니다. 학습을 많이 진행하다가 멈춘경우 좀 오래 기다려야 하더군요. 

런타임을 완전히 해제하기 전에는 체크포인트를 구글 드라이브에 옮겨두면 좋은 것 같습니다.  

학습을 재개할때 구글 드라이브의 경로를 바로 지정해도 됩니다. 

1. 코랩 저장공간에 있는 체크포인트 폴더를 확인합니다. 예시에서는 1970 스텝에서 저장된 체크포인터 폴더가 있습니다.

 

!ls outputs/

 

checkpoint-1970  runs

2. 체크포인터 폴더를 구글 드라이브로 복사합니다.

 

!cp -rf outputs '/content/drive/MyDrive/Colab Notebooks/'

3. 구글 드라이브에 있는 체크포인터 폴더를 확인합니다.

 

!ls '/content/drive/MyDrive/Colab Notebooks/outputs'

 

checkpoint-1970  runs

4. 구글 드라이브에 있는 체크포인터를 사용하여 학습을 이어서 진행합니다.

 

trainer.train(resume_from_checkpoint="/content/drive/MyDrive/Colab Notebooks/outputs/checkpoint-1970")

5. 10스텝 진행후 중지후 코랩 저장공간을 확인해보면 1980 스텝에서 저장된 체크포인터 폴더를 볼 수 있습니다.

 

!ls outputs/

 

checkpoint-1980  runs

참고

https://github.com/huggingface/transformers/issues/7198#issuecomment-694352941