แฟลช แอทเทนชัน: การปฏิวัติความมีประสิทธิภาพของทรานส์ฟอร์เมอร์

</pre>
import torch

<p>def standard_attention(Q, K, V):
# Q, K, V shape: (batch_size, seq_len, d_model)
d_k = K.size(-1)
scores = torch.matmul(Q, K.transpose(-2, -1)) / torch.sqrt(torch.tensor(d_k))
attention_weights = torch.softmax(scores, dim=-1)
return torch.matmul(attention_weights, V)</p>

การนำไปใช้นี้ แม้ว่าจะตรงไปตรงมา แต่ก็遭遇ปัญหาที่กล่าวถึงข้างต้น เทนเซอร์ scores ที่มีรูป (batch_size, seq_len, seq_len) สามารถใหญ่เกินไปสำหรับลำดับยาว

การเข้าสู่แฟลช แอทเทนชัน

แฟลช แอทเทนชัน ซึ่งถูกนำเสนอโดย Tri Dao และเพื่อนร่วมงาน ในปี 2022 เป็นวิธีการคำนวณการกระจายความสนใจที่ลดการใช้หน่วยความจำและปรับปรุงประสิทธิภาพการคำนวณอย่างมาก แนวคิดหลักที่อยู่เบื้องหลังแฟลช แอทเทนชันคือ

1. การแบ่งออกเป็นบล็อก: แบ่งเมทริกซ์การกระจายความสนใจที่ใหญ่ออกเป็นบล็อกที่เล็กกว่าที่สามารถใส่ในหน่วยความจำ SRAM ได้
2. การคำนวณใหม่: แทนที่จะจัดเก็บเมทริกซ์การกระจายความสนใจทั้งหมด การคำนวณบางส่วนของมันใหม่เมื่อจำเป็นในการย้อนกลับ
3. การนำไปใช้ที่ตระหนักถึง IO: ปรับอัลกอริทึมให้ลดการเคลื่อนย้ายข้อมูลระหว่างระดับต่างๆ ของลำดับหน่วยความจำของ GPU

อัลกอริทึมแฟลช แอทเทนชัน

ในแก่นแท้ แฟลช แอทเทนชันคือการคิดใหม่เกี่ยวกับวิธีการคำนวณการกระจายความสนใจ แทนที่จะคำนวณเมทริกซ์การกระจายความสนใจทั้งหมดในครั้งเดียว มันจะประมวลผลมันในบล็อก โดยใช้ลำดับหน่วยความจำของ GPU ในรูปแบบที่เหมาะสม

มาทำให้เห็นภาพกับอัลกอริทึมระดับสูง

1. อินพุต: เมทริกซ์ Q, K, V ใน HBM และ SRAM ขนาด M
2. ขนาดบล็อกคำนวณตาม SRAM ที่มีอยู่
3. การเริ่มต้นเมทริกซ์เอาต์พุต O และเวกเตอร์เสริม l และ m
4. การแบ่งเมทริกซ์อินพุตออกเป็นบล็อกเพื่อใส่ใน SRAM
5. ลูปที่ซ้อนกันสองชั้นประมวลผลบล็อกเหล่านี้

การประมวลผลแบบบล็อกนี้ทำให้แฟลช แอทเทนชันสามารถรักษาหน่วยความจำที่เล็กกว่ามากในขณะที่ยังคงคำนวณการกระจายความสนใจที่ถูกต้อง

คณิตศาสตร์เบื้องหลังแฟลช แอทเทนชัน

กุญแจสำคัญในการทำให้แฟลช แอทเทนชันทำงานคือเทคนิคคณิตศาสตร์ที่ทำให้สามารถคำนวณ softmax ในลักษณะบล็อกได้

1. การแยก softmax:
   softmax(x) = exp(x - m) / Σexp(x - m)

   โดยที่ m คือค่าสูงสุดใน x

2. การรวม softmax:
   softmax(x ∪ y) = softmax(softmax(x) * e^(m_x - m), softmax(y) * e^(m_y - m))

   โดยที่ m = max(m_x, m_y)

สูตรเหล่านี้ทำให้แฟลช แอทเทนชันสามารถคำนวณผลลัพธ์ softmax บางส่วนสำหรับแต่ละบล็อกและรวมเข้าด้วยกันเพื่อให้ได้ผลลัพธ์สุดท้าย

รายละเอียดการนำไปใช้

มาทำให้เห็นภาพกับการนำไปใช้แฟลช แอทเทนชันแบบง่ายๆ เพื่อแสดงแนวคิดหลัก

import torch

<p>def flash_attention(Q, K, V, block_size=256):
batch_size, seq_len, d_model = Q.shape</p>

<p># Initialize output and running statistics
O = torch.zeros_like(Q)
L = torch.zeros((batch_size, seq_len, 1))
M = torch.full((batch_size, seq_len, 1), float(&#039;-inf&#039;))</p>

<p>for i in range(0, seq_len, block_size):
Q_block = Q[:, i:i+block_size, :]</p>

<p>for j in range(0, seq_len, block_size):
K_block = K[:, j:j+block_size, :]
V_block = V[:, j:j+block_size, :]</p>

<p># Compute attention scores for this block
S_block = torch.matmul(Q_block, K_block.transpose(-2, -1)) / (d_model ** 0.5)</p>

<p># Update running max
M_new = torch.maximum(M[:, i:i+block_size], S_block.max(dim=-1, keepdim=True).values)</p>

<p># Compute exponentials
exp_S = torch.exp(S_block - M_new)
exp_M_diff = torch.exp(M[:, i:i+block_size] - M_new)</p>

<p># Update running sum
L_new = exp_M_diff * L[:, i:i+block_size] + exp_S.sum(dim=-1, keepdim=True)</p>

<p># Compute output for this block
O[:, i:i+block_size] = (
exp_M_diff * O[:, i:i+block_size] +
torch.matmul(exp_S, V_block)
) / L_new</p>

<p># Update running statistics
L[:, i:i+block_size] = L_new
M[:, i:i+block_size] = M_new</p>

return O

การนำไปใช้นี้ แม้ว่าจะถูกทำให้ง่าย แต่ก็แสดงถึงแก่นแท้ของแฟลช แอทเทนชัน

ผลกระทบของแฟลช แอทเทนชัน

การแนะนำแฟลช แอทเทนชันมีผลกระทบอย่างลึกซึ้งต่อสนามของการเรียนรู้ของเครื่อง โดยเฉพาะสำหรับโมเดลภาษาที่ใหญ่และแอปพลิเคชันที่มีบริบทยาว ผลประโยชน์หลักๆ ได้แก่

1. การใช้หน่วยความจำที่ลดลง: แฟลช แอทเทนชันลดความซับซ้อนของหน่วยความจำจาก O(N^2) เป็น O(N) โดยที่ N คือความยาวของลำดับ ทำให้สามารถประมวลผลลำดับที่ยาวกว่ามากด้วยฮาร์ดแวร์เดียวกัน
2. ความเร็วที่ดีขึ้น: โดยการลดการเคลื่อนย้ายข้อมูลและใช้ความสามารถในการคำนวณของ GPU ได้ดีขึ้น แฟลช แอทเทนชันสามารถเพิ่มความเร็วได้มากถึง 3 เท่าในการฝึกอบรม GPT-2 เมื่อเทียบกับการนำไปใช้มาตรฐาน
3. การคำนวณที่ถูกต้อง: ไม่เหมือนเทคนิคการเพิ่มประสิทธิภาพการกระจายความสนใจอื่นๆ แฟลช แอทเทนชันคำนวณการกระจายความสนใจที่ถูกต้อง ไม่ใช่การประมาณ
4. ความสามารถในการปรับขนาด: การลดการใช้หน่วยความจำทำให้สามารถปรับขนาดไปสู่ลำดับที่ยาวกว่ามากได้ โดยที่อาจมีหลายล้านโทเค็น

ผลกระทบในโลกแห่งความเป็นจริง

ผลกระทบของแฟลช แอทเทนชันไม่ได้จำกัดอยู่แค่ในงานวิจัยเท่านั้น แต่ยังได้รับการนำไปใช้ในหลายไลบรารีและโมเดลการเรียนรู้ของเครื่องยอดนิยม

• ทรานส์ฟอร์เมอร์ของ Hugging Face: ไลบรารีทรานส์ฟอร์เมอร์ที่ได้รับความนิยมสูงได้นำแฟลช แอทเทนชันมาใช้ ทำให้ผู้ใช้สามารถใช้ประโยชน์จากมันได้อย่างง่ายดาย
• GPT-4 และต่อๆ ไป: แม้ว่าจะไม่ได้รับการยืนยัน แต่มีการคาดเดาว่าโมเดลภาษาที่ล้ำสมัยอย่าง GPT-4 อาจใช้เทคนิคที่คล้ายกับแฟลช แอทเทนชันเพื่อจัดการบริบทที่ยาว
• โมเดลบริบทยาว: แฟลช แอทเทนชันทำให้เกิดรุ่นใหม่ของโมเดลที่สามารถจัดการบริบทที่ยาวมาก เช่น โมเดลที่สามารถประมวลผลหนังสือทั้งเล่มหรือวิดีโอที่ยาว

แฟลชแอทเทนชัน: การพัฒนาล่าสุด

การกระจายความสนใจมาตรฐานเทียบกับแฟลช แอทเทนชัน

แฟลชแอทเทนชัน-2

โดยการสร้างบนความสำเร็จของแฟลช แอทเทนชันเดิม ทีมเดียวกันได้นำเสนอแฟลชแอทเทนชัน-2 ในปี 2023 ซึ่งมีการปรับปรุงหลายประการ

1. การเพิ่มประสิทธิภาพเพิ่มเติม: แฟลชแอทเทนชัน-2 สามารถใช้ GPU ได้ดีขึ้น โดยสามารถเข้าถึงได้ถึง 70% ของความสามารถในการคำนวณสูงสุดบน GPU A100
2. การปรับปรุงการย้อนกลับ: การย้อนกลับได้รับการปรับปรุงให้เร็วขึ้น ทำให้การฝึกอบรมเร็วขึ้นอย่างมาก
3. การสนับสนุนรูปแบบการกระจายความสนใจที่แตกต่างกัน: แฟลชแอทเทนชัน-2 ขยายการสนับสนุนไปยังรูปแบบการกระจายความสนใจที่หลากหลาย รวมถึงการกระจายความสนใจแบบกลุ่มและแบบหลายคำถาม

แฟลชแอทเทนชัน-3

ที่เปิดตัวในปี 2024 แฟลชแอทเทนชัน-3 เป็นตัวแทนของความก้าวหน้าล่าสุดในซีรีส์นี้ โดยนำเสนอเทคนิคใหม่ๆ เพื่อปรับปรุงประสิทธิภาพ

1. การคำนวณแบบไม่สอดคล้องกัน: การใช้คำสั่ง GPU แบบไม่สอดคล้องกันเพื่อซ้อนทับคำนวณต่างๆ
2. การสนับสนุน FP8: การใช้การคำนวณ FP8 ที่มีความแม่นยำต่ำสำหรับการประมวลผลที่เร็วขึ้น
3. การประมวลผลที่ไม่สอดคล้องกัน: เทคนิคในการลดข้อผิดพลาดการปรับขนาดเมื่อใช้รูปแบบความแม่นยำต่ำ

ตัวอย่างง่ายๆ ของวิธีการที่แฟลชแอทเทนชัน-3 อาจใช้การคำนวณแบบไม่สอดคล้องกัน

import torch
from torch.cuda.amp import autocast

<p>def flash_attention_3(Q, K, V, block_size=256):
with autocast(dtype=torch.float8): # ใช้ FP8 สำหรับการคำนวณ
# ... (คล้ายกับการนำไปใช้ก่อนหน้า)</p>

<p># ตัวอย่างการคำนวณแบบไม่สอดคล้องกัน
with torch.cuda.stream(torch.cuda.Stream()):
# คำนวณ GEMM แบบไม่สอดคล้องกัน
S_block = torch.matmul(Q_block, K_block.transpose(-2, -1)) / (d_model ** 0.5)</p>

<p># ในขณะเดียวกันบนสตรีมหลัก:
# เตรียมการสำหรับการคำนวณ softmax</p>

<p># สynchronize สตรีม
torch.cuda.synchronize()</p>

<p># ต่อด้วยการคำนวณ softmax และเอาต์พุต
# ...</p>

return O

โค้ดนี้แสดงวิธีการที่แฟลชแอทเทนชัน-3 อาจใช้การคำนวณแบบไม่สอดคล้องกันและความแม่นยำ FP8

การนำแฟลช แอทเทนชันไปใช้ในโปรเจ็กต์ของคุณ

หากคุณตื่นเต้นกับการใช้แฟลช แอทเทนชันในโปรเจ็กต์ของคุณ คุณมีตัวเลือกหลายอย่าง

1. ใช้ไลบรารีที่มีอยู่: ไลบรารีหลายตัว เช่น ทรานส์ฟอร์เมอร์ของ Hugging Face ได้นำแฟลช แอทเทนชันมาใช้แล้ว การอัปเดตเป็นเวอร์ชันล่าสุดและเปิดใช้งานแฟลช แอทเทนชันอาจเป็นเพียงพอ
2. การนำไปใช้แบบกำหนดเอง: สำหรับการควบคุมหรือกรณีการใช้งานพิเศษ คุณอาจต้องการนำแฟลช แอทเทนชันไปใช้เอง ไลบรารี xformers ให้การนำไปใช้ที่ดีเป็นตัวอย่าง
3. การปรับให้เหมาะสมสำหรับฮาร์ดแวร์: หากคุณทำงานกับฮาร์ดแวร์เฉพาะ (เช่น GPU H100 ของ NVIDIA (NVDA )) คุณอาจต้องการใช้คุณสมบัติเฉพาะของฮาร์ดแวร์สำหรับการปรับให้เหมาะสมสูงสุด

ตัวอย่างของวิธีการใช้แฟลช แอทเทนชันกับไลบรารีทรานส์ฟอร์เมอร์ของ Hugging Face

from transformers import AutoModel, AutoConfig

<p># เปิดใช้แฟลช แอทเทนชัน
config = AutoConfig.from_pretrained(&quot;bert-base-uncased&quot;)
config.use_flash_attention = True</p>

<p># โหลดโมเดลพร้อมแฟลช แอทเทนชัน
model = AutoModel.from_pretrained(&quot;bert-base-uncased&quot;, config=config)</p>

<p># ใช้โมเดลตามปกติ
# ...

ความท้าทายและทิศทางในอนาคต

แม้ว่าแฟลช แอทเทนชันจะก้าวหน้าไปมากในเรื่องประสิทธิภาพของการกระจายความสนใจ แต่ก็ยังมีความท้าทายและพื้นที่สำหรับการวิจัยในอนาคต

1. ความเฉพาะเจาะจงของฮาร์ดแวร์: การนำไปใช้ในปัจจุบันมักจะถูกปรับให้เหมาะสมสำหรับฮาร์ดแวร์ GPU แบบเฉพาะ การทำให้การปรับให้เหมาะสมเหล่านี้ทำงานได้กับฮาร์ดแวร์ที่หลากหลายยังคงเป็นความท้าทาย
2. การรวมกับเทคนิคอื่นๆ: การรวมแฟลช แอทเทนชันกับเทคนิคการเพิ่มประสิทธิภาพอื่นๆ เช่น การตัดทอน การปรับขนาด และการบีบอัดโมเดล เป็นพื้นที่ที่กำลังได้รับการวิจัยอย่างต่อเนื่อง
3. การขยายไปยังโดเมนอื่นๆ: แม้ว่าแฟลช แอทเทนชันจะประสบความสำเร็จใน NLP แต่การขยายผลประโยชน์ของมันไปยังโดเมนอื่นๆ เช่น คอมพิวเตอร์วิชั่นและโมเดลหลายรูปแบบยังคงเป็นความพยายามอย่างต่อเนื่อง
4. ความเข้าใจเชิงทฤษฎี: การทำความเข้าใจแฟลช แอทเทนชันในระดับทฤษฎีเพิ่มเติมอาจนำไปสู่การเพิ่มประสิทธิภาพที่มีประสิทธิภาพมากขึ้น

สรุป

ด้วยการนำGPU Memory Hierarchy และเทคนิคคณิตศาสตร์มาใช้ แฟลช แอทเทนชันสามารถปรับปรุงความเร็วและลดการใช้หน่วยความจำได้อย่างมากโดยไม่สูญเสียความแม่นยำ

เมื่อเราสำรวจในบทความนี้ ผลกระทบของแฟลช แอทเทนชันไม่ได้จำกัดอยู่แค่การเพิ่มประสิทธิภาพเท่านั้น แต่ยังช่วยให้สามารถพัฒนาโมเดลที่มีประสิทธิภาพและใช้หน่วยความจำได้ดีขึ้น