একটি নতুন গবেষণা, যার নাম CODA: Rewriting Transformer Blocks as GEMM-Epilogue Programs, যেখানে ট্রান্সফরমার মডেল প্রশিক্ষণের দক্ষতা উন্নত করার লক্ষ্যে, বিশেষ করে যেসব প্রতিকূল কিন্তু সঞ্চিতভাবে সময়সাপেক্ষ “মেমোরি-ঘনত্বপূর্ণ” অপারেশনগুলির সমস্যা সমাধান করা হয়েছে।

লেখক, উৎস: মেশিন ইন্টেলিজেন্স

মে ২২ তারিখে, ট্রি ডাও সামাজিক মাধ্যমে হান গুওর একটি টুইট শেয়ার করেন। তিনি আরও লিখেছেন: "কিছু গাণিতিক পুনর্লিখনের পর, পাওয়া গেল যে Transformer-এর সমস্ত কিছুই একটি GEMM + epilogue (ম্যাট্রিক্স গুণন এবং এপিলগ) এর ধারাবাহিকতা। কিছু অপ্টিমাইজড প্রিমিটিভ প্রদান করলে, LLM (এবং নতুনরা) সমস্ত Transformer অপারেশনের জন্য আলোকবেগের কার্নেল লিখতে পারবে!"

ট্রি ডাও ফ্ল্যাশঅ্যাটেনশন সিরিজের একজন প্রধান লেখক, এবং এই টুইটটি তাদের ঐ দিন প্রকাশিত একটি পেপার: CODA-এর দিকে ইঙ্গিত করে।

• শিরোনাম: CODA: ট্রান্সফরমার ব্লকগুলিকে GEMM-এপিলগ প্রোগ্রাম হিসাবে পুনর্লিখন
• পেপারের ঠিকানা: https://arxiv.org/abs/2605.19269
• কোড ঠিকানা: https://github.com/HanGuo97/coda-kernels

এই নামটি পড়তে শোনায় "ফাইনাল" এবং উচ্চারণ করলে শোনায় "CUDA"। MIT, প্রিনসটন, Together AI এবং Meta-এর গবেষকরা ট্রান্সফরমার প্রশিক্ষণের সেই অল্প পরিচিত, কিন্তু সময় নষ্ট করে দেওয়া ছোট ছোট গণনাগুলিকে সিস্টেম্যাটিকভাবে শোষণ করার জন্য একটি নতুন প্রোগ্রামিং অ্যাবস্ট্রাকশন ব্যবহার করছেন।

বড় মডেল প্রশিক্ষণের "লাজুকি ট্যাক্স"

CODA কী সমস্যার সমাধান করছে তা বুঝতে, প্রথমে বুঝতে হবে যে বড় মডেল প্রশিক্ষণের সময় কোথায় যায়।

একটি নভিডিয়া H100 গ্রাফিক্স কার্ডে একটি 1B প্যারামিটার LLaMA-3 স্টাইল মডেল ট্রেন করার সময়, বেশিরভাগ মানুষ স্বাভাবিকভাবেই মনে করেন: সময়টা ম্যাট্রিক্স গুণন এবং অনুধাবন গণনায় ব্যয় হচ্ছে, কারণ এগুলোই 'আসল গণনা'। এই স্বাভাবিক ধারণা প্রায় সঠিক: ম্যাট্রিক্স গুণন (GEMM) এবং অনুধাবন প্রক্রিয়াটি প্রধান কম্পিউটেশনাল শক্তি নিয়ে নেয়।

কিন্তু যদি আপনি পারফরম্যান্স প্রোফাইলার খুলে সতর্কতার সাথে দেখেন, তবে দেখতে পাবেন যে একটি সেট ছোট অপারেটর নিঃশব্দে সময় খরচ করছে: নরমালাইজেশন (RMSNorm), এক্টিভেশন ফাংশন (SwiGLU, RoPE), রেসিডুয়াল যোগ, লেয়ার-ব্যাপী রিডাকশন... তাদের প্রতিটির গণনা বোঝা কম, কিন্তু তারা প্রায়শই বড় মধ্যবর্তী টেনসরগুলিকে জিপিইউ মেমোরির মধ্যে ঢুকিয়ে-বার করছে।

এটিকে বলা হয় “মেমোরি ব্যান্ডউইথ বটলনেক”: যেমন একজন অত্যন্ত দক্ষ রান্নাবান্না রাঁধুনি, যিনি প্রতিটি খাবার বানাতে উপকরণগুলি দূরের গুদাম থেকে আনতে হয়, ব্যবহারের পরে আবার ফিরিয়ে দিতে হয়, এবং হাতের কাছের টেবিলে রাখেন না। রান্নাবান্না করার গতি যতই দ্রুত হোক, পরিবহনের জন্য অপেক্ষা করার সময়টি বাস্তবিকভাবেই নষ্ট হয়।

আরও খারাপ হলো, নভিডিয়ার FP8, FP4 ইত্যাদি নিম্ন প্রেসিশন ফরম্যাটের কারণে ম্যাট্রিক্স গণনা আরও দ্রুত হচ্ছে, এই 'স্থানান্তর' অপারেশনগুলির আপেক্ষিক খরচ বাড়ছে: ম্যাট্রিক্স গুণন দ্রুত হচ্ছে, কিন্তু টেনসরগুলি ভিতরে-বাইরে স্থানান্তরের খরচ সমানুপাতিকভাবে কমছে না।

একটি গবেষণাপত্রে একটি স্পষ্ট ডেটা সেট রয়েছে: H100-এ TorchTitan ব্যবহার করে 1B প্যারামিটার মডেল প্রশিক্ষণের সময়, ম্যাট্রিক্স গুণন অপারেশন ছাড়া অন্যান্য অপারেশনগুলি এন্ড-টু-এন্ড রানটাইমের একটি উল্লেখযোগ্য অংশ অধিকার করে, এবং FP8 প্রিসিশনের প্রবর্তনের সাথে সাথে এই অনুপাতটি আরও বেড়ে যায়।

বর্তমান প্রোগ্রামিং ফ্রেমওয়ার্কগুলি এটির জন্য প্রায় অক্ষম। PyTorch Transformer-এর গণনাকে একটি অপারেটর সিরিজ হিসাবে প্রকাশ করে, যেখানে অপারেটরগুলির মধ্যে স্পষ্ট সীমানা রয়েছে। এই সীমানাগুলি অটোগ্র্যাড (autograd) এর জন্য অত্যন্ত বন্ধুত্বপূর্ণ, কিন্তু এগুলি ঠিক অপারেটরগুলির মধ্যে একীভূতকরণ অপ্টিমাইজেশনকে বাধা দেয়: প্রতিটি অপারেটর সীমানা, প্রায়শই একটি অপ্রয়োজনীয় VRAM রিটার্নকে নির্দেশ করে।

CODA: "শেষ" এর মধ্যে খনি লুকানো আছে

কোডা একটি সাধারণ পর্যবেক্ষণ থেকে শুরু হয়েছে।

একটি উচ্চ কার্যক্ষমতাবিশিষ্ট ম্যাট্রিক্স গুণন (GEMM) কোর GPU-এ দুটি অংশে বিভক্ত: মেইনলুপ (mainloop) মূল ম্যাট্রিক্স ব্লক গুণ এবং যোগের গণনা সম্পাদন করে, এবং এপিলোগ (epilogue) ফলাফল গ্রাফিক্স মেমরিতে লেখার আগে বাইয়াস যোগ, টাইপ কনভার্শন, সরল স্কেলিংয়ের মতো শেষ পরিশোধন কাজগুলি সম্পাদন করে।

শেষ পর্যায়ের অর্থ হলো, এই মুহূর্তে ম্যাট্রিক্স গুণনের আউটপুট এখনও চিপের রেজিস্টারে "জীবিত" আছে, এবং এখনও গ্লোবাল ভিডিও মেমোরিতে সংরক্ষিত হয়নি। এটি একটি সংক্ষিপ্ত সোনালি জানালা: যদি এই মুহূর্তে আরও কিছু গণনা করা যায়, তবে একবার ভিডিও মেমোরিতে লেখা এবং আবার পড়ার প্রয়োজনীয়তা সম্পূর্ণরূপে এড়ানো যাবে।

কোডার মূল অন্তর্দৃষ্টি হল: ট্রান্সফরমারের যেসব মেমোরি-ঘনত্বপূর্ণ অপারেশন রয়েছে, তাদের অনেকগুলিকে বীজগাণিতিকভাবে পুনঃপ্যারামিটাইজ করে এই «শেষ» উইন্ডোতে প্রক্রিয়া করা যায়।

এটি কিছুটা গাণিতিক দক্ষতা প্রয়োজন। সবচেয়ে সাধারণ GEMM-RMSNorm-GEMM প্যাটার্নটি নিয়ে কথা বলি: একটি ম্যাট্রিক্স গুণনের ফলাফল, যা রেসিডুয়াল যোগ, RMS নরমালাইজেশনের মাধ্যমে পাস হয়, এবং তারপর আরেকটি ম্যাট্রিক্স গুণন হয়। প্রচলিত পদ্ধতিতে, তিনটি স্বতন্ত্র অপারেটর ধারাবাহিকভাবে চালানো হয়, এবং মধ্যবর্তী ফলাফল দুইবার ডিসপ্লে মেমরিতে সংরক্ষিত হয়।

CODA টিম আবিষ্কার করেছে যে, RMS নরমালাইজেশনের প্রতিটি সারির স্কেলিং ফ্যাক্টর r, যেহেতু এটি প্রতিটি সারির জন্য শেয়ার করা স্কেলার, তাই এটি পরবর্তী ম্যাট্রিক্স গুণনের সাথে কমিউটেটিভ হয়: r-এর প্রয়োগকে «দ্বিতীয় GEMM-এর আগে» থেকে «দ্বিতীয় GEMM-এর শেষে» স্থানান্তরিত করা যায়। এই স্থানান্তরণের পর, প্রথম GEMM-এর শেষে শুধুমাত্র স্থানীয় «পার্টিয়াল RMS» গণনা করতে হয়, যা একটি অত্যন্ত হালকা সহায়ক রিডাকশন কার্নেল দ্বারা একীভূত হয়, এবং পূর্ণাঙ্গ RMSNorm গণনা অপসারণ করা হয়।

একই ধরনের পুনঃপ্যারামিটাইজেশন SwiGLU, RoPE (রোটেশনাল পজিশন এনকোডিং), ক্রস-এনট্রপি লস ইত্যাদি অপারেশনের জন্যও প্রযোজ্য, এবং ব্যাকওয়ার্ড প্রপাগেশনের জন্যও প্রযোজ্য। পেপারে একটি থিওরেম প্রমাণ করা হয়েছে: যদি ফরওয়ার্ড পাসের শেষাংশ "ব্লক-লোকাল" হয়, তবে ব্যাকওয়ার্ড প্রপাগেশন স্বয়ংক্রিয়ভাবে একই স্ট্রাকচারকে উত্তরাধিকারসূত্রে পায়। বিস্তারিত জানতে মূল পেপারটি দেখুন।

পাঁচটি "ব্লক" এবং একটি "লেগো ভাষা"

CODA একটি নির্দিষ্ট ফিউজড কার্নেল নয়, বরং একটি প্রোগ্রামিং অ্যাবস্ট্রাকশনের সেট।

এটি বিশেষজ্ঞদ্বারা অপ্টিমাইজ করা GEMM মূল লুপকে স্থির করে এবং শেষে পাঁচটি সংযোগযোগ্য মৌলিক প্রাইমিটিভকে প্রকাশ করে:

• প্রতিটি উপাদান পরিবর্তন (বাকি যোগ, সক্রিয়করণ ফাংশন, RoPE)
• ভেক্টর লোড এবং স্টোর (ব্রডকাস্ট RMSNorm ওজন)
• ম্যাট্রিক্স ব্লক লোড এবং স্টোর (ব্যাকপ্রোপাগেশনের জন্য মধ্যবর্তী এক্টিভেশন সংরক্ষণ করা)
• ব্লক রিডাকশন (লোকাল আরএমএস, ব্লক লগ-সাম-এক্সপ)
• স্টেটফুল ট্রান্সফর্মেশন (অনলাইন নরমালাইজেশনের জন্য ম্যাক্স এবং সাম-এক্সপ পরিসংখ্যান)

এই পাঁচ ধরনের ব্লক ব্যবহার করে, অনুসন্ধানের বাইরে একটি স্ট্যান্ডার্ড ট্রান্সফরমারের ফরওয়ার্ড এবং ব্যাকওয়ার্ড প্রপাগেশনের প্রায় সমস্ত অপারেশন কভার করা যায়।

আরও আকর্ষণীয় বিষয় হল এই বিমূর্ত পদ্ধতির কোড লেখার জন্য কারা দায়ী তা নিয়ে এর সহনশীলতা। গবেষণাপত্রটিতে দুটি বাস্তবায়ন পদ্ধতির মূল্যায়ন করা হয়েছে: একটি মানব প্রোগ্রামার দ্বারা লেখা, অন্যটি Claude Code দ্বারা তৈরি — CODA-এর প্রাইমিটিভ বর্ণনা, কিছু উদাহরণ এবং বাস্তবায়ন লগ প্রদান করে AI-এর দ্বারা প্রধান কোর কোডটি তৈরি করা হয়, মানবদের সামান্য পর্যবেক্ষণের সাথে।

উভয় মোডের পারফরম্যান্স উচ্চ স্তরে পৌঁছেছে। ট্রি ডাও টুইটে বলেছেন, "LLM এবং নতুনরাও লাইটনিং কার্নেল লিখতে পারে," যা প্রবন্ধের পরীক্ষার ফলাফলের বাস্তব প্রতিফলন।

পরীক্ষার ফলাফল

কোডা এর বেঞ্চমার্ক পছন্দ করেছে কঠোর প্রতিদ্বন্দ্বী: cuBLAS এবং torch.compile, এবং LLM-এর জন্য অপ্টিমাইজড Liger Kernel এবং FlashInfer।

প্রতিটি কোরের জন্য দুটি বাস্তবায়ন মূল্যায়ন করা হয়েছে: CODA (LLM), যা Claude Code দ্বারা উত্পাদিত, যেখানে গবেষকরা প্রাইমিটিভ বর্ণনা, কিছু উদাহরণ এবং একটি নিয়মিত আপডেট করা বাস্তবায়ন টিপস লগ প্রদান করেন, যাতে AI মূল কোডটি সম্পূর্ণ করে এবং মানুষ সামান্য পর্যবেক্ষণ করে; CODA (Human), যা মানুষের প্রোগ্রামারদের দ্বারা স্বতন্ত্রভাবে লেখা হয়েছে, যেখানে একই উচ্চস্তরের পুনঃপ্যারামিটারাইজেশন ধারণা ব্যবহার করা হয়েছে, কিন্তু CODA প্রাইমিটিভসের উপর নির্ভরশীল নয়। উভয় গ্রুপের ফলাফলকে cuBLAS + torch.compile, Liger Kernel, FlashInfer ইত্যাদি অপটিমাইজড লাইব্রেরির সাথে তুলনা করা হয়েছে।

একক অপারেটর স্তরে, GEMM-RMSNorm-GEMM এই প্রতিনিধিত্বকারী প্যাটার্নটির উদাহরণ হিসেবে, CODA 1B, 7B, এবং 70B মডেল সাইজগুলির জন্য লুকানো মাত্রাগুলিতে cuBLAS + PyTorch বেসলাইনকে ছাড়িয়ে গেছে। SwiGLU, RoPE, এনট্রপি ইত্যাদি শেষ কম্বিনেশনগুলিতেও অনুরূপ পারফরম্যান্স দেখা গেছে।

LLM দ্বারা তৈরি কর্ন অধিকাংশ বেঞ্চমার্কে ম্যানুয়ালি লেখা সংস্করণের সমান পারফরম্যান্স দেখায় এবং কিছু কনফিগারেশনে এটি সামান্য অতিক্রমও করে। এটি গতকাল থেকেই খুবই উচ্চ বাধা হিসেবে বিবেচিত GPU কর্ন অপ্টিমাইজেশনের ক্ষেত্রে একটি অত্যন্ত দুর্লভ উপলব্ধি।

রিভার্স প্রপাগেশনের উপকারিতা বিশেষভাবে উল্লেখযোগ্য: GEMM-Residual-PartialRMS-GEMM-এর রিভার্স কার্নেল বেসলাইনের তুলনায় 1.6 থেকে 1.8 গুণ পর্যন্ত গতি বৃদ্ধি পায়, এবং SwiGLU রিভার্সেও প্রায় 1.4 থেকে 1.6 গুণ উন্নতি দেখা যায়। এই দিকে, LLM এবং ম্যানুয়ালি বাস্তবায়িত মডেলের মধ্যে পার্থক্যও খুবই কম। এটি আশ্চর্যজনক নয়: রিভার্স প্রপাগেশন স্বাভাবিকভাবেই বেশি মধ্যবর্তী টেনসরের অ্যাক্সেসকে অন্তর্ভুক্ত করে, যার ফলে টেইল-ফিউশনের উপকারিতা বেশি; এবং CODA-এর প্রিমিটিভ ডিজাইন যথেষ্ট পরিষ্কার, যার ফলে AI মডেলগুলি সঠিকভাবে কম্বিনেশনটি সম্পন্ন করতে পারে।

পূর্ণ ট্রান্সফরমার স্তরের এন্ড-টু-এন্ড বেঞ্চমার্কে, কোডা এর ফরওয়ার্ড স্পিডআপ বিভিন্ন স্কেলে প্রায় 5% থেকে 20% এবং বড় মডেল সাইজে (70B স্কেলের হিডেন ডাইমেনশনের সাথে সম্পর্কিত) এটি আরও উল্লেখযোগ্য।

সংখ্যাগত সূক্ষ্মতার ক্ষেত্রে, CODA-এর পুনঃপ্যারামিটাইজেশন RMSNorm স্কেলিং ফ্যাক্টরের প্রয়োগের সময়কাল পরিবর্তন করে, তবে পরীক্ষা দেখায় যে এর সংখ্যাগত ত্রুটি PyTorch রেফারেন্স বাস্তবায়নের সমান, কিছু কনফিগারেশনে এটি এমনকি কম — GEMM মূল লুপটি নিজেই উচ্চতর সূক্ষ্মতার অ্যাকুমুলেটর বহন করার কারণে।

CODA কী করতে পারে: বড় দৃশ্যে প্রবেশ করার আগে CODA-এর ক্ষমতার সীমানা পরিষ্কার করুন।

• পরিসর: স্ট্যান্ডার্ড ট্রান্সফরমার (যেমন LLaMA আর্কিটেকচার) এর ফরওয়ার্ড এবং ব্যাকওয়ার্ড প্রপাগেশনে, অ্যাটেনশন এবং টোকেন এমবেডিং বাদে প্রায় সমস্ত গণনা, যা অন্তর্ভুক্ত করে RMSNorm, রেসিডুয়াল যোগ, SwiGLU এক্টিভেশন, RoPE রোটেশনাল পজিশন এনকোডিং, ক্রস-এনট্রপি লস, এবং উপরোক্ত অপারেশনগুলির ব্যাকওয়ার্ড গ্রেডিয়েন্ট গণনা।
• স্পিড বুস্ট: 1B থেকে 70B সাইজের হিডেন ডাইমেনশনের জন্য, cuBLAS + torch.compile বেসলাইনের তুলনায় একক অপারেটর লেভেলে বিভিন্ন মাত্রায় উন্নতি দেখা যায়, যার মধ্যে ব্যাকওয়ার্ড প্রপাগেশনের লাভ সবচেয়ে বেশি (কিছু কোর প্রায় 1.6x এর বেশি); পুরো Transformer লেয়ারের এন্ড-টু-এন্ড ফরওয়ার্ড স্পিড বুস্ট প্রায় 5% থেকে 20%, এবং বড় মডেল সাইজের ক্ষেত্রে এটি আরও উল্লেখযোগ্য।
• CODA, যা CuTeDSL (NVIDIA CUTLASS-এর Python DSL) এর উপর ভিত্তি করে তৈরি, মানব প্রোগ্রামার এবং AI মডেল উভয়ের জন্য কার্নেল লেখার সমর্থন করে, এবং উভয় পদ্ধতিতেই উচ্চ পারফরম্যান্স অর্জন করে।
• বর্তমান সীমাবদ্ধতা: বর্তমানে শুধুমাত্র একক GPU স্কেনারিও সমর্থিত, ডিস্ট্রিবিউটেড ট্রেনিং অন্তর্ভুক্ত নয়; পুনঃপ্যারামিটাইজেশন মূলত স্ট্যান্ডার্ড Transformer আর্কিটেকচারের জন্য, অন্যান্য আর্কিটেকচারের জন্য প্রযোজ্যতা এখনও যাচাইয়ের অপেক্ষায়।

শেষ কথা

CODA একা কাজ করে না। এটি একটি ধারণার বাস্তবায়ন: GPU-এ, প্রকৃত অপ্টিমাইজেশনের স্থান প্রায়শই "কী গণনা করবেন" নয়, বরং "কীভাবে স্থানান্তর করবেন"।

FlashAttention এটি মনোযোগ গণনাকে চিপের মেমোরিতে স্থানান্তরিত করে, আর CODA সাধারণীকরণ এবং এক্টিভেশন ফাংশনগুলিকেও সেখানে রাখার চেষ্টা করে। Triton কাস্টম কার্নেল লেখার বাধা কমিয়েছে, এবং ThunderKittens, TileLang ইত্যাদি বিভিন্ন স্তরে এই ক্ষেত্রটি আরও গভীরভাবে অন্বেষণ করছে। এই কাজগুলি একই দিকেই ইঙ্গিত করছে: PyTorch অপারেটর গ্রাফের সহজ এক্সপ্রেশনকে হস্তলিখিত CUDA-এর কাছাকাছি এক্সিকিউশন দক্ষতার সাথে একটি প্রোগ্রামযোগ্য ফ্রেমওয়ার্কে প্রকৃতপক্ষে একীভূত করা।

ট্রি ডাওয়ের টুইটের শেষ বাক্যটি আবার ভাবার মতো: 'LLM এবং নতুনরা সমস্ত Transformer অপারেশনের জন্য লাইট-স্পিড কার্নেল লিখতে পারে।' এর পিছনে একটি গভীর যুক্তি রয়েছে: যখন প্রোগ্রামিং অ্যাবস্ট্রাকশন যথেষ্ট ভালভাবে ডিজাইন করা হয়, তখন AI মডেলগুলি নিজেদের প্রশিক্ষণ অবকাঠামোর অপ্টিমাইজেশনে অংশগ্রহণ করতে পারে। এই চক্রটিই CODA-এর সবচেয়ে আকর্ষণীয় দিক।

এই দৃষ্টিকোণ থেকে, "CODA" নামটির সম্ভবত আরও গভীর অর্থ রয়েছে। ক্লাসিক্যাল সংগীতে, Coda হল সংগীতের শেষে সম্পূর্ণ রচনাকে সমাপ্ত করার অংশ। এখানে, এটি GEMM কোরের "শেষ" — এবং এই শেষটি ভালোভাবে লেখা, সম্ভবত Transformer প্রশিক্ষণ সিস্টেমের দক্ষতা বৃদ্ধির পরবর্তী গুরুত্বপূর্ণ অধ্যায়।