Εκπαιδεύστε γιγάντια μοντέλα με σχεδόν γραμμική κλιμάκωση χρησιμοποιώντας μοιρασμένο παραλληλισμό δεδομένων στο Amazon SageMaker

Για την επιδίωξη ανώτερης ακρίβειας, τα μοντέλα βαθιάς μάθησης σε τομείς όπως η επεξεργασία φυσικής γλώσσας και η όραση υπολογιστή έχουν αυξηθεί σημαντικά σε μέγεθος τα τελευταία χρόνια, μετρώντας συχνά σε δεκάδες έως εκατοντάδες δισεκατομμύρια παραμέτρους. Η εκπαίδευση αυτών των γιγαντιαίων μοντέλων είναι δύσκολη και απαιτεί πολύπλοκες στρατηγικές διανομής. Οι επιστήμονες δεδομένων και οι μηχανικοί μηχανικής μάθησης αναζητούν συνεχώς τον καλύτερο τρόπο για να βελτιστοποιήσουν τον υπολογισμό της εκπαίδευσής τους, ωστόσο αγωνίζονται με τα γενικά έξοδα επικοινωνίας που μπορεί να αυξηθούν μαζί με το συνολικό μέγεθος του συμπλέγματος.

Αυτός είναι ο λόγος που ξεκινήσαμε πρόσφατα sσκληρός παραλληλισμός δεδομένων on Amazon Sage Maker, μια νέα τεχνική κατανεμημένης προπόνησης εξοικονόμησης μνήμης στο Βιβλιοθήκη SageMaker Model Parallel (SMP).. Ο παραλληλισμός μοιρασμένων δεδομένων έχει σχεδιαστεί ειδικά για μοντέλα ακραίας κλίμακας και χρησιμοποιεί το Amazon εσωτερικά MiCS τεχνολογία κάτω από την κουκούλα, μια επιστημονική προσπάθεια ελαχιστοποίησης της κλίμακας επικοινωνίας μειώνοντας τα ακριβά έξοδα επικοινωνίας που βασίζονται στη συλλογή παραμέτρων και στο συγχρονισμό διαβάθμισης. Με ένα μοντέλο GPT-30 παραμέτρων 2Β με μήκος ακολουθίας 2048, αυτή η νέα δυνατότητα πέτυχε 141 TFLOP, ταχύτητα 39.7% σε σύγκριση με το DeepSpeed ​​ZeRO-3. Για ένα μοντέλο 10B GPT-2 με μήκος ακολουθίας 512, αυτή η νέα δυνατότητα πέτυχε επίσης 564 δείγματα ανά δευτερόλεπτο, μια ταχύτητα 13.9% σε σύγκριση με το Fully Sharded Data Parallel (FSDP) της PyTorch. Να θυμάστε ότι σε γιγαντιαία εκπαίδευση μοντέλων, κάθε ποσοστό επιτάχυνσης μεταφράζεται σε δολάρια που εξοικονομήθηκαν και σε παραγωγικότητα που αποκτήθηκε στην ομάδα σας.

Σε αυτήν την ανάρτηση ιστολογίου, θα ρίξουμε πρώτα μια πιο προσεκτική ματιά στους βασικούς παράγοντες διαφοροποίησης του παραλληλισμού μοιρασμένων δεδομένων και πότε να το χρησιμοποιήσουμε. Στη συνέχεια, θα μάθετε πώς να εκπαιδεύετε εύκολα ένα μοντέλο GPT-30 παραμέτρων 2Β στο SageMaker με αυτήν τη νέα δυνατότητα. Τέλος, θα συγκρίνουμε την απόδοση με άλλες επιλογές ανοιχτού κώδικα, κυρίως με καλύτερη απόδοση από το DeepSpeed ​​ZeRO έως και 39.7% σε 256 GPU.

Πώς λειτουργεί ο παραλληλισμός μοιρασμένων δεδομένων και πότε να χρησιμοποιείται

Πριν εισαγάγουμε τον παραλληλισμό των μοιρασμένων δεδομένων, ας δούμε την ευρύτερη οικογένεια τεχνικών του. Οι πρόσφατες κατανεμημένες προσεγγίσεις εκπαίδευσης για μεγάλα μοντέλα έχουν μετακινηθεί σε ένα παράδειγμα όπου οι παράμετροι του μοντέλου, οι διαβαθμίσεις και οι καταστάσεις βελτιστοποίησης μοιράζονται σε παράλληλους κόμβους δεδομένων. Σε αντίθεση με το Pipeline Parallelism που έχει την έμφυτη πολυπλοκότητα της επιλογής επιπέδων για διαχωρισμό σε συσκευές, ειδικά όταν το πλαίσιο σας δεν υποστηρίζει αυτοματοποιημένο διαχωρισμό μοντέλων, αυτό το παράδειγμα διατηρεί κομψά την απλότητα του παραλληλισμού δεδομένων, ενώ αφαιρεί τον περιορισμό του παραλληλισμού δεδομένων όπου ένα μοντέλο πρέπει να χωράει σε μια ενιαία GPU.

Σε υπάρχοντα πλαίσια που εμπίπτουν σε αυτό το παράδειγμα, ιδίως το DeepSpeed ​​ZeRO-3 και το FSDP της PyTorch που προέρχεται από το FairScale, οι καταστάσεις μοντέλων κατανέμονται όλοι GPU, μια στρατηγική που μειώνει την κατανάλωση μνήμης σε κάθε GPU με το κόστος της επιβάρυνσης μεγάλων επιβαρύνσεων επικοινωνίας που αυξάνονται με το μέγεθος του συμπλέγματος και ως εκ τούτου προκαλεί σημαντική πτώση της επεκτασιμότητας σε κλίμακα. Αντίθετα, ο παραλληλισμός μοιρασμένων δεδομένων στο μοντέλο κατατμήσεων βιβλιοθήκης SMP δηλώνει στο α με επίγνωση της κλίμακας τρόπο με την κατάτμηση κάθε αντίγραφου καταστάσεων μοντέλου μόνο εντός ένα υποσύνολο των GPU.

Ας δούμε πιο προσεκτικά το Διαμέριση μοντέλου με επίγνωση κλίμακας στο MiCS, η βασική τεχνολογία πίσω από την παράλληλη μοιρασιά δεδομένων. Η διαίσθηση πίσω από αυτόν τον σχεδιασμό είναι ότι η κατάτμηση καταστάσεων εκπαίδευσης σε ολόκληρη την παράλληλη ομάδα δεδομένων μπορεί να μην απαιτείται για την εκπαίδευση ενός μοντέλου με δεκάδες δισεκατομμύρια παραμέτρους. Για παράδειγμα, 8 GPU V100 (32 GB η καθεμία) επαρκούν για να κρατήσουν το αντίγραφο των καταστάσεων μοντέλου ενός μοντέλου παραμέτρων 10 Β το οποίο χρειάζεται περίπου 200 GB μνήμης κατά την προπόνηση με το Adam optimizer χρησιμοποιώντας μικτή ακρίβεια. Περιορίζοντας ένα πλήρες αντίγραφο καταστάσεων μοντέλου στο μικρότερο υποσύνολο των GPU, μπορούμε να μειώσουμε αποτελεσματικά την κλίμακα των επιβαρύνσεων επικοινωνίας σε σύγκριση με το DeepSpeed ​​και το PyTorch FSDP. Η παράλληλη κατανομή δεδομένων αξιοποιεί επίσης άλλες τεχνικές στο MiCS, όπως η Ιεραρχική Επικοινωνία και ο Συγχρονισμός Διαβάθμισης 2-Hop. Για περισσότερες πληροφορίες, ελέγξτε έξω Σχεδόν γραμμική κλιμάκωση της εκπαίδευσης με γιγάντια μοντέλα σε AWS or MiCS: Σχεδόν γραμμική κλιμάκωση για εκπαίδευση γιγαντιαίο μοντέλο σε δημόσιο σύννεφο.

Τώρα, πώς ξέρετε πότε να επιλέξετε μοιρασμένα δεδομένα παράλληλα σε σχέση με άλλες κατανεμημένες τεχνικές εκπαίδευσης; Ο γενικός κανόνας είναι ότι εάν το μοντέλο σας έχει λιγότερες από 1 δισεκατομμύριο παραμέτρους και μπορεί να χωρέσει στη μνήμη GPU, Παράλληλη βιβλιοθήκη δεδομένων SageMaker or Μεταγλωττιστής εκπαίδευσης SageMaker μπορεί να είναι επαρκής για εσάς. Εάν έχετε μεγαλύτερα μοντέλα γλώσσας ή όρασης υπολογιστή, η πρότασή μας είναι να τα εκπαιδεύσετε με την τεχνική του παραλληλισμού δεδομένων σε συνδυασμό με σημείο ελέγχου ενεργοποίησης και εκφόρτωση ενεργοποίησης στην παράλληλη βιβλιοθήκη του μοντέλου SageMaker πρώτα, πριν από άλλες τεχνικές όπως π.χ παραλληλισμός τανυστών ή παραλληλισμός αγωγών.

Χρησιμοποιώντας τον παραλληλισμό μοιρασμένων δεδομένων για την εκπαίδευση του GPT-2 στο Amazon SageMaker

Ας μάθουμε τώρα πώς να εκπαιδεύσετε ένα μοντέλο GPT-2 με μοιρασμένα δεδομένα παράλληλα, με το SMP να ενσωματώνει την πολυπλοκότητα για εσάς. Αυτό πλήρες εκπαιδευτικό σημειωματάριο σας καθοδηγεί σε όλη τη διαδικασία, από την επεξεργασία δεδομένων, τον καθορισμό και την υποβολή εργασιών κατάρτισης έως την παρακολούθηση των αρχείων καταγραφής εκπαίδευσης. Αυτό που ακολουθεί είναι μια σύντομη επισκόπηση που επισημαίνει τα βασικά βήματα για τη χρήση αυτής της δυνατότητας.

1. Ξεκίνα

Ο παραλληλισμός μοιρασμένων δεδομένων είναι διαθέσιμος στην έκδοση PyTorch 1.12.0+ και λειτουργεί τόσο με το FP16 όσο και με το BF16. Ο ευκολότερος τρόπος χρήσης της βιβλιοθήκης SMP είναι μέσω ενός προκατασκευασμένου AWS Deep Learning Container για το PyTorch. Ωστόσο, εάν θέλετε να φέρετε το δικό σας κοντέινερ Docker, μπορείτε να ανατρέξετε στο Δημιουργήστε το δικό σας Docker Container με την SageMaker Distributed Model Parallel Library. Για να ξεκινήσετε, ακολουθήστε Τροποποιήστε ένα εκπαιδευτικό σενάριο PyTorch για να προσαρμόσετε τα API των SMP στο σενάριο εκπαίδευσής σας. Σε αυτήν την ενότητα, αναφέρουμε μόνο μερικά βασικά βήματα με αποσπάσματα κώδικα από το έτοιμο προς χρήση σενάριο εκπαίδευσης train_gpt_simple.py. Μπορείτε να ακολουθήσετε τα σχόλια στο σενάριο και Έγγραφο API για να μάθετε περισσότερα σχετικά με το πού χρησιμοποιούνται τα SMP API.

Πρώτα, εισαγάγετε και αρχικοποιήστε τη βιβλιοθήκη καλώντας smdistributed.modelparallel.torch.init() στην αρχή του σεναρίου εκπαίδευσης:

import smdistributed.modelparallel.torch as smp

smp.init(smp_config)

Δεύτερον, τυλίξτε το μοντέλο με το οποίο θα χωριστεί smdistributed.modelparallel.torch.DistributedModel και χρησιμοποιήστε το επιστρεφόμενο DistributedModel αντικείμενο προς τα εμπρός:

from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_config(model_config)
model = smp.DistributedModel(model, trace_device="gpu", backward_passes_per_step=args.gradient_accumulation)

Τυλίξτε το βελτιστοποιητή με smdistributed.modelparallel.torch.DistributedOptimizer για αποθήκευση και φόρτωση καταστάσεων βελτιστοποίησης.

from torch import optim

optimizer = optim.Adam(
    param_groups, betas=(args.beta1, args.beta2), lr=args.lr, weight_decay=args.weight_decay
)

optimizer = smp.DistributedOptimizer(
        optimizer, 
        static_loss_scale=None, 
        dynamic_loss_scale=True,
        dynamic_loss_args={"scale_window": 1000, "min_scale": 1, "delayed_shift": 2},
        )

Βάλτε τη λογική προς τα εμπρός και προς τα πίσω σε μια συνάρτηση βήματος και διακοσμήστε τη με smdistributed.modelparallel.torch.step.  Οποιοσδήποτε υπολογισμός ορίζεται μέσα στο smp.step-decorated η λειτουργία εκτελείται με κατανεμημένο τρόπο.

@smp.step
def train_step(model, optimizer, input_ids, attention_mask, args):
    loss = model(input_ids=input_ids, attention_mask=attention_mask, labels=input_ids)["loss"]
    model.backward(loss)

    return loss

@smp.step
def test_step(model, input_ids, attention_mask):
    loss = model(input_ids=input_ids, attention_mask=attention_mask, labels=input_ids)["loss"]
    
    return loss

2. Προετοιμάστε το σύνολο δεδομένων

Χρησιμοποιούμε το openwebtext είναι το σύνολο δεδομένων που χρησιμοποιούμε σε αυτό το παράδειγμα. Το σημειωματάριο χρησιμοποιεί το σενάριο data_prep_512.py για λήψη και προεπεξεργασία του συνόλου δεδομένων. Μπορείτε επίσης να εκπαιδεύσετε με άλλα σύνολα δεδομένων τροποποιώντας data_pipeline.py. Όταν ασχολείστε με μεγάλο σύνολο δεδομένων και μοντέλο, μπορείτε να επιταχύνετε την εργασία εκπαίδευσης χρησιμοποιώντας δεδομένα που είναι αποθηκευμένα σε Amazon FSx για λάμψη, το οποίο παρέχει ένα σύστημα αρχείων υψηλής απόδοσης εγγενώς ενσωματωμένο με Απλή υπηρεσία αποθήκευσης Amazon (S3). Δείτε τις οδηγίες από Διαμορφώστε το κανάλι εισαγωγής δεδομένων για χρήση του Amazon FSx για Luster για καθοδήγηση σχετικά με τη ρύθμιση ενός συστήματος αρχείων FSx Luster ως καναλιού εισαγωγής δεδομένων.

3. Ξεκινήστε τις εργασίες εκπαίδευσης

Αυτό το βήμα προϋποθέτει ότι το έχετε ήδη κάνει τροποποίησε το σενάριο προπόνησής σας και ετοίμασε το σύνολο δεδομένων όπως αναφέρθηκε στις προηγούμενες ενότητες. Προς την ενεργοποιήστε τον παραλληλισμό μοιρασμένων δεδομένων, απλώς ορίστε το sharded_data_parallel_degree στο Εκτιμητής PyTorch. Σε αυτό το σεμινάριο, ορίσαμε sharded_data_parallel_degree=128 και  instace_count=32 για κόμβους p4d.24xlarge, γεγονός που υποδεικνύει ότι οι καταστάσεις του μοντέλου θα κατανεμηθούν σε 128 GPU μεταξύ των συνολικών 256 GPU. Με βάση αυτήν την επιλεγμένη τιμή, το SMP θα ορίσει αυτόματα τον βαθμό παράλληλων δεδομένων στο 2 (επειδή 256/128=2), που σημαίνει ότι θα έχουμε δύο αντίγραφα για παραλληλισμό δεδομένων. Ένας γενικός κανόνας για την επιλογή μιας ιδανικής τιμής sharded_data_parallel_degree είναι να προσθέσετε έναν ακόμη κόμβο στην ομάδα κοινής χρήσης για κάθε 3B παραμέτρων μοντέλου. Σε αυτό το σεμινάριο, το μέγεθος του μοντέλου μας είναι 30Β, επομένως θα πρέπει να χρησιμοποιήσουμε τουλάχιστον 10 κόμβους για διαμοιρασμό. Και επειδή οι 16 κόμβοι (128 GPU) είναι η μικρότερη ισχύς των 2 πάνω από το όριο, ορίσαμε sharded_data_parallel_degree=128.

Για σημεία ελέγχου, παρέχουμε επίσης ένα σύνολο βοηθητικών προγραμμάτων ελέγχου σημείων ελέγχου sharded_data_parallel_checkpoint.py , συμπεριλαμβανομένου ενός βοηθητικού προγράμματος για την ανακατασκευή του πλήρους state_dict για περιπτώσεις προηγμένης χρήσης. Τέλος, μπορούμε να ξεκινήσουμε μια κατανεμημένη εργασία εκπαίδευσης καλώντας τη fit() στον Εκτιμητή.

smp_estimator = PyTorch(
    entry_point="train_gpt_simple.py",
    instance_type="ml.p4d.24xlarge",
    source_dir=os.getcwd(),
    volume_size=500,
    instance_count=32,
    distribution={
        "mpi": {
            "enabled": True,
            "processes_per_host": processes_per_host,
            "custom_mpi_options": mpioptions,
        },
        "smdistributed": {
            "modelparallel": {
                "enabled": True,
                "parameters": {
                    "ddp": True,
                    "skip_tracing": True,
                    "delayed_parameter_initialization": True,
                    "offload_activations": True,
                    "activation_loading_horizon": 4,
                    # To enable sharded data parallelism.
                    # Here we shard model states across 128 GPUs. 
                    "sharded_data_parallel_degree": 128, 
                    "fp16": False,
                    "bf16": True,
                    # This is to disable pipeline parallelism.
                    "partitions": 1,
                },
            }
        },
    },
    framework_version="1.12",
    py_version="py38",
    hyperparameters=hyperparameters,
    checkpoint_s3_uri=checkpoint_s3_uri if not use_fsx else None,
    checkpoint_local_path=hyperparameters["checkpoint-dir"] if use_fsx else None,
    ...
)

smp_estimator.fit(inputs=data_channels)

4. Παρακολούθηση των εργασιών κατάρτισης

Μπορείτε να έχετε πρόσβαση στα αρχεία καταγραφής εκπαίδευσης και να παρακολουθείτε τη χρήση της GPU και της μνήμης amazoncloudwatch. Βεβαιωθείτε ότι έχετε δει τα αρχεία καταγραφής του "algo-1" επειδή αυτός είναι ο κύριος κόμβος του οποίου η ροή εξόδου έχει τα αρχεία καταγραφής εργασιών εκπαίδευσης από όλες τις περιπτώσεις.

Συγκριτική αξιολόγηση απόδοσης

Συγκρίναμε τον παραλληλισμό μοιρασμένων δεδομένων στη βιβλιοθήκη SMP και σε 16 και 32 p4d.24xlarge κόμβους για μήκος ακολουθίας 512 και 2048, αντίστοιχα. Το μοντέλο GPT30 παραμέτρων 2B έχει διαμορφωθεί ώστε να χρησιμοποιεί ένα κρυφό πλάτος 7168, 48 στρώσεων και 64 κεφαλών. Μπορείτε να υιοθετήσετε την ίδια ακριβώς διαμόρφωση όπου το μήκος ακολουθίας είναι 2048 με ρύθμιση model_config = "gpt2-30b" στο τετράδιο του φροντιστηρίου. Με αυτήν τη ρύθμιση, η SMP πέτυχε 73.52 δείγματα ανά δευτερόλεπτο, μια ταχύτητα 39.7% σε σύγκριση με το DeepSpeed ​​ZeRO-3. Εάν το μέγεθος του token σας είναι 500 δισεκατομμύρια, αυτή η επιτάχυνση σημαίνει σχεδόν 367 ώρες εξοικονόμησης σε κόμβους p4d.24x, που ισοδυναμεί με περισσότερους από 12,000 $ εξοικονόμηση προϋπολογισμού ανά εκπαίδευση! Ο παρακάτω πίνακας συνοψίζει τα αποτελέσματα αναφοράς μας.

1/ Για κάθε διαμόρφωση μοντέλου, δοκιμάσαμε διαφορετικά χαρακτηριστικά, στάδια και διαμορφώσεις στο DeepSpeed ​​ZeRO και επιλέξαμε αυτό που παρέχει την καλύτερη απόδοση ως γραμμή βάσης DeepSpeed. Το σημείο αναφοράς εφαρμόστηκε Amazon Elastic Compute Cloud (Amazon EC2). 2/ Αυτά τα αποτελέσματα βασίζονται σε βελτιωμένες συλλογικότητες επικοινωνίας βελτιστοποιημένες για AWS, οι οποίες θα είναι διαθέσιμες σύντομα. 3/ Ο χρόνος για την εκπαίδευση προβάλλεται από την ταχύτητα με βάση τον αριθμό των κουπονιών που υποβλήθηκαν σε επεξεργασία.

Συνοπτικά, παρατηρήσαμε σταθερά υψηλότερη απόδοση με διαμοιρασμένο παραλληλισμό δεδομένων στο SMP σε σύγκριση με το DeepSpeed ​​σε μια σειρά μοντέλων και διαμορφώσεων. Αυτή η νέα δυνατότητα έδειξε επίσης καλύτερη απόδοση μνήμης σε σύγκριση με το DeepSpeed, επιτρέποντας στο SMP να ταιριάζει σε μεγαλύτερο μέγεθος παρτίδας και να μειώσει το επίπεδο συσσώρευσης κλίσης που απαιτείται για να χωρέσει ένα συγκεκριμένο παγκόσμιο μέγεθος παρτίδας.

Συμπέρασμα

Σε αυτήν την ανάρτηση, εισαγάγαμε μια νέα τεχνική κατανεμημένης εκπαίδευσης - παραλληλισμός μοιρασμένων δεδομένων - και πώς επιταχύνει την εκπαίδευση γιγαντιαίων μοντέλων με σχεδόν γραμμική κλίμακα στο Amazon SageMaker. Περπατήσαμε επίσης στο πώς να εκπαιδεύσουμε ένα μοντέλο GPT-2 με τη νέα τεχνική που ακολουθεί αυτό πλήρες παράδειγμα. Μπορείτε να ακολουθήσετε το Παραδείγματα Amazon SageMaker repo GitHub για να παρακολουθήσετε όλα τα παράλληλα παραδείγματα μοντέλων SageMaker ή να παρακολουθήσετε το επόμενο διένειμε εκπαιδευτικά εργαστήρια. Για να μάθετε περισσότερα σχετικά με τον παραλληλισμό μοιρασμένων δεδομένων, ανατρέξτε στο τεκμηρίωση.

Σχετικά με τους συγγραφείς

Έμιλι Webber εντάχθηκε στο AWS αμέσως μετά την κυκλοφορία του SageMaker και από τότε προσπαθεί να το πει στον κόσμο! Εκτός από τη δημιουργία νέων εμπειριών ML για τους πελάτες, η Emily απολαμβάνει τον διαλογισμό και τη μελέτη του Θιβετιανού Βουδισμού.

Can Karakus είναι Ανώτερος Εφαρμοσμένος Επιστήμονας στο AWS, βελτιστοποιώντας μεγάλης κλίμακας κατανεμημένη βαθιά μάθηση στο AWS. Τα ερευνητικά του ενδιαφέροντα καλύπτουν τη βαθιά μάθηση, την κατανεμημένη βελτιστοποίηση, τα κατανεμημένα συστήματα και τη θεωρία πληροφοριών. Εκτός δουλειάς, του αρέσει το ποδήλατο, τα ταξίδια, το διάβασμα και η μάθηση.

Rahul Huilgol είναι Ανώτερος Μηχανικός Λογισμικού στην AWS. Εργάζεται σε κατανεμημένα συστήματα βαθιάς μάθησης, ώστε να είναι εύκολη και αποτελεσματική η εκπαίδευση μεγάλων μοντέλων βαθιάς μάθησης στο cloud. Στον ελεύθερο χρόνο του, του αρέσει η φωτογραφία, η ποδηλασία και η κηπουρική.

Suhit Kodgule είναι Μηχανικός Ανάπτυξης Λογισμικού με ομάδα τεχνητής νοημοσύνης AWS που εργάζεται σε πλαίσια βαθιάς μάθησης. Στον ελεύθερο χρόνο του, του αρέσει η πεζοπορία, τα ταξίδια και η μαγειρική.

Έριν Χο είναι Υπεύθυνος Προϊόντων για το AWS Deep Learning. Εργάζεται σε προϊόντα που διευκολύνουν τους πελάτες να εκπαιδεύσουν μοντέλα βαθιάς εκμάθησης στο AWS. Για διασκέδαση έξω από τη δουλειά, της αρέσει η πεζοπορία και το σκι.