FlashAttention کی پیمانہ بندی، کثیر GPU تربیت اور وزنوں کی کوانٹائزیشن

جدید AI بنیادی ڈھانچا بڑے ماڈلوں کو مؤثر انداز میں تربیت دینے اور چلانے کے لیے ہارڈویئر سے آگاہ اٹینشن کرنلز، متعدد GPU پر تقسیم شدہ متوازی عمل، اور کم بِٹ کوانٹائزیشن پر انحصار کرتا ہے۔ یہ تین تکنیکیں ماڈل کے نظریاتی پیمانے اور GPU کی طبعی حدود کے درمیان خلا پُر کرتی ہیں۔

1. Hopper (H100) GPUs پر FlashAttention-2 اور FlashAttention-3

اگرچہ اصل FlashAttention نے میموری بینڈوڈتھ کی رکاوٹ دور کر دی تھی، مگر جدید GPUs کی خام حسابی کارکردگی کا خاصا حصہ اب بھی غیر استعمال شدہ تھا۔ FlashAttention-2 (2023) اور FlashAttention-3 (2024) نے GPU تھریڈز اور ہارڈویئر انجنوں کے باہمی تال میل کا طریقہ بدل دیا۔

• FlashAttention-2 (کام کی بہتر تقسیم)۔

• غیر Matmul آپریشنز میں کمی۔ GPU کے Tensor Cores میٹرکس ضرب (matmul) کو اس رفتار سے تقریباً 16× تیز انجام دیتے ہیں جس سے معیاری CUDA کورز Softmax کی دوبارہ اسکیلنگ جیسے اسکیلر آپریشنز کرتے ہیں۔ FlashAttention-2 نے آن لائن Softmax کی ریاضی کو ازسرنو لکھا تاکہ دوبارہ اسکیلنگ اندرونی لوپ کے بالکل آخر تک مؤخر ہو، یوں سست غیر matmul ہدایات کم ہو گئیں۔
• سیکوئنس کی لمبائی میں متوازی عمل۔ اصل ورژن کام کو بنیادی طور پر بیچ سائز اور اٹینشن ہیڈز میں متوازی بناتا تھا، جس سے لمبے سیکوئنس اور چھوٹے بیچ سائز کی تربیت میں GPU کورز فارغ رہتے تھے۔ FlashAttention-2 سیکوئنس کی لمبائی والی جہت میں بھی متوازی عمل شامل کرتا ہے اور لوپوں کی ترتیب الٹ دیتا ہے (بیرونی لوپ Q بلاکس پر، اندرونی لوپ K اور V پر)، جس سے سست مشترکہ میموری میں لکھنا کم سے کم ہوتا ہے۔

• FlashAttention-3 (Hopper H100 ہارڈویئر کا استعمال)۔

• غیر ہم وقت ڈیٹا کی منتقلی (TMA)۔ NVIDIA کے Hopper آرکیٹیکچر نے Tensor Memory Accelerator (TMA) متعارف کرایا، جو ایک مخصوص ہارڈویئر یونٹ ہے اور حسابی تھریڈز کو مصروف کیے بغیر عالمی HBM اور چِپ پر موجود SRAM کے درمیان ڈیٹا نقل کرتا ہے۔
• وارپ کی تخصیص (پروڈیوسر-کنزیومر ماڈل)۔ FlashAttention-3 GPU وارپس (32 تھریڈز کے گروپ) کو مخصوص کرداروں میں بانٹتا ہے۔ پروڈیوسر وارپس TMA کے ذریعے K اور V کی آئندہ ٹائلیں پیشگی SRAM میں لاتے ہیں، جبکہ کنزیومر وارپس موجودہ ٹائلیں Tensor Cores کو دیتے ہیں (WGMMA ہدایات)۔ یوں میموری کی منتقلی، میٹرکس کی ریاضی اور Softmax کے حسابات بیک وقت ایک دوسرے کے ساتھ چلتے ہیں۔
• FP8 کم دقت کی معاونت۔ FlashAttention-3 Hopper کے FP8 Tensor Cores کو استعمال کے قابل بناتا ہے، جس سے FP16 کے مقابلے میں خام تھروپٹ دوگنا ہو جاتا ہے۔ FP8 میں غیر معمولی ایکٹیویشنز سے دقت کے نقصان کو روکنے کے لیے یہ بلاک کوانٹائزیشن (چھوٹی ٹائلوں کی الگ الگ اسکیلنگ) اور غیر مربوط پروسیسنگ (کوانٹائز کرنے سے پہلے انتہائی قدروں کو ہموار کرنے کے لیے Q اور K کو اتفاقی متعامد میٹرکس سے ضرب دینا) استعمال کرتا ہے۔

2. کثیر GPU پیمانہ بندی کی حکمت عملیاں (3D متوازی عمل)

جب کسی ماڈل کے پیرامیٹرز کی تعداد دسیوں یا سینکڑوں ارب ہو، تو اس کے وزن، گریڈینٹس اور آپٹمائزر کی حالتیں ایک GPU کی VRAM سے بڑھ جاتے ہیں۔ بڑے پیمانے کی تربیت متوازی بنانے کی تین الگ حکمت عملیوں کو یکجا کرتی ہے، جنہیں اکثر 3D متوازی عمل کہا جاتا ہے۔

• ڈیٹا متوازی عمل & ZeRO / FSDP۔ معیاری ڈیٹا متوازی عمل میں ہر GPU ماڈل کی مکمل نقل رکھتا ہے، جس سے میموری ضائع ہوتی ہے۔ جدید فریم ورکس ZeRO (Zero Redundancy Optimizer) یا PyTorch FSDP (Fully Sharded Data Parallel) استعمال کرتے ہیں، جو پیرامیٹرز، گریڈینٹس اور آپٹمائزر کی حالتوں کو تمام GPUs میں حصوں کی صورت بانٹتے ہیں اور ہر پرت کے وزن حساب کے عین وقت پر جمع کرتے ہیں۔

• ٹینسر متوازی عمل (پرت کے اندر)۔ بہت بڑی فیڈ فارورڈ میٹرکس ضرب Y = XW کو دہرانے کے بجائے ٹینسر متوازی عمل W کو کالموں کے لحاظ سے GPUs میں بانٹتا ہے (W = [W_1, W_2])۔ ہر GPU بیک وقت آؤٹ پٹ کے اپنے نصف کا حساب کرتا ہے اور تیز رفتار NVLink (H100s پر 900 GB/s) کے ذریعے جزوی نتائج کا تبادلہ کرتا ہے۔

• پائپ لائن متوازی عمل (پرتوں کے درمیان)۔ پرتیں 1–8، GPU 1 پر، پرتیں 9–16، GPU 2 پر ہوتی ہیں، اور اسی طرح آگے۔ بعد کے GPUs کو پہلے کی پرتیں مکمل ہونے کے انتظار میں فارغ بیٹھنے سے روکنے کے لیے ("پائپ لائن ببل") فریم ورکس ہر تربیتی بیچ کو چھوٹے مائیکرو بیچز میں بانٹتے ہیں، تاکہ تمام GPUs اسمبلی لائن کی طرح مصروف رہیں۔

3. کوانٹائزیشن (GPTQ، AWQ) اور موقع پر ڈی کوانٹائزیشن کرنے والے کرنلز

انفرنس کے دوران (ایک ایک کر کے ٹوکن بناتے وقت) بڑے لسانی ماڈلوں کی حد میموری بینڈوڈتھ ہوتی ہے، حسابی قوت نہیں۔ صرف ایک ٹوکن بنانے کے لیے VRAM سے 70 ارب 16-bit پیرامیٹرز (140 GB) لوڈ کرنا Tensor Cores کو میموری کی منتقلی کے انتظار میں رکھتا ہے۔ کوانٹائزیشن وزنوں کو 16-bit فلوٹس سے 4-bit صحیح اعداد (INT4) میں سکیڑتی ہے، جس سے میموری کا حجم اور منتقلی کا وقت تقریباً 4× کم ہوتے ہیں۔

• GPTQ (Hessian کیلیبریشن کے ذریعے تربیت کے بعد کوانٹائزیشن)۔

• 16-bit وزنوں کو سادہ انداز میں 4-bit صحیح اعداد پر گول کرنا ماڈل کی درستگی تباہ کر دیتا ہے۔ GPTQ ایک چھوٹا کیلیبریشن ڈیٹا سیٹ ماڈل میں سے گزارتا ہے اور دوسرے درجے کی ریاضی (معکوس Hessian میٹرکس) سے ناپتا ہے کہ ہر وزن کتنا حساس ہے۔
• یہ وزنوں کو کالم بہ کالم کوانٹائز کرتا ہے؛ جب بھی کسی وزن کو گول کرنے سے خطا پیدا ہوتی ہے، GPTQ اس قطار کے باقی غیر کوانٹائز شدہ وزنوں کو ریاضیاتی طور پر ایڈجسٹ کرتا ہے تاکہ خطا کا ازالہ ہو سکے۔

• AWQ (ایکٹیویشن سے آگاہ وزن کوانٹائزیشن)۔

• محققین نے دریافت کیا کہ وزنی چینلز کے تقریباً 0.1% سے 1% تک باقیوں سے کہیں زیادہ اہم ہوتے ہیں، کیونکہ وہ بڑی ایکٹیویشن قدروں سے ضرب کھاتے ہیں۔
• ان اہم وزنوں کو 16-bit میں رکھنے کے بجائے (جس سے بے ترتیب اور سست مخلوط دقت والے میموری لے آؤٹ بنتے ہیں)، AWQ کوانٹائزیشن سے پہلے ان نمایاں وزنی چینلز کا اسکیل بڑھاتا ہے اور رن ٹائم پر متعلقہ ان پٹ ایکٹیویشنز کا اسکیل گھٹاتا ہے، یوں 100% وزنوں کو یکساں 4-bit فارمیٹ میں رکھتے ہوئے درستگی محفوظ رہتی ہے۔

• حسب ضرورت CUDA کرنلز موقع پر ڈی کوانٹائزیشن کیسے کرتے ہیں (W4A16)۔

• GPUs کسی 16-bit ایکٹیویشن (FP16) کو براہ راست 4-bit صحیح عددی وزن (INT4) سے ضرب نہیں دے سکتے۔ اس کے بجائے مخصوص انفرنس کرنلز (جیسے Marlin یا ExLlama) W4A16 (4-bit وزن، 16-bit ایکٹیویشنز) عمل انجام دیتے ہیں:
• 1. سکیڑے ہوئے HBM کو پڑھنا: کرنل سست عالمی VRAM سے پیک شدہ 4-bit وزن (ایک 32-bit صحیح عدد میں پیک کیے گئے آٹھ INT4 اعداد) لوڈ کرتا ہے، جس سے میموری بینڈوڈتھ کا استعمال 4× کم ہو جاتا ہے۔
• 2. رجسٹر کی سطح پر کھولنا: GPU کے انتہائی تیز رجسٹرز کے اندر CUDA کرنل تیز بٹ وائز شفٹس (>>) اور بٹ ماسکس (&) سے INT4 صحیح اعداد کھولتا ہے اور محفوظ اسکیل/زیرو پوائنٹ عوامل لگا کر انہیں اسی وقت واپس FP16 قدروں میں بدلتا ہے۔
• 3. فوری Tensor Core حساب: ابھی کھولے گئے FP16 وزن، FP16 ایکٹیویشنز کے ساتھ براہ راست Tensor Cores کو دے دیے جاتے ہیں، اور پھیلائے گئے 16-bit وزن کبھی دوبارہ VRAM میں نہیں لکھے جاتے۔