Az adatelosztás optimalizálása a SageMaker elosztott adatokkal párhuzamosan

Ez egy háromrészes bejegyzés második része az elosztott képzés optimalizálása témájában. Az „első részben” rövid áttekintést adtunk az elosztott képzési algoritmusokról. Megjegyeztük, hogy minden algoritmusra jellemző, hogy több GPU között nagy sebességű kommunikációra támaszkodnak. Feltételeztük, hogy egy elosztott algoritmus, amely figyelembe veszi az alapul szolgáló példány topológiáját, különösen a GPU-párok közötti kommunikációs kapcsolatok különbségeit, jobban teljesít, mint az, amelyik nem.

A második részben bemutatjuk, hogyan használhatjuk az Amazon SageMaker „distributed data parallel” (SDP) könyvtárát az adatok olyan elosztására, amely megkülönbözteti a csomóponton belüli és a csomópontok közötti GPU-GPU kommunikációt.

Adatelosztás SageMakerrel párhuzamosan elosztott adatok

Az adatmegosztott oktatás során minden GPU fenntartja a modell saját példányát, és a másolatok közötti igazodás a gradiensmegosztás révén megmarad. Egy jó adatelosztási algoritmus úgy valósítja meg a gradiens megosztási mechanizmust, hogy korlátozza a betanítási teljesítményre gyakorolt ​​hatást. Egyes gradiensmegosztó algoritmusok egy vagy több központi paraméterszerverre támaszkodnak, amelyek összegyűjtik a gradiens frissítéseket az összes dolgozótól, majd az eredményeket visszaküldik a dolgozóknak. Mások a GPU-k közötti közvetlen peer-to-peer kommunikációra támaszkodnak. A színátmenetek megosztásának egyik népszerű algoritmusa a Ring-AllReduce, amelyben több üzenet halad át a dolgozók között egy egyirányú gyűrűn.Lásd itt, hogy nagyszerű vizuális áttekintést kapjon arról, hogyan Ring-AllReduce működik. A Ring-AllReducet a Horovod használja, amely egy népszerű keretrendszer az adatokon elosztott képzéshez.

Az Amazon SageMaker "elosztott adatok párhuzamos" (SDP) könyvtárának célja, hogy egyszerűsítse és felgyorsítsa az adatok elosztott képzését. Az SDP-vel kapcsolatos további részleteket a „bejelentésben”, az „Amazon SageMaker fejlesztői útmutatóban” és a kapcsolódó „fehér könyvben” találja. (Aki kíváncsi arra, hogy az általam választott borítókép mennyire releváns a bejegyzés tartalmához, csak a fehér könyv címét kell keresnie: Hering: A paraméterszerver újragondolása a felhő számára .) Az SDP gradiens megosztási algoritmusa számos ravasz technikára támaszkodik. Az ebben a bejegyzésben folytatott megbeszélésünk szempontjából leginkább releváns tulajdonság a csomóponton belüli GPU-GPU kommunikáció és a csomópontok közötti GPU-GPU kommunikáció közötti különbség. Ezt a különbséget az alábbi kép foglalja össze, amely egy kétszintű folyamatot mutat be, amelyben a csomóponton belüli GPU-k megosztják a gradienseket az NVLink segítségével, míg a különböző csomópontok GPU-i közötti kommunikációt az összes használt csomópont CPU-ján lévő szerverek közvetítik.

A „fehér könyv” bemutatja, hogy ez a fajta megközelítés, vagyis egy elosztott algoritmus, amely az alapul szolgáló képzési környezet topológiájához van szabva, hogyan képes felgyorsítani a nagy léptékű elosztott képzési feladatokat a szabványos Ring-AllReducehez képest. em> algoritmus.

Meg kell jegyeznünk, hogy az SDP által használt hierarchikus gradiens-eloszlási algoritmuson kívül számos további algoritmus és könyvtár kínál megoldásokat, amelyek figyelembe veszik az alapul szolgáló példánytopológiát. Például amellett, hogy népszerűsítette a Ring-AllReduce használatát, a Horovod támogatja a hierarchikus gradiensmegosztási algoritmust is. A projekt és a környezet részletei alapján vezérlőelemeket is elérhetővé tesz a gradiens áramlás "hangolásához". Ezenkívül a Horovod által használt „NCCL” mögöttes műveletek „fejlett hierarchikus technikákat” is tartalmaznak.

Az SDP egyik előnye a többi könyvtárhoz képest a példánytopológia felfedezése. A képzési környezet részletei, mint például a hálózati sávszélességek és a késleltetés fontos bemeneti adatok a legjobb gradiens megosztási algoritmus kiválasztásához. Mivel az SDP be van építve a SageMaker keretrendszerbe, azonnali és részletes ismeretekkel rendelkezik a példány topológiájáról. Más könyvtárak nem rendelkeznek ugyanolyan részletességgel, és arra kényszerülhetnek, hogy kitalálják a legjobb színátmenet-megosztási stratégiát, vagy megpróbálják felfedezni a hiányzó információkat, és ennek megfelelően beállítani.

Példa

Itt egy példát mutatunk be az SDP integrálására egy egyszerű TensorFlow (2.9) szkriptbe. A SageMaker dokumentációja számos „TensorFlow-példát” tartalmaz, amelyek bemutatják az SMP API-k meghívását. Példánkban két olyan hasznos technikát mutatunk be, amelyek jelenleg hiányoznak a SageMaker példákból:

  1. Hogyan programozza be szkriptjeit, hogy könnyen válthasson az SDP és a népszerű Horovod könyvtár között.
  2. Hogyan lehet kombinálni az SDP-t a TensorFlow magas szintű API-jával a modellképzéshez — tf.keras.Model.fit(). Míg a magas szintű API elrejti a TensorFlow GradientTape-hez való hozzáférést, az SDP megköveteli, hogy azt a tensorflow.DistributedGradientTape API-val csomagolják. Ezt a konfliktust a model.fit() hívás "képzési lépésének testreszabásával" oldjuk meg.

A példa két szkriptből áll. Az elosztott képzési munka indító szkriptje és a képzési szkript.

Az első szkript a SageMaker képzési munkamenet indító szkriptje. Ebben a példában úgy döntöttünk, hogy a beviteli módot "Fast File Mode"-ra állítjuk, amely egy Amazon SageMaker szolgáltatás, amely lehetővé teszi a bemeneti adatok közvetlen streamelését az "Amazon S3"-ból a képzési példányokba. Az adatok TFRecord fájlokban tárolódnak. Bár az elosztott képzések általában különösen nagy adatkészleteken futnak, ebben a példában a fájlokat a CIFAR-10 adatkészletből hoztam létre (ezt a szkriptet használva). A munkamenet négy p4d.24xlarge betanítópéldánnyal van példányosítva, és a elosztás beállítás a SageMaker adatelosztási könyvtár használatára van konfigurálva.

from sagemaker.tensorflow import TensorFlow
from sagemaker.session import TrainingInput
s3_input = TrainingInput(
          's3://'+S3_BUCKET_DATASET+'/cifar10-tfrecord/',
          input_mode='FastFile')
# Training using SMDataParallel Distributed Training Framework
distribution = {'smdistributed':
                    {'dataparallel':{'enabled': True}}
               }
tensorflow = TensorFlow(entry_point='train_tf.py',
                        role=<role>,
                        instance_type='ml.p4d.24xlarge',
                        instance_count=4,
                        framework_version='2.9.1',
                        py_version='py39',
                        distribution=distribution)
tensorflow.fit(s3_input, job_name='data-parallel-example')

Vegye figyelembe, hogy ugyanaz a szkript módosítható a Horovod-munkamenet elindításához, ha csak a terjesztési beállítást módosítja:

distribution = {
  'mpi': {
    'enabled': True,
    'processes_per_host': 8
  }
}

A második szkript tartalmazza a képzési ciklust. A szkriptet úgy alakítottuk ki, hogy bemutassuk, milyen egyszerű az átalakítás a Horovod adatelosztási keretrendszer és az Amazon SageMaker párhuzamos adattár között. A szkript egy run_hvd kapcsolóval kezdődik, amellyel váltani lehet a két lehetőség között. A következő if-elseblokk tartalmazza az egyetlen könyvtárspecifikus kódot. Amint azt fentebb tárgyaltuk, megvalósítottunk egy egyéni képzési lépést, amely a DistributedGradientTapeAPI-t használja.

import tensorflow as tf
# toggle flag to run Horovod
run_hvd = False
if run_hvd:
  import horovod.tensorflow as dist
  from horovod.tensorflow.keras.callbacks import \
                            BroadcastGlobalVariablesCallback
else:
  import smdistributed.dataparallel.tensorflow as dist
  from tensorflow.keras.callbacks import Callback
  class BroadcastGlobalVariablesCallback(Callback):
    def __init__(self, root_rank, *args):
      super(BroadcastGlobalVariablesCallback, self).
                                            __init__(*args)
      self.root_rank = root_rank
      self.broadcast_done = False
    def on_batch_end(self, batch, logs=None):
      if self.broadcast_done:
        return
      dist.broadcast_variables(self.model.variables, 
                               root_rank=self.root_rank)
      dist.broadcast_variables(self.model.optimizer.variables(),          
                               root_rank=self.root_rank)
      self.broadcast_done = True
# Create Custom Model that performs the train step using 
# DistributedGradientTape
from keras.engine import data_adapter
class CustomModel(tf.keras.Model):
  def train_step(self, data):
    x, y, w = data_adapter.unpack_x_y_sample_weight(data)
    with tf.GradientTape() as tape:
      y_pred = self(x, training=True)
      loss = self.compute_loss(x, y, y_pred, w)
    tape = dist.DistributedGradientTape(tape)
    self._validate_target_and_loss(y, loss)
    self.optimizer.minimize(loss, 
                            self.trainable_variables,
                            tape=tape)
    return self.compute_metrics(x, y, y_pred, w)
def get_dataset(batch_size, rank):
  def parse_sample(example_proto):
    image_feature_description = {
        'image': tf.io.FixedLenFeature([], tf.string),
        'label': tf.io.FixedLenFeature([], tf.int64)
    }
    features = tf.io.parse_single_example(example_proto, 
                                          image_feature_description)
    image = tf.io.decode_raw(features['image'], tf.uint8)
    image.set_shape([3 * 32 * 32])
    image = tf.reshape(image, [32, 32, 3])
    image = tf.cast(image, tf.float32)/255.
    label = tf.cast(features['label'], tf.int32)
    return image, label
  aut = tf.data.experimental.AUTOTUNE
  records = tf.data.Dataset.list_files(
                     os.environ.get("SM_CHANNEL_TRAINING")+'/*',
                     shuffle=True)
  ds = tf.data.TFRecordDataset(records, num_parallel_reads=aut)
  ds = ds.repeat()
  ds = ds.map(parse_sample, num_parallel_calls=aut)
  ds = ds.batch(batch_size)
  ds = ds.prefetch(aut)
  return ds
if __name__ == "__main__":
  import argparse, os
  parser = argparse.ArgumentParser(description="Train resnet")
  parser.add_argument("--model_dir", type=str, 
                      default="./model_keras_resnet")
  args = parser.parse_args()
  
  # init distribution lib
  dist.init()
  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[dist.local_rank()], 'GPU')
  input_shape = (32, 32, 3)
  classes = 10
  inputs = tf.keras.Input(shape=input_shape)
  outputs = tf.keras.applications.ResNet50(weights=None, 
                                          input_shape=input_shape, 
                                          classes=classes)(inputs)
  model = CustomModel(inputs, outputs)
  model.compile(loss=tf.losses.SparseCategoricalCrossentropy(),
                optimizer= tf.optimizers.Adam())
  dataset = get_dataset(batch_size = 1024, rank=dist.local_rank())
  cbs = [BroadcastGlobalVariablesCallback(0)]
  model.fit(dataset, steps_per_epoch=100, 
            epochs=10, callbacks=cbs, verbose=2)

Eredmények

Ebben a részben összehasonlítjuk az elosztott képzés végrehajtásának futásidejű eredményeit a Horovod adatelosztási keretrendszerrel és az Amazon SageMaker adatpárhuzamos könyvtárral. A futásidejű teljesítményt az edzési lépésenkénti átlagos másodpercek számával mértük.

Kísérletünk azt mutatja, hogy az SDP könyvtár környezettudatos elosztott tanítási algoritmusai nagyjából 20%-kal jobban teljesítenek, mint a Horovod könyvtár által használt algoritmusok. Vegye figyelembe, hogy az összehasonlító teljesítmény nagymértékben változhat a projekt részleteitől és a kiválasztott példány konfigurációjától függően. Még ugyanazon projekt és ugyanazon példánykonfiguráció esetén is eltérhetnek az eredmények attól függően, hogy a példányok pontosan hol helyezkednek el a képzési munkában. Amint fentebb megjegyeztük, a példányok különböző "klaszterelhelyezési csoportokba" helyezhetők, ami növelheti a várakozási időt és lassíthatja a képzést.

Kérjük, vegye figyelembe, hogy a cikk írásakor az SDP nem támogatott minden példánytípuson. A részletekért lásd a "dokumentációt".

Következő lépések

Bejegyzésünk "harmadik és utolsó részében" bemutatjuk, hogy az "Amazon SageMaker elosztott adat-párhuzamos könyvtára" hogyan támogatja az adatelosztást oly módon, hogy különbséget tesz a csomóponton belüli és a csomópontok közötti GPU-párok között.