Torchrun

Torchrun

PyTorch의 분산 훈련용 실행기(launcher)인 torchrun을 소개해 드릴게요. torch.distributed.run 모듈을 기반으로 여러 노드에 걸쳐 분산 훈련 프로세스를 자동으로 띄워 주는 도구예요. torchrun은 단일 노드에서 여러 프로세스를 실행하는 데도, 다중 노드로 확장하는 데도 모두 쓸 수 있어서 GPU/CPU 훈련 성능을 크게 끌어올릴 수 있어요. PyTorch 2.0 이후에는 --local-rank 인자 전달 방식도 바뀌었으니 참고하시면 좋아요.

출처: 문서

본문

torchrun이란?

torch.distributed.run은 각 훈련 노드에서 여러 개의 분산 훈련 프로세스를 띄우는 모듈이에요. torchrun은 이 메인 모듈을 가리키는 파이썬 콘솔 스크립트로, python -m torch.distributed.run을 실행하는 것과 같아요.

torchrun은 단일 노드 분산 훈련에 쓸 수 있는데, 노드마다 하나 이상의 프로세스를 띄워요. CPU 훈련에도 GPU 훈련에도 사용할 수 있고, GPU 훈련을 할 때는 각 분산 프로세스가 단일 GPU에서 동작해요. 다중 노드 분산 훈련에서도 노드마다 여러 프로세스를 띄워 성능을 높일 수 있어요. 특히 GPU를 직접 지원하는 Infiniband 인터페이스가 여러 개 있다면 통신 대역폭을 모두 활용할 수 있어서 특히 유리해요.

단일 노드든 다중 노드든 torchrun은 노드당 프로세스 수(--nproc-per-node)만큼 프로세스를 띄워요. GPU 훈련이라면 이 숫자는 시스템의 GPU 개수보다 작거나 같아야 하고, 각 프로세스는 GPU 0부터 GPU(nproc_per_node - 1)까지 하나씩 맡게 돼요.

버전 2.0.0부터 바뀐 점: torchrun은 스크립트에 --local-rank=<rank> 인자를 전달해요. PyTorch 2.0.0부터는 이전의 언더스코어 --local_rank보다 대시(-)로 구분한 --local-rank가 권장돼요. 이전 버전과의 호환을 위해 두 형태를 모두 파서에 넣어 두는 게 좋아요. --local_rank만 제공하면 "error: unrecognized arguments: –local-rank=" 오류가 나요.

>>> import argparse
>>> parser = argparse.ArgumentParser()
>>> parser.add_argument("--local-rank", "--local_rank", type=int)
>>> args = parser.parse_args()

사용 방법 (Usage)

단일 노드 다중 워커 (Single-node multi-worker)

보통은 --standalone과 함께 노드 수, 노드당 프로세스 수를 지정해요.

torchrun
    --standalone
    --nnodes=1
    --nproc-per-node=$NUM_TRAINERS
    YOUR_TRAINING_SCRIPT.py (--arg1 ... train script args...)

여기서 --nproc-per-node는 다음 값을 받을 수 있어요.

  • "gpu": GPU 하나당 프로세스 하나를 띄워요.
  • "cpu": CPU 하나당 프로세스 하나를 띄워요.
  • "xpu": XPU 하나당 프로세스 하나를 띄워요.
  • "auto": CUDA를 쓸 수 있으면 "gpu", XPU를 쓸 수 있으면 "xpu", 아니면 "cpu"로 동작해요.
  • 정수: 프로세스 개수를 직접 지정해요.

자세한 내용은 torch.distributed.run.determine_local_world_size 문서를 참고하세요.

같은 호스트에 여러 인스턴스 쌓기 (Stacked single-node multi-worker)

같은 호스트에서 단일 노드 다중 워커 작업을 여러 개(별도 잡)로 실행하려면 포트 충돌을 피해야 해요. --rdzv-backend=c10d로 실행하고 --rdzv-endpoint=localhost:$PORT_k로 포트를 다르게 지정하면 돼요. --nodes=1일 때는 localhost:0을 줘서 torchrun이 빈 무작위 포트를 자동으로 고르게 하는 게 편해요.

torchrun
    --rdzv-backend=c10d
    --rdzv-endpoint=localhost:0
    --nnodes=1
    --nproc-per-node=$NUM_TRAINERS
    YOUR_TRAINING_SCRIPT.py (--arg1 ... train script args...)

장애 허용 (Fault tolerant, 워커 수 고정, 3회 실패 허용)

torchrun
    --nnodes=$NUM_NODES
    --nproc-per-node=$NUM_TRAINERS
    --max-restarts=3
    --rdzv-id=$JOB_ID
    --rdzv-backend=c10d
    --rdzv-endpoint=$HOST_NODE_ADDR
    YOUR_TRAINING_SCRIPT.py (--arg1 ... train script args...)

HOST_NODE_ADDR는 <host>[:<port>] 형태예요 (예: node1.example.com:29400). C10d 렌더부(rzndezvous) 백엔드가 실행되고 호스트되는 노드와 포트를 지정해요. 클러스터의 아무 노드나 될 수 있지만, 대역폭이 높은 노드를 고르는 게 이상적이에요. 포트를 지정하지 않으면 기본값은 29400이에요.

탄력적 (Elastic, min=1, max=4, 멤버십 변경·실패 최대 3회 허용)

torchrun
    --nnodes=1:4
    --nproc-per-node=$NUM_TRAINERS
    --max-restarts=3
    --rdzv-id=$JOB_ID
    --rdzv-backend=c10d
    --rdzv-endpoint=$HOST_NODE_ADDR
    YOUR_TRAINING_SCRIPT.py (--arg1 ... train script args...)

--nnodes=1:4처럼 최소-최대 노드 수를 지정하면 탄력적으로 노드가 합류하거나 빠질 수 있어요.

셸 완성 (Shell completion)

torchrun은 bash, zsh, tcsh용 완성 스크립트를 만들 수 있어요. 스크립트는 인자 파서에서 생성되므로 위 옵션과 동기화돼요. 선택 패키지 shtab이 필요해요 (pip install shtab).

# zsh
torchrun --print-completion zsh > ~/.zsh/completions/_torchrun

# bash
torchrun --print-completion bash > ~/.local/share/bash-completion/completions/torchrun

완성 스크립트를 로드하는 디렉터리는 셸 문서를 참고하세요. 위 경로는 흔한 기본값이에요.

렌더부 백엔드에 대한 참고 (Note on rendezvous backend)

다중 노드 훈련에는 다음을 지정해야 해요.

  • --rdzv-id: 작업에 참여하는 모든 노드가 공유하는 고유한 작업 ID예요.
  • --rdzv-backend: torch.distributed.elastic.rendezvous.RendezvousHandler의 구현체예요.
  • --rdzv-endpoint: 렌더부 백엔드가 실행되는 엔드포인트로, 보통 host:port 형태예요.

현재 c10d(권장), etcd-v2, etcd(레거시) 렌더부 백엔드를 기본 제공해요. etcd-v2나 etcd를 쓰려면 v2 API를 활성화한 etcd 서버를 준비해야 해요 (예: --enable-v2).

⚠️ etcd-v2와 etcd 렌더부는 etcd API v2를 사용해요. etcd 서버에서 반드시 v2 API를 활성화해야 해요. 우리 테스트는 etcd v3.4.3을 사용해요.

⚠️ etcd 기반 렌더부는 etcd보다 etcd-v2를 권장해요. 기능적으로 동일하지만 구현이 개선됐고, etcd는 유지보수 모드라 향후 버전에서 제거될 예정이에요.

용어 정의 (Definitions)

  • Node: 물리 인스턴스 또는 컨테이너. 잡 매니저가 다루는 단위에 대응해요.
  • Worker: 분산 훈련 맥락에서의 워커예요.
  • WorkerGroup: 같은 함수(예: 트레이너)를 실행하는 워커 집합이에요.
  • LocalWorkerGroup: 같은 노드에서 실행되는 워커 그룹의 부분집합이에요.
  • RANK: 워커 그룹 안에서의 워커 순위예요.
  • WORLD_SIZE: 워커 그룹의 전체 워커 수예요.
  • LOCAL_RANK: 로컬 워커 그룹 안에서의 워커 순위예요.
  • LOCAL_WORLD_SIZE: 로컬 워커 그룹의 크기예요.
  • rdzv_id: 잡의 워커 그룹을 고유하게 식별하는 사용자 정의 ID예요.
  • rdzv_backend: 렌더부의 백엔드(예: c10d)로, 보통 강한 일관성을 가진 키-값 저장소예요.
  • rdzv_endpoint: 렌더부 백엔드 엔드포인트로, 보통 <host>:<port> 형태예요.

한 노드는 LOCAL_WORLD_SIZE개의 워커를 실행하며, 이것이 LocalWorkerGroup을 이뤄요. 잡의 모든 노드에 있는 LocalWorkerGroup을 합치면 WorkerGroup이 돼요.

환경 변수 (Environment Variables)

스크립트에서 다음 환경 변수를 사용할 수 있어요.

  • LOCAL_RANK: 로컬 순위.
  • RANK: 전역 순위.
  • GROUP_RANK: 워커 그룹의 순위. 0과 max_nnodes 사이의 숫자로, 노드당 한 워커 그룹을 실행하면 노드의 순위가 돼요.
  • ROLE_RANK: 같은 역할(role)을 가진 모든 워커 중에서의 순위.
  • LOCAL_WORLD_SIZE: 로컬 월드 크기. --nproc-per-node 값과 같아요.
  • WORLD_SIZE: 잡의 전체 워커 수.
  • ROLE_WORLD_SIZE: WorkerSpec에서 같은 역할로 실행된 전체 워커 수.
  • MASTER_ADDR: 순위 0 워커가 실행되는 호스트의 FQDN으로, Torch 분산 백엔드 초기화에 쓰여요.
  • MASTER_PORT: C10d TCP store를 호스팅할 수 있는 MASTER_ADDR의 포트.
  • TORCHELASTIC_RESTART_COUNT: 지금까지의 워커 그룹 재시작 횟수.
  • TORCHELASTIC_MAX_RESTARTS: 설정된 최대 재시작 횟수.
  • TORCHELASTIC_RUN_ID: 렌더부 run_id(고유 작업 ID)와 같아요.
  • PYTHON_EXEC: 시스템 실행 파일 오버라이드. 제공하면 사용자 파이썬 스크립트가 PYTHON_EXEC 값을 실행 파일로 사용해요. 기본값은 sys.executable이에요.

로깅 (Logging)

기본적으로 각 워커의 stdout/stderr는 콘솔에 그대로 출력돼서 모든 순위의 출력이 뒤섞여 구분하기 어려워요. --redirects는 콘솔 대신 --log-dir 아래 로그 파일로 스트림을 보내고, --tee는 로그 파일에 쓰면서 콘솔에도 함께 출력해요. 둘 다 같은 형식을 받아요: 단일 값(3은 두 스트림 모두, 1은 stdout, 2는 stderr)이나 로컬 순위별 매핑(0:1,1:2)을 지정할 수 있어요. 예를 들어 --tee 3은 모든 워커의 두 스트림을 모두 기록해요.

Tee된 콘솔 줄은 [${role_name}${local_rank}]: 접두사가 붙어요 (예: [default3]: foobar). --log-line-prefix-template으로 바꿀 수 있고, ${role_name}, ${local_rank}, ${rank}, ${hostname} 매크로가 워커별로 치환돼요.

torchrun --nnodes 2 --nproc-per-node 8 --tee 3              --log-line-prefix-template "${hostname}:${rank}: " train.py

r12i0n8:3: python: src/psm2_nccl_net.c:756: Assertion `r->used' failed.
r12i0n8:3: Fatal Python error: Segmentation fault

템플릿은 TORCHELASTIC_LOG_LINE_PREFIX_TEMPLATE 환경 변수로도 설정할 수 있고, 커맨드라인 옵션이 우선해요. --local-ranks-filter는 --redirects/--tee가 만드는 로그 파일에는 영향 없이 콘솔에 도달하는 순위만 제한해요.

배포 (Deployment)

  • (C10d 백엔드에는 불필요) 렌더부 백엔드 서버를 시작하고 엔드포인트를 얻어서 --rdzv-endpoint로 넘겨요.
  • 단일 노드 다중 워커: 호스트에서 torchrun을 시작하면 로컬 워커 그룹을 만들고 모니터링하는 에이전트 프로세스가 시작돼요.
  • 다중 노드 다중 워커: 훈련에 참여하는 모든 노드에서 같은 인자로 torchrun을 시작해요.

잡/클러스터 매니저를 쓸 때도 다중 노드 잡의 엔트리포인트 명령은 torchrun이어야 해요.

실패 모드 (Failure Modes)

  • 워커 실패: n개의 워커로 훈련하는 잡에서 k<=n개의 워커가 실패하면 모든 워커가 멈추고 max_restarts까지 재시작돼요.
  • 에이전트 실패: 에이전트 실패는 로컬 워커 그룹 실패로 이어져요. 잡 매니저가 전체 잡을 실패시키거나(gang semantics) 노드를 교체하는 건 잡 매니저의 몫이에요. 두 동작 모두 에이전트가 지원해요.
  • 노드 실패: 에이전트 실패와 같아요.

멤버십 변경 (Membership Changes)

  • 노드 이탈(축소): 에이전트가 이탈을 통지받고, 기존 워커를 모두 멈춘 뒤 새 WorkerGroup을 구성해 모든 워커를 새 RANK와 WORLD_SIZE로 시작해요.
  • 노드 합류(확장): 새 노드가 잡에 들어오면 기존 워커를 모두 멈추고 새 WorkerGroup을 구성해 모든 워커를 새 RANK와 WORLD_SIZE로 시작해요.

NUMA 바인딩 (NUMA Binding)

NUMA(Non-Uniform Memory Access) 구조의 다중 GPU 시스템에서는 워커 프로세스를 할당된 GPU 근처의 CPU에 바인딩해서 성능을 높일 수 있어요. --numa-binding 플래그를 쓰면 돼요.

torchrun --numa-binding=node --nproc-per-node=8 train.py

주요 유의사항 (Important Notices)

  • 이 유틸리티와 다중 프로세스 분산(단일·다중 노드) GPU 훈련은 현재 NCCL 분산 백엔드를 쓸 때 최고 성능을 내요. 따라서 GPU 훈련에는 NCCL 백엔드를 권장해요.
  • Torch 프로세스 그룹 초기화에 필요한 환경 변수는 이 모듈이 제공하므로 RANK를 직접 넘길 필요가 없어요. 훈련 스크립트에서 프로세스 그룹을 초기화하려면 다음처럼 해요.
>>> import torch.distributed as dist
>>> dist.init_process_group(backend="gloo|nccl")
  • 훈련 프로그램에서는 일반 분산 함수를 쓰거나 torch.nn.parallel.DistributedDataParallel() 모듈을 쓸 수 있어요. GPU 훈련에 DistributedDataParallel을 쓰려면 이렇게 설정해요.
local_rank = int(os.environ["LOCAL_RANK"])
model = torch.nn.parallel.DistributedDataParallel(
    model, device_ids=[local_rank], output_device=local_rank
)

device_ids는 코드가 다룰 유일한 GPU 장치 ID로 설정해야 해요. 보통 프로세스의 로컬 순위이고, device_ids는 [int(os.environ("LOCAL_RANK"))], output_device는 int(os.environ("LOCAL_RANK"))가 돼야 해요.

  • 실패나 멤버십 변경 시 살아남은 모든 워커는 즉시 종료돼요. 진행 상황을 반드시 체크포인트로 저장하세요. 체크포인트 빈도는 잡이 감당할 수 있는 손실에 맞춰 잡으세요.
  • 이 모듈은 동종(homogeneous) LOCAL_WORLD_SIZE만 지원해요. 즉 모든 노드가 역할당 같은 수의 로컬 워커를 실행한다고 가정해요.
  • RANK는 안정적이지 않아요. 재시작 사이에 노드의 로컬 워커가 이전과 다른 순위 범위를 받을 수 있어요. RANK의 안정성이거나 RANK와 LOCAL_RANK의 상관관계를 절대 가정하지 마세요.
  • 탄력성(min_size != max_size)을 쓸 때는 노드가 합류/이탈하면서 월드 크기가 변할 수 있으므로 WORLD_SIZE에 대한 가정을 코드에 하드코딩하지 마세요.

권장되는 스크립트 구조는 다음과 같아요.

def main():
    load_checkpoint(checkpoint_path)
    initialize()
    train()


def train():
    for batch in iter(dataset):
        train_step(batch)

        if should_checkpoint:
            save_checkpoint(checkpoint_path)
  • 워커 오류 발생 시 이 도구가 오류의 상세(시간, 순위, 호스트, pid, traceback 등)를 요약해 줘요. 각 노드에서 가장 먼저(타임스탬프 기준) 발생한 오류가 "Root Cause" 오류로 보고돼요. 오류 요약에 traceback을 포함하려면 메인 엔트리포인트 함수를 아래처럼 장식해야 해요. 장식하지 않으면 요약에 traceback이 빠지고 exit code만 포함돼요. 자세한 내용은 torchelastic 오류 처리 문서를 참고하세요.
from torch.distributed.elastic.multiprocessing.errors import record


@record
def main():
    # do train
    pass


if __name__ == "__main__":
    main()

더 알아보기 (Learn more)