Tipps und Fallstricke der SageMaker Distributed Model Parallelism Library

Die Bibliothek ist am effizientesten, wenn die Chargengröße erhöht wird. In Anwendungsfällen, in denen das Modell zwar in ein einzelnes Gerät passt, aber nur mit einer kleinen Chargengröße trainiert werden kann, kann und sollte die Chargengröße nach der Integration der Bibliothek erhöht werden. Modellparallelität spart Speicherplatz bei großen Modellen, sodass Sie mit Losgrößen trainieren können, die zuvor nicht in den Arbeitsspeicher passten.

Die Auswahl einer zu kleinen oder zu großen Anzahl von Mikrobatches kann zu Leistungseinbußen führen. Die Bibliothek führt jeden Mikrobatch sequentiell in jedem Gerät aus. Daher muss die Mikrobatch-Größe (Batchgröße geteilt durch die Anzahl der Mikrobatches) groß genug sein, um jede GPU voll auszunutzen. Gleichzeitig steigt die Effizienz der Pipeline mit der Anzahl der Mikrobatches, weshalb es wichtig ist, das richtige Gleichgewicht zu finden. In der Regel ist es ein guter Ausgangspunkt, 2 oder 4 Mikrobatches auszuprobieren, wobei die Chargengröße bis zur Speichergrenze erhöht wird, und dann mit größeren Chargengrößen und einer größeren Anzahl von Mikrobatches zu experimentieren. Wenn die Anzahl der Mikrobatches erhöht wird, könnten größere Chargengrößen realisierbar werden, wenn eine ineinander verschachtelte Pipeline verwendet wird.

Ihre Chargengröße muss immer durch die Anzahl der Mikrobatchen teilbar sein. Beachten Sie, dass je nach Größe des Datensatzes manchmal die letzte Charge jeder Epoche kleiner sein kann als die anderen, und diese kleinere Charge muss auch durch die Anzahl der Mikrobatches teilbar sein. Ist dies nicht der Fall, können Sie drop_remainder=True im -tf.Dataset.batch()Aufruf (in TensorFlow) oder drop_last=True in DataLoader (in PyTorch) festlegen, sodass dieser letzte kleine Batch nicht verwendet wird. Wenn Sie eine andere API für die Datenpipeline verwenden, müssen Sie den letzten Stapel möglicherweise manuell überspringen, wenn er nicht durch die Anzahl der Mikrobatches teilbar ist.

Wenn Sie die manuelle Partitionierung verwenden, sollten Sie die Parameter berücksichtigen, die für mehrere Operationen und Module in Ihrem Modell verwendet werden, z. B. für die Einbettungstabelle in Transformatorarchitekturen. Module, die denselben Parameter verwenden, müssen aus Gründen der Richtigkeit auf demselben Gerät platziert werden. Wenn die automatische Partitionierung verwendet wird, erzwingt die Bibliothek diese Einschränkung automatisch.

Wenn das Modell mehrere Eingaben benötigt, stellen Sie sicher, dass Sie die Zufallsoperationen in Ihrer Datenpipeline (z. B. Mischen) mit smp.dp_rank() starten. Wenn der Datensatz deterministisch auf datenparallele Geräte aufgeteilt wird, stellen Sie sicher, dass der Shard von smp.dp_rank() indiziert wird. Dadurch soll sichergestellt werden, dass die Reihenfolge der Daten auf allen Rängen, die eine Modellpartition bilden, konsistent ist.

Jeder Tensor, der von der Funktion smp.DistributedModel.call (für TensorFlow) oder smp.DistributedModel.forward (für PyTorch) zurückgegeben wird, wird an alle anderen Ränge übertragen, und zwar von dem Rang, der diesen bestimmten Tensor berechnet hat. Daher sollte jeder Tensor, der außerhalb der Call- und Forward-Methoden nicht benötigt wird (z. B. Zwischenaktivierungen), nicht zurückgegeben werden, da dies zu unnötiger Kommunikation und Speicheraufwand führt und die Leistung beeinträchtigt.

Wenn eine smp.step-dekorierte Funktion ein Tensor-Argument hat, das keine Stapeldimension hat, muss der Argumentname beim Aufruf von non_split_inputs in der smp.step-Liste angegeben werden. Dadurch wird verhindert, dass die Bibliothek versucht, den Tensor in Mikrobatches aufzuteilen. Weitere Informationen finden Sie unter smp.step in der API-Dokumentation.

Wenn Sie einen Seed für deterministische Ergebnisse verwenden, setzen Sie den Seed auf der Grundlage von smp.dp_rank() (z. B. torch.manual_seed(42 + smp.dp_rank())). Andernfalls werden verschiedene Partitionen eines nn.Parameter auf die gleiche Weise initialisiert, was die Konvergenz beeinträchtigt.

SageMakerDie Modellparallelitätsbibliothek von verwendet NCCL, um Kollektive zu implementieren, die für die Verteilung der Module erforderlich sind. Insbesondere bei kleineren Modellen kann die Speichernutzung aufgrund des zusätzlichen Speicherplatzes, den NCCL beansprucht, zunehmen, wenn zu viele NCCL-Aufrufe gleichzeitig auf der GPU geplant sind. Um dem entgegenzuwirken, drosselt smp die NCCL-Aufrufe, so dass die Anzahl der laufenden NCCL-Operationen zu einem bestimmten Zeitpunkt kleiner oder gleich einem bestimmten Grenzwert ist. Das Standardlimit ist 8, kann aber mithilfe der Umgebungsvariablen SMP_NCCL_THROTTLE_LIMIT angepasst werden. Wenn Sie bei der Verwendung der Tensorparallelität einen höheren Speicherverbrauch als erwartet feststellen, können Sie versuchen, diesen Grenzwert zu reduzieren. Wenn Sie jedoch ein zu kleines Limit wählen, kann dies zu Durchsatzverlusten führen. Um die Drosselung vollständig zu deaktivieren, können Sie SMP_NCCL_THROTTLE_LIMIT=-1 festlegen.

Die folgende Identität, die gilt, wenn der Grad der Tensorparallelität 1 ist, gilt nicht, wenn der Grad der Tensorparallelität größer als 1 ist: smp.mp_size() * smp.dp_size() == smp.size(). Dies liegt daran, dass die Tensorparallelitätsgruppe sowohl Teil der Modellparallelitätsgruppe als auch der Datenparallelitätsgruppe ist. Wenn Ihr Code bereits Verweise auf mp_rank, mp_size, MP_GROUP, usw. enthält und Sie nur mit der parallel Pipeline-Gruppe arbeiten möchten, müssen Sie die Verweise möglicherweise durch smp.pp_size() ersetzen. Die folgenden Identitäten sind immer wahr:

Da der smp.DistributedModel Wrapper die Modellparameter ändert, wenn die Tensorparallelität aktiviert ist, sollte der Optimierer nach dem Aufruf von smp.DistributedModel mit den verteilten Parametern erstellt werden. Zum Beispiel funktioniert das Folgende nicht:

## WRONG
model = MyModel()
optimizer = SomeOptimizer(model.parameters())
model = smp.DistributedModel(model)  # optimizer now has outdated parameters! 

Stattdessen sollte der Optimierer mit den folgenden Parametern von smp.DistributedModel erstellt werden:

## CORRECT
model = smp.DistributedModel(MyModel())
optimizer = SomeOptimizer(model.optimizers())

Wenn ein Modul durch Tensorparallelität durch sein verteiltes Gegenstück ersetzt wird, erbt das verteilte Modul seine Gewichte nicht vom ursprünglichen Modul und initialisiert neue Gewichte. Das bedeutet, dass, wenn die Gewichte in einem bestimmten Aufruf initialisiert werden müssen (z. B. durch einen load_state_dict-Aufruf), dies nach dem smp.DistributedModel-Aufruf geschehen muss, sobald die Modulverteilung stattfindet.

Beachten Sie beim direkten Zugriff auf die Parameter verteilter Module, dass das Gewicht nicht dieselbe Form wie das ursprüngliche Modul hat. Zum Beispiel, 

with smp.tensor_parallelism():
    linear = nn.Linear(60, 60)

# will pass
assert tuple(linear.weight.shape) == (60, 60)

distributed_linear = smp.DistributedModel(linear)

# will fail. the number of input channels will have been divided by smp.tp_size()
assert tuple(distributed_linear.module.weight.shape) == (60, 60)

Die Verwendung von torch.utils.data.distributed.DistributedSampler wird aus Gründen der Tensorparallelität dringend empfohlen. Dadurch wird sichergestellt, dass jeder parallel Datenrang die gleiche Anzahl von Datenproben empfängt, wodurch verhindert wird, dass verschiedene dp_ranks eine unterschiedliche Anzahl von Schritten ausführen.

Wenn Sie die joinAPI der PyTorchKlasse von verwenden, DistributedDataParallelum Fälle zu behandeln, in denen verschiedene parallele Datenränge eine unterschiedliche Anzahl von Batches haben, müssen Sie dennoch sicherstellen, dass Ränge, die sich in derselben befinden, dieselbe Anzahl von Batches TP_GROUPhaben. Andernfalls können die bei der verteilten Ausführung von Modulen verwendeten Kommunikationskollektive hängen bleiben. Ränge, die sich in unterschiedlichen TP_GROUP s befinden, können eine unterschiedliche Anzahl von Batches haben, solange die join API verwendet wird.

Wenn Sie Ihr Modell überprüfen und die Tensorparallelität verwenden möchten, sollten Sie Folgendes berücksichtigen:

Weitere Hinweise zum Checkpointing eines Modells mit Tensorparallelität finden Sie unter Überprüfung eines verteilten Modells.