AWS کے Neuron 2.33 کے ریلیز نوٹس میں TorchNeuron کے عوامی بیٹا کی تاریخ 19 اکتوبر دی گئی ہے۔ 19 اکتوبر ہی کی تاریخ والا PyTorch کا اعلان کہتا ہے کہ عوامی بیٹا اور اوپن سورس کا اجرا 20 اکتوبر کو دستیاب ہوگا۔ torch-neuronx پیکیج Trainium2 اور Trainium3 کے لیے PyTorch میں ایک neuron ڈیوائس رجسٹر کرتا ہے، جس سے ماڈل اور ٹینسر کی معاون کارروائیاں مانوس PyTorch API استعمال کر سکتی ہیں۔ ان چپس پر غور کرنے والے انجینئروں کے لیے مفید تبدیلی یہ ہے کہ موجودہ PyTorch پروگرام سے ڈیوائس پر آزمائش تک کا راستہ مختصر ہو جاتا ہے۔ پروگرام کا درست اور مؤثر طور پر چلنا اب بھی اس کی کارروائیوں، ٹینسر کی شکلوں اور سافٹ ویئر ماحول پر منحصر ہے۔

Trainium پر PyTorch کا پچھلا راستہ XLA پر مبنی torch-neuronx لائبریری اور torch_xla ڈیوائس کوڈ استعمال کرتا تھا۔ AWS کے مطابق نیٹو بیک اینڈ PyTorch کی PrivateUse1 ایکسٹینشن استعمال کرتا ہے اور xm.xla_device() کی جگہ neuron ڈیوائس پر تعیناتی لاتا ہے۔ یہ پہلے ہی Neuron استعمال کرنے والے منصوبوں کے لیے انٹرفیس میں ایک اہم تبدیلی ہے، اگرچہ پیکیج کا نام مانوس ہے۔ یہ ریلیز Trn2 اور Trn3 انسٹینسز پر تربیت اور انفرنس کے لیے ہے؛ معاونت کی جدول میں Trn1 درج نہیں۔

PyTorch پروگرام میں کیا بدلتا ہے

AWS کے ریلیز نوٹس میں ایک ماڈل کو .to("neuron") کے ذریعے منتقل کرنے اور ان پٹ ٹینسر کو device="neuron" پر بنانے کی مثال ہے۔ اس کے بعد torch.nn، autograd اور آپٹیمائزر کی معاون کارروائیاں بیک اینڈ کے ذریعے چلتی ہیں۔ PyTorch کی ڈیوائس سے غیر وابستہ torch.accelerator API کے لیے لکھا گیا کوڈ کسی خاص ڈیوائس کے لیے دوبارہ لکھے بغیر Neuron دریافت کر سکتا ہے، جیسا کہ TorchNeuron کے تعارف میں بتایا گیا ہے۔

عمل درآمد کے دو راستے ہیں۔ eager موڈ، Python کی طرف سے بلائے جانے پر PyTorch کی کارروائیاں انجام دیتا ہے اور معمول کے کنٹرول فلو، پرنٹ آؤٹ پٹ اور بریک پوائنٹس کو برقرار رکھتا ہے۔ ہر کارروائی پہلی بار استعمال پر کمپائل ہوتی ہے اور بعد کی مماثل کالز کیش استعمال کر سکتی ہیں۔ انجینئر کسی ماڈل یا منتخب حصے کو torch.compile(model, backend="neuron", dynamic=False) سے بھی کمپائل کر سکتے ہیں۔ نیٹو بیک اینڈ معاون فارورڈ اور بیک ورڈ گراف کمپائل کرتا ہے؛ کیپچر کیے گئے حصے سے باہر کا کوڈ eager موڈ میں چلتا رہتا ہے۔ PyTorch کے اعلان میں کمپائل کرنے سے پہلے ایک بار torch.compiler.set_default_backend("neuron") مقرر کرنے کا بھی ذکر ہے تاکہ بعد کی کمپائل کالز میں بیک اینڈ آرگیومنٹ چھوڑا جا سکے۔

ریلیز میں ایک neuron بیک اینڈ شامل ہے جو torch.distributed کے لیے ہے، اور اس کی collective اور point-to-point کارروائیاں دستاویزی صورت میں موجود ہیں۔ AWS نے معاون API میں DistributedDataParallel، FSDP2، DTensor، DeviceMesh اور پائپ لائن پیراللزم کو درج کیا ہے۔ تقسیم شدہ تربیت کے منصوبے میں ان راستوں کا جائزہ لینا مفید ہے، لیکن ریلیز نوٹس یہ ثابت نہیں کرتے کہ موجودہ کئی نوڈز والا کام بغیر تبدیلی چلے گا یا تھروپٹ کا کوئی خاص ہدف حاصل کرے گا۔

جائزے کا آغاز موزوں ماحول سے ہوتا ہے

AWS کی تنصیب کی فوری رہنما دستاویز کے مطابق ایک Trn2 یا Trn3 انسٹینس اور Neuron Base DLAMI یا معاون Ubuntu AMI درکار ہے۔ ریلیز کی جدول Ubuntu 24.04 یا 26.04 پر Linux x86_64 کو متعین کرتی ہے۔ اس ریلیز کے لیے Ubuntu 22.04 کی معاونت ختم ہو چکی ہے۔ اس میں Python 3.12، 3.13 اور 3.14 درج ہیں، اور PyTorch 2.13.0 کے لیے torch-neuronx 2.13.3.0 یا PyTorch 2.14.0 کے لیے torch-neuronx 2.14.3.0 بتایا گیا ہے۔ پیکیج کا خاندان PyTorch کے ورژن سے ملنا چاہیے۔

فوری رہنما Neuron اور PyTorch کے پیکیج انڈیکس سے torch-neuronx==2.14.* کے ساتھ neuronx-cc==2.* اور nki نصب کرتی ہے۔ پھر ڈیوائس کی دستیابی جانچتی ہے اور Neuron پر ایک چھوٹی میٹرکس کارروائی کا نتیجہ CPU کے نتیجے سے ملاتی ہے۔ یہ AWS کی مثال میں نصب شدہ ڈیوائس کا راستہ جانچتا ہے۔ اس سے قاری کے اپنے ماڈل کی توثیق نہیں ہوتی۔ اگلا معقول جائزہ یہ ہے کہ اصل ماڈل کے فارورڈ اور بیک ورڈ راستے چلائے جائیں، torch_neuronx.get_fallback_ops() سے ریکارڈ شدہ CPU fallback دیکھے جائیں، موجودہ ماحول سے آؤٹ پٹ کا موازنہ کیا جائے اور کمپائلیشن اور مستحکم حالت میں چلنے کے اوقات الگ ناپے جائیں۔ یہ آخری اقدامات BIG CHANGE کے تجویز کردہ جائزے کے سوالات ہیں، ہمارے کیے ہوئے ٹیسٹ نہیں۔

بیٹا اب بھی کہاں کام کو بدل سکتا ہے

eager موڈ میں غیر معاون کارروائی CPU پر واپس جا سکتی ہے اور انتباہ دے سکتی ہے۔ torch.compile کے تحت، گراف میں کیپچر کی گئی ایسی کارروائی جسے Neuron بیک اینڈ نچلی سطح کی صورت میں تبدیل نہ کر سکے، ڈیفالٹ طور پر کمپائلیشن ناکام کر دیتی ہے۔ ہر کارروائی کی معاونت جاننے کے لیے AWS ڈویلپرز کو اپنی آپریٹر مطابقت کی جدول کی طرف بھیجتا ہے۔ لہٰذا eager موڈ میں کسی رن کا مکمل ہو جانا بذات خود یہ ثابت نہیں کرتا کہ تمام عمل ڈیوائس پر ہوا یا کمپائل شدہ صورت میں مطابقت موجود ہے۔

ریلیز یہ بھی کہتی ہے کہ torch.compile کے لیے dynamic=False درکار ہے؛ ان پٹ کی نئی شکل ایک اور کمپائلیشن کا تقاضا کر سکتی ہے۔ AWS شکلوں میں تنوع کم کرنے کے لیے padding یا bucketing تجویز کرتا ہے۔ جب مماثل کمپائل شدہ آرٹیفیکٹ کیش میں موجود نہ ہو تو پہلے رن میں کمپائلیشن شامل ہوتی ہے۔ کسی ورک لوڈ کا وقت ناپتے ہوئے یہ تفصیلات اہم ہیں: کمپائلیشن کی تاخیر اور ان پٹ کی شکلوں کا انداز جائزے کا نتیجہ بدل سکتے ہیں۔

دیگر واضح حدود مطابقت کے دعوے کا دائرہ مزید محدود کرتی ہیں۔ بیٹا torch.sparse، torch.func، torch.nested یا torch.distributions کی معاونت نہیں کرتا۔ صرف کچھ torch.fft ٹرانسفارمز Neuron پر چلتے ہیں، اور float64 انتباہ کے ساتھ float32 کے طور پر چلتا ہے۔ اس ریلیز میں TorchAO انٹیگریشن اور مرکب MXFP8 ڈیٹا ٹائپ کی معاونت موجود نہیں۔ یہ سہولتیں استعمال کرنے والی ٹیم کو ڈیوائس پر تعیناتی کو منتقلی کا منصوبہ سمجھنے سے پہلے اپنے کوڈ کا راستہ جانچنا چاہیے۔

بڑی تبدیلی

TorchNeuron عوامی بیٹا میں Trainium کو PyTorch کے نیٹو ڈیوائس اور کمپائلیشن انٹرفیس کے پیچھے رکھتا ہے۔ اس سے معاون PyTorch پروگراموں کی پہلی آزمائش زیادہ براہ راست ہو جاتی ہے اور موجودہ Neuron صارفین کو اپنے XLA پر مبنی کوڈ سے مختلف راستہ ملتا ہے۔ اپنانے کے لیے فیصلہ کن شواہد ہر ورک لوڈ سے آئیں گے: کیا اس کے آپریٹر مطلوبہ ڈیوائس پر رہتے ہیں، کیا اس کے عددی نتائج قابل قبول رہتے ہیں، اور کیا ماپی گئی کارکردگی اور لاگت انسٹینس اور منتقلی کی محنت کا جواز بنتی ہیں۔ BIG CHANGE نے اس تجزیے کے لیے دستاویزات اور اعلانات کا جائزہ لیا؛ ہم نے TorchNeuron نہیں چلایا اور نہ AWS کے کارکردگی کے دعووں کی آزادانہ توثیق کی۔

ذرائع اور مزید مطالعہ

  • AWS Neuron 2.33 کے TorchNeuron ریلیز نوٹس: اجرا کی تاریخ، بیٹا کی حیثیت، ورژن اور آپریٹنگ سسٹم کی جدول، API اور معلوم حدود۔
  • AWS کا TorchNeuron تعارف: ڈیوائس انٹیگریشن، eager اور کمپائل شدہ عمل درآمد، CPU fallback کا رویہ اور Neuron سے گزرنے والا سافٹ ویئر راستہ۔
  • AWS کی تنصیب اور توثیق کی فوری رہنما: انسٹینس اور امیج کی پیشگی شرائط، پیکیج کی تنصیب اور دستاویزی ڈیوائس جانچ۔
  • PyTorch بلاگ: TorchNeuron کا تعارف: ڈیوائس بدلنے، کمپائلیشن اور تقسیم شدہ API کی PyTorch کی طرف سے وضاحت۔ یہ مضمون PyTorch بلاگ پر AWS کے معاونین کا اعلان ہے، آزاد بینچ مارک نہیں۔
  • Amazon کا اعلان: اہداف اور ابتدائی صارفین کے بارے میں AWS کا بیان۔ منتقلی، کارکردگی اور اقتصادی فائدے کے دعووں کو ہر ورک لوڈ کے لیے الگ جانچنا ضروری ہے۔