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

중요합니다

이 기능은 베타 버전으로 제공됩니다.

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

이 페이지의 주요 내용:

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

Note

이 페이지는 Databricks 노트북에서 Databricks AI 런타임 SDK를 사용한 분산 학습에 대해 다룹니다. 로컬 머신에서 분산 학습 워크로드를 제출하려면 퍼블릭 프리뷰 상태인 AI Runtime용 Databricks CLI 명령을 사용하세요. Databricks CLI를 AI Runtime과 함께 사용을(를) 참조하세요.

Quickstart

이 페이지의 예시들은 GPU 환경 버전 6 이상에 사전 설치된 Databricks AI 런타임 SDK를 사용합니다. 버전 5 이하에서는 대신 serverless_gpu 패키지를 사용하세요. serverless_gpu에서 마이그레이션을 참조하세요. 노트북을 가속기에 연결하려면 노트북에 연결을(를) 참조하세요.

훈련 함수에 @distributed 데코레이터를 적용한 다음, .distributed()로 호출하세요:

from databricks.air import distributed

# By default, @distributed uses every GPU on the node your notebook is connected
# to. To use fewer, pass num_accelerators (for example, num_accelerators=2).
@distributed()
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()

분산 학습에는 GPU가 8개인 단일 노드를 프로비전하는 8xH100 가속기가 필요합니다. 기본적으로 @distributed 8개의 GPU 전부를 사용합니다. num_accelerators를 더 적게 사용하도록 설정합니다. accelerator_type 이 파라미터도 선택 사항이며, 노트북이 연결된 가속기를 기반으로 자동으로 감지됩니다.

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

지원되는 프레임워크

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

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

각 라이브러리를 사용하는 실제 훈련 시나리오를 보려면 분산 훈련 예제를 참조하세요.

장식가의 @distributed 작동 원리

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

• 직렬화 및 팬아웃: 함수는 직렬화되어 선택한 각 GPU에서 실행됩니다. 모든 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 라이브러리를 기반으로 합니다.

serverless_gpu에서 마이그레이션

serverless_gpu 패키지는 더 이상 사용이 권장되지 않으며 새로운 기능도 제공되지 않습니다. GPU 환경 버전 6부터는 Databricks AI 런타임 SDK(databricks-sdk-air)가 이를 대체하며, databricks.air 네임스페이스 아래에서 인자명이 변경된 동등한 API를 제공합니다. 사용하는 패키지는 GPU 환경 버전에 따라 다릅니다:

• 버전 5 이하: serverless_gpu만 포함하세요. 이 버전들에서는 serverless_gpu를 계속 사용하세요.
• 버전 6: 두 패키지 모두 포함되지만, AI 런타임 SDK만 업데이트를 받습니다. 버전 6은 serverless_gpu을 포함하는 마지막 GPU 환경 버전입니다.

AI 런타임 SDK는 Python용 Databricks SDK와는 별도의 패키지입니다 (databricks-sdk).

	serverless_gpu(더 이상 사용되지 않음)	AI 런타임 SDK (databricks.air)
GPU 환경의 버전들	6세 이하	6세 이상
Import	from serverless_gpu import distributed, ray_init	from databricks.air import distributed, ray_init
GPU 개수 인수	gpus	num_accelerators
GPU 타입 인수	gpu_type(GPUType)	accelerator_type(AcceleratorType)
GPU 타입 값	"A10", "H100", "B300"	"GPU_1xA10", "GPU_1xH100", , "GPU_8xH100", "GPU_8xB300"
데이터 로딩 및 체크포인트 모듈	serverless_gpu.data	databricks.air.data; AI 런타임에서 데이터 로드를 참조하세요.

노트북을 마이그레이션하려면:

1. 노트북을 AI 런타임 SDK가 포함된 GPU 환경 버전 6 이상에 연결하세요.
2. serverless_gpu에서 databricks.air로 가져오기를 변경하세요.
3. gpus을 num_accelerators(으)로, gpu_type를 accelerator_type(으)로 이름을 바꾸고 새로운 가속기 유형 값을 사용하세요. 예전 인자 이름과 GPU 타입 값은 databricks.air에서 여전히 작동하지만 FutureWarning를 표시합니다. GPU 환경 버전 7에서 지원이 종료됩니다.

Tip

이를 위해 Genie Code에 다음을 수행하도록 지시합니다.

Migrate this notebook from the serverless_gpu SDK to the AI Runtime SDK.

databricks.air.data를 사용하는 Databricks CLI 워크로드의 경우, 워크로드의 databricks-sdk-air[data]에 environment.dependencies를 추가하세요.

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

Spark 클러스터에서 TorchDistributor를 사용해 분산 훈련을 실행하고 작업 부하가 단일 노드에 맞는 경우, AI 런타임에서는 @distributed API가 권장되는 대체 수단입니다. DDP, FSDP, 또는 DeepSpeed 모델 코드는 변하지 않습니다.

특징	@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 가속기에 연결하세요. @distributed 기본적으로 8개의 GPU 전부를 사용합니다.
• 장식된 함수 안에서 데이터 로딩을 이동하세요. 데이터 로딩을 참조하세요.
• 기존 DDP, FSDP, 또는 DeepSpeed 모델 코드를 유지하세요. 데코레이터는 세 가지 모두를 지원합니다.

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

전체 예시

다음 예시는 노트북에서 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. AI 런타임 SDK에서 distributed 데코레이터를 가져오세요.

   from databricks.air import distributed
   

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

   @distributed()
   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-parallel 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 databricks.air import distributed

# Don't do this: loading the dataset outside the function means it's captured
# and pickled with the function, which can raise a pickle error.
# dataset = get_dataset(file_path)

@distributed()
def run_train():
    # Load the dataset inside the decorated function instead.
    dataset = get_dataset(file_path)
    ...

Unity Catalog 볼륨에 저장된 파일 기반 데이터의 경우, 로컬 캐싱으로 파일을 스트리밍하고 이를 랭크와 워커에 자동으로 분할하는 UCVolumeDataset의 databricks.air.data를 사용합니다. 볼륨에 분산 학습의 검사점을 지정하려면 다음을 사용합니다 UCVolumeWriterUCVolumeReader. AI Runtime에서 데이터 로드 및 Distributed Checkpoint(DCP)를 사용한 체크포인트를 참조하세요.

Limitations

• @distributed 단일 노드만 지원합니다: 노트북이 연결된 노드의 GPU들에서 실행됩니다. 여러 노드에 걸쳐 학습하려면 AI 런타임 Databricks CLI를 사용하세요. Databricks CLI를 AI Runtime과 함께 사용을(를) 참조하세요.
• 만약 accelerator_type를 설정했다면, 노트북이 연결된 가속기와 일치해야 합니다(예: "GPU_8xH100"). 불일치로 인해 워크로드가 실패합니다. 이 매개변수는 선택 사항이며, 누락 시 자동으로 감지됩니다.
• AI 런타임은 GPU 환경 v4 이상을 지원합니다. 커스텀 타임아웃( timeout 파라미터)은 GPU 환경 v5 이상이어야 합니다.
• 장식가는 기본적으로 3시간 후에 타임아웃을 합니다. 변경하려면 초 단위로 timeout를 전달하거나, 비활성화하려면 timeout=None를 전달하세요.
• 실행은 노트북의 수명주기 내에서 진행됩니다. 노트북을 종료하면 실행 기간이 종료됩니다.

추가 리소스

• AcceleratorType 데코레이터, @distributed, Ray API에 대해서는 Databricks AI 런타임 SDK API 참조를 확인하세요.
• 훈련 파이프라인을 더 효율적이고 회복력 있게 만드는 패턴은 성과 및 회복력 가이드를 참조하세요.
• 종단 간 훈련 시나리오에 대해서는 분산 훈련 예제를 참조하세요.