A Diffusers könyvtár mostantól natívan képes betölteni Nunchaku-stílusú, SVDQuant alapú 4-bites (W4A4) checkpointokat a „Nunchaku Lite” úton. Ez azt jelenti, hogy nincs szükség külön Nunchaku inferencia motorra vagy helyi CUDA-fordításra: a modellek egyszerűen a megszokott from_pretrained() hívással kerülnek be, a szükséges CUDA-kernelfájlok pedig a Hugging Face Hubról töltődnek le a kernels csomagon keresztül.
Miért fontos ez?
Nagy diffúziós transzformerek BF16 precízióban gyakran 20–30 GB VRAM-ot igényelnek, ami sok fogyasztói GPU számára elérhetetlenné teszi őket. A SVDQuant megközelítés a súlyok és aktivációk egy részét 4 bitre kvantálja, és a maradék legnehezebb részt kis precíziójú alacsony rangú (16 bites) korrekcióval kezeli. Így csökken a memóriaigény, és a denoising ciklus is gyorsulhat — a Nunchaku Lite integrációval ezek az előnyök Diffusers-en belül elérhetők.
Hogyan működik a Nunchaku Lite?
A Nunchaku Lite nem követeli meg a modellstruktúra átalakítását a betöltés előtt; ehelyett a Diffusers betöltési folyamatánál futásidőben lecseréli a kompatibilis torch.nn.Linear modulokat SVDQ/AWQ típusú lineáris rétegekre (például SVDQW4A4Linear vagy AWQW4A16Linear), mielőtt a súlyokat beolvasná. A szükséges CUDA-kernelfájlok a kernels csomagon keresztül érkeznek a Hubról.
Két fő kernelcsalád létezik:
- svdq_w4a4: 4-bites súlyok és aktivációk SVDQuant alacsony rangú korrekcióval (INT4 és NVFP4 variánsok). Ezt a transzformer attention- és MLP-projekcióihoz használják.
- awq_w4a16: 4-bites súlyok és 16-bites aktivációk, modulációs és normalizációs projekciókhoz, ahol érzékenyebb a precizitás.
Az architektúra-független, „light” útvonallal nem érhető el az eredeti Nunchaku motor által kínált maximális gyorsulás (amely modell-specifikus, fused operátorokra épít), de a Nunchaku Lite így is körülbelül 30%-os gyorsulást és jelentős VRAM-megtakarítást ad.
Hardver- és precizitás-támogatás
A kernelváltozatok GPU-generációtól függenek:
- svdq_w4a4 (nvfp4): NVIDIA Blackwell (RTX 50-es sorozat, RTX PRO 6000, B200)
- svdq_w4a4 (int4): Turing / Ampere / Ada (RTX 30/40, A100, L40S)
- awq_w4a16 (int4): Turing / Ampere / Ada Volta és Hopper jelenleg nem támogatott a 4-bites kernelfuttatásokhoz. A kvantizáló ellenőrzi a GPU CUDA-képességét betöltéskor és világos hibát dob, ha inkompatibilitás van.
Kezdés: előre kvantált pipeline betöltése
A szükséges csomagok telepítése után (diffusers, transformers, accelerate, kernels, bitsandbytes) egy Nunchaku Lite pipeline példányosítása megegyezik egy szokványos Diffusers modéllal. Például egy NVFP4 ERNIE-Image-Turbo checkpoint 1024×1024 képet képes generálni körülbelül 1.7 másodperc alatt egy RTX 5090-en, csúcs VRAM-felhasználás ~12 GB mellett, szemben a BF16 pipeline ~24 GB-jával. Az NVFP4 variánsok Blackwell hardvert igényelnek; régebbi kártyákon az INT4 verziókat kell használni.
Teljesítmény és memória: mérések
RTX PRO 6000 (Blackwell) mérések 1024×1024 felbontáson a rootonchair/ERNIE-Image-Turbo-nunchaku-lite-int4-bnb4-text-encoder checkpointtal:
- BF16 alap: teljes pipeline 3.00 s, denoise loop 2.86 s, csúcs VRAM 31.1 GB (1.0×)
- Nunchaku Lite NVFP4: teljes pipeline 2.27 s, denoise loop 2.13 s, csúcs VRAM 20.6 GB (1.35× gyorsabb)
- Nunchaku Lite NVFP4 + torch.compile: teljes pipeline 1.68 s, denoise loop 1.53 s, csúcs VRAM 20.6 GB (1.8× gyorsabb)
- Nunchaku Lite NVFP4 + NF4 text encoder: teljes pipeline 2.29 s, denoise loop 2.13 s, csúcs VRAM 16.0 GB (1.35× gyorsabb)
A mérések azt mutatják, hogy Nunchaku Lite akár ~50%-kal is csökkentheti a csúcs VRAM-igényt, miközben a késleltetést kb. 30%-kal javítja; további gyorsulást hoz a torch.compile, amely csökkenti a kernelindítási overheadet.
Kombinálható optimalizációk
- torch.compile: a transzformer lefordítása teljesebb gyorsulást ad (például 1.35× → 1.8× a mérésekben).
- Kvantált text encoderek: bitsandbytes NF4 további VRAM-megtakarítást eredményezhet (a tesztben ~22%).
- Offloading: Diffusers meglévő CPU offload segédjei (enable_model_cpu_offload(), enable_sequential_cpu_offload()) továbbra is működnek.
Saját modell kvantizálása és közzététele
A diffuse-compressor eszközkészlet egy végponttól-végpontig SVDQuant munkafolyamatot nyújt: vizsgálat, kalibráció, kvantizálás, csomagolás és publikálás. A bemutatott példa a FLUX.2 Klein 4B kvantizálása, amely lépéseket tartalmaz a target-azonosítástól a checkpoint előállításán át a Diffusers pipeline csomagolásáig és a Hubra feltöltésig.
Fontos, hogy a generikus útvonal feltételezi: a modell kvantizálható struktúrális átírás nélkül. Egyes modelleken (például FLUX.1-dev) a Nunchaku eredeti útvonala struktúrális átalakításokat alkalmaz (például a külön Q/K/V projekciók egyesítése egy fused to_qkv modulba), amit a generikus szkenner nem tud automatikusan kitalálni. Ilyen esetekben target-config leírás és/vagy kicsi futásidő adapter szükséges a helyes betöltéshez. A diffuse-compressor és a rootonchair/nunchaku-lite repo-k tartalmaznak példákat és adaptereket ilyen esetekre.
Elérhető, kész checkpointok
Példák, amelyekkel azonnal kipróbálható a Nunchaku Lite:
- rootonchair/ERNIE-Image-Turbo-nunchaku-lite-int4-bnb4-text-encoder (INT4 + NF4 text encoder)
- rootonchair/ERNIE-Image-Turbo-nunchaku-lite-nvfp4-bnb4-text-encoder (NVFP4 + NF4 text encoder)
- OzzyGT/Krea_2_Turbo_nunchaku_lite_nvfp4
- lite-infer (gyűjtemény és további checkpointok)
Következtetés
A SVDQuant alapú Nunchaku kernelek hatékony módszert adnak a diffúziós transzformerek fogyasztói hardveren történő futtatására. A Diffusers-be integrált Nunchaku Lite lehetővé teszi előre kvantált checkpointok egyszerű betöltését a megszokott from_pretrained() API-val, míg a diffuse-compressor eszköz lépéseket nyújt új architektúrák kvantizálására és publikálására. A W4A4 útvonal csökkenti a memóriaigényt és javítja a denoising késleltetést úgy, hogy az eredményközeli képminőséget megőrzi.
Köszönet a Diffusers karbantartóinak, az MIT HAN Lab / Nunchaku csapatnak az SVDQuant munkáért, valamint a tesztelőknek és visszajelzőknek.
