Pelatihan terdistribusi dalam buku catatan

Penting

Fitur ini ada di Beta.

Dekorator @distributed dari Databricks AI Runtime SDK (databricks-sdk-air) adalah cara paling praktis untuk menjalankan pelatihan terdistribusi dari notebook Databricks. Tambahkan dekorator pada fungsi pelatihan Anda, lalu panggil fungsi tersebut, dan AI Runtime akan menjalankannya pada semua GPU di node tempat notebook Anda terhubung. Kode yang sama dapat diskalakan dari single-GPU ke multi-GPU tanpa perlu menyediakan cluster dan tanpa perlu mengonfigurasi launcher terdistribusi.

Poin-poin utama di halaman ini:

• Dekorator @distributed menjalankan fungsi pelatihan pada setiap GPU pada node Anda dari dalam notebook.
• Mendukung PyTorch DDP, FSDP, dan DeepSpeed, serta memindahkan kode GPU tunggal ke multi-GPU dengan perubahan minimal.
• Hubungkan notebook Anda ke akselerator 8xH100 untuk pelatihan multi-GPU penuh. Secara default, @distributed menggunakan semua GPU di node.

Note

Halaman ini membahas pelatihan terdistribusi dari notebook Databricks dengan Databricks AI Runtime SDK. Untuk mengirimkan beban pelatihan terdistribusi dari mesin lokal Anda, gunakan perintah Databricks CLI untuk AI Runtime, yang berada dalam Pratinjau Publik. Lihat Gunakan Databricks CLI dengan AI Runtime.

Quickstart

Contoh di halaman ini menggunakan Databricks AI Runtime SDK, yang sudah terpasang di lingkungan GPU versi 6 dan ke atas. Pada versi 5 ke bawah, gunakan serverless_gpu paket sebagai gantinya. Lihat Bermigrasi dari serverless_gpu. Untuk menghubungkan notebook ke akselerator, lihat Menghubungkan ke notebook.

Dekorasi fungsi latihan Anda dengan @distributed, lalu sebut dengan .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()

Pelatihan terdistribusi memerlukan akselerator 8xH100, yang menyediakan satu simpul dengan 8 GPU. Secara default, @distributed menggunakan semua delapan GPU. Atur num_accelerators agar menggunakan lebih sedikit. accelerator_type parameter ini juga opsional dan secara otomatis terdeteksi dari akselerator yang terhubung dengan notebook Anda.

Setiap pemanggilan .distributed() membuat run MLflow (atau run turunan bertingkat jika sudah ada run yang aktif) dan menampilkan tautan run pada output sel. Untuk panduan lengkap yang dapat dijalankan, lihat Contoh lengkap.

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.

Untuk skenario pelatihan nyata yang menggunakan setiap pustaka, lihat contoh pelatihan terdistribusi.

Cara kerja dekorator @distributed

Saat Anda memanggil fungsi yang didekorasi dengan .distributed(), AI Runtime menangani mekanik yang biasanya Anda konfigurasikan secara manual dengan peluncur terdistribusi:

• Serialisasi dan fan-out: Fungsi ini dijalankan secara serial dan diluncurkan pada setiap GPU yang dipilih. Setiap GPU menjalankan salinan fungsi dengan argumen yang sama.
• Sinkronisasi lingkungan: Lingkungan Python dan dependensi direplikasi di semua tingkatan, sehingga setiap proses menjalankan kode yang sama.
• Variabel lingkungan rank: Variabel standar seperti LOCAL_RANK ditetapkan untuk setiap proses. Bacalah semuanya dalam fungsi Anda untuk menempatkan model dan data pada perangkat yang tepat.
• Pengumpulan hasil: Nilai hasil dikumpulkan dari semua peringkat dan dikembalikan ke penelepon.
• Pelacakan MLflow: Setiap .distributed() panggilan membuat MLflow run, atau nested child run jika sudah aktif, sehingga metrik yang dicatat dari fungsi Anda akan berada di run yang sama.
• Siklus hidup dan batas waktu: Eksekusi terdistribusi berjalan selama siklus hidup buku catatan. Mengakhiri notebook akan mengakhiri proses berjalan. Dekorator memiliki waktu istirahat default selama 3 jam. Lewatkan timeout dalam hitungan detik untuk mengubahnya, atau timeout=None menonaktifkannya. Timeout khusus memerlukan lingkungan GPU versi 5 atau yang lebih baru.

API ini dibangun di atas pustaka PyTorch standar: Distributed Data Parallel (DDP), Fully Sharded Data Parallel (FSDP), dan DeepSpeed.

Migrasi dari serverless_gpu

Paket serverless_gpu sudah usang dan tidak lagi menerima fitur baru. Mulai dari lingkungan GPU versi 6, Databricks AI Runtime SDK (databricks-sdk-air) menggantikannya dan menyediakan API setara, dengan argumen yang diberi nama ulang, di bawah databricks.air namespace. Paket yang Anda gunakan tergantung pada versi lingkungan GPU:

• Versi 5 dan ke bawah: Sertakan hanya serverless_gpu. Lanjutkan penggunaan serverless_gpu pada versi-versi ini.
• Versi 6: Termasuk kedua paket, tetapi hanya AI Runtime SDK yang menerima pembaruan. Versi 6 adalah versi lingkungan GPU terakhir yang mencakup serverless_gpu.

AI Runtime SDK adalah paket terpisah dari Databricks SDK untuk Python (databricks-sdk).

	serverless_gpu (tidak digunakan lagi)	AI Runtime SDK (databricks.air)
Versi lingkungan GPU	6 ke bawah	6 tahun ke atas
Import	from serverless_gpu import distributed, ray_init	from databricks.air import distributed, ray_init
Parameter jumlah GPU	gpus	num_accelerators
Argumen tipe GPU	gpu_type(GPUType)	accelerator_type(AcceleratorType)
Nilai tipe GPU	"A10", "H100", "B300"	"GPU_1xA10","GPU_1xH100","GPU_8xH100","GPU_8xB300"
Modul pemuatan data dan checkpointing	serverless_gpu.data	databricks.air.data. Lihat Memuat data pada Runtime AI.

Untuk memigrasikan notebook:

1. Hubungkan notebook ke lingkungan GPU versi 6 atau lebih tinggi, yang mencakup AI Runtime SDK.
2. Ubah impor dari serverless_gpu ke databricks.air.
3. Ganti nama gpus menjadi num_accelerators dan gpu_type ke accelerator_type, lalu gunakan nilai tipe akselerator yang baru. Nama argumen lama dan nilai tipe GPU masih berfungsi di databricks.air tetapi menghasilkan sebuah FutureWarning Dukungan akan berakhir di lingkungan GPU versi 7.

Tip

Beri tahu Genie Code untuk melakukan ini untuk Anda:

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

Untuk beban kerja CLI Databricks yang menggunakan databricks.air.data, tambahkan databricks-sdk-air[data] ke environment.dependencies milik beban kerja.

Berasal dari TorchDistributor

Jika Anda menjalankan pelatihan terdistribusi menggunakan TorchDistributor pada cluster Spark dan beban kerja Anda muat pada satu node, @distributed API adalah pengganti yang direkomendasikan pada AI Runtime. Kode model DDP, FSDP, atau DeepSpeed Anda tidak berubah.

Feature	@distributed Antarmuka Pemrograman Aplikasi (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 berkas)	Melalui Spark atau sistem file

Untuk memigrasikan beban kerja satu node:

• Ganti pemanggilan TorchDistributor(...).run(train_fn, ...) dengan dekorator @distributed pada train_fn, lalu jalankan dengan train_fn.distributed(...).
• Hapus konfigurasi cluster Spark dan pekerja GPU. Sambungkan notebook Anda ke akselerator 8xH100 sebagai gantinya. @distributed menggunakan semua delapan GPU secara default.
• Pindahkan pemuatan data ke dalam fungsi yang didekorasi. Lihat Pemuatan Data.
• Pertahankan kode model DDP, FSDP, atau DeepSpeed yang sudah ada. Dekorator mendukung ketiganya.

@distributed berjalan pada satu node (lihat Keterbatasan), sehingga tidak menggantikan setiap beban kerja TorchDistributor. Simpan beban kerja yang bergantung pada integrasi Spark di TorchDistributor. Untuk menjalankan pelatihan terdistribusi dari mesin lokal Anda atau di beberapa node, gunakan AI Runtime CLI sebagai gantinya, yang berada di Public Preview. Lihat Gunakan Databricks CLI dengan AI Runtime.

Contoh lengkap

Contoh berikut melatih model multilayer perceptron (MLP) menggunakan 8 GPU H100 melalui 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():
       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. Impor distributed dekorator dari AI Runtime SDK.

   from databricks.air import distributed
   

3. Bungkus kode pelatihan model dalam fungsi dan dekori fungsi dengan dekorator @distributed. Fungsi yang didekorasi adalah titik masuk untuk eksekusi terdistribusi, jadi definisikan semua logika pelatihan, pemuatan data, dan inisialisasi model di dalamnya.

   @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. Jalankan pelatihan terdistribusi dengan memanggil fungsi terdistribusi menggunakan argumen yang ditentukan pengguna.

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

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

Pemuatan data

Tempatkan kode pemuatan data di dalam @distributed fungsi. Sebuah dataset dapat melebihi ukuran maksimum yang diizinkan oleh pickle, sehingga menghasilkan atau memuatnya di dalam dekorator menghindari kesalahan serialisasi:

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

Untuk data berbasis file yang disimpan dalam volume Unity Catalog, gunakan UCVolumeDataset dari databricks.air.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 Load data on AI Runtime dan Titik pemeriksaan dengan Distributed Checkpoint (DCP).

Limitations

• @distributed hanya mendukung satu node: ini berjalan menggunakan semua GPU pada node yang terhubung dengan notebook Anda. Untuk melatih lintas beberapa node, gunakan AI Runtime Databricks CLI. Lihat Gunakan Databricks CLI dengan AI Runtime.
• Jika Anda mengatur accelerator_type, maka harus sesuai dengan akselerator yang terhubung dengan notebook Anda (misalnya, "GPU_8xH100"). Ketidakcocokan menyebabkan beban kerja gagal. Parameter ini bersifat opsional dan otomatis terdeteksi jika dihilangkan.
• AI Runtime mendukung lingkungan GPU v4 dan ke atas. Batas waktu khusus (parameter timeout) memerlukan lingkungan GPU v5 atau yang lebih baru.
• Dekorator biasanya habis waktu setelah 3 jam secara default. Lewatkan timeout dalam hitungan detik untuk mengubahnya, atau timeout=None menonaktifkannya.
• Eksekusi berlangsung dalam siklus hidup notebook. Mengakhiri notebook akan mengakhiri proses berjalan.

Sumber daya tambahan

• Untuk dekorator @distributed, AcceleratorType, dan API Ray, lihat referensi API Databricks AI Runtime SDK.
• Untuk pola yang membuat pipeline pelatihan Anda lebih efisien dan tangguh, lihat panduan kinerja dan ketahanan.
• Untuk skenario pelatihan end-to-end, lihat contoh pelatihan terdistribusi.