Bagikan melalui

Beban kerja multi-GPU

Penting

Fitur ini ada di Beta. Admin ruang kerja dapat mengontrol akses ke fitur ini dari halaman Pratinjau . Lihat Kelola Pratinjau Azure Databricks.

Anda dapat meluncurkan beban kerja terdistribusi di beberapa GPU pada satu simpul menggunakan API Python GPU Tanpa Server. API menyediakan antarmuka terpadu sederhana yang mengabstraksi detail provisi GPU, penyiapan lingkungan, dan distribusi beban kerja. Dengan perubahan kode minimal, Anda dapat dengan mulus berpindah dari pelatihan GPU tunggal ke eksekusi terdistribusi multi-GPU dari notebook yang sama.

Kerangka kerja yang didukung

@distributed API terintegrasi dengan pustaka pelatihan terdistribusi utama:

  • PyTorch Distributed Data Parallel (DDP) — Paralelisme data multi-GPU standar.
  • Paralel Data Yang Dipecah Sepenuhnya (FSDP) — Pelatihan hemat memori untuk model besar.
  • DeepSpeed — Pustaka pengoptimalan Microsoft untuk pelatihan model besar.

serverless_gpu API vs. TorchDistributor

Tabel berikut membandingkan serverless_gpu@distributed API dengan TorchDistributor:

Feature serverless_gpu @distributed API TorchDistributor
Infrastruktur Sepenuhnya tanpa server, tidak ada manajemen kluster Memerlukan kluster Spark dengan pekerja GPU
Siapkan Dekorator tunggal, konfigurasi minimal Memerlukan kluster Spark dan pengaturan TorchDistributor
Dukungan kerangka kerja PyTorch DDP, FSDP, DeepSpeed Terutama PyTorch DDP
Pemuatan data Di dalam fungsi dekorator, menggunakan Unity Catalog Volumes Melalui Spark atau sistem file

serverless_gpu API adalah pendekatan yang direkomendasikan untuk beban kerja pembelajaran mendalam baru di Databricks. TorchDistributor tetap tersedia untuk beban kerja yang digabungkan erat dengan kluster Spark.

Cepat Mulai

API GPU tanpa server untuk pelatihan terdistribusi telah diinstal sebelumnya saat Anda terhubung ke GPU tanpa server dalam buku catatan dan pekerjaan Databricks. Kami merekomendasikan lingkungan GPU 4 ke atas. Untuk menggunakannya untuk pelatihan terdistribusi, impor dan gunakan dekorator distributed untuk mendistribusikan fungsi pelatihan Anda.

Bungkus kode pelatihan model dalam fungsi dan dekori fungsi dengan dekorator @distributed. Fungsi yang didekorasi menjadi titik masuk untuk eksekusi terdistribusi - semua logika pelatihan, pemuatan data, dan inisialisasi model harus ditentukan di dalam fungsi ini.

Peringatan

Parameter gpu_type di @distributed harus sesuai dengan jenis akselerator yang terhubung dengan laptop Anda. Misalnya, @distributed(gpus=8, gpu_type='H100') mengharuskan notebook Anda tersambung ke akselerator H100. Menggunakan jenis akselerator yang tidak cocok (seperti menyambungkan ke A10 saat menentukan H100) akan menyebabkan beban kerja gagal.

Cuplikan kode di bawah ini menunjukkan penggunaan @distributeddasar :

# Import the distributed decorator
from serverless_gpu import distributed

# Decorate your training function with @distributed and specify the number of GPUs and GPU type
@distributed(gpus=8, gpu_type='H100')
def run_train():
    ...

Di bawah ini adalah contoh lengkap yang melatih model multilayer perceptron (MLP) pada 8 GPU H100 dari notebook:

  1. Siapkan model Anda dan tentukan fungsi utilitas.

    
# Define the model
import os
import torch
import torch.distributed as dist
import torch.nn as nn

def setup():
    dist.init_process_group("nccl")
    torch.cuda.set_device(int(os.environ["LOCAL_RANK"]))

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. Impor pustaka serverless_gpu dan modul terdistribusi .

    import serverless_gpu
from serverless_gpu import distributed

  3. Bungkus kode pelatihan model dalam fungsi dan dekori fungsi dengan dekorator @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. Jalankan pelatihan terdistribusi dengan memanggil fungsi terdistribusi dengan argumen yang ditentukan pengguna.

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

  5. Saat dijalankan, tautan eksekusi MLflow dihasilkan dalam output sel buku catatan. Klik tautan eksekusi MLflow atau temukan di panel Eksperimen untuk melihat hasil eksekusi.

Rincian eksekusi terdistribusi

API GPU tanpa server terdiri dari beberapa komponen utama:

  • Manajer komputasi: Menangani alokasi dan manajemen sumber daya
  • Lingkungan runtime: Mengelola lingkungan dan dependensi Python
  • Peluncur: Mengatur eksekusi dan pemantauan pekerjaan

Saat berjalan dalam mode terdistribusi:

  • Fungsi ini diserialisasikan dan didistribusikan di seluruh jumlah GPU yang ditentukan
  • Setiap GPU menjalankan salinan fungsi dengan parameter yang sama
  • Lingkungan disinkronkan di semua GPU
  • Hasil dikumpulkan dan dikembalikan dari semua GPU

API mendukung pustaka pelatihan paralel populer seperti Distributed Data Parallel (DDP), Fully Sharded Data Parallel (FSDP), DeepSpeed.

Anda dapat menemukan skenario pelatihan terdistribusi yang lebih nyata menggunakan berbagai pustaka dalam contoh buku catatan.

FAQs

Di mana kode pemuatan data harus ditempatkan?

Saat menggunakan API GPU Tanpa Server untuk pelatihan terdistribusi, pindahkan kode pemuatan data di dalam dekorator @distributed . Ukuran himpunan data dapat melebihi ukuran maksimum yang diizinkan oleh pickle, sehingga disarankan untuk membuat himpunan data di dalam dekorator, seperti yang ditunjukkan di bawah ini:

from serverless_gpu import distributed

# this may cause pickle error
dataset = get_dataset(file_path)
@distributed(gpus=8, gpu_type='H100')
def run_train():
  # good practice
  dataset = get_dataset(file_path)
  ....

Pelajari lebih lanjut

Untuk referensi API, lihat dokumentasi API Python GPU Tanpa Server .

  • Last updated on