Pytorch torchtune

Colab에서 실행해 보기

torchtune은 대규모 언어 모델(LLM)의 작성, 파인튜닝, 실험 과정을 간소화하도록 설계된 PyTorch 기반 라이브러리입니다. 또한 torchtune은 W&B로 로깅을 기본 지원하여 트레이닝 과정을 더 효율적으로 추적하고 시각화할 수 있게 해줍니다.

torchtune를 사용한 Mistral 7B 파인튜닝에 대한 W&B 블로그 게시글을 확인하세요.

​

손쉽게 W&B 로깅 활용하기

명령줄

실행 시 사용할 명령줄 인수를 다음과 같이 재정의합니다:

tune run lora_finetune_single_device --config llama3/8B_lora_single_device \
  metric_logger._component_=torchtune.utils.metric_logging.WandBLogger \
  metric_logger.project="llama3_lora" \
  log_every_n_steps=5

Recipe

레시피의 config에서 W&B 로깅을 활성화합니다:

# inside llama3/8B_lora_single_device.yaml
metric_logger:
  _component_: torchtune.utils.metric_logging.WandBLogger
  project: llama3_lora
log_every_n_steps: 5

​

W&B 메트릭 로거 사용하기

레시피의 구성 파일에서 metric_logger 섹션을 수정하여 W&B 로깅을 활성화합니다. _component_를 torchtune.utils.metric_logging.WandBLogger 클래스로 변경하세요. 로깅 동작을 사용자 지정하기 위해 project 이름과 log_every_n_steps를 전달할 수도 있습니다.

또한 wandb.init() 메서드에 전달하듯이 다른 kwargs도 모두 전달할 수 있습니다. 예를 들어 팀에서 작업 중이라면, 팀 이름을 지정하기 위해 entity 인자를 WandBLogger 클래스에 전달할 수 있습니다.

Recipe

# llama3/8B_lora_single_device.yaml 내부
metric_logger:
  _component_: torchtune.utils.metric_logging.WandBLogger
  project: llama3_lora
  entity: my_project
  job_type: lora_finetune_single_device
  group: my_awesome_experiments
log_every_n_steps: 5

Command Line

tune run lora_finetune_single_device --config llama3/8B_lora_single_device \
  metric_logger._component_=torchtune.utils.metric_logging.WandBLogger \
  metric_logger.project="llama3_lora" \
  metric_logger.entity="my_project" \
  metric_logger.job_type="lora_finetune_single_device" \
  metric_logger.group="my_awesome_experiments" \
  log_every_n_steps=5

​

무엇이 기록되나요?

W&B 대시보드를 탐색하여 기록된 메트릭을 확인할 수 있습니다. 기본적으로 W&B는 설정 파일의 모든 하이퍼파라미터와 Launch 오버라이드 값을 기록합니다.

W&B는 Overview 탭에 최종 확정된 config를 기록합니다. 또한 Files 탭에 YAML 형식의 config를 저장합니다.

​

기록된 메트릭

각 레시피는 자체 트레이닝 루프를 갖습니다. 각 레시피별로 기록되는 메트릭을 확인하면, 기본적으로 아래 항목들이 포함됩니다:

Metric	Description
loss	모델의 loss 값
lr	learning rate 값
tokens_per_second	모델의 초당 토큰 수
grad_norm	모델의 gradient norm
global_step	현재 트레이닝 루프에서의 스텝에 해당합니다. 그래디언트 누적을 고려하며, 기본적으로 옵티마이저 스텝이 한 번 수행될 때마다 그래디언트가 누적되고, gradient_accumulation_steps마다 한 번씩 모델이 업데이트됩니다.

global_step은 트레이닝 스텝 수 자체와 동일하지 않습니다. 트레이닝 루프에서의 현재 스텝에 해당하며, 그래디언트 누적을 고려합니다. 기본적으로 옵티마이저 스텝이 수행될 때마다 global_step이 1씩 증가합니다. 예를 들어, 데이터로더에 배치가 10개 있고, gradient accumulation steps가 2이며 3 에포크를 실행하면, 옵티마이저는 총 15번 스텝을 수행합니다. 이 경우 global_step은 1에서 15까지의 값을 갖습니다.

torchtune의 간결한 설계 덕분에 사용자 정의 메트릭을 손쉽게 추가하거나 기존 메트릭을 수정할 수 있습니다. 해당 recipe 파일만 수정하면 되며, 예를 들어 전체 에포크 수에 대한 비율(%)로 current_epoch을 계산한 뒤, 다음과 같이 로깅할 수 있습니다:

# 레시피 파일의 `train.py` 함수 내부
self._metric_logger.log_dict(
    {"current_epoch": self.epochs * self.global_step / self._steps_per_epoch},
    step=self.global_step,
)

이 라이브러리는 빠르게 발전 중인 라이브러리로, 현재 제공되는 메트릭은 향후 변경될 수 있습니다. 커스텀 메트릭을 추가하려면 레시피를 수정한 뒤 해당 self._metric_logger.* 함수를 호출해야 합니다.

​

체크포인트 저장 및 로드

torchtune 라이브러리는 다양한 체크포인트 형식을 지원합니다. 사용 중인 모델의 출처에 따라 적절한 checkpointer 클래스를 사용해야 합니다.

모델 체크포인트를 W&B Artifacts에 저장하려면, 가장 간단한 방법은 해당 레시피 내부의 save_checkpoint 함수를 오버라이드하는 것입니다.

다음은 save_checkpoint 함수를 오버라이드하여 모델 체크포인트를 W&B Artifacts에 저장하는 예시입니다.

def save_checkpoint(self, epoch: int) -> None:
    ...
    ## 체크포인트를 W&B에 저장합니다
    ## Checkpointer Class에 따라 파일 이름이 달라집니다
    ## full_finetune의 경우 예시입니다
    checkpoint_file = Path.joinpath(
        self._checkpointer._output_dir, f"torchtune_model_{epoch}"
    ).with_suffix(".pt")
    wandb_artifact = wandb.Artifact(
        name=f"torchtune_model_{epoch}",
        type="model",
        # 모델 체크포인트 설명
        description="Model checkpoint",
        # dict 형태로 원하는 메타데이터를 추가할 수 있습니다
        metadata={
            utils.SEED_KEY: self.seed,
            utils.EPOCHS_KEY: self.epochs_run,
            utils.TOTAL_EPOCHS_KEY: self.total_epochs,
            utils.MAX_STEPS_KEY: self.max_steps_per_epoch,
        },
    )
    wandb_artifact.add_file(checkpoint_file)
    wandb.log_artifact(wandb_artifact)