Panduan Praktikal untuk PyTorch Lanjutan

Pengenalan

Menguasai PyTorch tidak lagi bermaksud sekadar mengenali ciri-cirinya; ia adalah tentang mengendalikan aliran kerja kejuruteraan yang boleh diulang di mana kod latihan kekal pantas, boleh diskala, dan boleh dipulihkan di bawah bebanan sebenar pengeluaran. Ciri-ciri seperti torch.compile, torch.profiler, DDP/FSDP, dan Pengarkiban Teragih adalah senjata berkuasa dalam kumpulan ML dan infrastruktur untuk memastikan latihan kekal pantas. Walau bagaimanapun, ia tidak bermakna jika tidak digunakan dalam urutan yang betul dan disahkan dengan teliti.

Dalam artikel ini, kami berkongsi laluan aliran kerja yang disyorkan: asas → kompil → profil → skala → pemeriksaan. Kami akan membincangkan apa yang perlu diukur sebelum mengoptimumkan, perangkap biasa yang perlu dielakkan dengan kompilator & profiler, peraturan keputusan yang perlu diikuti untuk DDP vs FSDP, dan cara untuk melaksanakan pemeriksaan yang tahan ralat untuk sesi latihan berbilang nod.

Perkara Penting

• Pendekatan penyetelan prestasi PyTorch sebagai proses kejuruteraan iteratif daripada senarai semak ciri: Menerapkan secara berulang langkah asas → kompil → profil → skala → titik semakan membawa kepada prestasi yang lebih mampan berbanding mengaktifkan pengoptimuman secara bersendirian.
• Kecuali anda mempunyai garis dasar single-GPU yang stabil dengan aliran data yang diketahui dan ketepatan yang disahkan, anda tidak dapat membandingkan prestasi atau menyelesaikan masalah dengan berkesan.
• Jangan sesekali mengoptimum sebelum menetapkan asas yang betul: Asas eager GPU tunggal yang stabil dengan kadar pemprosesan yang diukur dan ketepatan yang disahkan adalah penting untuk perbandingan prestasi yang bermakna dan penjejakan ralat.
• Gunakan torch.compile dengan sengaja, bukan secara membuta: Jejak pemutusan graf, uruskan bentuk, panaskan sebelum mengukur, dan sahkan bahawa prestasi dalam keadaan mantap benar-benar meningkat berbanding pelaksanaan eager.
• Profil untuk membimbing keputusan, bukan untuk mengesahkan andaian: Gunakan torch.profiler untuk membuat keputusan secara jelas, atasi kelewatan CPU, titik panas kernel, jejak semula bentuk, dan sahkan beban komunikasi dalam jalankan teragih.
• Reka bentuk titik semakan untuk kegagalan dari hari pertama: Pemeriksaan Titik Sebaran (dengan simpanan async pilihan) mesti menangkap keadaan latihan penuh, menyokong penyusunan semula merentas bilangan GPU, dan diuji secara rutin melalui latihan pemulihan untuk memastikan ketahanan terhadap kegagalan sebenar.

Garis Asas: Menetapkan Titik Rujukan

Mulakan dengan mempunyai contoh latihan GPU tunggal yang berfungsi. Ini akan menjadi titik rujukan anda untuk kedua-dua fungsi dan prestasi. Tetapkan model, pemuat data, dan gelung latihan. Pastikan semuanya berjalan dari awal hingga akhir dalam mod eager (tanpa penyusunan, satu proses). Berikut adalah contoh bagaimana gelung latihan itu mungkin kelihatan:

import torch
import torch.nn as nn
# Dummy model and data for illustration
model = nn.Sequential(nn.Linear(100, 50), nn.ReLU(), nn.Linear(50, 10))
data = torch.randn(32, 100)
targets = torch.randint(0, 10, (32,))
# Baseline forward + backward
optimizer = torch.optim.SGD(model.parameters(), lr=1e-3)
outputs = model(data)
loss = nn.CrossEntropyLoss()(outputs, targets)
loss.backward()
optimizer.step()

Laksanakan beberapa iterasi latihan dan ukur kadar pemprosesan (contohnya, sampel/saat). Mempunyai asas yang betul memastikan model anda konvergen dan menyediakan titik perbandingan bagi penambahbaikan masa depan. Pada ketika ini, ketepatan dan prestasi asas adalah keutamaan: pastikan GPU digunakan (lihat nvidia-smi atau log torch), dan tiada halangan besar yang jelas (loadgen tersekat pada data, operasi CPU berlebihan, dll.). Jangan terus ke langkah pengkompilasian dan penambahbaikan lain sehingga anda mempunyai asas yang stabil.

Senarai Semak Asas

• Ketepatan fungsi: Model dilatih dan menghasilkan keputusan yang dijangka pada satu GPU.
• Prestasi asas direkodkan: contohnya, masa setiap kumpulan atau peratus penggunaan GPU.
• Tiada kesesakan yang jelas: Saluran data memastikan GPU sentiasa sibuk (tiada masa menganggur yang lama dalam garis asas).

Sebaik sahaja anda telah melaksanakan asas yang baik, anda boleh mula menambah ciri-ciri lanjutan secara berperingkat, bermula dengan penyusun utama PyTorch 2.x untuk mempercepat latihan.

Kompilasi: Mempercepat dengan torch.compile dalam PyTorch

PyTorch 2 membolehkan pengguna untuk JIT kompilasi model bagi prestasi menggunakan torch.compile. Balut model atau fungsi anda dengan torch.compile, dan PyTorch akan menjana kod yang dioptimumkan di belakang tabir.

Lihat betapa mudahnya di bawah:

# Switch to compiled mode
model = torch.compile(model)  # uses default backend 'inductor'

Memanggil model(data) akan secara automatik menggunakan laluan kod yang dioptimumkan. Beberapa iterasi pertama akan menyusun model anda secara just-in-time; iterasi berikutnya akan menggunakan kernel yang dioptimumkan secara terus. PyTorch 2.9 secara automatik menyimpan hasil penyusunan untuk meningkatkan prestasi pada penggunaan seterusnya (bahkan antara proses).

Walau bagaimanapun, jika anda ingin mendapatkan prestasi terbaik daripada penyusunan, terdapat beberapa perkara yang perlu anda ketahui mengenai pecahan graf dan bentuk dinamik.

Graf terhenti

Gangguan graf berlaku setiap kali pengkompil tidak dapat menangkap sebahagian kod anda ke dalam satu graf (seperti aliran kawalan Python bergantung pada data). Secara lalai, torch.compile akan secara automatik kembali kepada pelaksanaan segera untuk bahagian yang tidak disokong (mengganggu graf) dan menyusun yang selebihnya. Ini memastikan kod anda tidak akan terhenti. Walau bagaimanapun, setiap gangguan graf menghalang sebahagian kod daripada dioptimumkan (bahagian itu dijalankan sepenuhnya dalam Python, bukan dalam graf yang digabungkan). torch.compile(fullgraph=True) berguna untuk debugging gangguan graf semasa pembangunan, kerana ia akan mengeluarkan ralat jika mana-mana bahagian model tidak dapat disusun. Sebagai contoh:

model = torch.compile(model, fullgraph=True)
try:
    model(data)
except Exception as e:
    print("Graph break:", e)
model = torch.compile(model, fullgraph=True)
try:
    model(data)
except Exception as e:
    print("Graph break:", e)

Ini akan menyebabkan ralat pada operasi pertama yang tidak disokong dan menunjukkan kepada anda bahagian yang perlu diubahsuai. Mesej ralat biasanya akan menyertakan beberapa petunjuk (atau URL) mengenai apa yang menyebabkan kerosakan dan cara mengelakkannya. Log PyTorch mungkin berguna: jika anda menjalankan skrip anda dengan pemboleh ubah persekitaran TORCH_LOGS=“graph_breaks”, ia akan mencetak sebab dan lokasi di mana gangguan graf berlaku. Gunakan maklumat yang disediakan untuk menulis semula/membuang operasi sisi Python yang menghalang pengkompilasian (seperti menggantikan operasi Python list() atau kod if/else bergantung data dengan operasi tensor, atau membuangnya sepenuhnya).

Bentuk dinamik

Secara lalai, torch.compile akan memfokuskan graf yang disusun pada bentuk yang dilihatnya. Jika model anda berjalan dengan input bersaiz berbeza, graf tersebut akan disusun semula setiap kali menemui bentuk baru. Penyusunan semula ini mempunyai sedikit overhead. Dalam PyTorch 2.9, kami menambah penanda dinamik. Menentukan dynamic=True terlebih dahulu cuba menjana satu kernel sahaja yang dapat mengendalikan pelbagai bentuk (bentuk simbolik). Contohnya:

model = torch.compile(model, dynamic=True)

Ini mengarahkan penyusun untuk menjejaki graf yang digeneralisasi, yang menukar semula penyusunan apabila saiz berubah (seperti panjang urutan pembolehubah). Sebaliknya, menetapkan dynamic=False tidak akan menghasilkan teras dinamik dan akan memaksa pengkhususan kepada bentuk tepat (yang boleh menghasilkan kod lebih pantas jika benar). Menetapkan dynamic=None adalah lalai bermula dari versi 2.9, dan memberitahu penyusun untuk mengesan sendiri jika terdapat beberapa penyusunan semula dan kemudian beralih menggunakan teras dinamik secara automatik.

• Biarkan dinamik pada tetapan lalai atau tetapkan kepada False jika bentuk input adalah tetap atau jarang berubah untuk kepakaran maksimum.
• Tetapkan dynamic=True jika anda kerap mengalami kompil semula disebabkan oleh variasi bentuk (model NLP yang menjalankan ayat dengan panjang berbeza, dsb.) dan ingin mengelakkan daripada mencapai had kompil semula (yang secara lalai adalah 8 sebelum menyerah dan kembali ke pelaksanaan segera).

Contoh Tersusun

Mari kita susun model ringkas dan tunjukkan cara mengendalikan gangguan graf dan bentuk dinamik:

import torch
# Sample model with a potential graph break (data-dependent control flow)
class ToyModel(torch.nn.Module):
    def forward(self, x):
        # Example: data-dependent branching (not traceable by Dynamo)
        if x.sum() > 0:
            return x * 2
        else:
            return x
model = ToyModel()
# Attempt full-graph compilation to catch breaks
try:
    model = torch.compile(model, fullgraph=True)
    output = model(torch.randn(4, 4))
except Exception as e:
    print("Graph break detected:", e)
    # Rewrite model or accept partial graph compilation
    model = torch.compile(model, fullgraph=False)  # fallback to allow breaks

Dalam kod di atas, if yang bergantung pada data akan menyebabkan ralat pecahan graf apabila menggunakan fullgraph=True. Kami menangkap ralat itu, mencetaknya, kemudian menyusun semula dengan tingkah laku lalai yang membenarkan pecahan graf. Dalam praktiknya, anda perlu membaiki model anda (menghapuskan percabangan bergantung pada data dalam kes ini) sehingga anda dapat menyusun dengan fullgraph=True tanpa sebarang pengecualian – yang bermaksud keseluruhan model boleh disusun secara statik sebagai satu graf untuk kelajuan maksimum.

Senarai Semak Susun Atur

• Balut model dengan torch.compile(): Perubahan kod kecil yang sering membawa kepada peningkatan kelajuan yang ketara. Pilih backend “inductor” secara lalai melainkan anda mempunyai sebab untuk tidak melakukannya.
• Pemanasan: Lakukan beberapa pusingan sebelum mengambil masa – larian pertama termasuk masa kompilasi.
• Semak pecahan graf: Semasa pembangunan, cuba jalankan dengan fullgraph=True dan tetapkan pembolehubah persekitaran TORCH_LOGS=“graph_breaks ” untuk mengenal pasti bahagian kod yang tidak disokong. Buang/ubah suai penggunaan tersebut.
• Laraskan bentuk dinamik: Jika anda mendapati diri anda sering mengalami banyak penyenaraian semula (cari “recompiling ” dalam stdout/stderr/logs), anda mungkin ingin mempertimbangkan dynamic=True untuk menghasilkan graf yang fleksibel bentuknya.
• Percepatan monitor: Ukur kelajuan pemprosesan berbanding asas anda. Iterasi pertama mungkin lebih perlahan, tetapi kelajuan keadaan tetap sepatutnya melebihi asas anda yang cepat. Jika tidak, jalankan dengan petunjuk prestasi diaktifkan (contohnya, tetapkan TORCH_LOGS=“perf_hints ”).

Selepas menyusun model kami dan mencapai prestasi yang lebih pantas, kita harus menganalisis pelaksanaannya untuk mengenal pasti sebarang kekangan yang masih ada dan menentukan kawasan lain untuk penalaan lanjut.

Profil: Mendiagnosis Kesesakan dengan torch.profiler

Kesesakan prestasi masih boleh wujud walaupun selepas penyusunan - sama ada GPU yang kurang digunakan, kesesakan I/O atau kernel yang tidak dioptimumkan. PyTorch dilengkapi dengan Profiler terbina dalam untuk mengumpul jejak pelaksanaan dan mengenal pasti kesesakan ini. Dalam PyTorch 2.9, torch.profiler boleh mengesan aktiviti CPU dan GPU, merekodkan bentuk dan menyepadukan dengan Chrome Trace Viewer atau TensorBoard untuk visualisasi. torch.profiler juga boleh menangkap jejak secara tak segerak. Dengan cara ini, anda boleh memprofilkan bahagian latihan tanpa mengganggu aliran program.

import torch
import torch.nn as nn
import torch.optim as optim
import torch.profiler
# -----------------------------
# Device setup
# -----------------------------
device = "cuda" if torch.cuda.is_available() else "cpu"
# -----------------------------
# Model definition
# -----------------------------
model = nn.Sequential(
    nn.Linear(100, 256),
    nn.ReLU(),
    nn.Linear(256, 10)
).to(device)
# -----------------------------
# Optimizer and loss
# -----------------------------
optimizer = optim.Adam(model.parameters(), lr=1e-3)
loss_fn = nn.CrossEntropyLoss()
# -----------------------------
# Data iterator (synthetic)
# -----------------------------
def data_generator(batch_size=32):
    while True:
        inputs = torch.randn(batch_size, 100, device=device)
        targets = torch.randint(0, 10, (batch_size,), device=device)
        yield inputs, targets
data_iter = data_generator()
# -----------------------------
# Training step
# -----------------------------
def train_step(batch):
    model.train()
    inputs, targets = batch
    optimizer.zero_grad(set_to_none=True)
    outputs = model(inputs)
    loss = loss_fn(outputs, targets)
    loss.backward()
    optimizer.step()
    return loss
# -----------------------------
# Warm-up phase
# -----------------------------
for _ in range(5):
    train_step(next(data_iter))
# -----------------------------
# Profiling phase
# -----------------------------
with torch.profiler.profile(
    activities=[
        torch.profiler.ProfilerActivity.CPU,
        torch.profiler.ProfilerActivity.CUDA,
    ],
    record_shapes=True,
    profile_memory=True,
) as prof:
    for _ in range(3):
        train_step(next(data_iter))
        prof.step()
# -----------------------------
# Export trace
# -----------------------------
prof.export_chrome_trace("trace.json")
print("Profiler trace saved to trace.json")

Skrip contoh ini menunjukkan cara untuk mentakrifkan gelung latihan PyTorch yang minimum (tetapi menyeluruh) dan menggunakan torch.profiler untuk menangkap prestasi CPU dan GPU dalam keadaan stabil. Ia membina rangkaian neural mainan, mencipta data sintetik dengan peniterasi, dan membungkus satu iterasi latihan ke dalam fungsi train_step().

Kami menjalankan beberapa iterasi pemanasan terlebih dahulu (supaya kami tidak mengukur beban inisialisasi/cache sekali sahaja). Kemudian kami menyemak beberapa langkah latihan, merekod bentuk operator dan penggunaan memori. Kami juga mengekspor jejak Chrome (trace.json), yang boleh dimuatkan di pelayar Chrome untuk memeriksa penggunaan GPU, pelancaran kernel, pertindihan CPU–GPU, dan potensi kesesakan prestasi. Kami telah mengaktifkan record_shapes=True dan profile_memory=True untuk menjejak bentuk input dan peruntukan memori bagi membantu menyahpepijat ralat OOM atau ketidakcekapan.

Petua: Buka chrome://tracing dalam Google Chrome dan muatkan trace.json untuk mendapatkan paparan garis masa.

Untuk pengesanan leher botol automatik, anda boleh menggunakan torch.utils.bottleneck atau torch.profiler.schedule untuk menangkap snapshoot. Sebagai contoh, melakukan profiling beberapa langkah setiap epoch menggunakan penjadual:

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, TensorDataset
import torch.profiler
# -----------------------------
# Device
# -----------------------------
device = "cuda" if torch.cuda.is_available() else "cpu"
# -----------------------------
# Model
# -----------------------------
model = nn.Sequential(
    nn.Linear(100, 256),
    nn.ReLU(),
    nn.Linear(256, 10),
).to(device)
# -----------------------------
# Optimizer and loss
# -----------------------------
optimizer = optim.Adam(model.parameters(), lr=1e-3)
loss_fn = nn.CrossEntropyLoss()
# -----------------------------
# Dataset and DataLoader
# -----------------------------
num_samples = 10_000
batch_size = 32
inputs = torch.randn(num_samples, 100)
targets = torch.randint(0, 10, (num_samples,))
dataset = TensorDataset(inputs, targets)
data_loader = DataLoader(
    dataset,
    batch_size=batch_size,
    shuffle=True,
    num_workers=2,
)
# -----------------------------
# Training step
# -----------------------------
def train_step(batch):
    model.train()
    inputs, targets = batch
    inputs = inputs.to(device, non_blocking=True)
    targets = targets.to(device, non_blocking=True)
    optimizer.zero_grad(set_to_none=True)
    outputs = model(inputs)
    loss = loss_fn(outputs, targets)
    loss.backward()
    optimizer.step()
    return loss
# -----------------------------
# Warm-up (outside profiler)
# -----------------------------
for i, batch in enumerate(data_loader):
    train_step(batch)
    if i >= 5:
        break
# -----------------------------
# Profiling with schedule
# -----------------------------
with torch.profiler.profile(
    activities=[
        torch.profiler.ProfilerActivity.CPU,
        torch.profiler.ProfilerActivity.CUDA,
    ],
    schedule=torch.profiler.schedule(
        wait=1,
        warmup=1,
        active=3,
        repeat=2,
    ),
    on_trace_ready=torch.profiler.tensorboard_trace_handler("./prof_log"),
    record_shapes=True,
    profile_memory=True,
    with_stack=True,
) as prof:
    for step, batch in enumerate(data_loader):
        loss = train_step(batch)
        prof.step()

        if step >= 50:
            break

Ini akan memulakan dua tempoh jejak secara tak segerak (setiap 3 langkah panjang) selepas melangkau tempoh menunggu/pemanasan. Panggilan balik on_trace_ready kemudian mengeluarkan jejak yang boleh anda lihat dalam tab profil TensorBoard tanpa menghentikan latihan. Tangkap jejak secara tak segerak membolehkan anda membuat profil secara berselang-seli (contohnya, setiap 3 langkah dalam tetingkap 50 langkah) untuk mengurangkan beban sambil masih menangkap data prestasi.

Selepas pemprofilan, tafsirkan hasilnya untuk mengambil tindakan:

• Terikat CPU (banyak kelewatan apabila GPU sedang menunggu): Pertimbangkan sama ada anda boleh menyalurkan lebih banyak kerja ke GPU, atau persilangkan pemuatan data dengan pengiraan menggunakan pemuatan data tak segerak (contohnya, menghantar `workers` ke DataLoader atau pra-pemprosesan di GPU). Juga, periksa bahawa tiada gelung Python yang melakukan kerja semasa latihan; jika ini berlaku, cuba untuk menghapuskan atau `kompilkan` mereka.
• Teras GPU adalah leher botol: Anda boleh mengoptimumkan operasi tersebut (operasi gabungan, ketepatan lebih rendah, dll.). Panduan prestasi PyTorch mencadangkan untuk terus menggunakan operasi tahap tinggi (seperti panggilan torch.nn.functional) supaya perpustakaan dapat mengoptimumkannya. Walau bagaimanapun, jika anda perasan operasi tertentu yang jauh lebih perlahan, tentukan sama ada ia dijangka (memuat FFT yang besar boleh menjadi perlahan) atau jika terdapat cara alternatif untuk mengiranya.
• Pelbagai bentuk menyebabkan penjejakan semula: pertimbangkan untuk menyusun input mengikut saiz, atau gunakan dynamic=True seperti yang disebut di atas.
• Pastikan tiada kesesakan luaran: Sebagai contoh, jika menggunakan latihan teragih, pastikan komunikasi (all-reduce) bukan menjadi kesesakan - ini akan ditunjukkan oleh kernel NCCL atau CPU yang menunggu rangkaian.

Senarai Semak Profil

• Panaskan diri sebelum membuat profiling: Elakkan mengukur beban kompilasi atau inisialisasi malas.
• Tangkap kedua-dua CPU & GPU: Tangkap kedua-dua CPU & GPU (aktiviti=[CPU, CUDA]) untuk memahami interaksi mereka.
• Rekod bentuk & memori: Membantu mengenal pasti masalah berkaitan bentuk dan lonjakan penggunaan memori.
• Gunakan penjadualan untuk larian panjang: Gunakan profiler.schedule untuk menangkap tetingkap masa yang singkat secara berkala. Ini mengurangkan beban profiler dan masih memberikan keterlihatan terhadap tingkah laku program.
• Analisis penggunaan: Pastikan GPU anda digunakan sepenuhnya. Jika tidak, kenal pasti sama ada CPU atau I/O menjadi penghalang.
• Kenal pasti operasi utama: Profiler boleh menghasilkan senarai operasi yang disusun mengikut metrik penting seperti masa sendiri. Tentukan operasi yang mengambil masa paling lama dan optimakan ia (algoritma yang lebih baik, tingkatkan saiz kumpulan jika GPU tidak digunakan sepenuhnya).

Dengan aplikasi yang dioptimumkan pada satu mesin, cabaran seterusnya adalah mengembangkan skala untuk memanfaatkan berbilang GPU atau nod, yang memperkenalkan paradigma latihan teragih PyTorch.

Skala: Latihan Teragih melalui DDP atau FSDP

Apabila beban kerja atau model anda berkembang melebihi satu GPU, anda perlu meningkatkan latihan ke beberapa GPU. PyTorch menyediakan dua kaedah utama untuk melatih model merentas GPU: Distributed Data Parallel dan Fully Sharded Data Parallel. Walaupun kedua-dua kaedah menggunakan paralelisme data di bawahnya (setiap proses melatih pada subset data yang berbeza), mereka mempunyai pendekatan yang berbeza untuk menguruskan parameter model dalam memori. Kita akan meneroka bila untuk memilih setiap kaedah dan bagaimana untuk mengkonfigurasikannya untuk latihan satu-nod dan multi-nod dengan torchrun.

Data Sejajar Teragih

DDP mengekalkan salinan keseluruhan model pada setiap GPU dan menyelaraskan salinan tersebut dengan mengurangkan semua gradien selepas setiap langkah. Sesuai digunakan untuk model yang mudah muat dalam memori satu GPU, DDP adalah mudah dari segi konsep: Mulakan kumpulan proses, kemudian bungkus model dengan torch.nn.parallel.DistributedDataParallel.

import os
import torch
import torch.nn as nn
import torch.optim as optim
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.utils.data import DataLoader, TensorDataset
from torch.utils.data.distributed import DistributedSampler
# -----------------------------
# Setup distributed (safe)
# -----------------------------
def setup_distributed():
    world_size = int(os.environ.get("WORLD_SIZE", "1"))
    if world_size > 1:
        dist.init_process_group(backend="nccl")
        local_rank = int(os.environ["LOCAL_RANK"])
        torch.cuda.set_device(local_rank)
        device = torch.device(f"cuda:{local_rank}")
        distributed = True
    else:
        local_rank = 0
        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        distributed = False

    return distributed, local_rank, device

def cleanup_distributed(distributed):
    if distributed:
        dist.destroy_process_group()

# -----------------------------
# Main training
# -----------------------------
def main():
    distributed, local_rank, device = setup_distributed()

    # -----------------------------
    # Model
    # -----------------------------
    model = nn.Sequential(
        nn.Linear(100, 256),
        nn.ReLU(),
        nn.Linear(256, 10),
    ).to(device)
    if distributed:
        model = DDP(model, device_ids=[local_rank])
    # -----------------------------
    # Optimizer and loss
    # -----------------------------
    optimizer = optim.Adam(model.parameters(), lr=1e-3)
    loss_fn = nn.CrossEntropyLoss()
    # -----------------------------
    # Dataset
    # -----------------------------
    num_samples = 20_000
    batch_size = 32
    inputs = torch.randn(num_samples, 100)
    targets = torch.randint(0, 10, (num_samples,))
    dataset = TensorDataset(inputs, targets)
    if distributed:
        sampler = DistributedSampler(dataset, shuffle=True)
    else:
        sampler = None
    train_loader = DataLoader(
        dataset,
        batch_size=batch_size,
        sampler=sampler,
        shuffle=(sampler is None),
        num_workers=2,
        pin_memory=True,
    )

    # -----------------------------
    # Training loop
    # -----------------------------
    epochs = 3
    for epoch in range(epochs):
        if distributed:
            sampler.set_epoch(epoch)

        for step, (x, y) in enumerate(train_loader):
            x = x.to(device, non_blocking=True)
            y = y.to(device, non_blocking=True)

            optimizer.zero_grad(set_to_none=True)
            outputs = model(x)
            loss = loss_fn(outputs, y)
            loss.backward()
            optimizer.step()
            if step % 50 == 0 and local_rank == 0:
                print(
                    f"[Epoch {epoch}] Step {step} | "
                    f"Loss {loss.item():.4f}"
                )
    cleanup_distributed(distributed)
# -----------------------------
# Entry point
# -----------------------------
if __name__ == "__main__":
    main()

Melancarkan kerja DDP dilakukan dengan torchrun. Sebagai contoh, jika anda ingin melancarkan kerja pada satu nod dengan 4 GPU:

torchrun --nproc_per_node=4 train.py

Perintah ini akan menjana 4 proses (satu untuk setiap GPU) dan menetapkan pembolehubah persekitaran yang sesuai (pangkat, saiz dunia, dll). Menentukan pelbagai nod memerlukan anda untuk memberikan --nnodes dan butiran rangkaian:

# On node 0:
torchrun --nnodes=2 --nproc_per_node=4 --node_rank=0 --master_addr="<IP of node0>" --master_port=12345 train.py
# On node 1:
torchrun --nnodes=2 --nproc_per_node=4 --node_rank=1 --master_addr="<IP of node0>" --master_port=12345 train.py

Paralel Data Berpecah Sepenuhnya (FSDP)

FSDP melangkah lebih jauh dengan membahagikan parameter model dan keadaan pengoptimum di antara GPU, dan bukannya menyalin sepenuhnya pada setiap GPU. Menggunakan FSDP dalam PyTorch 2.9 adalah mudah dengan pembungkus FullyShardedDataParallel.

Perkara utama

• Bungkus model (atau submodul) dengan FSDP sebelum penciptaan pengoptimum.
• Gunakan torch.cuda.set_device(rank) dan pindahkan model ke GPU seperti dengan DDP.
• Opsional: sediakan polisi pembalutan automatik jika anda ingin membahagikan pada tahap modul (contohnya, membalut setiap blok transformer secara berasingan). Secara lalai, FSDP(model) akan membahagikan keseluruhan model sebagai satu unit shard (secara teknikal satu unit FSDP), yang menjadikannya kurang optimum untuk model yang sangat dalam.

Ini akan membahagikan secara rekursif sebarang submodul yang mempunyai lebih daripada 100k parameter ke dalam pecahan FSDP mereka sendiri. Jika anda mempunyai terlalu sedikit pecahan (contohnya, keseluruhan model adalah satu pecahan), anda tidak akan mendapat penjimatan memori antara lapisan. Namun, terlalu banyak pecahan meningkatkan beban komunikasi. Mari kita ringkaskan perbezaan antara DDP dan FSDP:

Sebagai contoh, anda mungkin mula dengan DDP. Jika anda menghadapi masalah kekurangan memori atau memerlukan model yang lebih besar, pertimbangkan untuk menggunakan FSDP. Perlu diambil perhatian bahawa FSDP dalam PyTorch 2.9 telah berkembang; secara lalai, ia "memecahkan semuanya" (ZeRO Tahap 3) dan menggunakan tetapan lalai lain, seperti limit_all_gathers=True, untuk mengelakkan lonjakan memori yang tidak dijangka.

Senarai Semak Skala

• Pastikan determinisme: Pastikan anda menetapkan benih rawak yang sama pada semua pangkat (torch.manual_seed(seed untuk offset rawak) jika anda memerlukan kebolehulangan.
• Gunakan DistributedSampler: Bahagikan data anda kepada beberapa peringkat dan panggil set_epoch() pada setiap epoch untuk menjadikannya berbeza setiap kali diacak.
• DDP: Gradien secara automatik dijumlahkan merentasi GPU selepas memanggil backward() pada model anda yang dibungkus dalam DDP. Jangan lupa untuk meningkatkan kadar pembelajaran atau saiz batch apabila menambah bilangan GPU!
• FSDP: Balut modul anda sebelum membuat pengoptimum. Cuba bermain dengan auto_wrap_policy untuk mengelakkan satu shard yang besar (kemungkinan jika anda mempunyai model yang sangat mendalam). Perhatikan GPU– pastikan saiz shard benar-benar muat.
• Backend komunikasi: Untuk pelbagai nod, pastikan anda mempunyai set rangkaian (biasanya NCCL). Tetapkan persekitaran NCCL_P2P_LEVEL=NVL jika anda mempunyai NVLink, dan lain-lain, untuk mengoptimumkan NCCL.
• Pengumpulan kecerunan: Jika digunakan, pastikan dengan FSDP, anda memanggil model.no_sync() jika melangkau penyegerakan kecerunan pada beberapa iterasi tertentu (sama seperti penggunaan DDP).

Checkpoint: Pulihkan Latihan dengan Boleh Dipercayai menggunakan Checkpoint Teragih

Dalam sesi latihan jangka panjang atau berskala besar, adalah penting untuk menyimpan checkpoint. Anda mungkin ingin menyambung semula latihan selepas berlaku kerosakan, menyelaras dari tengah-tengah latihan, atau memeriksa keputusan sementara. PyTorch 2.9 termasuk primitif yang mudah untuk checkpoint pengedaran (DCP). Ini membolehkan pengguna menyimpan/memuat dengan lebih cekap dan mantap berbanding pendekatan tradisional dengan torch.save. Mari kita bandingkan kedua-dua pendekatan sebelum menyelami pendekatan yang disyorkan.

Obor tradisional.save/torch.load

Kaedah tradisional dalam skrip latihan GPU tunggal atau bahkan DDP akan menyimpan state_dict() model ke dalam fail (biasanya hanya pada rank 0, kerana semua proses akan mempunyai akses kepada fail ini). Contohnya:

# On rank 0 only:
torch.save(model.state_dict(), "checkpoint.pt")
and to load:
model.load_state_dict(torch.load("checkpoint.pt"))

Walaupun ini sesuai untuk model/config kecil di mana satu proses boleh menyimpan keseluruhan keadaan model dalam memori (biasanya rank 0), terdapat masalah dengan menggunakan pendekatan ini untuk latihan teragih di mana model mungkin dibahagikan merentasi proses.

Checkpoint Teragih (DCP): modul torch.distributed.checkpoint dalam PyTorch (ditambah sejak v2.1+, matang pada 2.9) menangani isu-isu ini. Ia memparalelkan penyimpanan: setiap kedudukan menulis sebahagian daripada keadaan modelnya, menghasilkan beberapa fail (sekurang-kurangnya satu bagi setiap kedudukan) yang membentuk checkpoint. Ia juga melaksanakan penstrukturan semula semasa pemuatan: anda boleh menyimpan pada N kedudukan dan memuat pada M kedudukan, dan DCP akan secara automatik mengumpul/mengagihkan semula mengikut kesesuaian.

Corak penggunaan biasa melibatkan membungkus model anda (dan pengoptimum, jika perlu) dalam bekas berkeadaan yang menyediakan state_dict. Anda boleh melakukan ini menggunakan fungsi utiliti yang disediakan get_state_dict dan set_state_dict, yang menyedari FSDP:

import torch.distributed.checkpoint as dcp
from torch.distributed.checkpoint.state_dict import get_state_dict, set_state_dict
from torch.distributed.checkpoint.stateful import Stateful

# Define a stateful container for model & optimizer
class AppState(Stateful):
    def __init__(self, model, optimizer=None):
        self.model = model
        self.optimizer = optimizer
    def state_dict(self):
        model_sd, opt_sd = get_state_dict(self.model, self.optimizer)
        return {"model": model_sd, "optim": opt_sd}
    def load_state_dict(self, state):
        set_state_dict(self.model, self.optimizer, state["model"], state.get("optim"))

AppState di sini melaksanakan antara muka Stateful, jadi DCP tahu cara mendapatkan dan menetapkan keadaan. Perlu diingat bahawa get_state_dict secara automatik menggunakan kamus keadaan berpecah FSDP jika model dipecahkan (setiap pangkat hanya mendapat bahagian model mereka).

Penjimatan dengan DCP

Semua pangkat laksanakan:

app_state = AppState(model, optimizer)
dcp.save(app_state, checkpoint_id="mycheckpoint")

Ini akan menghasilkan direktori yang mengandungi fail yang diawali dengan rentetan "mycheckpoint" (mycheckpoint/mpirank_0_0.pdparams). Setiap pangkat menulis shardnya secara selari. Ini boleh mengurangkan masa dengan ketara berbanding jika hanya satu pangkat yang melakukannya.

Memuat dengan DCP

Untuk menyambung semula dari titik semak, anda memerlukan model dengan topologi yang sama (mungkin pada bilangan pangkat yang berbeza). Bina model anda (dan pengoptimum), bungkus dalam AppState, dan kemudian:

app_state = AppState(model, optimizer)
dcp.load(app_state, checkpoint_id="mycheckpoint");

Dengan syarat bilangan pangkat berbeza daripada semasa anda menyimpan, DCP akan menguruskan rombakan semula serpihan secara dalaman (contohnya, beralih daripada 8 kepada 4 GPU bermakna setiap pangkat akan memuatkan dua serpihan yang disimpan, atau jika meningkatkan pangkat, ia akan diagihkan dengan sewajarnya). Langkah perhimpunan semula masa beban ini menghapuskan keperluan untuk sebarang penukaran pusat pemeriksaan manual.

Pemeriksaan titik henti secara tak segerak

Salah satu ciri paling menarik dalam PyTorch 2.9 ialah penyimpanan pemeriksaan asin, yang boleh menumpang penulisan pemeriksaan dengan latihan. Menggunakan dcp.async_save, API memberikan masa depan dan menjalankan I/O berat dalam benang latar belakang. Idea asasnya:

save_future = dcp.async_save(app_state, checkpoint_id="chkpt_epoch10")
# ... training can continue immediately ...
save_future.wait()  # later, wait for completion (or periodically check)

async_save mula-mula peringkat data secara tempatan pada CPU (menyalin keadaan model/pengoptimum ke penimbal memori yang disematkan), kemudian menulis penimbal tersebut ke cakera/rangkaian secara tak segerak. GPU boleh menyambung semula kerja pengiraan (iaitu, latihan) semasa pusat pemeriksaan sedang disiram.

Penyimpanan tak segerak sementara memerlukan dua kali ganda keperluan memori untuk menyimpan keadaan model anda. Setiap pangkat memperuntukkan satu penampan di CPU yang saiznya hampir sama dengan bahagian sharnya dari titik semak. Jadi menyimpan model yang sangat besar secara tak segerak merentasi banyak pangkat boleh meningkatkan penggunaan memori dengan ketara (RAM CPU + penggunaan memori terpin ~ saiz_titik_semak_setiap_pangkat * bilangan_pangkat, dalam kes terburuk).

Sekarang, mari kita buat perbandingan kaedah penanda aras dengan panduan tentang bila untuk menggunakan setiap satu:

Senarai Semak Titik Periksa

• Sertakan semua keadaan yang diperlukan: Pastikan menjimatkan berat model dan keadaan pengoptimum (dan penjadual, keadaan RNG, dsb, jika perlu) untuk menyambung semula latihan sepenuhnya kemudian. Jika menggunakan kaedah Stateful DCP, anda boleh memasukkan pengoptimum dengan mudah melalui get_state_dict(..., model, pengoptimum).
• Ujian pemulihan: Sebaik sahaja anda menyimpan titik semak, muatkan semula (sebaiknya pada jumlah GPU yang sama dan berbeza jika diedarkan) dalam persekitaran yang bersih untuk memastikan semuanya dipulihkan seperti yang dijangkakan. Ini akan membantu anda menemui kesilapan, seperti terlupa menyimpan sebahagian daripada keadaan latihan pada awalnya.
• Uruskan storan: Titik semakan boleh menjadi sangat besar dengan cepat (terutamanya dengan titik semakan teragih yang dibahagikan). Laksanakan beberapa strategi pembersihan (simpan k titik semakan terakhir, dsb.) untuk mengelakkan kehabisan cakera.
• Penggunaan Async: async_save adalah ciri yang berguna, tetapi anda mesti berhati-hati untuk tidak mempunyai beberapa simpanan async yang sedang berjalan pada masa yang sama (melainkan anda sangat berhati-hati dengan cara ia menggunakan memori dan membenarkan pertindihan). Akhirnya, anda mesti memanggil .wait() pada future yang terhasil untuk menaikkan sebarang pengecualian atau mengesahkan kejayaan.
• Konsistensi: Apabila menyambung semula dari titik semak, pastikan semua pangkat memanggil dcp.load. yang sama. Ini sangat penting untuk FSDP, kerana ia perlu menyelaras pemuatan dalaman. Selain itu, pastikan model berada pada peranti yang dijangka sebelum memuat jika menggunakan pemindahan beban CPU atau mekanisme lain.

Kesimpulan

Dengan mengikuti aliran kerja ini, anda akan membangunkan kod latihan PyTorch yang berprestasi baik dan tahan terhadap kegagalan:

• Mulakan dari asas yang berfungsi.
• Kompil untuk kelajuan maksimum.
• Profil untuk mengenal pasti peluang untuk pengoptimuman.
• Kembangkan dengan strategi selari data atau model yang sesuai.
• Titik semak berkesan untuk membolehkan pemulihan daripada kegagalan.

Terdapat buku panduan yang tersedia untuk setiap langkah. Gunakannya secara berulang (contohnya, profil semula selepas langkah 4 mengubah ciri prestasi anda, atau semak semula tetapan kompilasi dalam langkah 2 untuk menyesuaikan dengan persekitaran baru anda). Oleh kerana PyTorch telah berkembang pesat, terdapat keupayaan untuk membantu di setiap langkah. Pelajari bila dan bagaimana menggunakan setiap satu, dan anda akan bersedia untuk melatih model besar dengan pantas dan berskala, dengan strategi pemulihan yang baik.

Soalan Lazim

Mengapa menetapkan garis asas begitu penting sebelum pengoptimuman?

Anda perlu mempunyai garis dasar GPU tunggal yang betul sebelum mengoptimumkan supaya anda mempunyai sesuatu yang anda percayai untuk ketepatan dan prestasi. Jika tidak, pengoptimuman mungkin menyembunyikan pepijat yang menyebabkan hasil yang salah. Hanya dengan garis dasar yang betul anda boleh mengaitkan pecutan atau regresi kepada perubahan kod tertentu dan bukannya pepijat samar atau bunyi pengukuran.

Bila saya patut menggunakan torch.compile, dan apa yang perlu saya beri perhatian?

torch.compile perlu digunakan selepas anda mempunyai asas yang kukuh untuk mempercepatkan pelaksanaan pada keadaan mantap. Pantau pemutusan graf, panaskan sebelum pengukuran, dan pastikan bentuk dinamik diurus dengan sengaja untuk mengelakkan penyusunan semula yang berlebihan.

Bagaimana torch.profiler membantu selain mengesahkan andaian?

torch.profiler mendedahkan kesesakan sebenar anda - sama ada itu adalah kelewatan CPU, kernel GPU yang tidak efisien, penjejakan bentuk semula, atau beban komunikasi -. Ini membolehkan anda menumpukan usaha pengoptimuman di tempat yang penting berdasarkan bukti dan bukan intuisi.

Bagaimana saya memilih antara DDP dan FSDP untuk latihan diedarkan?

Jika model anda muat dengan selesa pada setiap GPU yang anda ingin gunakan, gunakan DDP kerana kesederhanaannya dan kelajuannya yang relatif. Jika anda terhad oleh memori GPU atau bekerja dengan model yang sangat besar, FSDP akan memberikan kebolehan skala yang lebih tinggi dengan kos kerumitan tambahan.

Mengapa Penanda Henti Tersebar lebih digemari berbanding torch.save dalam jalankan berskala besar?

Checkpointing Teragih membolehkan menyimpan secara selari merentasi ranking, menyokong resharing semasa pemuatan, dan membolehkan checkpointing tak segerak, menjadikan pemulihan boleh dipercayai dan cekap untuk tugasan latihan besar, multi-GPU atau multi-node.

Rujukan

• Panduan utama untuk pustaka PyTorch dalam Python
• Melaras Latihan Model Merentasi Berbilang GPU: Strategi Berkesan dengan PyTorch DDP dan FSDP
• DataSelariBerpecahPenuh
• Penanda Aras Teragih Torch
• Tutorial torch.compile

Terima kasih kerana belajar bersama Komuniti Linux-Console.net.