Beban kerja multi-GPU dan multi-simpul

Anda dapat meluncurkan beban kerja terdistribusi di beberapa GPU -- baik dalam satu simpul atau di beberapa 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 berpindah dari pelatihan GPU tunggal ke eksekusi terdistribusi di seluruh GPU jarak jauh dari notebook yang sama.

Memulai dengan Cepat

API GPU tanpa server untuk pelatihan terdistribusi telah diinstal sebelumnya di lingkungan komputasi GPU tanpa server untuk notebook Databricks. Kami merekomendasikan lingkungan GPU 4 ke atas. Untuk menggunakannya untuk pelatihan terdistribusi, impor dan gunakan dekorator distributed untuk mendistribusikan fungsi pelatihan Anda.

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, the GPU type,
# and whether or not the GPUs are remote
@distributed(gpus=8, gpu_type='A10', remote=True)
def run_train():
    ...

Di bawah ini adalah contoh lengkap yang melatih model multilayer perceptron (MLP) pada 8 node GPU A10 melalui 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():
    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)

Impor pustaka serverless_gpu dan modul terdistribusi .

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='A10', remote=True)
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 node 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 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 simpul
Hasil dikumpulkan dan dikembalikan dari semua GPU

Jika remote diatur ke True, beban kerja didistribusikan pada GPU jarak jauh. Jika remote diatur ke False, beban kerja berjalan pada simpul GPU tunggal yang tersambung ke buku catatan saat ini. Jika node memiliki beberapa chip GPU, semuanya akan digunakan.

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

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

Luncurkan dengan Ray

API gpu tanpa server juga mendukung peluncuran pelatihan terdistribusi menggunakan Ray menggunakan @ray_launch dekorator, yang berlapis di atas @distributed. Setiap ray_launch tugas pertama-tama memulai rendezvous terdistribusi Torch untuk menentukan pekerja utama Ray dan mengumpulkan IP. Peringkat-nol dimulai ray start --head (dengan ekspor metrik jika diaktifkan), mengatur RAY_ADDRESS, dan menjalankan fungsi yang didekorasi sebagai driver Ray. Simpul lain bergabung melalui ray start --address dan menunggu hingga driver menulis penanda penyelesaian.

Detail konfigurasi tambahan:

Untuk mengaktifkan pengumpulan metrik sistem Ray pada setiap simpul, gunakan RayMetricsMonitor dengan remote=True.
Tentukan opsi runtime Ray (aktor, himpunan data, grup penempatan, dan penjadwalan) di dalam fungsi yang didekorasi menggunakan API Ray standar.
Mengelola kontrol seluruh kluster (jumlah dan jenis GPU, mode jarak jauh vs. lokal, perilaku asinkron, dan variabel lingkungan kumpulan Databricks) di luar fungsi dalam argumen dekorator atau lingkungan notebook.

Contoh di bawah ini menunjukkan cara menggunakan @ray_launch:

from serverless_gpu.ray import ray_launch
@ray_launch(gpus=16, remote=True, gpu_type='A10')
def foo():
    import os
    import ray
    print(ray.state.available_resources_per_node())
    return 1
foo.distributed()

Untuk contoh lengkapnya, lihat notebook ini, yang meluncurkan Ray untuk melatih jaringan saraf Resnet18 pada beberapa GPU A10.

Anda juga dapat menggunakan API ini untuk memanggil Ray Data, pustaka pemrosesan data yang dapat diskalakan untuk beban kerja AI, untuk menjalankan inferensi batch terdistribusi pada LLM. Lihat contoh vllm dan sglang .

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, remote=True)
def run_train():
  # good practice
  dataset = get_dataset(file_path)
  ....

Dapatkah saya menggunakan kumpulan GPU yang dipesan?

Jika kumpulan GPU yang dipesan tersedia (silakan periksa dengan admin Anda) di ruang kerja Anda dan Anda tentukan remoteTrue di @distributed dekorator, beban kerja akan diluncurkan pada kumpulan GPU yang dipesan secara default. Jika Anda ingin menggunakan kumpulan GPU sesuai permintaan, atur variabel DATABRICKS_USE_RESERVED_GPU_POOL lingkungan ke False sebelum memanggil fungsi terdistribusi, seperti yang ditunjukkan di bawah ini:

import os
os.environ['DATABRICKS_USE_RESERVED_GPU_POOL'] = 'False'
@distributed(gpus=8, remote=True)
def run_train():
    ...

Pelajari lebih lanjut

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

Saran dan Komentar

Apakah halaman ini membantu?

Last updated on 2026-01-10