Upanuzi wa FlashAttention, ufundishaji wa GPU nyingi na quantization ya uzani

Miundombinu ya kisasa ya AI inategemea kernels za attention zinazojali vifaa, usambambaji wa GPU nyingi uliosambazwa, na quantization ya biti chache ili kufundisha na kutumia miundo mikubwa kwa ufanisi. Mbinu hizi tatu huziba pengo kati ya kiwango cha kinadharia cha modeli na mipaka halisi ya GPU.

1. FlashAttention-2 na FlashAttention-3 kwenye GPU za Hopper (H100)

Ingawa FlashAttention ya awali ilitatua tatizo la kizibwaji cha memory bandwidth, bado iliacha sehemu kubwa ya nguvu halisi ya kompyuta bila kutumika kwenye GPU za kisasa. FlashAttention-2 (2023) na FlashAttention-3 (2024) zilibuni upya jinsi thread za GPU na injini za vifaa zinavyoshirikiana.

• FlashAttention-2 (Ugawaji Bora wa Kazi).

• Kupunguza Operesheni Zisizo za Matmul. Tensor Cores za GPU hutekeleza kuzidisha matrici (matmul) kwa kasi ya takribani mara 16 kuliko jinsi CUDA cores za kawaida zinavyotekeleza operesheni za scalar kama upimaji upya wa Softmax. FlashAttention-2 iliandika upya hisabati ya Softmax ya mtandaoni ili kuahirisha upimaji upya hadi mwishoni kabisa mwa kitanzi cha ndani, ikapunguza maagizo mabaya yasiyo ya matmul.
• Kufanya Sambamba kwa Urefu wa Mlolongo. Toleo la awali lilifanya kazi sambamba hasa kwa ukubwa wa batch na vichwa vya attention, likiacha cores za GPU bila kazi wakati wa kufundisha mialango mirefu na batch ndogo. FlashAttention-2 inaongeza usambambaji katika kipimo cha urefu wa mlolongo na kugeuza mpangilio wa vitanzi (kitanzi cha nje juu ya vitalu vya Q, kitanzi cha ndani juu ya K na V), kupunguza uandishi wa polepole wa kumbukumbu ya pamoja.

• FlashAttention-3 (Matumizi ya Vifaa vya Hopper H100).

• Uhamishaji wa Data Usiolingana (TMA). Usanifu wa Hopper wa NVIDIA ulitambulisha Tensor Memory Accelerator (TMA), kipengo maalum cha vifaa kinachonakili data kati ya HBM ya jumla na SRAM ya ndani ya chip bila kufunga thread za kompyuta.
• Utaalamishaji wa Warp (Muundo wa Mzalishaji-Mteja). FlashAttention-3 hugawanya warp za GPU (vikundi vya thread 32) katika majukumu maalum. Warp za Mzalishaji hutumia TMA kupakua awali vitalu vijavyo vya K na V ndani ya SRAM, huku warp za Mteja zikilisha vitalu vya sasa kwenye Tensor Cores (maagizo ya WGMMA). Hii hupachika uhamishaji wa kumbukumbu na hisabati za matrici na mahesabu ya Softmax kwa wakati mmoja.
• Msaada wa FP8 wa Usahihi wa Chini. FlashAttention-3 hufungua Tensor Cores za FP8 za Hopper, ikiongeza mara mbili throughput halisi ikilinganishwa na FP16. Ili kuzuia upotevu wa usahihi unaosababishwa na activations za kupindukia katika FP8, hutumia block quantization (kupima vitalu vidogo kivyake kila kimoja) na incoherent processing (kuzidisha Q na K kwa matrici za nasibu za orthogonal ili kusawazisha thamani za kupindukia kabla ya quantization).

2. Mikakati ya Kupanua Multi-GPU (3D Parallelism)

Wakati modeli ina vigezo mabilioni kadhaa au mamia ya mabilioni, uzani wake, gradients, na hali za optimizer huzidi VRAM ya GPU moja. Ufundishaji wa kiwango kikubwa huunganisha mikakati mitatu tofauti ya usambambaji—inayoitwa mara nyingi 3D Parallelism.

• Data Parallelism & ZeRO / FSDP. Katika Data Parallelism ya kawaida, kila GPU hubeba nakala kamili ya modeli, ikihujumu kumbukumbu. Mifumo ya kisasa hutumia ZeRO (Zero Redundancy Optimizer) au PyTorch FSDP (Fully Sharded Data Parallel), ambayo hugawanya vigezo, gradients, na hali za optimizer kati ya GPU zote, ikikusanya uzani wa kila tabaka kwa wakati sahihi wa kompyuta.

• Tensor Parallelism (Ndani ya Tabaka). Badala ya kurudia kuzidisha matrici kubwa ya feed-forward Y = XW, Tensor Parallelism hugawanya W kwa nguzo kati ya GPU (W = [W_1, W_2]). Kila GPU huhesabu nusu yake ya pato kwa wakati mmoja na kubadilishana matokeo ya sehemu kupitia NVLink ya kasi kubwa (900 GB/s kwenye H100).

• Pipeline Parallelism (Kati ya Tabaka). Tabaka 1–8 huwekwa kwenye GPU 1, Tabaka 9–16 kwenye GPU 2, na kadhalika. Ili kuzuia GPU za chini zisikae bila kazi wakati zikisubiri tabaka za awali zimalize ("pengo la pipeline"), mifumo hukatakata kila batch ya ufundishaji kuwa micro-batches ndogo ili GPU zote ziendelee kufanya kazi kwa mtindo wa laini ya kukusanyia.

3. Quantization (GPTQ, AWQ) na Kernels za Dequantization ya Papo Hapo

Wakati wa inference (kuzalisha token moja moja), Miundo Mikubwa ya Lugha ni memory-bandwidth bound, si compute bound. Kupakia vigezo bilioni 70 vya biti 16 (140 GB) kutoka VRAM ili kuzalisha token moja tu huacha Tensor Cores zikisubiri uhamishaji wa kumbukumbu. Quantization hupunguza uzani kutoka floats za biti 16 hadi integers za biti 4 (INT4), ikikata ukubwa wa kumbukumbu na muda wa uhamishaji kwa karibu mara 4.

• GPTQ (Quantization ya Baada ya Ufundishaji kupitia Urekebishaji wa Hessian).

• Kukata kwa kawaida uzani wa biti 16 kuwa integers za biti 4 huharibu usahihi wa modeli. GPTQ hupitisha seti ndogo ya data ya urekebishaji kupitia modeli na kutumia hisabati ya daraja la pili (matrici ya inverse Hessian) kupima jinsi kila uzani ulivyo nyeti.
• Hufanya quantization ya uzani nguzo kwa nguzo; kila kukata uzani mmoja kunapoingiza hitilafu, GPTQ hurekebisha kihesabati uzani uliobaki ambao bado haujafanyiwa quantization katika safu hiyo ili kuifuta hitilafu hiyo.

• AWQ (Quantization ya Uzani Inayojali Activation).

• Watafiti waligundua kuwa takribani 0.1% hadi 1% ya njia za uzani ni muhimu zaidi kuliko zingine kwa sababu huzidisha thamani kubwa za activation.
• Badala ya kuweka uzani huo muhimu katika biti 16 (jambo linaloumba mpangilio wa kumbukumbu wa usahihi mchanganyiko uliotatizika na wa polepole), AWQ huongeza kiwango cha njia hizo maalum za uzani kabla ya quantization na kupunguza kiwango cha activations za ingizo zinazolingana wakati wa utekelezaji, ikihifadhi usahihi huku ikitunza 100% ya uzani katika umbizo sare la biti 4.

• Jinsi Kernels Maalum za CUDA Zinavyofanya Dequantization Papo Hapo (W4A16).

• GPU haziwezi kuzidisha moja kwa moja activation ya biti 16 (FP16) kwa uzani wa integer ya biti 4 (INT4). Badala yake, kernels maalum za inference (kama Marlin au ExLlama) hutekeleza utekelezaji wa W4A16 (Uzani wa biti 4, Activations za biti 16):
• 1. Kusoma HBM Iliyobanwa: Kernel hupakia uzani wa biti 4 uliopakiwa pamoja (thamani nane za INT4 zilizobanwa ndani ya integer moja ya biti 32) kutoka VRAM ya jumla ya polepole, ikikata matumizi ya memory bandwidth mara 4.
• 2. Kufungua Kifurushi Katika Registers: Ndani ya registers za GPU zenye kasi kubwa, kernel ya CUDA hutumia mabadiliko ya haraka ya bitwise (>>) na bitmasks (&) kufungua integers za INT4 na kutumia vigezo vilivyohifadhiwa vya scale/zero-point kuzirudisha kuwa thamani za FP16 papo hapo.
• 3. Hisabati za Papo Hapo za Tensor Core: Uzani wa FP16 uliofunguliwa sasa hivi hulishwa moja kwa moja kwenye Tensor Cores pamoja na activations za FP16 bila kuandika uzani uliopanuliwa wa biti 16 tena kwenye VRAM.