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(**kwargs) 실행하고 기본 파일 경로인 경우 CLI 명령을 torchrun main *args 실행합니다.

요구 사항

• Spark 3.4
• Databricks Runtime 13.0 ML 이상

Notebook에 대한 개발 워크플로

모델 만들기 및 학습 프로세스가 로컬 컴퓨터의 Notebook 또는 Databricks Notebook에서 완전히 발생하는 경우 코드를 분산 학습에 대비하기 위해 약간만 변경하면 됩니다.

1. 단일 노드 코드 준비: PyTorch, PyTorch Lightning 또는 PyTorch/PyTorch Lightning을 기반으로 하는 다른 프레임워크(예: HuggingFace Trainer API)를 사용하여 단일 노드 코드를 준비하고 테스트합니다.

2. 표준 분산 학습을 위한 코드 준비: 단일 프로세스 학습을 분산 학습으로 변환해야 합니다. 이 분산 코드는 모두 함께 사용할 수 있는 하나의 학습 함수 내에 포함되도록 TorchDistributor합니다.

3. 학습 함수 내에서 가져오기 이동: 학습 함수 내에 필요한 가져오기(예: import torch가져오기)를 추가합니다. 이렇게 하면 일반적인 피클링 오류를 방지할 수 있습니다. 또한 device_id 모델과 데이터가 연결되는 것은 다음을 통해 결정됩니다.

   device_id = int(os.environ["LOCAL_RANK"])
   

4. 분산 학습 시작: 원하는 매개 변수를 사용하여 인스턴스화 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로 쉽게 마이그레이션할 수 있습니다.

1. 리포지토리 가져오기: 외부 리포지토리를 Databricks Git 폴더로 가져옵니다.
2. 새 Notebook 을 만들어 리포지토리 내에서 새 Azure Databricks Notebook을 초기화합니다.
3. 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: ) 을 포함합니다.

예제 Notebook

다음 Notebook 예제에서는 PyTorch를 사용하여 분산 학습을 수행하는 방법을 보여 줍니다.

Databricks Notebook에 대한 엔드 투 엔드 분산 학습

전자 필기장 가져오기

Hugging Face 모델 Notebook 분산 미세 조정

전자 필기장 가져오기

PyTorch 파일 Notebook에 대한 분산 학습

전자 필기장 가져오기

PyTorch Lightning Notebook을 사용한 분산 학습

전자 필기장 가져오기

Petastorm Notebook을 사용하여 분산 데이터 로드

전자 필기장 가져오기