Az NVIDIA a Nemotron család nyílt modelljeihez kvantálás‑érzékeny distillációt (quantization‑aware distillation, QAD) alkalmazott, hogy a Nemotron 3.5 Lightning modellt NVFP4 formátumban telepíthetővé tegye: a checkpoint memóriamérete a BF16‑hoz képest jelentősen csökkent, miközben a QAD visszaállítja a szigorú kvantálás okozta pontosságveszteség nagy részét. A munkafolyamat két lépésből áll: először post‑training quantization (PTQ), majd QAD a Nemotron 3.5 Lightning BF16 tanítómodelljétől való logit‑distillációval.
Mi történt és miért fontos
- A Nemotron 3.5 Lightning NVFP4 checkpoint például körülbelül 22 GB‑ra tömöríthető a 66 GB‑os (vagy részletesebb mérés szerint ~65.85 GB) BF16 alapú checkpointról azáltal, hogy sok súlyt 4 bites formátumba kvantálnak (W4 / NVFP4 jelölés). A publikációban említett PTQ és QAD végleges PTQ/ QAD checkpointok jellemző mérete 21.19 GB volt, szemben a BF16 ~65.85 GB‑jával. Ez akár 3–4× gyorsabb áteresztőképességet tehet lehetővé telepítéskor.
- A PTQ önmagában egyszerű és gyors módszer modell tömörítésére, de ha nagyon agresszív kvantálást választunk (pl. Mamba lineáris rétegek W4A16‑ra), a teljesítmény romolhat. Ebben a helyzetben a QAD lehetővé teszi, hogy a kvantált „student” modell megtanuljon alkalmazkodni a kvantálás zajához a BF16 „teacher” logitjeiből származó distillációs jel segítségével, így visszanyerve a pontosságot.
A QAD módszere röviden
- Kétlépcsős eljárás:
- PTQ: a teljes pontosságú BF16 modellt kvantálják előállítva egy W4A16/NVFP4 formátumú student checkpointot. Célszerű itt inkább agresszívan kvantálni, ha QAD követi, mert a QAD képes a későbbi helyrehozatalra.
- QAD: a student minden előremeneti futtatása során szimulált kvantáláson megy át; a tanulás KL‑divergerencia alapú distillációs veszteséggel történik, ahol a BF16 teacher marad fagyasztva.
Részletek: PTQ‑receptetek és skálázás
- Több PTQ receptet teszteltek, amelyek különböztek a súlyok kalibrálásában és abban, mennyire kvantálják a Mamba projekciókat és a KV cache‑t. Öt kiemelt receptet vizsgáltak, mindegyik W4A16/NVFP4 súlyformátumot használ, de eltért például a dinamikus (max‑kalibrált) és MSE‑alapú (statikus) skálázásban.
- A szerzők azt javasolják, hogy a PTQ során érdemes agresszívebb beállításokat választani, ha később QAD követi: például Mamba lineáris rétegeket W4A16‑ra vinni (helyett biztonságosabb FP8‑nak) — a QAD a kieső pontosság nagy részét vissza tudja adni.
QAD tréningbeállítások (konkrét paraméterek)
- A distillációhoz használt munkamenet NVIDIA Model Optimizer reléjén fut: a tanítás során ugyanaz a batch lefut a BF16 teacheren és a kvantált studenten, és a student a logitokra alkalmazott KL veszteséggel illeszkedik a teacherhez.
- Példa hyperparaméterek, amiket a cikk megad:
- konstans tanulási ráta: 5e‑6 (nincs warmup)
- dropout kikapcsolva
- gradient clipping: 1.0
- hardver: két node × 8 GPU (TP=2, EP=4) a bemutatott futtatáshoz
- train samples: 6,400 (≈400 iteráció) a bemutatott kis futtatáshoz; nagyobb végső futtatásnál ~524K token hosszúságot (524k) is használtak
- Sorrend: először PTQ kalibráció 32K vagy 256K/524K sorhosszakkal, majd QAD futtatás (például 256K előkísérletek, végső QAD 524K sorhosszal).
Skálakezelés QAD alatt
- Két stratégia:
- Dynamic scale QAD: a PTQ recept „max” vagy „mamba_fp8_max” esetén a skálák minden lépésben újraszámolódnak, így a súlyok és skálák együtt adaptálódnak.
- Frozen scale QAD: MSE alapú PTQ receptnél a skálákat a PTQ során keresik és lefagyasztják; a QAD alatt csak a súlyok frissülnek.
- A választott skálastratégia függ a PTQ recepttől; érdemes mindkettőt kipróbálni különböző PTQ variánsokkal.
Eredmények: köztes és végső checkpointok
- Méretek: agresszív W4A16 PTQ és QAD checkpointok 21.19 GB‑ot foglaltak, míg a BF16 alap ~65.85 GB volt.
- Checkpoint A (SFT intermediát): PTQ csak 96.33% medián score‑visszanyerést ért el; QAD után 99.72% lett, azaz a QAD 3.39 százalékpontnyi medián javulást hozott, és 10/11 benchmarkon javulás volt.
- Példák: AIME 2025 PTQ‑nál −3.70 pont, QAD után csak −0.57 pont a BF16‑hoz képest.
- Checkpoint B (RL intermediát): PTQ 95.84% mediánról QAD‑dal 98.53%‑ra emelkedett (+2.69). AA v4.1 Index például 20.03‑ról 23.48‑ra nőtt QAD után (BF16: 24.81).
- Végső NVFP4 checkpoint (konzervatív PTQ beállítással optimalizálva a pontosságot): PTQ medián visszanyerés 99.24%, QAD 98.97% — itt a különbségek kisebbek, mert PTQ maga is közel van a BF16‑hoz; QAD viszont jobb teljesítményt mutat több agentikus és kódolási benchmarkon (pl. Terminal‑Bench v2.1, SWE‑Bench Multilingual, PinchBench, HLE).
Reprodukálhatóság és eszközök
- A teljes QAD recept elérhető az NVIDIA Model Optimizer csomagjában, egyetlen launcher fájllal (megatron_lm_qad.yaml) az end‑to‑end pipeline futtatásához. Van Megatron‑Bridge változat is (mbridge_qad.yaml) PyTorch‑native futtatáshoz.
- A végső NVFP4 checkpoint közzétéve: Hugging Face‑en publikálták.
- A cikk bemutat kódrészleteket és parancsokat a modelopt könyvtár mtq.quantize használatára, valamint a launcher példákat a kvantizáció és distilláció futtatására.
Következtetés
A kvantizáció‑érzékeny distilláció lehetővé teszi, hogy a Nemotron 3.5 Lightninghez hasonló, sűrűn aktivált, kompakt modelleknél agresszív kvantálást alkalmazzanak anélkül, hogy tartósan elvesztenék a pontosság nagy részét. A PTQ‑t követő QAD módszer különösen hasznos, ha memóriára és áteresztőképességre szigorú korlátok vonatkoznak: a PTQ által kiszökő hibát a QAD nagy részben vissza tudja hozni. A teljes recept és a szükséges eszközök nyíltan elérhetők az NVIDIA Model Optimizer repójában és a kapcsolódó Hugging Face kiadásoknál.
Elismerések
A közlemény köszönetet mond a Nemotron és az NVIDIA Model Optimizer csapatának, valamint név szerint felsorolja a hozzájáruló kutatókat és mérnököket, akiket a forrásanyag megemlített (Asma Kuriparambil Thekkumpate, Carlo del Mundo, Chenjie Luo, Daniel Lo, Daria Levy, Frank Sun, Hung‑Yueh Chiang, James Shen, Jinhang Choi, Sweta Priyadarshi, Konstantinos Krommydas, Meng Xin, Rohan Joshi, Wei‑Ming Chen, Victor Cui, Trenton Starkey, Yaniv Galron).
