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 সমর্থিত নয়। Neuron-এ শুধু কিছু torch.fft ট্রান্সফর্ম চলে, এবং 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-এর লক্ষ্য ও প্রাথমিক ব্যবহারকারীদের বর্ণনা। এর মাইগ্রেশন, কর্মক্ষমতা ও অর্থনৈতিক দাবিগুলো প্রতিটি কাজের ভিত্তিতে যাচাই করা দরকার।