Βελτιστοποίηση Μνήμης για Εφαρμογές Σύνθεσης Γλώσσας και Προσαρμογή

Οι μεγάλες μοντέλα γλώσσας (LLM) όπως το GPT-4, το Bloom και το LLaMA έχουν επιτύχει αξιοσημείωτες ικανότητες αυξάνοντας το μέγεθός τους σε δισεκατομμύρια παραμέτρους. Ωστόσο, η ανάπτυξη αυτών των τεράστιων μοντέλων για εφαρμογές σύνθεσης ή προσαρμογής είναι μια πρόκληση λόγω των τεράστιων απαιτήσεων μνήμης. Σε αυτό το τεχνικό blog, θα εξερευνήσουμε τεχνικές για την εκτίμηση και την βελτιστοποίηση της κατανάλωσης μνήμης κατά τη διάρκεια της σύνθεσης και της προσαρμογής των LLM σε διάφορες ρυθμίσεις υλικού.

Κατανόηση των Απαιτήσεων Μνήμης

Η μνήμη που απαιτείται για να φορτωθεί ένα LLM καθορίζεται κυρίως από τον αριθμό των παραμέτρων και την αριθμητική ακρίβεια που χρησιμοποιείται για την αποθήκευση των παραμέτρων. Ένας απλός κανόνας είναι:

• Το φορτώνωμα ενός μοντέλου με X δισεκατομμύρια παραμέτρους απαιτεί περίπου 4X GB VRAM σε 32-bit ακρίβεια float
• Το φορτώνωμα ενός μοντέλου με X δισεκατομμύρια παραμέτρους απαιτεί περίπου 2X GB VRAM σε 16-bit bfloat16/float16 ακρίβεια

Για παράδειγμα, το φορτώνωμα του μοντέλου GPT-3 με 175 δισεκατομμύρια παραμέτρους θα απαιτούσε περίπου 350 GB VRAM σε bfloat16 ακρίβεια. Μέχρι σήμερα, οι μεγαλύτερες εμπορικά διαθέσιμες GPU όπως το NVIDIA A100 και το H100 προσφέρουν μόνο 80 GB VRAM, καθιστώντας αναγκαίες τις τεχνικές tensor parallelism και model parallelism.

Κατά τη διάρκεια της σύνθεσης, η μνήμη που χρησιμοποιείται κυριαρχείται από τις παραμέτρους του μοντέλου και τις προσωρινές ενεργοποιήσεις που παράγονται. Μια υψηλού επιπέδου εκτίμηση για τη μέγιστη χρήση μνήμης κατά τη διάρκεια της σύνθεσης είναι το άθροισμα της μνήμης που απαιτείται για να φορτωθεί το μοντέλο και της μνήμης για τις ενεργοποιήσεις.

Ποσοτικοποίηση της Μνήμης Σύνθεσης

Ας ποσοτικοποιήσουμε τις απαιτήσεις μνήμης για τη σύνθεση χρησιμοποιώντας το μοντέλο OctoCode, το οποίο έχει περίπου 15 δισεκατομμύρια παραμέτρους σε μορφή bfloat16 (~ 31 GB). Θα χρησιμοποιήσουμε τη βιβλιοθήκη Transformers για να φορτώσουμε το μοντέλο και να генνήσουμε κείμενο:

from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
import torch

<p>model = AutoModelForCausalLM.from_pretrained(&quot;bigcode/octocoder&quot;,
torch_dtype=torch.bfloat16,
device_map=&quot;auto&quot;,
pad_token_id=0)
tokenizer = AutoTokenizer.from_pretrained(&quot;bigcode/octocoder&quot;)
pipe = pipeline(&quot;text-generation&quot;, model=model, tokenizer=tokenizer)</p>

<p>prompt = &quot;Ερώτηση: Παρακαλώ γράψτε μια συνάρτηση Python για τη μετατροπή byte σε gigabyte.\n\nΑπάντηση:&quot;
result = pipe(prompt, max_new_tokens=60)[0][&quot;generated_text&quot;][len(prompt):]</p>

<p>def bytes_to_gigabytes(bytes):
return bytes / 1024 / 1024 / 1024</p>

<p>bytes_to_gigabytes(torch.cuda.max_memory_allocated())

Εξοδος:

Η μέγιστη χρήση μνήμης GPU είναι περίπου 29 GB, που συμφωνεί με την εκτίμησή μας των 31 GB για το φορτώνωμα του μοντέλου σε μορφή bfloat16.

Βελτιστοποίηση Μνήμης Σύνθεσης με Quantization

Ενώ η bfloat16 είναι η κοινή ακρίβεια που χρησιμοποιείται για την εκπαίδευση των LLM, οι ερευνητές έχουν βρει ότι η quantization των βαρών του μοντέλου σε χαμηλότερες ακρίβειες όπως 8-bit ακέραιοι (int8) ή 4-bit ακέραιοι μπορεί να μειώσει σημαντικά την χρήση μνήμης με ελάχιστη απώλεια ακρίβειας για εργασίες σύνθεσης όπως η γεννήθηκε κειμένου.

Ας δούμε τις αποταμιεύσεις μνήμης από την quantization 8-bit και 4-bit του μοντέλου OctoCode:

</div>
# 8-bit quantization
model = AutoModelForCausalLM.from_pretrained("bigcode/octocoder", load_in_8bit=True,
pad_token_id=0)
pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)
result = pipe(prompt, max_new_tokens=60)[0]["generated_text"][len(prompt):]
bytes_to_gigabytes(torch.cuda.max_memory_allocated())</pre>

Εξοδος:

# 4-bit quantization
model = AutoModelForCausalLM.from_pretrained("bigcode/octocoder", load_in_4bit=True,
low_cpu_mem_usage=True, pad_token_id=0)
pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)
result = pipe(prompt, max_new_tokens=60)[0]["generated_text"][len(prompt):]
bytes_to_gigabytes(torch.cuda.max_memory_allocated())

Εξοδος:

Με την quantization 8-bit, η απαιτούμενη μνήμη μειώνεται από 31 GB σε 15 GB, ενώ η quantization 4-bit μειώνει περαιτέρω την μνήμη σε μόλις 9,5 GB! Αυτό επιτρέπει την εκτέλεση του μοντέλου OctoCode 15B παραμέτρων σε καταναλωτικές GPU όπως η RTX 3090 (24 GB VRAM).

Ωστόσο, σημειώστε ότι μια πιο激 quantization όπως η 4-bit μπορεί να οδηγήσει σε απώλεια ακρίβειας σε σύγκριση με την 8-bit ή bfloat16 ακρίβεια. Υπάρχει ένα εύρος μεταξύ των αποταμιεύσεων μνήμης και της ακρίβειας που οι χρήστες πρέπει να αξιολογήσουν για την περίπτωσή τους.

Η quantization είναι μια ισχυρή τεχνική που μπορεί να επιτρέψει την ανάπτυξη των LLM σε περιβάλλοντα με περιορισμένα ресурσία όπως οι cloud instances, τα edge devices ή ακόμη και τα κινητά τηλέφωνα, μειώνοντας δραστικά την αποτύπωση μνήμης.

Εκτίμηση Μνήμης για Προσαρμογή

Ενώ η quantization χρησιμοποιείται κυρίως για την αποτελεσματική σύνθεση, τεχνικές όπως η tensor parallelism και η model parallelism είναι κρίσιμες για τη διαχείριση των απαιτήσεων μνήμης κατά τη διάρκεια της εκπαίδευσης ή προσαρμογής των μεγάλων μοντέλων γλώσσας.

Η μέγιστη κατανάλωση μνήμης κατά τη διάρκεια της προσαρμογής είναι συνήθως 3-4 φορές υψηλότερη από τη σύνθεση λόγω των πρόσθετων απαιτήσεων μνήμης για:

• Γραδίες
• Καταστάσεις βελτιστοποιητή
• Ενεργοποιήσεις από το forward pass που αποθηκεύονται για backpropagation

Μια συντηρητική εκτίμηση είναι ότι η προσαρμογή ενός LLM με X δισεκατομμύρια παραμέτρους απαιτεί περίπου 4 * (2X) = 8X GB VRAM σε bfloat16 ακρίβεια.

Για παράδειγμα, η προσαρμογή του μοντέλου LLaMA με 7 δισεκατομμύρια παραμέτρους θα απαιτούσε περίπου 7 * 8 = 56 GB VRAM ανά GPU σε bfloat16 ακρίβεια. Αυτό υπερβαίνει την ικανότητα μνήμης των τρεχουσών GPU, καθιστώντας αναγκαίες τις τεχνικές distributed fine-tuning.

Διανεμημένες Τεχνικές Προσαρμογής

Πολλές διανεμημένες μεθόδοι προσαρμογής έχουν προταθεί για την υπέρβαση των περιορισμών μνήμης GPU για μεγάλα μοντέλα:

1. Διανεμημένη Παραλληλία: Η κλασική προσέγγιση διανεμημένης παραλληλίας αναπαράγει το ολόκληρο μοντέλο σε πολλαπλά GPU ενώ διαιρεί και διανέμει τις δόσεις εκπαίδευσης. Αυτό μειώνει τον χρόνο εκπαίδευσης γραμμικά με τον αριθμό των GPU αλλά δεν μειώνει την απαιτούμενη μνήμη σε κάθε GPU.
2. ZeRO Stage 3: Μια προηγμένη μορφή διανεμημένης παραλληλίας που διαμερίζει τις παραμέτρους του μοντέλου, τις γραδίες και τις καταστάσεις βελτιστοποιητή σε GPU. Μειώνει την απαιτούμενη μνήμη σε σύγκριση με την κλασική διανεμημένη παραλληλία, κρατώντας μόνο τα απαραίτητα διαμερισμένα δεδομένα σε κάθε GPU κατά τη διάρκεια διαφόρων φάσεων της εκπαίδευσης.
3. Tensor Parallelism: Αντί να αναπαράγουμε το μοντέλο, η tensor parallelism διαιρεί τις παραμέτρους του μοντέλου σε γραμμές ή στήλες και τις διανέμει σε GPU. Κάθε GPU λειτουργεί σε ένα διαμερισμένο σύνολο παραμέτρων, γραδιών και καταστάσεων βελτιστοποιητή, οδηγώντας σε σημαντικές αποταμιεύσεις μνήμης.
4. Pipeline Parallelism: Αυτή η τεχνική διαμερίζει τους.layers του μοντέλου σε διαφορετικά GPU/εργαζόμενους, με κάθε συσκευή να εκτελεί ένα υποσύνολο των layer. Οι ενεργοποιήσεις μεταφέρονται μεταξύ των εργαζομένων, μειώνοντας την απαιτούμενη μνήμη αλλά αυξάνοντας το κόστος επικοινωνίας.

Η εκτίμηση της χρήσης μνήμης για αυτές τις διανεμημένες μεθόδοιες δεν είναι εύκολη, καθώς η διανομή των παραμέτρων, γραδιών, ενεργοποιήσεων και καταστάσεων βελτιστοποιητή ποικίλλει μεταξύ των τεχνικών. Επιπλέον, διαφορετικά компоненты όπως το σώμα του μετασχηματιστή και το κεφάλι της μονάδας γλώσσας μπορεί να εμφανίζουν διαφορετικές συμπεριφορές απομείωσης μνήμης.

Η Λύση LLMem

Ερευνητές πρόσφατα πρότειναν την LLMem, μια λύση που εκτιμά με ακρίβεια την κατανάλωση μνήμης GPU όταν εφαρμόζονται διανεμημένες μεθόδοι προσαρμογής σε LLM σε πολλαπλά GPU.

Εκτίμηση Χρήσης Μνήμης GPU για Προσαρμογή Προ-Εκπαιδευμένων LLM

Το LLMem λαμβάνει υπόψη παράγοντες όπως η ανασύνθεση παραμέτρων πριν από την επεξεργασία (ZeRO Stage 3), η συλλογή εξόδου στην αντίστροφη διαδικασία (tensor parallelism) και οι διαφορετικές στρατηγικές απομείωσης μνήμης για το σώμα του μετασχηματιστή και το κεφάλι της μονάδας γλώσσας.

Πειραματικά αποτελέσματα δείχνουν ότι το LLMem μπορεί να εκτιμήσει την μέγιστη χρήση μνήμης GPU για την προσαρμογή των LLM σε ένα單 GPU με ποσοστό λάθους μέχρι 1,6%, υπερβαίνοντας το μέσο ποσοστό λάθους του DNNMem 42,6%. Όταν εφαρμόζονται διανεμημένες μεθόδοι προσαρμογής σε LLM με πάνω από ένα δισεκατομμύριο παραμέτρους σε πολλαπλά GPU, το LLMem επιτυγχάνει ένα εντυπωσιακό μέσο ποσοστό λάθους 3,0%.