Pelatihan terdistribusi dalam buku catatan

Penting

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

Nota

Halaman ini membahas pelatihan terdistribusi dari notebook Databricks dengan API Python GPU nirserver. Untuk mengirimkan beban kerja pelatihan terdistribusi dari komputer lokal Anda, gunakan CLI Runtime AI, yang ada di Pratinjau Umum. Lihat AI Runtime CLI.

Anda dapat meluncurkan beban kerja terdistribusi di beberapa GPU pada satu simpul menggunakan GPU Serverless Python API. 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.

Nota

Pelatihan terdistribusi memerlukan akselerator 8xH100, yang menyediakan satu simpul dengan 8 GPU. Saat menggunakan @distributed dekorator, atur gpus=8. Parameter gpu_type bersifat opsional dan terdeteksi secara otomatis dari akselerator yang tersambung dengan buku catatan Anda.

Kerangka kerja yang didukung

@distributed API terintegrasi dengan pustaka pelatihan terdistribusi utama:

PyTorch Distributed Data Parallel (DDP): Paralelisme data multi-GPU standar.
Fully Sharded Data Parallel (FSDP): Pelatihan hemat memori untuk model besar.
DeepSpeed: pustaka pengoptimalan Microsoft untuk pelatihan model besar.

API serverless_gpu 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	Dalam dekorator, menggunakan Volume Unity Catalog (`UCVolumeDataset` untuk streaming data file)	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, sehingga semua logika pelatihan, pemuatan data, dan inisialisasi model harus ditentukan di dalam fungsi ini.

Untuk meluncurkan eksekusi terdistribusi, panggil fungsi yang didekorasi menggunakan train_function.distributed(). Setiap pemanggilan secara otomatis membuat run eksperimen MLflow, atau run turunan bertingkat jika sudah ada run yang aktif.

Peringatan

Jika Anda menetapkan gpu_type di @distributed, pastikan nilainya sesuai dengan jenis akselerator yang terhubung ke notebook Anda ("H100" atau "A10"). Menentukan jenis akselerator yang salah akan menyebabkan beban kerja gagal.

Cuplikan kode di bawah ini menunjukkan penggunaan @distributeddasar :

from serverless_gpu import distributed

# Decorate your training function with @distributed and specify the number of GPUs.
# gpu_type='H100' is optional and will be auto-detected if not set.
@distributed(gpus=8, gpu_type='H100')
def run_train():
    ...

run_train.distributed()

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

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():
    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)

Impor pustaka serverless_gpu dan modul distributed.

import serverless_gpu
from serverless_gpu import distributed

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()

Jalankan pelatihan terdistribusi dengan memanggil fungsi terdistribusi dengan argumen yang ditentukan pengguna.
```
run_train.distributed(num_epochs=3, batch_size=1)
```
Saat dijalankan, tautan proses MLflow dibuat pada output sel notebook. Klik tautan eksekusi MLflow atau temukan di panel Eksperimen untuk melihat hasil eksekusi. Untuk detail tentang menyesuaikan nama eksperimen, melacak metrik, dan melanjutkan eksekusi, lihat Pelacakan dan pengamatan eksperimen.

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
Manajemen siklus hidup: Eksekusi terdistribusi berjalan dalam siklus hidup buku catatan. Saat notebook berhenti, eksekusinya juga berhenti. Dekorator @distributed memiliki batas waktu default selama 3 jam. Untuk mengatur batas waktu kustom, lewati timeout dalam detik, atau timeout=None untuk menonaktifkannya. Mengatur batas waktu tersedia di lingkungan GPU v5 ke atas.

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.

Untuk saran dan pola agar jalur pelatihan Anda lebih efisien dan tangguh, lihat panduan kinerja dan ketahanan.

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)
  ....

Untuk data berbasis file yang disimpan dalam volume Unity Catalog, gunakan UCVolumeDataset dari serverless_gpu.data, yang melakukan streaming file dengan cache lokal dan membagi file-file tersebut ke seluruh rank dan pekerja secara otomatis. Untuk memeriksa pelatihan terdistribusi ke volume, gunakan UCVolumeWriter dan UCVolumeReader. Lihat Memuat data pada Runtime AI dan Titik pemeriksaan Model.

Pelajari lebih lanjut

Untuk referensi API, lihat dokumentasi Serverless GPU Python API.

Saran dan Komentar

Apakah halaman ini membantu?

Last updated on 2026-07-31