노트북 환경에서의 분산 학습

중요합니다

이 기능은 베타 버전으로 제공됩니다. 작업 영역 관리자는 미리 보기 페이지에서 이 기능에 대한 액세스를 제어할 수 있습니다. Azure Databricks 미리 보기 관리를 참조하세요.

@distributed Serverless GPU Python API의 데코레이터는 Databricks 노트북에서 분산 학습을 실행하는 가장 편리한 방법입니다. 훈련 기능을 꾸미고 호출하면 AI 런타임이 노트북이 연결된 노드의 모든 GPU에 걸쳐 실행됩니다. 동일한 코드가 클러스터 프로비저닝이나 분산 런처 설정 없이 단일 GPU에서 다중 GPU로 확장됩니다.

Tip

• @distributed 데코레이터는 노트북 내에서 사용 중인 노드의 모든 GPU에서 학습 함수를 실행합니다.
• PyTorch DDP, FSDP, DeepSpeed를 지원하며, 단일 GPU 코드를 최소한의 변경으로 멀티 GPU로 전환합니다.
• 노트북을 8xH100 가속기에 연결하고 멀티 GPU 훈련을 시작 gpus=8 하세요.

Quickstart

serverless_gpu 노트북이 서버리스 GPU에 연결되어 있을 때 패키지가 미리 설치되어 있습니다. 훈련 함수에 @distributed 데코레이터를 적용한 다음, .distributed()로 호출하세요:

from serverless_gpu import distributed

# gpus is the number of GPUs on the node. gpu_type is optional and
# auto-detected from the accelerator your notebook is connected to.
@distributed(gpus=8, gpu_type="H100")
def train():
    import os
    import torch
    import torch.distributed as dist

    # Bind this process to its own GPU before training.
    local_rank = int(os.environ["LOCAL_RANK"])
    torch.cuda.set_device(local_rank)
    device = torch.device(f"cuda:{local_rank}")
    dist.init_process_group("nccl")
    # ... build the model and data on `device`, then run your training loop ...
    dist.destroy_process_group()

train.distributed()

각 호출은 .distributed() MLflow 실행(또는 이미 활성화된 경우 중첩된 자식 실행)을 생성하고, 셀 출력에 실행 링크를 출력합니다. 완전하고 실행 가능한 공략은 전체 예제를 참조하세요.

지원되는 프레임워크

API는 @distributed 주요 분산 학습 라이브러리와 통합됩니다.

• PyTorch DDP(분산 데이터 병렬): 표준 다중 GPU 데이터 병렬 처리입니다.
• FSDP(완전 분할 데이터 병렬) : 대형 모델에 대한 메모리 효율적인 학습입니다.
• DeepSpeed: 대규모 모델 학습을 위한 Microsoft 최적화 라이브러리입니다.

각 라이브러리를 활용한 실제 훈련 시나리오는 노트북 예제를 참고하세요.

장식가의 @distributed 작동 원리

.distributed()로 데코레이터가 적용된 함수를 호출하면 AI Runtime이 분산 런처를 사용해 직접 구성해야 하는 관련 작업을 처리합니다:

• 직렬화 및 팬아웃: 함수는 요청하는 각 gpus 지점에서 직렬화되어 실행됩니다. 모든 GPU는 동일한 인자를 가진 함수 복사본을 실행합니다.
• 환경 동기화: Python 환경과 의존성이 모든 랭크에 복제되어 모든 프로세스가 동일한 코드를 실행합니다.
• 랭크 환경 변수: LOCAL_RANK와 같은 표준 변수는 각 프로세스마다 설정됩니다. 함수에서 모델과 데이터를 올바른 장치에 배치하기 위해 그것들을 읽으세요.
• 결과 수집: 모든 랭크에서 반환 값을 수집하여 호출자에게 반환합니다.
• MLflow 추적: 각 호출은 .distributed() MLflow 런을 생성하거나, 이미 활성화된 경우 중첩된 자식 런을 만들어 함수에서 기록된 메트릭이 같은 런에 포함되도록 합니다.
• 라이프사이클과 타임아웃: 분산 실행은 노트북의 라이프사이클 내에서 실행됩니다. 노트북을 종료하면 실행 기간이 종료됩니다. 장식가의 기본 타임아웃은 3시간입니다. 변경하려면 초 단위로 timeout를 전달하거나, 비활성화하려면 timeout=None를 전달하세요. 커스텀 타임아웃은 GPU 환경 v5 이상이 필요합니다.

이 API는 Distributed Data Parallel (DDP), Fully Sharded Data Parallel (FSDP), DeepSpeed라는 표준 PyTorch 라이브러리를 기반으로 합니다.

토치디스트리뷰터에서 온 이야기입니다

오늘날 Spark에서 TorchDistributor 와 함께 분산 PyTorch를 실행하고 작업 부하가 단일 노드에 맞는다면, serverless_gpu@distributed API가 새로운 딥러닝 작업을 대체하는 데 권장됩니다. Spark 클러스터를 제거하고 단일 GPU에서 다중 GPU로 동일한 코드 경로를 제공합니다.

특징	serverless_gpu @distributed API	토치디스트리뷰터
인프라	완전 서버리스, 클러스터 관리 없음	GPU 작업자가 있는 Spark 클러스터 필요
설치	단일 데코레이터, 최소 구성	Spark 클러스터 및 TorchDistributor 설정 필요
프레임워크 지원	파이토치 DDP, FSDP, 딥스피드	주로 PyTorch DDP
데이터 로드	데코레이터 내부에서는 Unity 카탈로그 볼륨(UCVolumeDataset 스트리밍 파일 데이터용)을 사용합니다.	Spark 또는 파일 시스템을 통해

단일 노드 워크로드를 마이그레이션하려면:

• TorchDistributor(...).run(train_fn, ...)에서 @distributed 호출을 train_fn 데코레이터로 바꾼 다음, train_fn.distributed(...)로 실행하세요.
• Spark 클러스터와 GPU 워커 구성을 제거하세요. 노트북을 8xH100 가속기에 연결하고 대신 gpus=8로 설정하세요.
• 장식된 함수 안에서 데이터 로딩을 이동하세요. 데이터 로딩을 참조하세요.
• 기존 DDP, FSDP, 또는 DeepSpeed 모델 코드를 유지하세요. 데코레이터는 세 가지 모두를 지원합니다.

@distributed 단일 노드에서 실행되므로( 제한사항 참조), 모든 TorchDistributor 작업을 대체하지는 않습니다. Spark 통합에 의존하는 작업은 TorchDistributor에 보관하세요. 로컬 머신이나 여러 노드에서 분산 학습을 실행하려면 공개 프리뷰에 있는 AI 런타임 CLI를 사용하세요. AI 런타임 CLI를 참조하세요.

전체 예시

다음 예시는 노트북에서 8개의 H100 GPU로 다층 퍼셉트론(MLP) 모델을 학습시킵니다.

1. 모델을 설정하고 유틸리티 함수를 정의합니다.

   
   # Define the model
   import os
   import torch
   import torch.distributed as dist
   import torch.nn as nn
   
   def setup():
       torch.cuda.set_device(int(os.environ["LOCAL_RANK"]))
       dist.init_process_group("nccl")
   
   def cleanup():
       dist.destroy_process_group()
   
   class SimpleMLP(nn.Module):
       def __init__(self, input_dim=10, hidden_dim=64, output_dim=1):
           super().__init__()
           self.net = nn.Sequential(
               nn.Linear(input_dim, hidden_dim),
               nn.ReLU(),
               nn.Dropout(0.2),
               nn.Linear(hidden_dim, hidden_dim),
               nn.ReLU(),
               nn.Dropout(0.2),
               nn.Linear(hidden_dim, output_dim)
           )
   
       def forward(self, x):
           return self.net(x)
   

2. serverless_gpu 라이브러리와 distributed 모듈을 가져오세요.

   import serverless_gpu
   from serverless_gpu import distributed
   

3. 모델 학습 코드를 함수로 감싸고 @distributed 데코레이터를 사용하여 함수를 장식합니다. 장식된 함수는 분산 실행의 진입점이므로 모든 학습 논리, 데이터 로딩, 모델 초기화 기능을 그 안에 정의하세요.

   @distributed(gpus=8, gpu_type='H100')
   def run_train(num_epochs: int, batch_size: int) -> None:
       import mlflow
       import torch.optim as optim
       from torch.nn.parallel import DistributedDataParallel as DDP
       from torch.utils.data import DataLoader, DistributedSampler, TensorDataset
   
       # 1. Set up multi-GPU environment
       setup()
       device = torch.device(f"cuda:{int(os.environ['LOCAL_RANK'])}")
   
       # 2. Apply the Torch distributed data parallel (DDP) library for data-parellel training.
       model = SimpleMLP().to(device)
       model = DDP(model, device_ids=[device])
   
       # 3. Create and load dataset.
       x = torch.randn(5000, 10)
       y = torch.randn(5000, 1)
   
       dataset = TensorDataset(x, y)
       sampler = DistributedSampler(dataset)
       dataloader = DataLoader(dataset, sampler=sampler, batch_size=batch_size)
   
       # 4. Define the training loop.
       optimizer = optim.Adam(model.parameters(), lr=0.001)
       loss_fn = nn.MSELoss()
   
       for epoch in range(num_epochs):
           sampler.set_epoch(epoch)
           model.train()
           total_loss = 0.0
           for step, (xb, yb) in enumerate(dataloader):
               xb, yb = xb.to(device), yb.to(device)
               optimizer.zero_grad()
               loss = loss_fn(model(xb), yb)
               # Log loss to MLflow metric
               mlflow.log_metric("loss", loss.item(), step=step)
   
               loss.backward()
               optimizer.step()
               total_loss += loss.item() * xb.size(0)
   
           mlflow.log_metric("total_loss", total_loss)
           print(f"Total loss for epoch {epoch}: {total_loss}")
   
       cleanup()
   

4. 분산 학습을 실행하려면 사용자 정의 인수로 분산 함수를 호출합니다.

   run_train.distributed(num_epochs=3, batch_size=1)
   

5. 실행되면 Notebook 셀 출력에 MLflow 실행 링크가 생성됩니다. MLflow 실행 링크를 클릭하거나 실험 패널에서 찾아서 실행 결과를 확인합니다. 실험 이름 사용자 지정, 메트릭 추적 및 실행 재개에 대한 자세한 내용은 실험 추적 및 관찰 가능성을 참조하세요.

데이터 로딩

함수 안에 @distributed 데이터 로딩 코드를 넣으세요. 데이터셋은 가 허용하는 pickle최대 크기를 초과할 수 있으므로, 데코레이터 내에서 생성하거나 로드하면 직렬화 오류를 피할 수 있습니다:

from serverless_gpu import distributed

# This may cause a pickle error because the dataset is captured by the function.
dataset = get_dataset(file_path)

@distributed(gpus=8, gpu_type='H100')
def run_train():
    # Load the dataset inside the decorated function instead.
    dataset = get_dataset(file_path)
    ...

Unity Catalog 볼륨에 저장된 파일 기반 데이터의 경우, 로컬 캐싱으로 파일을 스트리밍하고 이를 랭크와 워커에 자동으로 분할하는 UCVolumeDataset의 serverless_gpu.data를 사용합니다. 볼륨에 분산 학습의 검사점을 지정하려면 다음을 사용합니다 UCVolumeWriterUCVolumeReader. AI 런타임 및 모델 검사점에서 데이터 로드를 참조하세요.

Limitations

• 분산 학습은 노트북이 연결된 단일 노드의 GPU를 통해 실행됩니다. 완전한 멀티 GPU 훈련을 원한다면, 8개의 H100 가속기를 연결해 하나의 노드에 8개의 GPU를 배치하고, 를 설정 gpus=8하세요.
• 가속기 종류가 일치해야 합니다. gpu_type에서 @distributed를 설정하는 경우, 노트북이 연결된 가속기("H100" 또는 "A10")와 일치해야 합니다. 불일치로 인해 워크로드가 실패합니다. 이 매개변수는 선택 사항이며, 누락 시 자동으로 감지됩니다.
• AI 런타임은 GPU 환경 v4 이상을 권장합니다. 커스텀 타임아웃( timeout 파라미터)은 GPU 환경 v5 이상이어야 합니다.
• 장식가는 기본적으로 3시간 후에 타임아웃을 합니다. 변경하려면 초 단위로 timeout를 전달하거나, 비활성화하려면 timeout=None를 전달하세요.
• 실행은 노트북의 수명주기 내에서 진행됩니다. 노트북을 종료하면 실행 기간이 종료됩니다.

자세히 알아보기

• 데코레이터, GPUTypeRay API에 대해서는 @distributedServerless GPU Python API 참고 문서를 참조하세요.
• 훈련 파이프라인을 더 효율적이고 회복력 있게 만드는 패턴은 성과 및 회복력 가이드를 참조하세요.
• 엔드 투 엔드 학습 시나리오는 노트북 예제를 참조하세요.