Mengonfigurasi pelatihan terdistribusi untuk PyTorch

Dokumen ini menjelaskan cara membuat tugas pelatihan PyTorch terdistribusi. Saat Anda membuat tugas pelatihan terdistribusi, AI Platform Training akan menjalankan kode Anda di cluster instance virtual machine (VM), yang juga dikenal sebagai node, dengan variabel lingkungan yang mendukung pelatihan PyTorch terdistribusi. Hal ini dapat membantu menskalakan tugas pelatihan Anda untuk menangani data dalam jumlah besar.

Panduan ini mengasumsikan bahwa Anda menggunakan penampung PyTorch bawaan untuk pelatihan, seperti yang dijelaskan dalam Memulai PyTorch. Menyesuaikan kode PyTorch untuk pelatihan terdistribusi memerlukan perubahan minimal.

Menentukan struktur cluster pelatihan

Untuk pelatihan PyTorch terdistribusi, konfigurasikan tugas Anda untuk menggunakan satu node master worker dan satu atau beberapa node worker. Peran ini memiliki perilaku berikut:

  • Pekerja master: VM dengan peringkat 0. Node ini menyiapkan koneksi antar-node dalam cluster.
  • Worker: Node yang tersisa di cluster. Setiap node melakukan sebagian pelatihan, seperti yang ditentukan oleh kode aplikasi pelatihan Anda.

Untuk mempelajari cara menentukan node pekerja master dan node pekerja untuk cluster pelatihan, baca Menentukan jenis mesin atau tingkat skala.

Menentukan image container

Saat membuat tugas pelatihan, tentukan image container Docker untuk pekerja master yang akan digunakan di kolom trainingInput.masterConfig.imageUri, dan tentukan image container Docker untuk setiap pekerja yang akan digunakan di kolom trainingInput.workerConfig.imageUri. Lihat daftar penampung PyTorch yang telah dibuat sebelumnya.

Jika menggunakan perintah gcloud ai-platform jobs submit training untuk membuat tugas pelatihan, Anda dapat menentukan kolom ini dengan flag --master-image-uri dan --worker-image-uri.

Namun, jika Anda tidak menentukan kolom trainingInput.workerConfig.imageUri, nilainya akan ditetapkan secara default ke nilai trainingInput.masterConfig.imageUri. Sering kali lebih baik menggunakan container PyTorch bawaan yang sama di setiap node.

Memperbarui kode pelatihan

Dalam aplikasi pelatihan, tambahkan kode berikut untuk melakukan inisialisasi cluster pelatihan:

import torch

torch.distributed.init_process_group(
    backend='BACKEND',
    init_method='env://'
)

Ganti BACKEND dengan salah satu backend pelatihan terdistribusi yang didukung yang dijelaskan di bagian berikut. Argumen kata kunci init_method='env://' memberi tahu PyTorch untuk menggunakan variabel lingkungan guna melakukan inisialisasi komunikasi dalam cluster. Pelajari lebih lanjut di Bagian variabel lingkungan dalam panduan ini.

Selain itu, perbarui kode pelatihan Anda untuk menggunakan class torch.nn.parallel.DistributedDataParallel. Misalnya, jika Anda telah membuat modul PyTorch yang disebut model dalam kode, tambahkan baris berikut:

model = torch.nn.parallel.DistributedDataParallel(model)

Untuk mempelajari lebih lanjut cara mengonfigurasi pelatihan terdistribusi, baca panduan dokumentasi PyTorch untuk pelatihan terdistribusi.

Backend pelatihan terdistribusi

AI Platform Training mendukung backend berikut untuk pelatihan PyTorch terdistribusi:

  • gloo: direkomendasikan untuk tugas pelatihan CPU
  • nccl: direkomendasikan untuk tugas pelatihan GPU

Baca tentang perbedaan antara backend.

Variabel lingkungan

Saat Anda membuat tugas pelatihan PyTorch terdistribusi, AI Platform Training akan menetapkan variabel lingkungan berikut di setiap node:

  • WORLD_SIZE: Jumlah total node dalam cluster. Variabel ini memiliki nilai yang sama di setiap node.
  • RANK: ID unik untuk setiap node. Pada pekerja master, nilai ini ditetapkan ke 0. Pada setiap pekerja, nilai ini ditetapkan ke nilai yang berbeda dari 1 ke WORLD_SIZE - 1.
  • MASTER_ADDR: Nama host node pekerja master. Variabel ini memiliki nilai yang sama di setiap node.
  • MASTER_PORT: Port yang digunakan node pekerja master untuk berkomunikasi. Variabel ini memiliki nilai yang sama di setiap node.

PyTorch menggunakan variabel lingkungan ini untuk melakukan inisialisasi cluster.

Langkah selanjutnya