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 অপারেশন কমানো। সাধারণ CUDA কোর যে গতিতে Softmax পুনঃস্কেলিংয়ের মতো স্কেলার অপারেশন চালায়, GPU Tensor Cores তার প্রায় 16× দ্রুত ম্যাট্রিক্স গুণ (matmul) করে। 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 সমান্তরালতা বলা হয়।
| কৌশল | কীভাবে কাজ ভাগ করে | যোগাযোগের বাধা | সবচেয়ে উপযুক্ত ব্যবহার |
|---|---|---|---|
| ডেটা সমান্তরালতা (DP / FSDP) | GPUs-এ মডেলের প্রতিলিপি রাখে; প্রতিটি GPU ডেটার আলাদা ব্যাচ প্রক্রিয়া করে। | GPUs-এর মধ্যে গ্রেডিয়েন্ট ও ভাগ করা ওয়েটের All-Reduce / All-Gather। | স্তরগুলি একটি ডিভাইসে ধরলে শত শত GPUs-এ প্রশিক্ষণের থ্রুপুট বাড়ানো। |
| টেনসর সমান্তরালতা (TP) | পৃথক ওয়েট ম্যাট্রিক্স (স্তর) GPUs-এ উল্লম্ব বা অনুভূমিকভাবে ভাগ করে। | প্রতিটি Transformer স্তরের পরে All-Reduce (অতি দ্রুত NVLink প্রয়োজন)। | একটি GPU-তে না ধরা বিশাল স্তর, যেগুলি কঠোরভাবে একটি 8-GPU সার্ভার নোডের মধ্যেই রাখা হয়। |
| পাইপলাইন সমান্তরালতা (PP) | মডেলকে গভীরতা অনুযায়ী ভাগ করে, ধারাবাহিক স্তরগুচ্ছ আলাদা GPUs-এ রাখে। | পাইপলাইনের ধাপগুলির মধ্যে Point-to-Point (P2P) অ্যাক্টিভেশন স্থানান্তর। | ধীরতর InfiniBand/Ethernet নেটওয়ার্কে যুক্ত একাধিক সার্ভার নোডে স্কেল করা। |
- ডেটা সমান্তরালতা & 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-বিট প্যারামিটার (140 GB) লোড করলে Tensor Cores-কে মেমরি স্থানান্তরের অপেক্ষায় থাকতে হয়। কোয়ান্টাইজেশন ওয়েটকে 16-বিট ফ্লোট থেকে 4-বিট পূর্ণসংখ্যায় (INT4) সংকুচিত করে, ফলে মেমরির আকার ও স্থানান্তরের সময় প্রায় 4× কমে।
- GPTQ (Hessian ক্যালিব্রেশনের মাধ্যমে প্রশিক্ষণ-পরবর্তী কোয়ান্টাইজেশন)।
- 16-বিট ওয়েটকে সরলভাবে 4-বিট পূর্ণসংখ্যায় রাউন্ড করলে মডেলের নির্ভুলতা নষ্ট হয়। GPTQ মডেলের মধ্য দিয়ে একটি ছোট ক্যালিব্রেশন ডেটাসেট চালায় এবং প্রতিটি ওয়েট কতটা সংবেদনশীল তা মাপতে দ্বিতীয়-ক্রমের গণিত (বিপরীত Hessian ম্যাট্রিক্স) ব্যবহার করে।
- এটি কলাম ধরে ওয়েট কোয়ান্টাইজ করে। কোনো ওয়েট রাউন্ড করার ফলে ত্রুটি হলে GPTQ সেই সারির বাকি অকোয়ান্টাইজড ওয়েট গাণিতিকভাবে সমন্বয় করে, যাতে ত্রুটি বাতিল হয়।
- AWQ (অ্যাক্টিভেশন-সচেতন ওয়েট কোয়ান্টাইজেশন)।
- গবেষকেরা দেখেছেন, প্রায় 0.1% থেকে 1% ওয়েট চ্যানেল বাকিগুলির তুলনায় অনেক বেশি গুরুত্বপূর্ণ, কারণ সেগুলি বড় অ্যাক্টিভেশন মানের সঙ্গে গুণ হয়।
- ওই গুরুত্বপূর্ণ ওয়েটকে 16-বিটে রাখার বদলে (যা জটিল ও ধীর মিশ্র-নির্ভুলতার মেমরি বিন্যাস তৈরি করে), AWQ কোয়ান্টাইজেশনের আগে ওই উল্লেখযোগ্য ওয়েট চ্যানেল বড় করে স্কেল করে এবং রানটাইমে সংশ্লিষ্ট ইনপুট অ্যাক্টিভেশন ছোট করে স্কেল করে। এতে নির্ভুলতা বজায় থাকে, আবার 100% ওয়েট একই 4-বিট বিন্যাসে থাকে।
- কাস্টম CUDA কার্নেল কীভাবে তাৎক্ষণিক ডিকোয়ান্টাইজ করে (W4A16)।
- GPUs সরাসরি 16-বিট অ্যাক্টিভেশন (FP16)-কে 4-বিট পূর্ণসংখ্যার ওয়েট (INT4) দিয়ে গুণ করতে পারে না। পরিবর্তে বিশেষ ইনফারেন্স কার্নেল (যেমন Marlin বা ExLlama) W4A16 (4-বিট ওয়েট, 16-বিট অ্যাক্টিভেশন) চালায়:
- 1. সংকুচিত HBM রিড: কার্নেল ধীর গ্লোবাল VRAM থেকে প্যাক করা 4-বিট ওয়েট লোড করে (একটি 32-বিট পূর্ণসংখ্যার মধ্যে আটটি INT4 মান প্যাক করা), ফলে মেমরি ব্যান্ডউইডথের ব্যবহার 4× কমে।
- 2. রেজিস্টার-স্তরের আনপ্যাকিং: GPU-র অতি দ্রুত রেজিস্টারের মধ্যে CUDA কার্নেল দ্রুত বিটওয়াইজ শিফট (>>) ও বিটমাস্ক (&) দিয়ে INT4 পূর্ণসংখ্যা আনপ্যাক করে। এরপর সংরক্ষিত স্কেল/জিরো-পয়েন্ট গুণক প্রয়োগ করে তাৎক্ষণিকভাবে সেগুলিকে FP16 মানে ফেরায়।
- 3. তাৎক্ষণিক Tensor Core গণিত: সদ্য আনপ্যাক করা FP16 ওয়েট সরাসরি FP16 অ্যাক্টিভেশনের সঙ্গে Tensor Cores-এ পাঠানো হয়। সম্প্রসারিত 16-বিট ওয়েট কখনোই আবার VRAM-এ লেখা হয় না।