파이토치 분산 학습 도구 비교

1. 개요

• torch.distributed.launch와 torchrun은 둘 다 PyTorch에서 분산 학습을 실행하기 위한 도구이다.
• torch.distributed.launch는 이전부터 사용되어 온 유틸리티이며, torchrun은 더 최근에 등장한 도구로 사용 편의성을 개선했다.

2. 주요 차이점

	torch.distributed.launch	torchrun
사용 편의성	--local_rank 인자를 수동으로 스크립트에 전달해야 함	--local_rank 자동 처리
기능 범위	기본적인 분산 학습 실행 기능 제공	torch.distributed.launch 기능 + 탄력성(elasticity), 오류 처리 등 추가 기능 제공
오류 처리	worker 실패 시 자동 재시작 기능 없음	worker 실패 시 자동 재시작 기능 제공
탄력성 (Elasticity)	지원하지 않음	최소/최대 노드 수 사이에서 노드 수 변경 가능

3. 간단한 예제

torch.distributed.launch

python -m torch.distributed.launch --nproc_per_node=2 \
    your_training_script.py --arg1 --arg2 --local_rank=$LOCAL_RANK

• --nproc_per_node: 각 노드에서 실행할 프로세스 수
• your_training_script.py: 학습 스크립트 파일
• --arg1 --arg2: 학습 스크립트에 전달할 추가 인자
• --local_rank=$LOCAL_RANK: 각 프로세스의 로컬 랭크를 환경 변수로 전달

torchrun

torchrun --nproc_per_node=2 your_training_script.py --arg1 --arg2

• --nproc_per_node: 각 노드에서 실행할 프로세스 수
• your_training_script.py: 학습 스크립트 파일
• --arg1 --arg2: 학습 스크립트에 전달할 추가 인자

4. 결론

• torchrun은 torch.distributed.launch보다 사용하기 쉽고, 더 많은 기능을 제공한다.
• 새로운 프로젝트에서는 torchrun을 사용하는 것이 권장된다.
• 기존 프로젝트에서 torch.distributed.launch를 사용하고 있다면, torchrun으로 마이그레이션하는 것을 고려해 볼 수 있다.