PyTorch, Hugging Face Trainer를 활용한 분산 환경 학습

 

LLM 모델 분산 학습의 구성요소

1. PyTorch

• 개요:
  PyTorch는 Facebook AI Research(FAIR)에서 개발한 오픈 소스 딥러닝 프레임워크로, 연구와 산업계에서 널리 사용됩니다.
• 특징:
  • 동적 계산 그래프: 모델을 개발하고 디버깅할 때 유연성이 뛰어납니다.
  • GPU 가속: 손쉽게 GPU를 활용하여 연산을 가속할 수 있습니다.
  • 커뮤니티와 생태계: 다양한 라이브러리와 도구들이 활발히 개발되고 있으며, 풍부한 문서와 튜토리얼이 제공됩니다.

---

2. Hugging Face Trainer

• 개요:
  Hugging Face의 Transformers 라이브러리에서 제공하는 고수준 학습 인터페이스로, 복잡한 학습 루프를 단순화시켜 줍니다.
• 특징:
  • 자동화: 학습 루프, 평가, 로깅, 체크포인트 저장 등 반복되는 작업들을 자동으로 처리합니다.
  • 사용 편의성: 복잡한 딥러닝 모델도 몇 줄의 코드로 학습시킬 수 있어, 연구자나 개발자가 모델 실험에 집중할 수 있습니다.
  • 분산 학습 지원: PyTorch의 분산 학습 기능과 연동되어 대규모 모델 학습도 지원합니다.

---

3. rdzv (Rendezvous)

• 개요:
  rdzv는 “rendezvous”의 약자로, 분산 학습 환경에서 여러 프로세스나 노드가 서로를 인식하고 동기화하는 초기 단계에서 사용되는 메커니즘입니다.
• 역할:
  • 프로세스 동기화: 분산 학습을 시작할 때, 각 학습 프로세스가 서로 연결되어 통신할 수 있도록 초기 설정을 합니다.
  • 노드 연결: rdzv를 통해 모든 참여 프로세스가 공통의 “만남의 장소”(endpoint)를 사용하여 서로를 찾고, 이후에 작업을 분배받게 됩니다.
• Hugging Face Trainer와의 관계:
  • Trainer는 내부적으로 PyTorch의 분산 학습 기능을 사용하며, 이때 rdzv 설정(예: rdzv_endpoint, rdzv_backend 등)을 활용하여 다수의 프로세스가 원활히 협력할 수 있도록 합니다.
  • 이를 통해 대규모 모델을 여러 GPU나 노드에 분산하여 효율적으로 학습할 수 있습니다.

---

요약

• PyTorch: 강력하고 유연한 딥러닝 프레임워크로, 동적 계산 그래프와 GPU 가속을 지원합니다.
• Hugging Face Trainer: PyTorch 기반의 모델 학습을 단순화하는 고수준 API로, 자동화된 학습 루프와 분산 학습 기능을 제공합니다.
• rdzv (Rendezvous): 분산 학습 환경에서 여러 프로세스가 초기 동기화 및 연결을 할 수 있도록 하는 메커니즘으로, 대규모 모델 학습 시 필수적인 역할을 합니다.

 

코드 예시

train.py

import os
import torch
from transformers import Trainer, TrainingArguments, BertForSequenceClassification, BertTokenizerFast
from datasets import load_dataset

def main():
    # 데이터셋 불러오기 (예시로 GLUE의 MRPC 사용)
    dataset = load_dataset("glue", "mrpc")
    tokenizer = BertTokenizerFast.from_pretrained("bert-base-uncased")
    
    def tokenize_function(examples):
        return tokenizer(examples["sentence1"], examples["sentence2"], truncation=True)
    
    tokenized_datasets = dataset.map(tokenize_function, batched=True)
    # 불필요한 컬럼 제거 및 텐서 포맷 지정
    tokenized_datasets = tokenized_datasets.remove_columns(["sentence1", "sentence2", "idx"])
    tokenized_datasets.set_format("torch")

    # 모델 초기화
    model = BertForSequenceClassification.from_pretrained("bert-base-uncased", num_labels=2)

    # TrainingArguments: torchrun 환경변수(예: LOCAL_RANK, RANK, WORLD_SIZE)를 활용해 분산 학습이 자동 인식됩니다.
    training_args = TrainingArguments(
        output_dir="./results",
        num_train_epochs=3,
        per_device_train_batch_size=8,
        per_device_eval_batch_size=8,
        evaluation_strategy="epoch",
        save_strategy="epoch",
        logging_dir="./logs",
        logging_steps=10,
    )

    trainer = Trainer(
        model=model,
        args=training_args,
        train_dataset=tokenized_datasets["train"],
        eval_dataset=tokenized_datasets["validation"],
    )

    trainer.train()

if __name__ == "__main__":
    # torchrun으로 실행하면 환경변수에 따라 분산 학습 설정이 전달됩니다.
    # 이미 초기화되지 않은 경우, 아래와 같이 명시적으로 프로세스 그룹을 초기화할 수 있습니다.
    if "RANK" in os.environ and "WORLD_SIZE" in os.environ:
        backend = "nccl" if torch.cuda.is_available() else "gloo"
        torch.distributed.init_process_group(backend=backend)
    main()

 

실행 방법 (3개 노드 분산 학습)

각 노드에서는 아래와 같이 torchrun 명령어를 사용하여 실행합니다. 예를 들어, 마스터 노드의 IP가 192.168.1.1이고 포트 29500을 사용한다고 가정합니다.

 

노드 0 (node_rank=0):

torchrun --nproc_per_node=1 --nnodes=3 --node_rank=0 --rdzv_backend=c10d --rdzv_endpoint=192.168.1.1:29500 train.py

 

노드 1 (node_rank=1):

torchrun --nproc_per_node=1 --nnodes=3 --node_rank=1 --rdzv_backend=c10d --rdzv_endpoint=192.168.1.1:29500 train.py

 

노드 2 (node_rank=2):

torchrun --nproc_per_node=1 --nnodes=3 --node_rank=2 --rdzv_backend=c10d --rdzv_endpoint=192.168.1.1:29500 train.py

 

각 노드에서 --nproc_per_node 옵션은 해당 노드에서 사용할 GPU(혹은 프로세스) 수를 지정합니다. 이 예제에서는 각 노드당 1개의 프로세스를 사용합니다.

---

내부 통신 및 리소스 할당 과정

1. Rendezvous 단계 (rdzv):
   • 초기 등록:
     각 노드는 torchrun 명령어를 통해 실행되며, --rdzv_endpoint로 지정된 마스터(또는 공용) IP와 포트에 접속하여 자신의 존재를 알립니다.
   • 환경 변수 설정:
     torchrun은 각 노드에 RANK, WORLD_SIZE, LOCAL_RANK 등의 환경 변수를 설정하여 전체 분산 환경 정보를 전달합니다.
   • 프로세스 그룹 초기화:
     코드 내에서 torch.distributed.init_process_group를 호출하여, 각 노드가 rdzv를 통해 서로를 확인하고 동일한 프로세스 그룹에 가입합니다.
2. 통신 방식:
   • 백엔드 선택:
     GPU가 사용 가능하면 nccl, 그렇지 않으면 gloo 백엔드를 사용해 통신합니다. 이 백엔드는 TCP/IP를 기반으로 각 프로세스 간의 데이터(예: gradient) 교환 및 동기화를 담당합니다.
   • All-Reduce 연산:
     학습 과정에서 각 프로세스가 계산한 gradient는 All-Reduce 방식으로 집계되어, 모든 프로세스가 동일한 모델 파라미터 업데이트를 받게 됩니다.
3. 리소스 할당:
   • 노드별 할당:
     각 노드는 실행 시 LOCAL_RANK에 따라 특정 GPU에 할당됩니다. Hugging Face Trainer 및 PyTorch는 이를 자동으로 인식하여, 각 프로세스가 자신에게 할당된 GPU에서 연산을 수행하도록 합니다.
   • 동기화 및 학습:
     모든 노드가 초기 rdzv 과정을 마치면, 각 노드는 동기화된 상태에서 학습 루프를 시작합니다. 모델의 파라미터 업데이트와 gradient 계산은 분산 환경에서 통신 및 동기화를 통해 이루어집니다.

---

요약

• rdzv (Rendezvous):
  각 노드가 중앙의 rendezvous 엔드포인트를 통해 서로를 인식하고, 환경 변수(RANK, WORLD_SIZE 등)를 기반으로 분산 프로세스 그룹을 형성합니다.
• 통신:
  선택된 백엔드(nccl/gloo)를 통해 각 노드 간에 gradient와 모델 업데이트 정보를 All-Reduce 등의 집계 연산으로 동기화합니다.
• 리소스 할당 및 학습 진행:
  각 노드는 자신에게 할당된 GPU에서 학습을 진행하며, 모든 노드가 동기화되어 동일한 모델 파라미터를 업데이트하면서 분산 학습이 효율적으로 수행됩니다.

이와 같이 Hugging Face Trainer는 torchrun과 rdzv 메커니즘을 활용해 여러 노드에서 분산 학습을 원활하게 수행할 수 있도록 도와줍니다.