FlashAttention-2/3, 3D Parallelism, Quantization: நவீன AI உள்கட்டமைப்பின் மூன்று தூண்கள்

நவீன AI உள்கட்டமைப்பு hardware-ஐ அறிந்த attention kernel-கள், விநியோகிக்கப்பட்ட பல-GPU parallelism, குறைந்த-bit quantization ஆகியவற்றைச் சார்ந்து, பெரிய மாதிரிகளைத் திறமையாக train செய்து serve செய்கிறது. கோட்பாட்டு மாதிரி அளவுக்கும் உடல்ரீதியான GPU வரம்புகளுக்கும் இடையிலான இடைவெளியை இந்த மூன்று நுட்பங்கள் இணைக்கின்றன.

1. Hopper (H100) GPU-களில் FlashAttention-2 மற்றும் FlashAttention-3

அசல் FlashAttention நினைவக bandwidth நெரிசலைத் தீர்த்தாலும், நவீன GPU-களில் கணிசமான மூல கணிசச் செயல்திறனை (compute) பயன்படுத்தாமலே விட்டது. FlashAttention-2 (2023) மற்றும் FlashAttention-3 (2024) GPU thread-களும் hardware engine-களும் எவ்வாறு ஒருங்கிணைகின்றன என்பதை மறுவடிவமைப்பு செய்தன.

• FlashAttention-2 (சிறந்த பணிப் பிரிவு).

• Non-Matmul செயல்களைக் குறைத்தல். GPU Tensor Cores matrix multiplication-களை (matmul) சாதாரண CUDA core-கள் Softmax rescaling போன்ற scalar செயல்களைச் செய்வதை விட சுமார் 16× மடங்கு வேகமாகச் செய்கின்றன. FlashAttention-2, rescaling-ஐ உள் loop-இன் மிகக் கடைசிவரை ஒத்திவைக்க online Softmax கணிதத்தை மீளெழுதியது; இது மெதுவான non-matmul கட்டளைகளைக் குறைத்தது.
• தொடர் நீளத்தில் (Sequence Length) இணையாக்கம். அசல் பதிப்பு பணியை முக்கியமாக batch size மற்றும் attention head-கள் மட்டுமே இணையாக்கியது; நீளமான தொடர்கள் மற்றும் சிறிய batch size-களுடன் train செய்யும்போது GPU core-கள் வெற்றே இருந்தன. FlashAttention-2 sequence length பரிமாணத்திலும் parallelism-ஐச் சேர்த்து, loop ஒழுங்கையும் மாற்றியது (வெளி loop Q block-கள் மீது, உள் loop K மற்றும் V மீது); இது மெதுவான shared-memory எழுத்துகளைக் குறைத்தது.

• FlashAttention-3 (Hopper H100 Hardware பயன்பாடு).

• ஒத்திசைவற்ற தரவு நகர்வு (TMA). NVIDIA-இன் Hopper கட்டமைப்பு Tensor Memory Accelerator (TMA)-ஐ அறிமுகப்படுத்தியது; இது compute thread-களைப் பிடிக்காமல் global HBM-க்கும் on-chip SRAM-க்கும் இடையே தரவை நகலெடுக்கும் அர்ப்பணிக்கப்பட்ட hardware அலகு.
• Warp Specialization (Producer-Consumer மாதிரி). FlashAttention-3 GPU warp-களை (32 thread-களின் குழுக்கள்) அர்ப்பணிக்கப்பட்ட பணிகளாகப் பிரிக்கிறது. Producer warp-கள் TMA-ஐப் பயன்படுத்தி K, V-யின் எதிர்கால tile-களை SRAM-க்கு முன்கூட்டியே ஏற்றுகின்றன (prefetch); Consumer warp-கள் தற்போதைய tile-களை Tensor Cores-க்கு (WGMMA கட்டளைகள்) அளிக்கின்றன. இது நினைவகப் பரிமாற்றங்களை matrix கணிதத்துடனும் Softmax கணக்கீடுகளுடனும் ஒரே நேரத்தில் மேலடுக்குச் (overlap) செய்கிறது.
• FP8 குறைந்த-துல்லிய ஆதரவு. FlashAttention-3 Hopper-இன் FP8 Tensor Cores-ஐத் திறந்து, FP16-ஐ விட மூல throughput-ஐ இரட்டிப்பாக்குகிறது. FP8-இல் outlier activation-களால் துல்லிய இழப்பு ஏற்படாமல் தடுக்க, அது block quantization-ஐ (சிறிய tile-களைத் தனித்தனியாக அளவிடல்) மற்றும் incoherent processing-ஐ (quantize செய்வதற்கு முன் தீவிர மதிப்புகளைச் சமன் செய்ய Q, K-ஐ சீரற்ற நேர்மையான matrix-களால் (random orthogonal matrices) பெருக்கல்) பயன்படுத்துகிறது.

2. பல-GPU அளவீட்டு உத்திகள் (3D Parallelism)

ஒரு மாதிரி பல்டஜன் அல்லது நூற்றுக்கணக்கான பில்லியன் parameter-களைக் கொண்டிருக்கும்போது, அதன் weight-கள், gradient-கள், optimizer நிலைகள் ஒரே GPU-இன் VRAM-ஐ மீறிவிடுகின்றன. பெரிய-அளவு training மூன்று வெவ்வேறு parallelization உத்திகளை இணைக்கிறது — இது பெரும்பாலும் 3D Parallelism எனப் பேசப்படுகிறது.

• Data Parallelism மற்றும் ZeRO / FSDP. சாதாரண Data Parallelism-இல், ஒவ்வொரு GPU-ம் மாதிரியின் முழு நகலையும் வைத்திருக்கும்; இது நினைவகத்தை வீணடிக்கிறது. நவீன framework-கள் ZeRO (Zero Redundancy Optimizer) அல்லது PyTorch FSDP (Fully Sharded Data Parallel)-ஐப் பயன்படுத்துகின்றன; அவை parameter-கள், gradient-கள், optimizer நிலைகளை எல்லா GPU-களிலும் துண்டுபிரித்து (shard), ஒவ்வொரு layer-இன் weight-களையும் கணக்கீட்டிற்குச் சரியான நேரத்தில் (just in time) சேகரிக்கின்றன.

• Tensor Parallelism (Layer-க்குள்). மாபெரும் feed-forward matrix பெருக்கல் Y = XW-ஐ நகல் செய்வதற்குப் பதிலாக, Tensor Parallelism W-ஐ GPU-களில் நிரல்-வாரியாகப் பிரிக்கிறது (W = [W_1, W_2]). ஒவ்வொரு GPU-ம் output-இன் தனது பாதியை ஒரே நேரத்தில் கணக்கிட்டு, அதிவேக NVLink (H100-களில் 900 GB/s) வழியாகப் பகுதி முடிவுகளைப் பரிமாறிக்கொள்கிறது.

• Pipeline Parallelism (Layer-களுக்கிடையே). Layer 1–8 GPU 1-இல், Layer 9–16 GPU 2-இல் என அமைகின்றன. முந்தைய layer-கள் முடியும் வரை காத்திருக்கும் பின்னடை GPU-கள் வெற்றே இருப்பதைத் (pipeline bubble) தடுக்க, framework-கள் ஒவ்வொரு training batch-ஐயும் சிறிய micro-batch-களாக வெட்டி, எல்லா GPU-களும் அசெம்பிளி-லைன் பாணியில் பிஸியாக இருக்கச் செய்கின்றன.

3. Quantization (GPTQ, AWQ) மற்றும் உடனுக்குடன் Dequantize செய்யும் Kernel-கள்

Inference-இன் போது (token-களை ஒன்றன் பின் ஒன்றாக உருவாக்குதல்), Large Language Model-கள் கணிச நெரிசலாக (compute bound) அல்லாமல் நினைவக-bandwidth நெரிசலாக (memory-bandwidth bound) இருக்கின்றன. ஒரே ஒரு token-ஐ உருவாக்க 70 பில்லியன் 16-bit parameter-களை (140 GB) VRAM-இலிருந்து ஏற்றுவது Tensor Cores-ஐ நினைவகப் பரிமாற்றங்களுக்காகக் காத்திருக்க வைக்கிறது. Quantization weight-களை 16-bit float-களிலிருந்து 4-bit integer-களாக (INT4) சுருக்கி, நினைவக அளவையும் பரிமாற்ற நேரத்தையும் கிட்டத்தட்ட 4×-ஆகக் குறைக்கிறது.

• GPTQ (Hessian Calibration மூலம் Post-Training Quantization).

• 16-bit weight-களை 4-bit integer-களாக நேரடியாக (naive) முழுமைப்படுத்துவது (rounding) மாதிரியின் துல்லியத்தை அழித்துவிடும். GPTQ ஒரு சிறிய calibration dataset-ஐ மாதிரி வழியாக அனுப்பி, இரண்டாம்-வரிசை கணிதத்தை (inverse Hessian matrix) பயன்படுத்தி ஒவ்வொரு weight-ம் எவ்வளவு உணர்திறன் உள்ளது என்பதை அளக்கிறது.
• அது weight-களை நிரல்-வாரியாக (column-by-column) quantize செய்கிறது; ஒரு weight-ஐ முழுமைப்படுத்தும்போது பிழை ஏற்பட்டால், அந்த row-இல் மீதமுள்ள quantize செய்யப்படாத weight-களைக் கணிதமாகச் சரிசெய்து பிழையை நீக்குகிறது.

• AWQ (Activation-அறிந்த Weight Quantization).

• ஆராய்ச்சியாளர்கள், weight channel-களில் சுமார் 0.1% முதல் 1% மீதமுள்ளவற்றை விட மிக முக்கியமானவை என்பதைக் கண்டறிந்தனர்; ஏனெனில் அவை பெரிய activation மதிப்புகளைப் பெருக்குகின்றன.
• அந்த முக்கியமான weight-களை 16-bit-இல் வைத்திருப்பதற்குப் (இது குழப்பமான, மெதுவான mixed-precision நினைவக அமைப்புகளை உருவாக்குகிறது) பதிலாக, AWQ quantization-க்கு முன் அந்த முக்கிய (salient) weight channel-களை மேலே அளவிட்டு (scale up), runtime-இல் தொடர்புடைய input activation-களைக் கீழே அளவிடுகிறது (scale down); இது 100% weight-களையும் ஒருசீர் 4-bit வடிவத்தில் வைத்திருந்தபடியே துல்லியத்தைப் பாதுகாக்கிறது.

• Custom CUDA Kernel-கள் உடனுக்குடன் எவ்வாறு Dequantize செய்கின்றன (W4A16).

• GPU-கள் ஒரு 16-bit activation-ஐ (FP16) நேரடியாக ஒரு 4-bit integer weight-ஆல் (INT4) பெருக்க முடியாது. அதற்குப் பதிலாக, சிறப்பு inference kernel-கள் (Marlin அல்லது ExLlama போன்றவை) W4A16 (4-bit Weights, 16-bit Activations) செயலாக்கத்தை நடத்துகின்றன:
• 1. சுருக்கப்பட்ட HBM வாசிப்பு: Kernel, packing செய்யப்பட்ட 4-bit weight-களை (ஒரே 32-bit integer-க்குள் packing செய்யப்பட்ட எட்டு INT4 மதிப்புகள்) மெதுவான global VRAM-இலிருந்து ஏற்றி, நினைவக bandwidth பயன்பாட்டை 4×-ஆகக் குறைக்கிறது.
• 2. Register-மட்ட Unpacking: GPU-இன் மிக வேகமான register-களுக்குள், CUDA kernel வேகமான bitwise shift-கள் (>>) மற்றும் bitmask-கள் (&) பயன்படுத்தி INT4 integer-களை unpack செய்து, சேமிக்கப்பட்ட scale/zero-point காரணிகளைப் பயன்படுத்தி அவற்றை உடனுக்குடன் FP16 மதிப்புகளாக மீள் மாற்றுகிறது.
• 3. உடனடி Tensor Core கணிதம்: புதிதாக unpack செய்யப்பட்ட FP16 weight-கள், விரிவாக்கப்பட்ட 16-bit weight-களை ஒருபோதும் VRAM-இல் மீண்டும் எழுதாமல், FP16 activation-களுடன் நேரடியாக Tensor Cores-க்கு அனுப்பப்படுகின்றன.