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 pada cluster instance virtual machine (VM), yang juga dikenal sebagai node, dengan variabel lingkungan yang mendukung pelatihan PyTorch terdistribusi. Hal ini dapat membantu skala tugas pelatihan Anda untuk menangani data dalam jumlah besar.

Panduan ini mengasumsikan bahwa Anda menggunakan container PyTorch yang telah dibangun sebelumnya untuk pelatihan, seperti yang dijelaskan dalam Mulai menggunakan PyTorch. Mengadaptasi kode PyTorch untuk pelatihan terdistribusi memerlukan sedikit perubahan.

Menentukan struktur cluster pelatihan

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

  • Pekerja master: VM dengan peringkat 0. Node ini menyiapkan koneksi antara node dalam cluster.
  • Pekerja: 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 Anda membuat tugas pelatihan, tentukan image container Docker yang akan digunakan pekerja master di kolom trainingInput.masterConfig.imageUri, lalu tentukan image container Docker yang akan digunakan oleh setiap pekerja di kolom trainingInput.workerConfig.imageUri. Lihat daftar container PyTorch yang telah dibangun 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. Sebaiknya gunakan container PyTorch yang telah dibangun sebelumnya yang sama pada setiap node.

Memperbarui kode pelatihan Anda

Dalam aplikasi pelatihan Anda, tambahkan kode berikut untuk menginisialisasi cluster pelatihan:

import torch

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

Ganti BACKEND dengan salah satu backend pelatihan terdistribusi yang didukung dan dijelaskan di bagian berikut. Argumen kata kunci init_method='env://' memberi tahu PyTorch agar menggunakan variabel lingkungan untuk menginisialisasi komunikasi di 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 Anda, 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 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 pada setiap {i>node<i}.
  • RANK: ID unik untuk setiap node. Pada pekerja master, nilai ini disetel ke 0. Pada setiap pekerja, nilai ini ditetapkan ke nilai yang berbeda dari 1 hingga WORLD_SIZE - 1.
  • MASTER_ADDR: Nama host node pekerja master. Variabel ini memiliki nilai yang sama pada setiap {i>node<i}.
  • MASTER_PORT: Port tempat node pekerja master berkomunikasi. Variabel ini memiliki nilai yang sama di setiap node.

PyTorch menggunakan variabel lingkungan ini untuk menginisialisasi cluster.

Langkah selanjutnya