Google کے Gemma-4 31B ماڈل کو AWS Inferentia2 inf2.24xlarge پر پورٹ کرنے کے نتیجے میں CPU ریفرنس کے مقابلے میں بالکل درست ٹوکن-فار-ٹوکن میچ حاصل ہوا—لیکن تیار کردہ ہر جملہ بے معنی (gibberish) تھا۔ "میچنگ" اور "کام کرنے" کے درمیان یہ فرق اب ان تمام لوگوں کے لیے ایک انتباہ ہے جو بڑے LLMs کو Amazon کے کسٹم انفرنس چپس پر چلانے کی کوشش کر رہے ہیں۔

ٹوکن-بائی-ٹوکن میچ کیوں کافی نہیں ہے

ڈویلپر نے Inferentia ڈیوائس سے حاصل ہونے والے ہر آؤٹ پٹ ٹوکن کا موازنہ ماڈل کے CPU رن سے حاصل ہونے والے ٹوکن سے کیا۔ دونوں اسٹریمز بالکل ایک جیسی تھیں، اس لیے ایسا لگا کہ ہارڈ ویئر نے ریفرنس امپلیمنٹیشن کو بالکل درست طریقے سے دوبارہ تخلیق کیا ہے۔ حقیقت میں، دونوں اسٹریمز نے ایک ایسے ماڈل میں غلط فارمیٹ شدہ پرامپٹ (malformed prompt) بھیجا جس کا چیٹ ٹیمپلیٹ (chat template) ہٹا دیا گیا تھا اور اسے غلط ٹرن مارکرز (turn markers) فراہم کیے گئے تھے۔ ٹیمپلیٹ کی کمی نے ماڈل کو ایک انفینٹ لوپ (infinite loop) میں ڈال دیا، جس سے بے معنی نتائج نکلنے لگے۔ ہارڈ ویئر نے اپنا کام کیا—اس نے اس بگ (bug) کو دوبارہ تخلیق کیا جو ریفرنس کوڈ میں پہلے سے موجود تھا۔

سبق سادہ ہے: SEQ_MATCH (سیکوئینشل ٹوکن برابری) کا مطلب درستگی نہیں ہے۔ اگر ریفرنس امپلیمنٹیشن خراب ہے، تو ہارڈ ویئر کا وفادار نقل بھی اسی ناکامی کو اپنا لے گا۔ ویلیڈیشن کو ٹوکن لیول کی برابری سے آگے بڑھنا چاہیے؛ اس کے لیے درست فارمیٹ شدہ ان پٹس کے ساتھ اینڈ-ٹو-اینڈ فنکشنل چیکس کی ضرورت ہے۔

پیرامیٹرز کے روپ میں چھپے ہوئے بفرز (Buffers)

لوڈنگ کے مرحلے کے دوران، ماڈل لوڈر نے layer_scalar نامی ایک جزو کو چھوڑ دیا۔ کوڈ نے PyTorch ماڈل کی تعریف میں اس آبجیکٹ کو پیرامیٹر کے بجائے ایک بفر (buffer) کے طور پر رجسٹر کیا تھا۔ بفرز اسٹیٹک ٹینسرز (static tensors) ہوتے ہیں جنہیں ٹریننگ کے دوران اپ ڈیٹ نہیں کیا جاتا، اور بہت سے لوڈرز Neuron-compatible فارمیٹس میں تبدیل کرتے وقت انہیں نظر انداز کر دیتے ہیں۔ اسے چھوڑ دینے سے کئی تہوں (layers) کے اسکیلنگ فیکٹرز (scaling factors) ان کی ڈیفالٹ ویلیوز پر رہ گئے، جس سے پورے نیٹ ورک میں ریاضیاتی حساب کتاب بگڑ گیا۔ کوئی ایرر (error) سامنے نہیں آیا؛ ماڈل کمپائل ہو گیا، اور انفرنس پائپ لائن بھی چل پڑی، لیکن عددی نتائج غلط تھے۔

جو کوئی بھی بڑے ماڈلز کو Inferentia پر منتقل کر رہا ہے، وہ ہر نان-پیرامیٹر ٹینسر (non-parameter tensor) کا آڈٹ کرے۔ اگرچہ کسی ٹینسر کو سیکھنے (learn) کے لیے نہیں بنایا گیا، پھر بھی وہ درست فارورڈ-پاس کمپیوٹیشن کے لیے ضروری ہو سکتا ہے۔ بفر کے شامل ہونے کی دستی طور پر تصدیق کرنے سے خاموش اسکیلنگ کی غلطیوں (silent scale errors) سے بچا جا سکتا ہے جن کی تشخیص کرنا مشکل ہوتا ہے۔

اسپاٹ انسٹنس کی غیر یقینی صورتحال اور 39 منٹ کا کمپائل

اسپاٹ انسٹنس (spot instance) پر 31 بلین پیرامیٹر والا ماڈل چلانا سستا لگتا ہے، لیکن اس بچت کے ساتھ غیر متوقع ریکلیم ایونٹس (reclaim events) کا خطرہ بھی ہوتا ہے۔ ڈویلپر کا کمپائل ٹائم—ماڈل کو Neuron-compatible کوڈ میں تبدیل کرنے کے لیے تقریباً 39 منٹ—تب ختم ہو گیا جب AWS نے انسٹنس واپس لے لیا۔ رکاوٹوں سے بچنے کے لیے انہوں نے تین پہلوؤں پر مشتمل ایک حفاظتی جال بنایا:

  • ModelBuilder نے میموری کے استعمال کو 384 GB ہوسٹ کی حد کے اندر رکھا، تاکہ کریشز سے بچا جا سکے جو دوبارہ اسٹارٹ کرنے پر مجبور کر دیتے ہیں۔
  • خام وزن کی فائلوں (raw weight files) اور کمپائل شدہ "neffs" (Neuron executable files) دونوں کی فوری S3 mirroring نے ایک نئے انسٹنس کو بالکل وہیں سے شروع کرنے کی اجازت دی جہاں پچھلا انسٹنس رکا تھا۔
  • ایک multi-region poller نے دستیاب اسپاٹ کیپیسٹی کے لیے AWS ریجنز کو اسکین کیا اور جیسے ہی کوئی دستیاب ہوا، نیا انسٹنس لانچ کر دیا۔

ان اقدامات نے ایک نازک، سنگل پوائنٹ کمپائل کو ایک لچکدار پائپ لائن میں بدل دیا جو اسپاٹ مارکیٹ کے اتار چڑھاؤ میں بھی برقرار رہتی ہے۔

مکسڈ اٹینشن لے آؤٹس کے ساتھ شارڈنگ (Sharding) کے خطرات

Gemma-4 31B دو اٹینشن کنفیگریشنز استعمال کرتا ہے۔ کچھ تہیں (layers) چار کی-وی (KV) ہیڈز استعمال کرتی ہیں، جبکہ دیگر کی تعداد مختلف ہوتی ہے۔ ماڈل کو آٹھ متوازی رینکس (parallel ranks) میں برابر تقسیم کرنا اس وقت ناکام ہو جاتا ہے جب کسی تہہ کے KV ہیڈ کی تعداد آسانی سے تقسیم نہ ہو۔ ایک 4-ہیڈ والی تہہ کو آٹھ رینکس میں تقسیم کرنے کی کوشش کرنے سے ہر رینک کو آدھے ہیڈ کو سنبھالنا پڑے گا—جو کہ ایک ریاضیاتی طور پر ناممکن کام ہے اور اس سے شیپ مِس میچ (shape mismatches) اور رن ٹائم ایررز پیدا ہوتے ہیں۔

اس کا حل یہ تھا کہ گلوبلی-شارڈڈ تہوں (وہ تہیں جن کے ہیڈ کاؤنٹ مطابقت رکھتے ہوں) کو تمام رینکس پر ریپلیکیٹ (replicate) کیا جائے اور صرف ان "سلائیڈنگ" تہوں کو شارڈ کیا جائے جن کے ہیڈ کاؤنٹ کی تقسیم ممکن ہو۔ اس ہائبرڈ حکمت عملی نے KV ہیڈز کی غیر قانونی تقسیم سے بچتے ہوئے ٹینسر-پیرل ایفیشنسی (tensor-parallel efficiency) کو برقرار رکھا، اور ان ٹینسر-پیرلائزیشن ایررز کو ختم کر دیا جنہوں نے پچھلی کوششوں کو ناکام بنایا تھا۔

خلاصہ (Takeaway)

ایک بڑے LLM کو Inferentia پر پورٹ کرنا محض "کمپائل اور رن" کرنے کا عمل نہیں ہے۔ اس کے لیے ٹوکن برابری سے ہٹ کر سخت فنکشنل ٹیسٹنگ، اس بات کی باریک بینی سے تصدیق کہ ہر ٹینسر—چاہے وہ پیرامیٹر ہو یا بفر—درست طریقے سے ہینڈل کیا گیا ہے، اور ایک ایسی ڈیپلائمنٹ حکمت عملی کی ضرورت ہے جو اسپاٹ انسٹنس کی واپسی (reclamation) کا اندازہ لگا سکے۔ آخر میں، شارڈنگ کو ماڈل کی اندرونی اٹینشن جیومیٹری کا احترام کرنا چاہیے؛ ورنہ، وہ پیرا للیزم (parallelism) جو رفتار کا وعدہ کرتی ہے، خاموش ناکامی کا ذریعہ بن جاتی ہے۔