TorchDistributor를 사용하여 분산 학습
이 문서에서는 TorchDistributor를 사용하여 PyTorch ML 모델에서 분산 학습을 수행하는 방법을 설명합니다.
TorchDistributor는 사용자가 Spark 클러스터에서 PyTorch를 사용하여 분산 학습을 수행하는 데 도움이 되는 PySpark의 오픈 소스 모듈이므로 PyTorch 학습 작업을 Spark 작업으로 시작할 수 있습니다. 내부적으로는 작업자 간의 환경 및 통신 채널을 초기화하고 CLI 명령 torch.distributed.run을 사용하여 작업자 노드 간에 분산 학습을 실행합니다.
TorchDistributor API는 다음 표에 표시된 메서드를 지원합니다.
| 메서드 및 서명 | 설명 |
|---|---|
init(self, num_processes, local_mode, use_gpu) |
TorchDistributor의 인스턴스를 만듭니다. |
run(self, main, *args) |
main이 함수인 경우 torchrun main *args를 호출하여 분산 학습을 실행하고, main이 파일 경로인 경우 CLI 명령 main(**kwargs)을 실행합니다. |
요구 사항
- Spark 3.4
- Databricks Runtime 13.0 ML 이상
Notebook에 대한 개발 워크플로
모델 만들기 및 학습 프로세스가 로컬 컴퓨터의 Notebook 또는 Databricks Notebook에서 완전히 발생하는 경우 코드를 분산 학습에 대비하기 위해 약간만 변경하면 됩니다.
단일 노드 코드 준비: PyTorch, PyTorch Lightning 또는 PyTorch/PyTorch Lightning을 기반으로 하는 다른 프레임워크(예: HuggingFace Trainer API)를 사용하여 단일 노드 코드를 준비하고 테스트합니다.
표준 분산 학습을 위한 코드 준비: 단일 프로세스 학습을 분산 학습으로 변환해야 합니다. 이 분산 코드를 모두
TorchDistributor와 함께 사용할 수 있는 하나의 교육 함수로 포함시킵니다.학습 함수 내에서 가져오기를 이동: 학습 함수 내에 필요한 가져오기(예:
import torch)를 추가합니다. 이렇게 하면 일반적인 피클링 오류를 방지할 수 있습니다. 또한 모델과 데이터가 연결되는device_id은 다음을 통해 결정됩니다.device_id = int(os.environ["LOCAL_RANK"])분산 학습 시작: 원하는 매개변수로
TorchDistributor을 인스턴스화하고 학습 시작을 위해.run(*args)을 호출합니다.
다음은 학습 코드의 예제입니다.
from pyspark.ml.torch.distributor import TorchDistributor
def train(learning_rate, use_gpu):
import torch
import torch.distributed as dist
import torch.nn.parallel.DistributedDataParallel as DDP
from torch.utils.data import DistributedSampler, DataLoader
backend = "nccl" if use_gpu else "gloo"
dist.init_process_group(backend)
device = int(os.environ["LOCAL_RANK"]) if use_gpu else "cpu"
model = DDP(createModel(), **kwargs)
sampler = DistributedSampler(dataset)
loader = DataLoader(dataset, sampler=sampler)
output = train(model, loader, learning_rate)
dist.cleanup()
return output
distributor = TorchDistributor(num_processes=2, local_mode=False, use_gpu=True)
distributor.run(train, 1e-3, True)
외부 리포지토리에서 학습 마이그레이션
외부 리포지토리에 저장된 기존 분산 학습 프로시저가 있는 경우 다음을 수행하여 Azure Databricks로 쉽게 마이그레이션할 수 있습니다.
- 리포지토리 가져오기: 외부 리포지토리를 Databricks Git 폴더로 가져옵니다.
- 새 Notebook을 만들기 리포지토리 내에서 새 Azure Databricks Notebook을 초기화합니다.
- 분산 학습 시작 Notebook 셀에서 다음과 같은
TorchDistributor를 호출합니다.
from pyspark.ml.torch.distributor import TorchDistributor
train_file = "/path/to/train.py"
args = ["--learning_rate=0.001", "--batch_size=16"]
distributor = TorchDistributor(num_processes=2, local_mode=False, use_gpu=True)
distributor.run(train_file, *args)
문제 해결
Notebook 워크플로의 일반적인 오류는 분산 학습을 실행할 때 개체를 찾거나 절일 수 없다는 것입니다. 라이브러리 가져오기 문이 다른 실행기에 배포되지 않는 경우에 발생할 수 있습니다.
이 문제를 피하려면 TorchDistributor(...).run(<func>)로 호출되는 학습 함수의 맨 위와 학습 메서드에서 호출되는 다른 사용자 정의 함수 내부에 모든 import 문(예: import torch)을 포함합니다.
NCCL 실패: ncclInternalError: Internal check failed.
다중 노드 학습 중에 이 오류가 발생하면 일반적으로 GPU 간의 네트워크 통신에 문제가 있음을 나타냅니다. 이 문제는 NCCL(NVIDIA 집단 통신 라이브러리)이 GPU 통신에 특정 네트워크 인터페이스를 사용할 수 없는 경우에 발생합니다.
이 오류를 해결하려면 학습 코드에 다음 코드 조각을 추가하여 기본 네트워크 인터페이스를 사용합니다.
import os
os.environ["NCCL_SOCKET_IFNAME"] = "eth0"
예제 Notebook
다음 Notebook 예제에서는 PyTorch를 사용하여 분산 학습을 수행하는 방법을 보여 줍니다.