სამი მატრიცა თითოეულ თავზე, ან ერთი დიდი მატრიცა

თავთავობრივი ხედი

თითოეული თავი საჭიროებს საკუთარ კითხვის პროექციას, გასაღების პროექციას და მნიშვნელობის პროექციას. თავ h-სთვის:

Q_h = X · W_Q^h    სადაც W_Q^h-ს აქვს ფორმა [d_model, head_dim]
K_h = X · W_K^h    სადაც W_K^h აქვს ფორმა [d_model, head_dim]
V_h = X · W_V^h    სადაც W_V^h აქვს ფორმა [d_model, head_dim]

X შეიცავს შეყვანის ფორმას [batch, seq_len, d_model]. პროექციის შემდეგ Q_h, K_h, V_h თითოეული შეიცავს ფორმას [batch, seq_len, head_dim].

შერწყმული ხედი

თავისთვის მატრიცები მეხსიერებაში ერთმანეთის გვერდით დგანან. ერთი შერწყმული მატრიცა W_Q ფორმით [d_model, d_model] ქმნის ყველა თავს ერთდროულად:

Q_fused = X · W_Q              # [batch, seq_len, d_model]
Q_per_head = reshape(Q_fused)  # [batch, n_head, seq_len, head_dim]

შერწყმული მატრიცული გამრავლება ერთ BLAS ზარს ახორციელებს 12-ის ნაცვლად. CUDA tensor cores იღებენ მაქსიმალურ შესრულებას ამ ზომის მატრიცულ გამრავლებებზე; თავშეკის მიხედვით გამრავლებები ნაკლებად გამოიყენებდნენ აპარატს.

პარამეტრების რაოდენობა

სამი შერწყმული მატრიცა W_Q, W_K, W_V, თითოეული d_model × d_model. პლუს გამომავალი პროექციის W_O, ასევე d_model × d_model. ANDREA-120M-ისთვის:

პარამეტრები თითოეული ფენის ატენშენზე = 4 × 768² = 2,359,296 ≈ 2.36M
პარამეტრები 12 ფენაზე      = 12 × 2.36M ≈ 28.3M

ANDREA-120M-ის მთლიანი პარამეტრების დაახლოებით მეოთხედი განთავსებულია ყურადღების პროექციებში. დარჩენილი სამეოთხედი მდებარეობს MLP ქვაშრეში და ემბედინგებში.

პროექციების დასახელება

თორმეტი ვექტორი ერთად იქცევა

თითოეული Head-ის გამოთვლის შემდეგ

თითოეული head გამოიმუშავებს გამომავალ ტენსორს ფორმით [batch, seq_len, head_dim]. თორმეტი head გამოიმუშავებს თორმეტ ასეთ ტენსორს. კონკატენაცია ფუნქციის განზომილება მzd მისი გამოყენება მათ ხელახლა ერთად stack-ს:

concat_output = concat(head_1, head_2, ..., head_12)
shape         = [batch, seq_len, n_head × head_dim]
= [batch, seq_len, 768]    # for ANDREA-120M

Concat ხდის split-ის შემობრუნებას. საერთო ფიჩერის განზომილება ბრუნდება d_model-ში. არანაირი ინფორმაციის დაკარგვა განზომილებებში; განსხვავება იმაშია, თუ რას შეიცავს თითოეული ჩანქი: head 1-ის ჩანქი ასახავს head 1-ის შეკვეთილ attention pattern-ს.

გამოყოფის პროექცია W_O

კონკატენაცია მარტო ტოვებს თავებს იზოლირებულად: head 4-ის გამოყოფა მდებარეობს head 7-ის გამოყოფის გვერდით, არც ერთი არ იცის მეორეზე. გამოყოფის პროექცია W_O ფორმით [d_model, d_model] აურევს მათ:

attention_output = concat_output · W_O
shape            = [batch, seq_len, d_model]

W_O-ის შემდეგ, ყველა გამომავალი განზომილება შეიცავს ყველა თორმეტი თავის შეკვეთილი ხაზოვანი კომბინაციას. ინფორმაცია თავისუფლად მიედინება თავებს შორის ამ ერთი მატრიცის გამრავლების მეშვეობით.

რატომ სპეციალიზდება თავები

არქიტექტურაში არაფერი აიძულებს თავს 4 შეძენილი ზმნის დროს ან თავს 9 შეძენილი წყვილი პუნქტუაციის შესწავლას. სპეციალიზაცია ჩნდება გრადიენტის ზეწოლისგან: ტრენინგის დროს, თავები რომლებიც რედუნდანტულად წვლილს შეაქვთ იღებენ უფრო მცირე გრადიენტებს ვიდრე თავები რომლებიც უნიკალურად წვლილს შაქვთ. ათასობით ნაბიჯის შემდეგ, თითოეული თავი რჩება ნიშაში რომელიც ყველაზე ეფექტურად ამცირებს საერთო ზარალს.

ემპირიულად, გაწვრთნილი ტრანსფორმერები აჩვენებენ თავებს რომლებიც მართავენ: პოზიციურ პატერნებს (თავი უყურებს წინა ტოკენს), სინტაქსურ პატერნებს (თავი უყურებს შესაბამის დახურულ ხაზგასმას), სემანტიკურ პატერნებს (თავი უყურებს ყველაზე უახლეს სახელწოდებულ ერთეულს). არანაირი ეტიკეტები არ ტრენინგავს ამ სპეციალიზაციას. მხოლოდ ტრენინგის სიგნალი, რომელიც გავრცელებულია W_O-ის მეშვეობით, სორტირებს თავებს.

რატომ ორმოცდაექვსი თავი, არა ერთი უფრო ფართო თავი

როგორ ინახავს CUDA თავებს

ერთი ტენზორი, ჩამოყალიბებული

ANDREA-ს სწავლების ძრავი microgpt_cuda.cu არ გამოყოფს ცალკეულ ბუფერებს თხავიანთა სათავეებისთვის. ის გამოყოფს ერთ შერწყმულ ტენზორს და თავის განზომილებას მკურნალ ნიმუშად მოეხმარება:

// Q = X · W_Q-ის შემდეგ (ერთი მატრიცული გამრავლება, შერწყმული თავებზე)
// Q-ს ფორმა აქვს [batch, seq_len, d_model]
// გადაკეთება    [batch, seq_len, n_head, head_dim]-ზე
// ტრანსპონირება  [batch, n_head, seq_len, head_dim]-ზე
// თითოეული head ახლა კონტიგური არის შიდა ორ ზომაში

ტრანსპოზიცია გადააადგილებს n_head-ს seq_len-ის წინ. რატომ? იმიტომ, რომ შემდეგი ოპერაცია (Q_h · K_h^T) საჭიროებს თითოეული head-ის seq_len × head_dim ნაჭერის კონტიგურობას მეხსიერებაში. CUDA matmuls უფრო სწრაფად მუშაობს კონტიგურ ტენზორებზე.

ერთი Kernel, მრავალი Head

ერთი attention CUDA კერნელი მუშაობს ყველა თავის (head) მსგავსად პარალელურად. თითოეული thread block-ი მართავს ერთ (batch, head) წყვილს; ბლოკის შიგნით thread-ები თანამშრომლობენ seq_len × head_dim ტილაზე. კერნელი არასდროს იცის, რომ ამუშავებს მრავალ თავს; გაშვების გრიდი უზრუნველყოფს პარალელიზმს.

კონფიგურაცია ასახავს აპარატურას

ANDREA-120M-ის არჩევანი n_head=12, head_dim=64 შესაბამისია RTX 4090 tensor cores-თან, რომლებიც ურჩევნიათ matmul ტილები 16-ის ჯერადებზე. head_dim=64 = 4 × 16 ზუსტად შეხვდება ტილის ფორმას. head_dim=32 (ANDREA-12M) ასევე შეხვდება, მაგრამ ნაკლებად გამოიყენებს ტილას. head_dim=72 არ შეხვდებოდა და გამოიწვევდა fallback კერნელებს.

საბოლოო სურათი

ხუთი ნაბიჯი. სამი მატრიცის გამრავლება ეხება შეყვანას (Q, K, V პროექციები). ერთი მატრიცის გამრავლება ეხება შერწყმულ თავებს (W_O). ერთი ყურადღების კერნელი მართავს ყველა თავს პარალელურად. ANDREA-120M ასრულებს ყველა ხუთ ნაბიჯს ერთხელ თითოეულ ფენაზე, თხრილობა შვიდიანი ფენა, ყოველ წინსვლის გავლაში.

პაიპლაინის წაკითხვა