Gunakan SMDDP pustaka dalam skrip TensorFlow latihan Anda (tidak digunakan lagi) - Amazon SageMaker

Terjemahan disediakan oleh mesin penerjemah. Jika konten terjemahan yang diberikan bertentangan dengan versi bahasa Inggris aslinya, utamakan versi bahasa Inggris.

Gunakan SMDDP pustaka dalam skrip TensorFlow latihan Anda (tidak digunakan lagi)

penting

SMDDPPustaka menghentikan dukungan untuk TensorFlow dan tidak lagi tersedia di TensorFlow lebih dari DLCs v2.11.0. Untuk menemukan sebelumnya TensorFlow DLCs dengan SMDDP perpustakaan yang diinstal, lihatKerangka kerja yang didukung.

Langkah-langkah berikut menunjukkan cara memodifikasi skrip TensorFlow pelatihan untuk memanfaatkan SageMaker perpustakaan paralel data terdistribusi. 

Perpustakaan APIs dirancang agar mirip dengan HorovodAPIs. Untuk detail tambahan tentang setiap perpustakaan API yang ditawarkan TensorFlow, lihat TensorFlow APIdokumentasi paralel data SageMaker terdistribusi.

catatan

SageMaker distributed data parallel dapat disesuaikan dengan skrip TensorFlow pelatihan yang terdiri dari modul tf inti kecuali tf.keras modul. SageMaker distributed data parallel tidak mendukung TensorFlow implementasi Keras.

catatan

Pustaka paralelisme data SageMaker terdistribusi mendukung Automatic Mixed Precision (AMP) di luar kotak. Tidak ada tindakan tambahan yang diperlukan untuk mengaktifkan AMP selain modifikasi tingkat kerangka kerja pada skrip pelatihan Anda. Jika gradien masukFP16, pustaka paralelisme SageMaker data menjalankan operasinya AllReduce di. FP16 Untuk informasi selengkapnya tentang penerapan AMP APIs skrip pelatihan Anda, lihat sumber daya berikut:

  1. Impor TensorFlow klien perpustakaan dan inisialisasi.

    import smdistributed.dataparallel.tensorflow as sdp 
sdp.init()

  2. Pin masing-masing GPU ke satu smdistributed.dataparallel proses dengan local_rank —ini mengacu pada peringkat relatif dari proses dalam node tertentu. sdp.tensorflow.local_rank()APIIni memberi Anda peringkat lokal perangkat. Node pemimpin adalah peringkat 0, dan node pekerja adalah peringkat 1, 2, 3, dan seterusnya. Ini dipanggil dalam blok kode berikut sebagaisdp.local_rank(). set_memory_growthtidak terkait langsung dengan SageMaker didistribusikan, tetapi harus ditetapkan untuk pelatihan terdistribusi dengan TensorFlow. 

    gpus = tf.config.experimental.list_physical_devices('GPU')
for gpu in gpus:
    tf.config.experimental.set_memory_growth(gpu, True)
if gpus:
    tf.config.experimental.set_visible_devices(gpus[sdp.local_rank()], 'GPU')

  3. Skala tingkat pembelajaran dengan jumlah pekerja. sdp.tensorflow.size()APIIni memberi Anda jumlah pekerja di cluster. Ini dipanggil dalam blok kode berikut sebagaisdp.size(). 

    learning_rate = learning_rate * sdp.size()

  4. Gunakan perpustakaan DistributedGradientTape untuk mengoptimalkan AllReduce operasi selama pelatihan. Ini membungkustf.GradientTape.  

    with tf.GradientTape() as tape:
      output = model(input)
      loss_value = loss(label, output)
    
# SageMaker data parallel: Wrap tf.GradientTape with the library's DistributedGradientTape
tape = sdp.DistributedGradientTape(tape)

  5. Siarkan variabel model awal dari node pemimpin (peringkat 0) ke semua node pekerja (peringkat 1 hingga n). Ini diperlukan untuk memastikan inisialisasi yang konsisten di semua peringkat pekerja. Gunakan sdp.tensorflow.broadcast_variables API setelah model dan variabel pengoptimal diinisialisasi. Ini dipanggil dalam blok kode berikut sebagaisdp.broadcast_variables(). 

    sdp.broadcast_variables(model.variables, root_rank=0)
sdp.broadcast_variables(opt.variables(), root_rank=0)

  6. Terakhir, modifikasi skrip Anda untuk menyimpan pos pemeriksaan hanya pada node pemimpin. Node pemimpin memiliki model yang disinkronkan. Ini juga menghindari node pekerja yang menimpa pos pemeriksaan dan mungkin merusak pos pemeriksaan. 

    if sdp.rank() == 0:
    checkpoint.save(checkpoint_dir)

Berikut ini adalah contoh skrip TensorFlow pelatihan untuk pelatihan terdistribusi dengan perpustakaan.

import tensorflow as tf

# SageMaker data parallel: Import the library TF API
import smdistributed.dataparallel.tensorflow as sdp

# SageMaker data parallel: Initialize the library
sdp.init()

gpus = tf.config.experimental.list_physical_devices('GPU')
for gpu in gpus:
    tf.config.experimental.set_memory_growth(gpu, True)
if gpus:
    # SageMaker data parallel: Pin GPUs to a single library process
    tf.config.experimental.set_visible_devices(gpus[sdp.local_rank()], 'GPU')

# Prepare Dataset
dataset = tf.data.Dataset.from_tensor_slices(...)

# Define Model
mnist_model = tf.keras.Sequential(...)
loss = tf.losses.SparseCategoricalCrossentropy()

# SageMaker data parallel: Scale Learning Rate
# LR for 8 node run : 0.000125
# LR for single node run : 0.001
opt = tf.optimizers.Adam(0.000125 * sdp.size())

@tf.function
def training_step(images, labels, first_batch):
    with tf.GradientTape() as tape:
        probs = mnist_model(images, training=True)
        loss_value = loss(labels, probs)

    # SageMaker data parallel: Wrap tf.GradientTape with the library's DistributedGradientTape
    tape = sdp.DistributedGradientTape(tape)

    grads = tape.gradient(loss_value, mnist_model.trainable_variables)
    opt.apply_gradients(zip(grads, mnist_model.trainable_variables))

    if first_batch:
       # SageMaker data parallel: Broadcast model and optimizer variables
       sdp.broadcast_variables(mnist_model.variables, root_rank=0)
       sdp.broadcast_variables(opt.variables(), root_rank=0)

    return loss_value

...

# SageMaker data parallel: Save checkpoints only from master node.
if sdp.rank() == 0:
    checkpoint.save(checkpoint_dir)

Setelah Anda selesai mengadaptasi skrip pelatihan Anda, lanjutkan keMeluncurkan pekerjaan pelatihan terdistribusi dengan SMDDP menggunakan SageMaker Python SDK.