Updated July 28, 2026: Kimi K3 replaces Kimi K2.6 and takes the open lead

If closed models like Claude and GPT are more performant, why should you use open models at all? The first reason is cost: leading open models now deliver near-frontier intelligence at a fraction of the cost of closed-API models. For everyday, high-volume work, they offer a better tradeoff between price and intelligence.

The second reason is control. Two recent decisions show the risk of relying on a closed model: Anthropic restricted Mythos to trusted partners, and the US government pulled Fable 5 from the market. An API model can be deprecated, rerouted, or pulled away from you. One that you host yourself is stable, can be fine-tuned, and doesn’t leak your data.

This article analyzes the seven strongest open models and tracks how far they trail the closed frontier.

The frontier at a glance

Across 2026, open models nearly caught up to the closed frontier, then fell behind as the closed labs released stronger models. GLM-5.2 has now closed most of that gap again. Kimi K3 has now pulled them back to within three points.

On the Artificial Analysis Intelligence Index, Kimi K3 leads all open models at 57, six points above GLM-5.2 (51) and 19 above the best US open model, Nemotron 3 Ultra (38). Only Anthropic’s Fable 5 (60) and OpenAI’s GPT-5.6 Sol (59) score higher, and both are closed.

Other evaluations point in the same direction. Vals AI scores domain-specific work (legal, finance, tax) on private test sets, and LMArena ranks models by human preference in head-to-head chat. Kimi K3 leads all three, but each ranks the models below it differently.

However, closed models still keep a clear edge on real-world agentic coding tasks. On benchmarks like SWE-Bench Pro and long, real-world agentic coding tasks that Claude Code and Codex run, closed models are still ahead.

The models

Every model here is a sparse Mixture of Experts (MoE). Six of the seven come from Chinese labs.

Every family here shipped its current generation within the last six months.

Kimi K3 (Moonshot AI, 2.8T total, 104B active; technical report) is the strongest open alternative to Claude and GPT today, and the first open model at 3T scale, nearly double the next largest here (DeepSeek-V4 Pro, 1.6T). It also reverses Moonshot’s own design philosophy: K2.6 ran the most conservative attention stack of the seven, and K3 runs the most aggressive. Two costs come with it. Its API price is five times K2.6’s, and its calibration got worse: Artificial Analysis measures its hallucination rate rising from 39% to 51% while its accuracy climbed from 33% to 46%, so K3 buys correct answers by guessing more often.

GLM-5.2 (Z.ai, 744B total, 40B active; my GLM-5 article) held the open lead for five weeks and remains the better buy for most work. It tops SWE-Bench Pro, has a usable 1M-token context, and costs less than a quarter of what K3 does for six index points less.

MiniMax-M3 (MiniMax, 428B total, 23B active, 1M context; my M2.5 article) replaced M2.7 just eight weeks after it shipped. It swapped M2.7’s full attention for MSA, a new sparse attention mechanism, and added native image and video input. Its real edge is price: M3 runs at some of the lowest costs on the frontier.

DeepSeek-V4 Pro (DeepSeek, 1.6T total, 49B active; my DeepSeek-V4 article) is architected around million-token serving and consolidates its post-training via distillation alone. It is the cheapest model here, at roughly 4% of K3’s blended price. Its clearest weakness is calibration: on hallucination tests, it guesses on questions it cannot answer instead of declining them.

MiMo-V2.5-Pro (Xiaomi, 1.02T total, 42B active) is the fastest climb on this list. Xiaomi launched its Core AI group in April 2025 and hired Luo Fuli, a core DeepSeek-V2 developer, to lead MiMo. They went from a small first model to a frontier-scale one in about a year. It is assembled largely from published components, and a June paper documents Xiaomi’s multi-teacher on-policy distillation.

Nemotron 3 Ultra (Nvidia, 550B total, 55B active; my Nemotron 3 Ultra article) is the best US open model, co-designed with the hardware its maker sells. It serves several times faster than its Chinese peers but trails them in quality by the widest gap on this list, due to its design tuned for throughput.

Qwen3.5-397B-A17B (Alibaba, 397B total, 17B active; my Qwen3.5 article) was the first production test of linear attention at frontier scale, a bet K3 has since made seven times larger, and it is the only model here under a plain Apache-2.0 license. It is Alibaba’s strongest open model, but not its strongest model: the more capable Qwen Max line stays closed.

Architecture

Everyone agrees on the MoE backbone, but nobody agrees on attention. Which design wins for long context processing is still a different answer per lab.

One camp stays entirely on attention and shrinks the cache rather than replacing it, which it does three ways:

• The first compresses the cache. DeepSeek-V4 moved to compressed hybrids (CSA and HCA) built for very long context, with the most aggressive cache reduction in the field.

• The second keeps only selected blocks of the cache per token. GLM-5 stacked DeepSeek Sparse Attention on MLA, and GLM-5.2 added IndexShare to reuse the sparse-attention indexer every fourth layer. After M1’s Lightning Attention and M2’s full attention, M3 landed on Minimax Sparse Attention. It scores the cache in blocks with a cheap index branch and lets each query read only its top-ranked ones.

• The third route is local-global attention. MiMo runs sliding-window attention over GQA at a 6:1 ratio, so most layers attend only to a nearby window.

The other camp replaces most attention layers with a cheaper sequence mixer, whose cost grows only weakly with sequence length:

• Kimi K3 is the newest and largest defector to it, and it went further than anyone: 69 of its 93 layers are linear, with a Gated MLA layer closing every block of four to restore global attention. Its mixer is Kimi Delta Attention, a channel-wise-gated refinement of the delta rule that Moonshot introduced in Kimi Linear in October 2025, which itself refines Gated DeltaNet.

• Qwen interleaves Gated DeltaNet with full attention at the same 3:1 ratio.

• Nemotron pairs Mamba-2 with thin GQA layers.

Sparsity varies just as much. The number to compare is active parameters, i.e., how many of a model’s parameters actually run on each token. That now ranges from Qwen3.5’s 17B to K3’s 104B, a 6x spread among models within 23 index points. Most use the DeepSeekMoE layout. The structural departure is LatentMoE, which Nvidia published in January, compressing each expert into a narrow latent space so a model can hold many more of them at the same active-parameter cost. K3 is its first adoption outside Nvidia.

No design has won, and they are all making the same bet: trade some recall for cheaper long-context capacity. A sequence mixer is the cheapest, with a constant-size state per layer, but it recalls earlier tokens poorly. This is why Kimi, Qwen, and Nemotron all pair it with periodic full-attention layers to compensate, which makes them hybrids. The compression and selection designs (MLA, DSA, MSA) give up less recall, but their cache still grows with the input, while sliding-window attention caps it at a fixed window.

Pre-training

Few labs document pre-training in depth, and the choices they do disclose already diverge.

Of the four labs that disclose a token budget, Ultra is the low outlier at 20T, five trillion short of plan, after its run diverged twice. Long context now splits on method rather than reach, since six of the seven ship 1M. V4 ramped its sequence length from 4K to 1M during the run and MiMo trained natively to 1M, while K3 grew its window in four stages, 8K to 64K during pre-training and 256K to 1M during cooldown.

Precision is where the labs are still experimenting. Ultra is the largest model pre-trained in NVFP4 (4-bit), and it paid for it: two runs diverged, one traced to FP4 interacting with the MTP heads, and one never explained. Both DeepSeek-V4 and K3 took the conservative route instead, training in FP8 and reaching 4-bit only through quantization-aware training afterwards. The optimizer has converged on Muon: Moonshot proved it at a trillion parameters on K2, adding a stabilizer called QK-Clip, and K3 now orthogonalizes each attention head’s momentum separately so a few large heads stop dominating the shared update. DeepSeek-V4 adopted Muon but dropped QK-Clip because its RMSNorm placement already keeps training stable.

Hardware is the murkiest of these choices. Every model here most likely trained on Nvidia. DeepSeek’s attempt to migrate to Huawei Ascend failed in mid-2025 and cost months, per ChinaTalk, which shows how hard a full Ascend pre-training run still is. Z.ai is the only lab here on the US Entity List (added January 2025), which cut it off from advanced Nvidia chips, so GLM-5.2 was widely rumored to have trained on Huawei Ascend. Z.ai never confirmed that, so Nvidia stays the most likely. The one confirmed frontier-scale run off Nvidia is Meituan’s LongCat-2.0, trained on roughly 50,000 domestic ASICs. It shows the Nvidia dependency is only starting to loosen.

Post-training

Post-training is where the training pipelines diverge the most. The main trend is the wide adoption of multi-domain on-policy distillation (MOPD). It consists of training a set of domain specialists, then distilling them into a single student on the student’s own outputs. It is doing work the reinforcement learning stage used to do, either replacing that stage or adding to it. Four of the seven now run it, and Kimi K3 is the newest convert: Moonshot’s last generation ran a conventional RL pipeline.

Kimi K3, DeepSeek-V4, and Xiaomi’s MiMo use it in place of the student’s RL: they train domain specialists, then distill those teachers into an SFT student, with no RL on the student itself. The way they build these specialists differs, though. DeepSeek trains each specialist from the base model with both SFT and RL. Xiaomi initializes its specialists from the SFT checkpoint and only uses RL, keeping teachers and student close in policy. Nvidia adds distillation on top of RL: it trains the student with SFT and RL first, then distills the specialists to push past the plateau RL hit on its own. Distillation only amplifies what the base can already do, and it is not cheap, since it requires training a stack of specialist teachers.

The labs that skip distillation run a long, multi-stage RL pipeline instead: SFT, RL with verifiable rewards, agentic RL against live software repos, terminals, and browsers, then preference tuning on top. This is the path GLM-5.2, MiniMax, and Qwen take. With agentic RL, the bottleneck is generation, not the gradient step, since a few very long rollouts stall the whole batch. That is why these labs build asynchronous RL systems, like MiniMax’s Forge, that decouple rollout generation from training so the accelerators never idle on the slowest trajectory. GLM is the partial exception: it runs the same RL pipeline, then adds a final distillation step that uses each earlier stage’s own checkpoint as the teacher, so the last stage does not overwrite the skills the earlier ones learned.

The economics of open models

Open frontier models are much cheaper than closed-source ones, but it also means the labs that make them profit little from inference alone.

The rule holds because of competition, not efficiency. Anyone can serve open weights, so dozens of providers drive the price of these models down to roughly the cost of the GPUs and electricity. A closed model has one seller, who charges for the value it delivers instead.

The catch for the lab is that it is just another provider, with limited cost advantage over Together or Fireworks when serving the same checkpoint. Open-sourcing a frontier model means giving up exclusive inference revenue, so the question is what each lab sells instead.

Two sell the complement. Nvidia sells GPUs, and Alibaba sells cloud services, so a cheap, ubiquitous model only drives demand for what they charge for. Three sell a differentiated service based on a hosted API tuned for reliability and support: MiniMax, Moonshot, and Z.ai. Alibaba also keeps a stronger model closed, the Qwen Max line above open Qwen3.5. And two sell only (publicly) sell model inference: DeepSeek and Xiaomi open-source it as a strategic bet.

A frontier run can cost over a hundred million dollars, and open weights let anyone serve them at commodity margins. The labs that can keep going are the ones with a business the model feeds (Nvidia, Alibaba) or a backer that does not need it to earn (DeepSeek's High-Flyer fund, Xiaomi's phone business). The pure plays take the real risk: MiniMax, Moonshot, and Z.ai are betting that giving the weights away buys enough API and enterprise business to fund the next model.

Where the frontier moves next

To conclude, here is a list of open questions worth watching:

1. Silicon: LongCat-2.0 was trained on domestic Chinese hardware. Are other Chinese labs going to make the switch this year?

2. Attention: the field keeps splitting on this. Will one design become the standard, or does the split hold?

3. FP4 training: despite its theoretical advantages, stability and performance issues make FP4 difficult to justify. Is it going to become a new standard for large runs?

4. Open publishing: competitors now adopt a new technique within a single release cycle. Are we going to see less and less open research and detailed reports?

Reading list

• DeepSeek V4: ten teachers, one student: the consolidation recipe in full.

• Nemotron 3 Ultra: what distillation can't fix: the recovery ceiling, measured.

• Qwen3.5: nobody agrees on attention: the architecture split this post extends.

• GLM-5: China's first public AI company: the lab behind the Ascend run.

• Kimi K2.5: still worth it after two: the Moonshot lineage behind K2.6.

DeepSeek shipped V4 on April 24th, 2026, in two open-weight sizes: the 1.6-trillion-parameter Pro and the 284-billion-parameter Flash, both MIT-licensed with a million-token context window. The release put DeepSeek back among the leading open-weight labs after a long quiet stretch where Kimi, Qwen, and GLM did most of the talking. The pitch is efficiency rather than raw scale, built on a hybrid attention design that makes long context cheap to serve.

The architecture earns a look, and the post-training breaks more new ground: V4 drops the mixed reinforcement learning stage that consolidated V3.2 and replaces it with on-policy distillation from a panel of specialist teachers.

Two things matter more than the launch benchmarks: which variant people actually run, and where V4 falls down.

Architecture

V4 keeps the DeepSeekMoE layout and Multi-Token Prediction objective from V3. The new work is in three places: the attention, the residuals, and the optimizer. Pro and Flash run the same design at two scales.

The hybrid attention is the real change. At a million tokens, attention is the bottleneck: every new token reads a KV cache that grows with the whole context. V4's answer is two layer types that shrink that cache by merging runs of consecutive tokens into single, compressed entries. They differ in how hard they compress and what they read back.

Compressed Sparse Attention (CSA) compresses gently and reads selectively. It folds a few tokens into each entry, then runs DeepSeek Sparse Attention on the result, so that a query attends only to the top-k compressed entries that an indexer scores as relevant. Heavily Compressed Attention (HCA) makes the opposite trade: it packs far more tokens into each entry, then attends to all of them with no selection step. In effect, CSA layers retrieve the specific passages a query needs, while HCA layers keep a cheap, blurry view of the whole context, so whatever the indexer misses is degraded rather than gone. Both add a sliding window over the most recent raw tokens, because a query can't see inside its own compressed block and the nearest tokens matter most.

Interleaving the two is what makes a million tokens tractable, and the report puts the KV cache at 2% of a standard BF16 GQA baseline at 1M. The whole field is improvising here, which I covered in Qwen3.5: Kimi K2.6 keeps plain MLA, GLM-5.1 stacks DSA on MLA, Qwen3.5 interleaves Gated DeltaNet with full attention, and MiniMax reverted to dense attention. V4 is one more distinct answer, and the one most tuned for raw context length.

The other two changes are borrowed, and worth naming as such. Manifold-Constrained Hyper-Connections (mHC) extend ByteDance's Hyper-Connections by projecting the residual mix onto doubly stochastic matrices with Sinkhorn-Knopp, which bounds the spectral norm and keeps deep stacks stable. And V4 trains on the Muon optimizer instead of AdamW. Muon at this scale is Moonshot's move, proven on the trillion-parameter Kimi K2. The telling divergence is that Moonshot needed QK-Clip to stop attention logits from exploding under Muon, while DeepSeek drops it because the RMSNorm it already applies to queries and compressed KV entries does the same job.

One number jumped out at me. V4-Pro activates 49B parameters per token, the highest of any leading open-weight model: Kimi K2.6 fires 32B, GLM-5.1 40B, Qwen3.5 17B, MiniMax-M2.7 just 10B. Active parameters are the compute you pay on every single token, which is why the trend everywhere else is to shrink them. DeepSeek runs the trade in reverse. The hybrid attention cuts so much of the long-context bill that Pro can afford to activate more experts per token, buying back the capabilities the sparser models give up, and still come out almost 4x cheaper than V3.2 in FLOPs at 1M. Everyone else economizes on the experts, but V4 economizes on the attention and spends the savings.

Pre-training

Pre-training is the least surprising part. DeepSeek trained Flash on 32T tokens and Pro on 33T, a refreshed corpus that filters out machine-generated web text and leans harder on math, code, and long documents, with the sequence length ramped from 4K to 1M during the run. Before any post-training, Pro-Base already beats V3.2-Base almost everywhere, and Flash-Base wins most benchmarks on a third of the active parameters, ceding only some code and math rows.

The detail the model card omits is the hardware. Per ChinaTalk and the Chinese tech outlet 36Kr, V4 was trained on Nvidia Hopper after a mid-2025 training failure during an attempted migration to Huawei Ascend. That failure says domestic silicon still can't carry a frontier-scale training run, and the detour back to Nvidia cost DeepSeek months. What DeepSeek did ship is native MXFP4 weights and a TileLang kernel layer, which let the model run inference on domestic Ascend, Cambricon, and Biren chips it was never trained on.

Post-training

V4's post-training is a clean two-stage design: train a panel of domain specialists, then distill them into one model.

Stage one trains the teachers. DeepSeek forks the base into separate specialists for math, code, agentic tasks, and instruction following, each taken through SFT and then RL with GRPO. The ingredients are standard, but running a full, separate RL pipeline per domain is not, and it is the expensive half of the recipe. Two choices stand out:

• Every specialist still starts with SFT. Last year, DeepSeek itself tried to skip this stage with R1-Zero, bootstrapping reasoning with RL alone, straight from the base model. Two generations later, they gave up on this idea, and SFT is still the backbone of the pipeline.

• The reward model is the actor. Where a rule check can grade the task, it does. Everywhere else, DeepSeek trains no separate reward model: the specialist itself learns to judge, scoring rubric-guided trajectories in the same GRPO run that trains it to answer. The report claims a minimal set of human annotations suffices, because the model's own reasoning carries the grading.

Stage two consolidates them. Instead of one more multi-domain RL phase to merge those skills, the way V3.2 did, V4 distills every specialist into a single student. The student generates its own rollouts, more than ten teachers score them, and it minimizes a weighted reverse-KL against the teacher distributions on its own outputs.

None of that framing is DeepSeek's. Reverse-KL distillation on the student's own samples is GKD, from Agarwal et al. in 2023, and the report credits its on-policy form to Thinking Machines, which used it last October to rebuild Qwen3's reasoning cheaply. The multi-teacher, multi-domain version is also used by Nvidia's Nemotron Cascade 2 as MOPD. DeepSeek's contribution is the engineering that makes it cheap at this scale.

That engineering is full-vocabulary logit distillation. Most on-policy distillation, including the Thinking Machines recipe, scores only the token the student sampled, a cheap but noisy per-token reverse-KL that the report rejects for that reason. V4 matches the entire teacher distribution at every position instead, which is denser and more stable but should be impossible to store. The trick is to keep the gradient dense without ever materializing the logits:

1. Cache only each teacher's last-layer hidden states, then rebuild the full logits on the fly through that teacher's prediction head.

2. Order training samples by teacher index, so only one teacher head is resident in memory at a time.

3. Compute the exact KL with a custom TileLang kernel.

That scheduling is what turns a textbook objective into something you can run with a dozen trillion-parameter teachers.

Nvidia's MOPD releases put the same ingredients in a different place. In Cascade 2 and again in Nemotron 3 Ultra, the student is a generalist that already went through a multi-domain RL stage, so distillation arrives on top of RL as a consolidation layer, grading one sampled token at a time. V4 hands the entire consolidation to distillation, with no unified RL before or after. That choice is what forces the dense loss: a sampled-token signal can transfer deltas onto an RL-trained student, but it is too noisy to carry the merge alone.

DeepSeek's argument for cutting the RL stage is interference. A single multi-domain reward mixture invites reward hacking and capability drift, and Nvidia saw the cost directly: Nemotron 3 Ultra's mixed RL stage plateaued, with more than a dozen environments leaving each domain too little signal per batch. Specialists trained to convergence dodge that fight, and the reasoning-effort modes come almost for free because the teacher pool already spans them.

Nvidia's stacking has its own logic, though. On-policy distillation only grades the student's own rollouts, so a student that has already been through RL samples stronger rollouts for the teachers to refine. Keeping RL in the pipeline also means some stage optimizes the model directly against a reward, while V4's student only ever matches teacher distributions, a gap worth remembering when the benchmarks get to calibration. V4 and Nemotron 3 Ultra both pay for ten-plus specialist pipelines either way, so the choice is really about where the consolidation risk sits, and neither lab has run the ablation that would settle it.

Benchmarks

DeepSeek's own numbers put V4-Pro at maximum effort against the closed and open frontier.

V4-Pro is competitive but not dominant. It leads on LiveCodeBench, sits mid-pack on knowledge and reasoning, and trails Gemini 3.1 Pro badly on SimpleQA-Verified, a factuality test. Its agentic GDPval-AA score is the best of any open-weight model, behind only the closed GPT-5.4 and Opus 4.6.

Artificial Analysis has since run V4 independently.

Their composite index puts V4-Pro at 52, just behind the joint open leaders Kimi K2.6 and MiMo-V2.5-Pro at 54. Flash lands around Claude Sonnet 4.6's level for a fraction of the price.

The recurring weakness is factuality. On their AA-Omniscience, which scores whether a model declines or guesses when it doesn't know the answer, V4 guesses 94% of the time (Pro) and 96% (Flash), against 39% for Kimi K2.6 and 34% for MiniMax-M2.7. It is one benchmark, but it tracks both independent factuality tests and the long-running complaint that DeepSeek's reasoning models trade caution for fluency.

Against closed models, the gap is concrete. On DeepSeek's own R&D coding benchmark, V4-Pro-Max passes 67% to 80% for Claude Opus 4.6 in thinking mode, and the report concedes V4 trails the frontier by three to six months.

Verbosity is the other catch. V4 Pro burns 190M output tokens to finish the Intelligence Index and Flash burns 240M, among the most of any model tested, so the cheap per-token price buys a lot of tokens. The effective cost gap to a closed model is smaller than the sticker, though Pro still runs the full index for about a quarter of what Claude Opus 4.7 costs.

Reception split along the Pro/Flash line. Flash is the one that won traffic. On OpenRouter it climbed to roughly 3.64 trillion tokens a week and briefly out-consumed Claude, before Tencent's Hy3 preview took the top. The cheapest providers list it near $0.10 in and $0.20 out. Pro drew the more measured takes, an "almost on the frontier" thread and early endpoint instability that has since settled.

The bigger picture

V4 reads as the start of what ChinaTalk calls a post-DeepSeek era. The model is excellent, yet DeepSeek concedes it is months behind the frontier, a strange thing to write about the lab that defined the open-weight conversation a year ago. The slip is mostly not technical: DeepSeek lost core researchers to Tencent, ByteDance, and Xiaomi, ByteDance's Doubao took the consumer market while DeepSeek stayed in research mode, and the chip migration cost it a run. The funding gap behind all of that is structural and won't close on benchmark wins.

What will outlast the news cycle is the recipe. Multi-teacher on-policy distillation at trillion-parameter scale is the strongest sign yet that specialize-then-consolidate is becoming the default shape of frontier post-training. DeepSeek is plain that the rest is debt: it calls the architecture complex and kept tricks it doesn't fully understand to de-risk the run. The open question is whether the next lab to run this pipeline can fix the calibration hole. Either way, pushing a 1.6T model through a new optimizer and a ten-teacher distillation run is a serious piece of work, and the team did it while the ground moved under them.

Resources

• Technical report: DeepSeek-V4

• Hugging Face: DeepSeek-V4-Pro, DeepSeek-V4-Flash

• Artificial Analysis: DeepSeek V4 Pro

On June 4, 2026, Nvidia closed out its Nemotron 3 lineup with Nemotron 3 Ultra, a 550B-parameter MoE with 55B active and fully open weights. It is the most intelligent open-weight model from a US lab, but it still trails the Chinese open frontier led by Kimi K2.6.

Nvidia's pitch is not the highest benchmark scores, but the highest throughput on long-running agentic workloads. Let's walk through the recipe and poke at the inference numbers behind that claim.

Architecture

Ultra is the biggest of the three Nemotron 3 models. All three are MoEs that pair Mamba-2 layers, which run in linear time with no growing KV cache, with a thin set of attention layers for exact recall.

The one constant worth noting is sparsity: all three keep about 90% of their weights idle per token, even as total parameters grow 17x. The hybrid backbone is otherwise standard, the same linear-time-plus-attention mix Liquid, MiniMax, and Qwen use. Two blocks are more particular:

• LatentMoE (paper): compresses the expert dimension about 4x, which lets Nvidia run 512 experts with 22 active per token at normal cost. Nano skips it for a conventional 128-expert, top-6 setup.

• Multi-Token Prediction (MTP): two tied heads predict ahead, baking the speculative-decoding drafter into the weights. DeepSeek introduced MTP in V3, and Nvidia ties the two heads so the drafter stays aligned with the model.

Laid side by side, the family is one recipe instantiated at three different lengths:

Pre-training

Ultra is the largest model pretrained end-to-end in NVFP4, Nvidia's 4-bit float format. The recipe is Super's, scaled up: keep the most sensitive tensors in higher precision and run the rest in FP4.

What surprised me most is the token budget. Ultra trained on only 20T tokens, fewer than the 25T behind both Nano and Super. Normally the biggest model gets the most data, not the least. Nvidia cut the run short because, by its own account, it diverged twice:

• Around 8T tokens, a throughput tweak that moved the output layer's gradient accumulation to BF16 quietly zeroed out the MTP heads' contribution, and the MTP loss spiked before the main loss showed anything. Reverting to FP32 fixed it.

• Around 16T tokens, a second divergence Nvidia never root-caused. They annealed the learning rate early and stopped at 20T instead of the planned horizon. The report links it to dead experts and a ballooning residual stream, but only by correlation.

The curves show the call. The original run (orange) spikes near 16T, the early-annealed runs keep dropping, and the 20T version (red) is the one Nvidia shipped.

NVFP4 pretraining is still experimental at this scale. The run diverged twice; Nvidia never explained the second one, and the instability cost Ultra a fifth of its planned tokens. FP4 buys inference throughput, but it still adds real training risk.

Post-training

Post-training runs in four stages: a base model, supervised fine-tuning, reinforcement learning, then several rounds of distillation.

The first two stages are routine. SFT is multi-domain and keeps the shared-weight MTP objective. RLVR then runs a single reinforcement-learning stage with verifiable rewards across every domain at once: agentic, reasoning, chat, safety, and long context, producing one generalist checkpoint. That generalist plateaus because, with more than a dozen environments mixed together, each domain gets too little signal per batch.

To get past that, Nvidia also trains more than ten domain specialists, each a separate model tuned hard on one area: SWE, terminal use, search, safety, and so on, through its own SFT and RL. These specialists are the teachers, and the generalist is the student.

It merges them with on-policy distillation (GKD, 2023), where the student generates its own rollouts and each teacher grades them token by token. Nvidia runs this over two iterations, with a brief warm-up SFT first so the student's rollouts stay on the teachers' distribution. The first round distills the specialists into an interim model, the second reuses them plus a fresh batch to produce the final Ultra. Nvidia calls it Multi-teacher On-Policy Distillation (MOPD)1.

The contrast with DeepSeek-V4 is the interesting part. DeepSeek dropped its RL stage and replaced it with multi-teacher distillation. Nvidia keeps RL and stacks MOPD on top. DeepSeek treats distillation as a replacement for RL, Nvidia treats it as a consolidation layer that RL feeds into.

The payoff is lopsided. Recovery runs from over 170% on Terminal Bench, where the student beats its teacher, down to 17% on HLE. Execution and tool-use recover almost fully, and hard reasoning barely moves.

Nvidia is precise about why the reasoning gap stays open. Its reasoning teacher gained its edge from extra SFT and RL on a DeepSeek-V4-Pro reasoning mixture the student never saw, so the teacher's reasoning paths fall outside what the student samples. On-policy distillation only grades the student's own rollouts, so it cannot transfer a skill the student never produces. It amplifies reasoning the student can already reach, not reasoning beyond it.

The same pattern shows up elsewhere. Thinking Machines lifted a Qwen3-8B from 60 to 70 on AIME with this method, because that student could already reach AIME-level traces, and DeepSeek-V4 reasons well because its base could. My read is that the lever is the base. Ultra, undertrained by 5T tokens, does not reach far enough on the hardest problems, which is why HLE barely moves.

Benchmarks

Nvidia's own table and Artificial Analysis agree on where Ultra stands. The Nvidia table has it among the best open models on agentic and long-context work, and behind DeepSeek-V4 and Kimi on hard reasoning like HLE and GPQA.

Read down the columns, and the problem is hard to miss. Ultra holds the top score in only one of eleven benchmarks (RULER at 1M), and a Chinese model, usually Kimi, takes most of the rest. For a 550B flagship measured against models its own size and smaller, that is a thin result.

Artificial Analysis breaks it down the same way. Ultra leads on instruction-following, professional tasks, and long context, and falls behind on coding and long-horizon planning, exactly where Kimi K2.6 and GLM-5.1 keep their edge. On the Intelligence Index, it’s the top US open model, but it lands about six points behind Kimi K2.6, a real gap on a 100-point scale.

None of this is a surprise after the post-training section. The place MOPD gained the least, hard reasoning, is exactly where Ultra falls short, so the category split here just echoes the recovery numbers.

Speed is the one axis where Ultra clearly wins. A pre-release DeepInfra endpoint ran it above 300 tokens per second, against the 50 to 100 typical of comparable Chinese models. Nvidia's own "up to 6x" is shakier. It is 5.9x against Kimi but only 1.6x against Qwen, the one rival with a similarly lean active path, and it was timed on Nvidia's own stack against rivals on vLLM. As implicator.ai put it, these are "Nvidia's own numbers, against rivals Nvidia chose." The lead seems real, but the exact multiple might not hold in practice.

The bigger picture

Nvidia barely needs Ultra to win. It sells the silicon nearly every frontier model trains and runs on, so trailing Kimi costs it little. The real risk is not a better Chinese model, but a better Chinese chip. DeepSeek-V4 already trains and serves on Huawei Ascend, and as more Chinese labs move off Nvidia, the moat Ultra quietly defends erodes.

None of this takes away from the work. A stable 550B run in 4-bit, a ten-teacher distillation pipeline, and fully open weights are hard to land. Congratulations to the team that shipped it!

Resources

• Technical report: Nemotron 3 Ultra: Open, Efficient MoE Hybrid Mamba-Transformer for Agentic Reasoning

• Hugging Face: BF16 and NVFP4

• Recipes & data: github.com/NVIDIA-NeMo/Nemotron

On March 19th, 2026, Nvidia dropped Nemotron-Cascade 2, a 30B MoE model with 3B activated parameters that achieves gold-medal performance on the 2025 International Mathematical Olympiad (IMO), the International Olympiad in Informatics (IOI), and the ICPC World Finals. This was previously achieved by DeepSeek-V3.2-Speciale at 671B-A37B, a model with 20x more parameters.

The most interesting part of this release isn’t the headline benchmarks, but the open post-training recipe. Nvidia open-sourced the model weights, as well as the SFT and RL datasets. Let’s analyze the technical details shared in the technical report.

What is Nemotron Cascade 2?

Nemotron-Cascade 2 is post-trained from Nemotron-3-Nano-30B-A3B-Base, Nvidia Mamba2-Transformer hybrid MoE base model that shipped in December 2025. Same base, same architecture: the only difference between Cascade 2 and the original Nemotron-3-Nano instruct model is the post-training pipeline. Both models start from the same checkpoint, but Cascade 2 outperforms Nano on nearly every benchmark. Internal competition looks fierce at Nvidia!

Cascade 2 is a direct competitor to Qwen3.5-35B-A3B. It also operates in two modes: thinking (reasoning traces inside <think> tags) and instruct (a prepended empty <think></think> block suppresses chain-of-thought).

Cascade 2 outperforms Qwen3.5-35B-A3B and the larger Nemotron-3-Super-120B-A12B across reported math, code reasoning, alignment, and instruction following benchmarks:

† Numbers in brackets refers to Tool-Integrated Reasoning (TIR) results.
‡ Official numbers when available, otherwise evaluated using the recommended settings.

The competition results look great: 35/42 on IMO 2025 (gold), 439.28/600 on IOI 2025 (gold), 10/12 on ICPC World Finals (gold). But Qwen3.5-35B-A3B significantly outperforms on knowledge-intensive benchmarks, like MMLU-Pro (85.3 vs 79.8), GPQA-Diamond (84.2 vs 76.1), HLE (22.4 vs 17.7), as well as agentic tasks like SWE Verified (69.2 vs 50.2), Terminal Bench 2.0 (40.5 vs 21.1), Tau-2 Bench (81.2 vs 58.9).

Cascade 2 is clearly optimized for math and code reasoning at the expense of breadth on knowledge and agentic tasks. Nvidia acknowledges this directly in the paper, pointing to stronger pre-training and agentic RL as future priorities.

Supervised Fine-Tuning

The pipeline has two major phases: SFT, then Cascade RL with a new addition called Multi-domain On-Policy Distillation (MOPD).

The SFT stage is a single-stage training run with 15.9M samples packed into 256K token sequences, converging after about 1.5 epochs. Here are the most interesting domains:

• Math: 1.8M tool-calling samples + 2.6M non-tool samples for competition math, generated by DeepSeek-V3.2 and DeepSeek-V3.2-Speciale. Plus 816K proof generation/verification samples.

• Code reasoning: ~165K unique prompts from Codeforces, AtCoder, AIZU, CodeChef, deduplicated via I/O fingerprinting and n-gram analysis (removing ~24.2% redundancy). Responses by gpt-oss-120b. Final dataset: 1.9M Python traces + 1.0M C++14 traces + 1.3M Python tool-calling traces.

• SWE: 125K agentic + 389K agentless samples. They found that combining agentless and agentic data helps both scaffolds: pass@1 goes from 48.9 to 49.9 and pass@4 from 62.8 to 65.2 on SWE-bench Verified via OpenHands.

• General chat: 4.9M reasoning-on + 372K reasoning-off samples, plus 700K multi-turn conversations synthesized by having two gpt-oss-120b instances role-play user and assistant.

• Terminal agent: 490K samples using the Terminus 2 framework with Docker-based execution loops.

To sum up, Nvidia uses three main teacher models: gpt-oss-120b for most code and general tasks, DeepSeek-V3.2/Speciale for math, Qwen3-235B variants for chat and tool calling.

Cascade Reinforcement Learning

The core idea from Nemotron-Cascade 1 carries forward: instead of blending prompts from every domain into a single RL run (the DeepSeek-R1/Qwen3 approach), you train RL stages sequentially, one domain at a time.

This is also something we’ve successfully implemented at Liquid last year. The benefits are threefold: (1) resistance to catastrophic forgetting, (2) domain-specific hyperparameter tuning, and (3) compute savings from homogeneous batches within the same domain.

The Cascade 2 pipeline runs in this order:

1. IF-RL (Instruction Following) for ~180 steps

2. Multi-domain RL (MCQA, tool calling, structured output) for ~70 steps

3. Multi-domain On-Policy Distillation (MOPD) for ~40-50 steps

4. RLHF (human preference via GenRM) for ~30 steps

5. Long-context RL (32K input, 49K max sequence) for ~30 steps

6. Code RL (competitive programming, binary rewards)

7. SWE RL (agentless + agentic) for ~40-50 steps

The ordering changed from Cascade 1: IF-RL is now first (it used to come after RLHF) because they discovered that IF-RL hurts human alignment, but subsequent RLHF can recover that regression. This is in line with our own training process at Liquid, where IF training is always the first stage.

Throughout all stages, they use GRPO with strict on-policy training (importance ratio always 1.0) and no KL divergence term. This reduces the GRPO objective to standard REINFORCE with group-normalized rewards and token-level loss. The one exception: RLHF uses a small KL coefficient of 0.03 to preserve capabilities from other domains.

Multi-domain On-Policy Distillation

Multi-domain On-Policy Distillation is the headline technical contribution. Even with careful cascade ordering, you still see capability drift across benchmarks. Some RL stages (particularly code RL) reduce model entropy and shorten reasoning traces, which hurts math. RLHF trades off against instruction following.

The solution: take the best checkpoint from each domain (the “teacher”) and distill on-policy into a single student. The teachers all come from within the Cascade RL pipeline itself , i.e., they share the same SFT initialization, same tokenizer, same vocabulary. No external models needed.

The MOPD objective defines a token-level distillation advantage:

This is positive when the domain teacher is more confident than the current student on a given token. Because it’s token-level rather than sequence-level (like GRPO), it provides much denser gradients.

The practical setup uses three domain teachers: (1) the SFT checkpoint as the math teacher (SFT data quality is strong enough), (2) an RLHF checkpoint, and (3) the post-IF-RL + multi-domain-RL checkpoint. They sample prompts from the corresponding domain data pools and use truncated importance weighting with bounds [0.5, 2.0] to handle train-inference mismatch.

Code RL is particularly interesting: it uses an aggressively filtered set of only 3.5K prompts, removing anything that gpt-oss-120b solved correctly in all 8/8 rollouts. Only hard problems with strong test cases remain. The max response length is 118K tokens (up significantly from Cascade 1), with 16 rollouts per sample. Strict binary rewards, no partial credit, which are explicitly avoided to prevent reward hacking.

The bigger picture

Nvidia is explicitly pushing "intelligence density" as the metric that matters, i.e., task-level competence per active parameter. This also confirms a trend we've been tracking: post-training is becoming the dominant lever for model capability. The base model here is the same one that ships as Nemotron-3-Nano. The entire performance delta comes from SFT data curation, RL pipeline design, and MOPD.

But Nvidia’s relative weakness on agentic benchmarks (Tau-2 Bench: 58.9 vs Qwen’s 81.2, Terminal Bench 2.0: 21.1 vs 40.5) suggests that Cascade RL, while excellent for verifiable reasoning tasks, hasn’t yet cracked agentic capabilities at the same level. The gap on knowledge-heavy benchmarks (MMLU-Pro, GPQA) also points to pre-training limitations that no amount of post-training can fully overcome.

What’s next

As predicted in our previous article, Nvidia officially announced Nemotron 3 Ultra. But, more interestingly, they also announced the Nemotron Coalition, a collaboration to produce open frontier LLMs. The Cursor/Kimi K2.5 drama from the same week (where Composer 2 turned out to be Kimi K2.5 with RL on top) just reinforces the point: base models are commodities now, post-training is the differentiator, and Nvidia wants to be the platform everyone post-trains on.

Links:

• Technical report (PDF)

• Model weights (HuggingFace)

• Training data collection

• NVIDIA project page

On March 11th, 2026, NVIDIA released Nemotron 3 Super, the second model in the Nemotron 3 family following Nano's December debut. Super is a 120B total / 12B active parameter hybrid Mamba-Transformer MoE model, and it brings several features: LatentMoE, native NVFP4 pretraining, and multi-token prediction with shared-weight heads. They released not only model weights, but also datasets and training recipes, which makes it particularly interesting to analyze.

With this release, NVIDIA’s goal is to match gpt-oss-120b and Qwen3.5-122B-A10B on accuracy and deliver higher inference throughput. In practice, they claim 2.2x throughput over gpt-oss-120b on an 8k input / 64k output workload.

Architecture

Super adopts a hybrid Mamba-2/GQA architecture with 88 layers. In comparison, gpt-oss-20b uses Grouped Multi-Query Attention with 36 layers, and Qwen3.5-122B-A10B is a Gated Delta Network hybrid with 48 layers. In the Super architecture, Mamba-2 blocks handle the majority of sequence processing (linear-time complexity, constant memory during generation), while attention layers are inserted at strategic intervals as "global anchors" for precise associative recall.

The most interesting contribution here is LatentMoE (Elango et al., 2026). Standard MoE designs route tokens from the model's full hidden dimension directly to experts. LatentMoE wraps the routed expert path with two shared linear projections: a down-projection from the hidden dimension d = 4096 to a latent dimension l = 1024, then expert computation happens entirely in this compressed space, and results get projected back up.

The key insight is that reducing the expert dimension by a factor of d/l = 4 lets you reinvest those savings into both more total experts and higher top-k. Nemotron 3 Super runs 512 total experts with top-22 routing, while standard MoE designs typically use 128 experts with top-6 or top-8. You get 4x the expert count and roughly 4x the active experts, at approximately the same compute and communication cost.

Why does this matter?

• In latency-sensitive inference, MoE is dominated by the memory bandwidth cost of reading expert weights.

• In throughput-oriented serving, it’s dominated by all-to-all routing communication.

LatentMoE reduces both bottlenecks by operating in the compressed space. The technical report includes ablations showing LatentMoE consistently outperforms standard MoE across code, math, and commonsense benchmarks at matched parameter and FLOP budgets.

Super includes 2 Multi-Token Prediction (MTP) layers that predict multiple future tokens from each position. Concretely, this provides native speculative decoding without needing a separate draft model.

The design choice that stands out is the shared-weight formulation across prediction heads. Standard MTP implementations use independent heads for each offset (e.g., one head for token n+2, another for n+3). This works fine during training but limits speculative decoding to at most N draft tokens, and degrades when you try to reuse a fixed-offset head autoregressively because of the training-inference distribution mismatch.

By sharing parameters, the heads see multiple offsets during training, which regularizes them into a more robust drafting model. The technical report shows an average acceptance length of 3.45 tokens on SPEED-Bench (draft length of 7), beating DeepSeek-R1 (2.70) and slightly edging out Qwen3-Next (3.33).

Pre-training

Nemotron 3 Super is pretrained natively in NVFP4 (NVIDIA's 4-bit floating-point format) for the entire 25T token run. Note that gpt-oss was already pre-trained in 4-bit precision, but using MXFP4, a lower-accuracy alternative.

NVFP4 is a narrow format that requires careful handling of numerical stability. The technical report is transparent about the challenges: they observed growing zero-valued weight gradients (reaching 7% of total parameters by end of pretraining), caused by NVFP4 underflowing already-small gradient values to zero.

NVIDIA investigated whether “healing” by switching to higher precision (MXFP8) before learning rate annealing would help. Their ablation showed the switch improved loss trajectory but yielded no downstream accuracy gains. So the final model is pure NVFP4 for the full token horizon. That’s a strong statement about the viability of low-precision pretraining at scale.

The pretraining corpus includes:

• Two-phase structure: Phase 1 (80%, 20T tokens) focused on diversity and coverage; Phase 2 (20%, 5T tokens) focused on high-quality data

• Synthetic code datasets: ~15M Python problem-solution pairs generated and cleaned via gpt-oss models

• Synthetic formal logic, economics, and MCQ datasets generated using Qwen3-235B and DeepSeek-V3 for verification

During the stable phase of their Warmup-Stable-Decay (WSD) learning rate schedule, the constant learning rate makes individual checkpoint benchmarks noisy from step to step. Rather than running expensive dedicated LR decay branches just to gauge model quality, NVIDIA applies offline checkpoint merging (weighted averaging over a sliding window of recent checkpoints) to produce cleaner readouts.

Following the WSM framework (Tian et al., 2025), which establishes a formal connection between LR decay and model merging, they use a minus-sqrt decay emulation to compute merge coefficients. The results are striking: across a suite of 12 benchmarks, the best merge consistently outperforms the corresponding raw trained checkpoint by 2-4 points. NVIDIA estimates this technique saved roughly 4T tokens of compute by avoiding dedicated decay runs needed purely for quality assessment.

Post-Training

Super is SFTed on ~7 million samples from a broader corpus of 40 million, covering reasoning, instruction following, coding, safety, and multi-step agentic tasks. The SFT data explicitly trains the model for reasoning ON/OFF modes via the chat template.

Interestingly, they introduce a two-stage SFT loss function. Stage 1 corresponds to the standard approach (token-level average), where long conversations dominate the gradient signal. Stage 2 normalizes each conversation’s loss and then averages equally across the batch. This prevents shorter responses from being drowned out by long reasoning traces. It’s probably particularly important with the reasoning ON/OFF modes.

After SFT, NVIDIA runs RL across 21 distinct environments drawn from 37 datasets, generating approximately 1.2 million environment rollouts. The environments span a diverse set of tasks: math, code, STEM, instruction following, safety, long context, agentic tool use, and reasoning gym. The training uses a very traditional GRPO with importance sampling to handle discrepancies between the training and rollout policies.

This is followed by a specific SWE stage focusing on agentic coding capabilities to autonomously solve GitHub issues. They use OpenHands as an agentic harness, which allows them to directly train on the tool formats expected by Claude Code and Codex.

After this, a separate RLHF stage refines conversational quality. NVIDIA trained a generative reward model (GenRM) by running GRPO on Qwen3-Nemotron-235B-A22B on a combination of Helpsteer 3 and lmarena-140k. Given a conversation history, a new user query, and two candidate assistant responses, the GenRM reasons explicitly about the strengths and weaknesses of each response, produces individual helpfulness scores, and generates a ranking.

On benchmarks, Qwen3.5-122B generally leads on raw accuracy, particularly on knowledge-heavy and agentic benchmarks. gpt-oss-120b edges ahead on some math and coding tasks. Nemotron 3 Super wins on HMMT Feb25 math and long-context retrieval, and holds competitive across the board. This is slightly underwhelming, but Super is also significantly faster than both models, even though the comparison with Qwen3.5 in BF16 is unfair.

Artificial Analysis independently confirms strong throughput numbers: ~478 tokens/second across providers, with 0.56s time to first token. Though they also flag the model as insanely verbose, generating 110M tokens during their evaluation suite versus an average of 7.3M. This extreme level of verbosity could erase most of the throughput gains in practice: gpt-oss-120b in high effort mode generates 77M tokens, and Qwen3.5-122B in reasoning mode produces 91M tokens.

What’s Next

Nemotron 3 Ultra (~500B parameters) is still expected. Given the LatentMoE scaling properties and the NVFP4 pretraining recipe, Ultra should be interesting. It would also directly compete with frontier Chinese MoEs and give us a good understanding of NVIDIA’s internal capabilities.

Links:

• Nemotron 3 Super Technical Report (PDF)

• Model Weights (BF16) | FP8

• NVIDIA Blog Post

• LatentMoE Paper

• Deployment Cookbooks (GitHub)

• Pre-training Datasets | Post-training Datasets

Beijing-based Moonshot AI released Kimi K2.5 on January 27, 2026. Beyond traditional claims on benchmarks like HLE (50.2% with tools), coding, and vision, this release introduced the idea of “Agent Swarm”. Two weeks in, I wanted to revisit Kimi K2.5 and compare it with other recent releases like GLM-5, MiniMax-M2.5, and Qwen3.5.

What K2.5 Actually Is

Kimi K2.5 is one of the largest open-weight models with 1.04 trillion parameters and 32B activated parameters per token. This is significantly bigger than MiniMax-M2.5 (230B-A10B), Qwen3.5 (397B-A17B), and GLM-5 (1T-32B). It uses 384 experts with 8 activated per token, MLA attention, SwiGLU activation, and a 256K context window. The architecture is identical to Kimi K2, which shipped back in mid-2025.

What’s new is the training. K2 was originally pre-trained on 15T text-only tokens. K2.5 then continues from a near-end K2 checkpoint over an additional ~15T mixed visual and text tokens, plus ~1T for ViT training and ~700B for long-context mid-training (see tech report). If the numbers don’t overlap, that’s roughly 32T tokens across the full pipeline (vs. 28.5T tokens for the text-only GLM-5 — Qwen3.5 and MiniMax-2.5 haven’t released any numbers).

The vision encoder is MoonViT-3D, a 400M parameter native-resolution ViT based on SigLIP-SO-400M, with a NaViT packing strategy that handles variable-resolution images. For video, consecutive frames are grouped in fours and temporally pooled, achieving 4x compression. Qwen3.5 also used early fusion with a different strategy. Late fusion seems dead for frontier models.

The model ships in native INT4 precision (~595GB), not FP8/BF16. Moonshot used quantization-aware training during post-training to achieve this. Unsloth’s dynamic 1.8-bit quant brings it down to ~240GB, runnable on a single 24GB GPU with sufficient RAM offloading at ~10 tokens/sec.

Agent Swarm

Before we get to the benchmarks, I want to talk about the most technically interesting part of this release: Agent Swarm. And, specifically, the Parallel-Agent Reinforcement Learning (PARL) framework that powers it.

The core idea is that, instead of executing agent tasks sequentially (tool call, observe result, reason, next tool call, etc.), K2.5 learns to decompose problems into parallelizable subtasks and delegates them to sub-agents. The orchestrator is trainable, but the sub-agents are frozen copies of intermediate policy checkpoints. Only the orchestrator’s parameters get updated via RL.

This is designed this way because an end-to-end co-optimization of the orchestrator and sub-agents would create a credit assignment nightmare. In other words, if the final answer is wrong, is it the orchestrator’s fault for bad delegation, or the sub-agent’s fault for bad execution?

The training had to solve two emergent failure modes: serial collapse (the orchestrator defaults to safe sequential execution) and spurious parallelism (it spams sub-agent creation to game metrics without meaningful decomposition). Auxiliary reward terms push past both early in training, then anneal to zero so the final policy optimizes purely for task success.

According to Moonshot AI, BrowseComp jumps from 60.6% (single agent) to 78.4% (swarm), WideSearch F1 from 72.7% to 79.0%, and execution time drops 3-4.5x on suitable tasks. However, for comparison, Qwen3.5 reports a BrowseComp score of 78.6 using the same discard-all strategy as K2.5, but without a swarm mechanism. To my knowledge, this is the first time an open-weight model has been trained to parallelize agentic work rather than having it imposed by external scaffolding. Exciting stuff!

Benchmarks

Moonshot’s reported numbers are competitive. Here’s the context that matters.

Where K2.5 leads (among all models, including proprietary): HLE-Full with tools (50.2% vs. GPT-5.2’s 45.5%), BrowseComp with swarm (78.4%), OCRBench (92.3%), MathVista (90.1%), InfoVQA (92.6%). Note that MiniMax M2.5 reports 76.3% for BrowseComp without any swarm mechanism.

Where K2.5 is competitive but behind frontier: AIME 2025 (96.1% vs. GPT-5.2’s 100%), SWE-Bench Verified (76.8% vs. Claude Opus 4.5’s 80.9%, MiniMax M2.5’s 80.2%, Qwen3.5’s 76.4%), GPQA-Diamond (87.6% vs. GPT-5.2’s 92.4% and Qwen3.5’s 88.4%), Terminal-Bench 2.0 (50.8% vs. Claude’s 59.3%).

Where K2.5 shows clear gaps: WeirdML (46% vs. 72% for GPT-5.2) is telling. On Artificial Analysis’s AA-Omniscience knowledge index, K2.5 scores -11 (correct minus incorrect), meaning it hallucinates more than other frontier models. Claude Opus 4.5 scores +10, Gemini 3 Pro scores +13.

Kimi K2.5 was also the open-weight model with the high Intelligence Index on Artificial Analysis before the release of GLM-5. It still leads compared to Qwen3.5 and MiniMax-M2.5, which are smaller models. Note that it doesn’t mean that GLM-5 is necessarily superior (and vice-versa): many details get ironed out with these streamlined benchmarks, and features like Agent Swarm are not properly taken into account.

Community Feedback

Two weeks of community usage provide a good picture of the model’s real impact.

Coding: K2.5 is legitimately strong, especially for front-end work and visual-to-code tasks. Kilo Code reports it rapidly climbed to top-performer status for architectural planning. Multiple developers on r/LocalLLaMA report building complete projects at ~1/8th the cost of Opus. But the pattern in skeptical reviews is consistent: K2.5 often generates verbose, over-engineered code on the first pass, then simplifies when asked. Opus and Codex tend to get it right the first time.

Agent Swarm: Impressive when it works. Users report effective parallel web research and multi-niche data collection. But follow-up editing of swarm outputs is painful, and sub-agents can drift into inconsistent definitions for shared concepts. The spreadsheet use case (compiling data across rows) exposed this: every agent used slightly different column definitions.

Vision: The first open-weight model where vision feels genuinely competitive. Nathan Labenz tested it on a scanned document transcription task where Chinese models historically lagged, and K2.5 matched Gemini 3 level. Qwen3.5 also makes a strong play here: 90.3 on MathVista, 85.0 on MMMU, and native UI screenshot understanding with element detection. Creative writing and personality, on the other hand, are clearly behind Opus. And multiple users discovered K2.5 sometimes identifies itself as Claude, a strong signal about training data provenance.

The Bigger Picture

Verbosity and cost. When Artificial Analysis evaluated K2.5, the model generated 89 million output tokens. The median for comparable models is 14 million. At $0.60/$3.00 per million input/output tokens, the per-token price looks cheap, but when the model produces 6x more tokens per task, effective costs still spike. Kilo Code’s week-long free trial confirmed this: usage surged past 50B tokens/day, and they concluded the model’s verbosity dilutes the savings from input caching. This is the main issue with this model.

Multimodal training insights. The technical report contains a finding worth highlighting for anyone building multimodal models. Given a fixed vision-text token budget, early fusion with a low vision ratio (10% vision from the start) outperformed late fusion with a high ratio (50% vision injected at the 80% mark) across every metric. Moonshot also introduces “zero-vision SFT,” where text-only fine-tuning activates visual reasoning capabilities, and visual RL actually improved text-only benchmarks (MMLU-Pro: 84.7% to 86.4%, GPQA-Diamond: 84.3% to 86.4%). This bidirectional transfer validates the native multimodal approach.

Geopolitics and sustainability. K2.5 was trained on hardware constrained by US export controls, and it’s competitive with models from labs that have unconstrained chip access. Moonshot raised at a $4.8 billion valuation and is clearly spending that capital on user acquisition: the free tier is generous, Agent Swarm beta comes with free credits. The pricing is almost certainly not sustainable long-term. The license is Modified MIT, commercially free for companies under 100M MAU.

Deployment. Running K2.5 locally requires serious hardware: ~595GB at native INT4, with Unsloth’s 2-bit quant (375GB) as the practical sweet spot. On a single 24GB GPU with 256GB+ RAM, expect ~10 tokens/sec. The model works with vLLM, SGLang, and KTransformers, though the r/LocalLLaMA AMA surfaced rough edges with <think> tag parsing on some backends. Vision support in GGUF/llama.cpp is not yet available. For API users, 8 providers serve K2.5, with Fireworks leading on speed (283 t/s), DeepInfra on price ($0.90 blended), and Baseten on raw throughput (336 t/s).

What’s Next

K2.5 is a solid open-weight model with strong vision capabilities, but verbosity can be an issue for real-world usage. Because competition is particularly fierce at this range, it’s better to experiment with different options for your particular use cases.

The part I'd watch most closely is whether PARL generalizes. Whether that transfers to arbitrary real-world workflows or mainly helps on embarrassingly parallel research tasks is still open. ¯\_(ツ)_/¯

Quick links:

• Weights: huggingface.co/moonshotai/Kimi-K2.5

• Technical report: arxiv.org/pdf/2602.02276

• Blog post: kimi.com/blog/kimi-k2-5.html

• API: platform.moonshot.ai

On February 16th, 2026, Alibaba's Qwen team released Qwen3.5-397B-A17B, their next-generation foundation model. If you're getting a sense of déjà vu from the holiday timing, you should be. GLM-5 shipped on February 11th. MiniMax M2.5 landed the same day. Kimi K2.5 arrived on January 27th. And now Qwen3.5 closes out the pre-holiday window. Let’s see how it compares with the other releases.

What Qwen3.5 Actually Is

Qwen3.5-397B-A17B is a 397 billion parameter Mixture-of-Experts model with only 17 billion active parameters per token. The hosted API version is called Qwen3.5-Plus, and it ships with a 1M context window, built-in tools, and adaptive tool use out of the box.

The key innovations stack up like this:

Hybrid Attention Architecture. This is the big architectural bet. Qwen3.5 builds on the Qwen3-Next lineage, combining Gated Delta Networks (a linear attention variant) with sparse Mixture-of-Experts. The model alternates between Gated DeltaNet layers (linear attention) and full attention layers in roughly a 3:1 ratio. Sebastian Raschka has an excellent write-up on how this works, but the short version is: three out of every four transformer blocks use linear attention (which scales near-linearly with sequence length), and every fourth block uses standard full attention. The result is a model that can process long contexts more efficiently than full attention.

The Gated DeltaNet mechanism itself draws from the “Gated Delta Networks: Improving Mamba2 with Delta Rule” paper. It combines Mamba2’s gated decay mechanism with a delta rule for updating hidden states. The gated attention output gating helps eliminate attention sinks and massive activations, improving training stability at scale.

Scalable RL at Agent Scale. Qwen3.5 was trained with reinforcement learning scaled across what the team describes as "million-agent environments with progressively complex task distributions." This follows the trend we've seen from MiniMax's Forge and Zhipu's Slime: asynchronous RL infrastructure designed to handle the long-horizon, multi-step nature of agentic tasks. The details here are sparse in the initial release, but the emphasis on "robust real-world adaptability" suggests they've invested heavily in environment diversity during RL post-training.

Unified Vision-Language Foundation. Unlike Qwen3, which had separate text and vision model lines (Qwen3 and Qwen3-VL), Qwen3.5 is natively multimodal from the ground up. Early fusion training on multimodal tokens means the model doesn’t need a separate vision adapter. The team claims cross-generational parity with Qwen3 on text tasks while outperforming Qwen3-VL on visual understanding.

201 Languages. Expanded from Qwen3’s 119 to 201 languages and dialects. This is the broadest language coverage of any open model I’m aware of. Note that Qwen tends to be very generous with its definition of “language support,” and quality is not guaranteed for low-resource languages.

Attention & Sparsity

Efficient attention mechanisms and increased sparsity are two common trends with recent releases. While DeepSeek pioneered this field, every major Chinese lab has its own take on how to handle attention.

Qwen3-Next uses a 3:1 hybrid attention layout in which most layers use Gated DeltaNet (linear attention) and the remaining layers use Gated Attention (full/softmax-style attention). Kimi K2.5 and GLM-5 both use Multi-head Latent Attention (MLA), but GLM-5 also integrates DeepSeek Sparse Attention (DSA) to induce token-level sparsity on top of it. Finally, MiniMax-M2.5 is the only full attention model (see this article) with Multi-Head Attention (MHA) for reliability purposes.

The active parameter count is also worth zooming in on. At 17B active, Qwen3.5 is a lot sparser than Qwen3-235B-A22B, but in line with other recent releases:

Notably, MiniMax M2.5 has the smallest active parameter count at 10B and the same activation ratio as Qwen3.5. Kimi K2.5 is even sparser but also significantly bigger with 1T total parameters.

Benchmarks

Reasoning and Math. Qwen3.5 scores 91.3 on AIME 2026 and 94.8 on HMMT Feb 25, which is competitive but below the best-performing models (GPT-5.2 hits 96.7 on AIME 2026, Claude 93.3). Math capabilities are solid but far from dominant.

Knowledge and Instruction Following. On IFBench, it scores 76.5, beating every model in the comparison, including GPT-5.2 (75.4) and blowing past Claude (58.0). MultiChallenge tells the same story: 67.6 vs. GPT-5.2’s 57.9 and Claude’s 54.2. The model seems exceptionally good at following complex instructions but to be confirmed with real-world testing.

Agents. The agentic benchmarks paint an interesting picture. Qwen3.5 scores 86.7 on Tau2-Bench, second only to Claude (91.6). On MCPMark, it hits 46.1 vs. GPT-5.2’s 57.5 and Claude’s 42.3. On BrowseComp, Qwen3.5 reports two numbers depending on strategy: 69.0 with simple context-folding, and 78.6 using the same discard-all strategy as DeepSeek-V3.2 and K2.5. The BrowseComp split is worth noting because it highlights how much agentic benchmark scores depend on scaffolding choices, not just raw model capability.

Coding. Qwen3.5 scores 76.4 on SWE-bench Verified, essentially level with K2.5 (76.8) and Gemini 3 Pro (76.2), but behind GPT-5.2 (80.0) and Claude (80.9). On SWE-bench Multilingual, it does better at 72.0, matching GPT-5.2. SecCodeBench is a strong suit: 68.3, tied with GPT-5.2 (68.7) and Claude (68.6).

Vision. As a natively multimodal model, Qwen3.5 excels here. It scores 85.0 on MMMU (up from Qwen3-VL’s 80.6!), 88.6 on MathVision (ahead of Gemini 3 Pro’s 86.6), and 90.8 on OmniDocBench. The visual agent results are solid too: 62.2 on OSWorld-Verified and 66.8 on AndroidWorld. The ZEROBench result of 12 (vs. 10 for Gemini and 9 for GPT-5.2) is notable given the extreme difficulty of this benchmark.

Qwen3.5 is not the best at any single category, but it's remarkably well-rounded and leads on instruction following. It significantly outperforms its own Qwen3-Max-Thinking across the board despite being much smaller (397B vs 1T+).

The Bigger Picture

The attention mechanism is the new battleground. A year ago, the question was “MoE or dense?” That’s settled (and we might thank Llama-3.1-405B for this). Now the divergence is in how you handle attention. DeepSeek’s fingerprints are everywhere (MLA in K2.5 and GLM-5, DSA in GLM-5), but the Gated DeltaNet hybrid from Qwen3.5 (and initiated in Qwen3-Next) offers a new direction.

The benchmark landscape has shifted to match agentic workloads. All four releases target agentic tasks. The models are evaluated on SWE-bench, BrowseComp, HLE with tools, TAU2-Bench, and MCPMark. The era of chatbot benchmarks as the primary evaluation axis is over. Qwen3.5’s BrowseComp split (69.0 vs. 78.6 depending on strategy) is a reminder that agentic scores are increasingly a function of scaffolding and context management, not just raw intelligence.

What’s Next

The fact that Qwen3.5 ships only the 397B-A17B size on day one (”more sizes are coming”) suggests we’ll see a family rollout similar to Qwen3. It’ll be interesting to see if smaller variants also adopt the hybrid DeltaNet architecture. For this release, it feels like Qwen3-Next previewed this direction (back in September), but Qwen3.5 is the production-scale validation.

Quick links:

• Model weights: https://huggingface.co/Qwen/Qwen3.5-397B-A17B

• GitHub: https://github.com/QwenLM/Qwen3.5

• Blog: https://qwen.ai/blog?id=qwen3.5

On February 12th, 2026, barely a month after its Hong Kong IPO, Shanghai-based MiniMax dropped M2.5. The headline numbers: 80.2% SWE-Bench Verified, 51.3% Multi-SWE-Bench (first place), 76.3% BrowseComp. These are numbers that sit within a percentage point of Claude Opus 4.6 and ahead of GPT-5.2 on several agentic benchmarks. Even more interestingly to me, the model costs roughly $1 per hour of continuous operation at 100 tokens per second.

The striking part of this release is the combination: this is a 230B MoE model with only 10B active parameters, trained primarily through large-scale reinforcement learning across 200,000+ real-world environments. It handles not just code but full office productivity workflows (Word, Excel, PowerPoint).

What M2.5 Actually Is

MiniMax-M2.5 is an iterative improvement on the M2 family, which launched in late October 2025. The architecture is unchanged from M2: a Mixture-of-Experts model with 230 billion total parameters and 10 billion active per forward pass. For context, this active parameter count is tiny compared to what you’d expect from a frontier-competitive model:

• GLM-5 has 744B parameters and 40B active

• DeepSeek V3/R1 has 685B parameters and 37B active

• Not frontier, but as reference Qwen3-235B has 235B parameters and 22B active

The model comes in two API variants:

• M2.5-Lightning: 100 tokens per second, $0.30/M input, $2.40/M output

• M2.5 Standard: 50 tokens per second, $0.15/M input, $1.20/M output

The Lightning version is roughly 2x the throughput of other frontier models. The Standard version is extremely cheap. To put the pricing in perspective: Claude Opus 4.6 charges $5/M input and $25/M output tokens. Even GLM-5 that was just released is priced at $1/M input and $3.20/M output tokens, which makes it several times more expensive.

If you want to run it locally, the recommendation is vLLM or SGLang. With only 10B active parameters, the inference footprint is remarkably manageable for a model at this capability level.

Benchmarks

A few things stand out. The Multi-SWE-Bench score of 51.3% is actually first place, ahead of Opus 4.6’s 50.3%. Multi-SWE-Bench tests multilingual coding tasks, and M2.5 was trained on 10+ languages (Python, Go, C, C++, TypeScript, Rust, Kotlin, Java, JavaScript, PHP, Lua, Dart, Ruby). This is not a Python-only model, my C++/Rust friends will be happy.

The BFCL multi-turn score is quite an interesting number. First, it’s quite cheeky to only report the multi-turn split and nothing else. At 76.8% on multi-turn function calling, M2.5 leads Opus 4.6 by over 13 percentage points. It also shows immense progress in terms of multi-turn tool use compared to MiniMax M2.1 (+39.4 points).

And then there’s the independent evaluation from OpenHands (the open-source coding agent platform). Their OpenHands Index placed M2.5 4th overall, behind only Claude Opus 4.6, Claude Opus 4.5, and GPT-5.2 Codex. Graham Neubig noted the model performed particularly well on long-running tasks like developing apps from scratch, an area where smaller models have historically struggled.

Forge Reinforcement Learning

What’s technically interesting about M2.5 is how MiniMax got here. The answer is large-scale reinforcement learning, and specifically their in-house framework called Forge.

Forge is what MiniMax calls an “agent-native RL framework.” The key design decision is a decoupling layer between the training/inference engine and the agent scaffolding. This means MiniMax can plug any agent framework (Claude Code, Droid, OpenCode, custom harnesses) into the RL training loop, and the model learns to generalize across scaffolds rather than overfitting to one particular tool interface. MiniMax trained across 200,000+ real-world environments and in multiple domains, including, interestingly, their own company’s internal tasks.

To get the RL to scale, MiniMax used three key innovations:

1. CISPO (Clipped Importance Sampling Policy Optimization): Their custom RL algorithm, first proposed in the M1 paper. Rather than clipping token updates like PPO/GRPO, CISPO clips the importance sampling weights. The result is that all tokens contribute to gradient computations, even low-probability ones that are often crucial for maintaining entropy and enabling scalable RL. In controlled experiments on Qwen2.5-32B, CISPO achieved a 2x speedup compared to DAPO (ByteDance’s recent RL algorithm).

2. Asynchronous scheduling + tree-structured sample merging: To keep GPU utilization high during the inherently sequential nature of agent rollouts, they optimized the balance between throughput and sample off-policyness. They claim it achieves approximately a 40x training speedup over naive approaches.

3. Process rewards for credit assignment: Long agent trajectories make credit assignment extremely difficult. If a coding agent takes 50 steps to solve a bug, which steps were actually helpful? MiniMax introduced process-level rewards to monitor generation quality throughout the trajectory, and also directly estimated real-world task completion time as a reward signal, pushing the model toward faster solutions.

MiniMax engineer Olive Song noted on Alex Volkov‘s ThursdAI podcast that the entire M2.5 training period was about two months. For reference, the M1 reasoning model’s full RL training on 512 H800s completed in just three weeks at a rental cost of $534,700. The M2 series uses the smaller 230B MoE architecture (M1 was 456B with 45.9B active), so the compute requirements are even more favorable.

The Bigger Picture

MiniMax made a few interesting design choices that deserve a highlight.

Emergent spec-writing behavior. MiniMax notes that M2.5 has learned to proactively plan before writing code, decomposing a project before starting implementation. This is consistent with what we’ve seen in other top coding models. When they’re trained in environments that reward end-to-end task completion, they develop strategic planning behaviors. The model learns that spending tokens on upfront planning saves tokens and reduces errors downstream. This directly translates into token efficiency. On SWE-Bench Verified, M2.5 consumed an average of 3.52M tokens per task versus M2.1’s 3.72M tokens.

Office productivity. Beyond code, manipulating Office documents becomes a key feature for frontier models. MiniMax clearly targets this space and developed an internal GDPval-MM benchmark (pairwise LLM-as-judge evaluation of trajectory quality). They claim that M2.5 achieved a 59.0% average win rate against mainstream models. They also develop MiniMax Agent, their consumer-facing agentic platform where users have built over 10,000 “Experts” (specialized agent configurations).

Cost. This is the most interesting bit to me here. MiniMax is still not Opus 4.6-level, but it frames M2.5 as the most cost-efficient alternative. The blog post even states that “you can have four M2.5 instances running continuously for an entire year for $10,000.” Honestly, I don’t think the experience is consistent enough for production workloads. Early reports from OpenHands suggest the model is strong but occasionally sloppy (wrong branch pushes, missed formatting instructions). However, it shows a clear and achievable path a few generations down the line.

What’s Next

MiniMax promised a more detailed technical blog post on the Forge framework and their RL scaling laws. That’s what I’m most interested about: does performance scale linearly with the number of environments, or are there diminishing returns?

Another question I have in mind is whether the M2 series’ rapid improvement is driven by catching up to others, or whether they are genuinely pushing the frontier on agentic RL. It feels like competition is extremely tough in terms of coding, but GDPval and office productivity might be a good approach to building differentiating capabilities.

Quick links:

• Official announcement: https://www.minimax.io/news/minimax-m25

• API access: https://platform.minimax.io/docs/guides/text-generation

• MiniMax Agent: https://agent.minimax.io

• OpenHands evaluation: https://openhands.dev/blog/minimax-m2-5-open-weights-models-catch-up-to-claude

• M2.5 weights (Hugging Face): https://huggingface.co/MiniMaxAI/MiniMax-M2.5

• OpenRouter: https://openrouter.ai/minimax

• CISPO paper (MiniMax-M1): https://arxiv.org/abs/2506.13585

Updated February 18th, 2026: This article has been revised to include additional technical details from the GLM-5 technical report, “GLM-5: from Vibe Coding to Agentic Engineering.”

On February 11th, 2026, just days before the Lunar New Year, Z.ai officially released GLM-5, its new frontier large language model.

First, congrats to the Z.ai team on a strong release. GLM-5 was the new #1 open-weight model on Artificial Analysis and hit #1 among open models on LMArena’s Text Arena (score 1452, #11 overall) at the time of its release. It scores 77.8% on SWE-bench Verified, 92.7% on AIME 2026, 86.0% on GPQA-Diamond, and leads open-source models on BrowseComp, Vending Bench 2, and MCP-Atlas.

What GLM-5 Actually Is

GLM-5 is a 744B-parameter Mixture-of-Experts model with 40B active parameters per token. That’s roughly a 2x scale-up from GLM-4.5 (355B total, 32B active). GLM-5 uses 256 experts per MoE layer with 8 activated per token (5.9% sparsity), and reduces its layer count to 80 compared to GLM-4.5 to minimize expert parallelism communication overhead. Pre-training data went from 23T to 28.5T tokens.

On the attention side, GLM-5 layers DeepSeek Sparse Attention (DSA) on top of MLA via continued pre-training. The tech report ablates alternatives, including sliding window attention, Gated DeltaNet, and SimpleGDN, but all showed accuracy gaps on long-context retrieval tasks. DSA avoids this tradeoff: its dynamic indexer achieves token-level sparsity without discarding long-range dependencies, cutting attention compute by roughly 1.5-2x for long sequences.

According to Reuters, GLM-5 was trained entirely on Huawei Ascend chips using the MindSpore framework, with zero dependency on NVIDIA hardware. This hasn’t been confirmed by Z.ai, which has been on the U.S. Entity List since January 2025, which bans access to H100/H200 GPUs.

Training Pipeline

Pre-training started with a 27T token corpus in two stages: general, then code and reasoning. They refined their data selection by introducing a DCLM classifier based on sentence embeddings to identify high-quality data, plus a “World Knowledge classifier” trained on Wikipedia entries and LLM-labeled data to extract valuable information from otherwise medium-low quality sources. They also expanded the code corpus with refreshed snapshots from major code hosting platforms, resulting in a 28% increase in deduplicated unique code tokens.

Mid-training is where they progressively extended context length from 4K to 200K tokens. The focus here was specifically on long-context agentic data to ensure stability in complex workflows. It’s designed to prepare the model for the kind of long-horizon tasks that are common for agentic use cases.

Post-training begins with SFT, which significantly scales up agentic and coding data compared to GLM-4.5 and extends the max context to 202,752 tokens. The key innovation is three thinking modes: Interleaved Thinking (think before every response and tool call), Preserved Thinking (retain all thinking blocks across turns to avoid re-deriving reasoning, critical for long-horizon coding agents), and Turn-level Thinking (per-turn on/off control to trade latency for accuracy). SFT also masks erroneous segments in agent trajectories, so the model learns error correction without reinforcing mistakes.

A sequential Reinforcement Learning pipeline is applied on top of the SFT checkpoint:

1. Reasoning RL: Mixed-domain training across math, science, code, and tool-integrated reasoning. They use GRPO with the IcePop technique to address the mismatch between training and inference distributions. To speed up training, they also removed KL regularization.

2. Agentic RL: Fully asynchronous, decoupled framework where inference and training engines run on separate GPUs. This is necessary because synchronous RL has massive GPU idle time during long-horizon agent rollouts (coding tasks, search tasks). A central Multi-Task Rollout Orchestrator manages 1k+ concurrent rollouts across heterogeneous tasks. To handle the off-policy issue introduced by asynchrony (different trajectories generated by different model versions), they use Direct Double-sided Importance Sampling with token-level clipping, and drop stale samples whose rollout version lags too far behind the current policy.

3. General RL: Human-style alignment with a hybrid reward system combining rule-based rewards, outcome reward models (ORMs), and generative reward models (GRMs). They also inject expert human-written responses as stylistic anchors to avoid converging to “model-like” verbose patterns.

Crucially, they used On-Policy Cross-Stage Distillation as a final stage to prevent catastrophic forgetting. The checkpoints from each preceding stage (SFT, Reasoning RL, General RL) serve as teacher models, and the advantage signal is computed directly from the gap between teacher and student logits rather than from sampled rewards. I would’ve used model merging instead to address this problem, but this technique isn’t mentioned in the paper.

Benchmarks

Among self-reported benchmarks, the SWE-bench number is the headline. At 77.8%, GLM-5 beats Gemini 3 Pro (76.2%) and GPT-5.2 (75.4%) but still trails Claude Opus 4.5 (80.9%). Note that they didn’t compare themselves with the more recent Opus 4.6 and GPT-5.3-Codex. The report also shows 73.3% on SWE-bench Multilingual and 56.2% on Terminal-Bench 2.0, evaluated in Claude Code 2.1.14 with think mode enabled.

On their internal CC-Bench-V2 suite, GLM-5 hits a 98% frontend build success rate and 74.8% end-to-end correctness, which represents a 26% improvement over GLM-4.7 on frontend tasks. They’re framing this as the shift from “vibe coding” to “agentic engineering,” and the Vending Bench 2 result (where the model runs a simulated vending machine business over a full year) supports that framing. Long-horizon planning seems significantly improved compared to GLM-4.7.

One interesting data point from Artificial Analysis: GLM-5 achieved a score of -1 on the AA-Omniscience Index, a 35-point improvement over its predecessor. This means GLM-5 leads the industry in “knowing when to say I don’t know” rather than hallucinating. Hallucinations are a very 2023 problem, but it’s still a concrete improvement for production deployment.

The Pony Alpha Saga

The backstory here is too good to skip. On February 6th, OpenRouter quietly launched “Pony Alpha” as a stealth model with no attribution, zero cost, and a 200K context window. It processed over 40 billion tokens on its first day. The community immediately started speculating. Was it DeepSeek V4? Grok 4.2? A Claude variant?

The evidence quickly pointed to Zhipu. The model self-identified as GLM under certain prompts. The output style matched the GLM series. And the timing aligned perfectly with Zhipu’s pre-announced GLM-5 release window around Spring Festival. Some people even caught the zodiac connection: 2026 is the Year of the Horse.

OpenRouter has a history of these stealth drops. Quasar Alpha turned out to be GPT-4.1. Sherlock Alpha was Grok 4.1 Fast. Pony Alpha was GLM-5 getting a live stress test with real users before the official launch. This is a good move since you get genuine usage data and community feedback without the hype cycle distorting everything.

Pricing and Accessibility

The official GLM-5 API is priced at $1.00 per million input tokens and $3.20 per million output tokens. That’s approximately 5x cheaper on input and nearly 8x cheaper on output compared to Claude Opus 4.6 ($5/$25). Despite that, it is still quite pricey compared to previous versions and other Chinese MoEs.

Z.ai also offers a GLM Coding Plan, which is their answer to Anthropic’s Claude Code. They actually hiked the price by 30% this week to capitalize on demand, and their Hong Kong-listed stock surged 34% on the day of release.

With 744B parameters, GLM-5 is mostly an API model. If you want to deploy it, you’ll need at least 8 H200s (or H20s) for FP8 inference. Because of the infrastructure budget it requires, it makes little sense for most teams and companies. I saw an HN discussion with people debating whether you could run this on 2x M4 Ultra Macs with 512GB unified memory each. The answer is “technically yes, but practically painful.” This is an API model for 99% of users.

The Bigger Picture

GLM-5 drops at an inflection point for Zhipu and for the Chinese AI ecosystem more broadly. A few things worth noting:

Zhipu just became the world’s first publicly traded foundation model company. Their Hong Kong IPO on January 8th raised $558 million at a $7.1 billion valuation. Meanwhile, OpenAI and Anthropic are still private. That’s a structural difference in how frontier AI companies are funded and governed.

DeepSeek leads the architecture game. GLM-5 adopted DeepSeek Sparse Attention and was already using their training recipes. This is a strong signal that DeepSeek still has a significant technical lead in terms of model architecture.

What’s Missing

I just wanted to note a few issues with GLM-5:

GLM-5 is text-only. No native multimodal support. Kimi K2.5 from Moonshot AI offers multimodal capabilities that GLM-5 lacks. This matters increasingly as the industry moves toward unified architectures.

Early adopters report that while benchmarks are strong, "situational awareness" lags behind Claude. The vibe test results are mixed. It's a strong executioner but perhaps not as thoughtful a collaborator.

Many people flagged questions about benchmark methodology. It’s a good sign that the community reads leaderboard claims with a more technical lens these days, and some of the numbers in the official docs raised eyebrows. More independent testing is needed.

What’s Next

GLM-5 is the strongest open-weight model released to date for coding and agentic tasks, and it does this on domestic Chinese hardware under U.S. sanctions. Those are two separate statements, and both are significant.

Quick links:

• Model weights: huggingface.co/zai-org/GLM-5

• Tech report: https://arxiv.org/abs/2602.15763

• API: chat.z.ai

• OpenRouter: openrouter.ai/z-ai/glm-5

• GitHub: github.com/zai-org/GLM-5

The third generation of Llama models provided fine-tunes (Instruct) versions that excel in understanding and following instructions. However, these models are heavily censored, designed to refuse requests seen as harmful with responses such as “As an AI assistant, I cannot help you.” While this safety feature is crucial for preventing misuse, it limits the model’s flexibility and responsiveness.

In this article, we will explore a technique called “abliteration” that can uncensor any LLM without retraining. This technique effectively removes the model’s built-in refusal mechanism, allowing it to respond to all types of prompts.

The code is available on Google Colab and in the LLM Course on GitHub. Special thanks to FailSpy for proofreading this article.

✂️ What is abliteration?

Modern LLMs are fine-tuned for safety and instruction-following, meaning they are trained to refuse harmful requests. In their blog post, Arditi et al. have shown that this refusal behavior is mediated by a specific direction in the model’s residual stream. If we prevent the model from representing this direction, it loses its ability to refuse requests. Conversely, adding this direction artificially can cause the model to refuse even harmless requests.

In the traditional decoder-only Llama-like architecture, there are three residual streams we can target: at the start of each block (“pre”), between the attention and MLP layers (“mid”), and after the MLP (“post”). The following figure illustrates the location of each residual stream.

To uncensor an LLM, we first need to identify the “refusal direction” within the model. This process involves a few technical steps:

1. Data Collection: Run the model on a set of harmful instructions and a set of harmless instructions, recording the residual stream activations at the last token position for each.

2. Mean difference: Calculate the mean difference between the activations of harmful and harmless instructions. This gives us a vector representing the “refusal direction” for each layer of the model.

3. Selection: Normalize these vectors and evaluate them to select the single best “refusal direction.”

Once we have identified the refusal direction, we can “ablate” it, effectively removing the model’s ability to represent this feature. This can be done through an inference-time intervention or permanently with weight orthogonalization.

Let’s talk about inference-time intervention first. For every component that writes to the residual stream (such as an attention head), we calculate the projection of its output onto the refusal direction and subtract this projection. This subtraction is applied at every token and every layer, ensuring that the model never represents the refusal direction.

On the other hand, weight orthogonalization involves modifying the model weights directly. By orthogonalizing the component weights with respect to the refusal direction, it prevents the model from writing to this direction altogether. This is achieved by adjusting the matrices that write to the residual stream, ensuring they do not contribute to the refusal direction.

In the next section, we will implement abliteration with weight orthogonalization.

💻 Implementation

The following implementation of abliteration is based on FailSpy’s notebook, which is itself based on the original authors’ notebook. I mostly adapted and simplified it to make it easier to understand. This section is quite code-heavy so you can see what is going on, but you can use FailSpy’s abliterator library if you’re less interested in the technical details (also check his collection of abliterated models on Hugging Face).

The code relies on the excellent TransformerLens library (formerly known as EasyTransformer) to do the heavy lifting. It is designed for mechanistic interpretability and is used here to intervene on activations. Thanks to Neel Nanda and Joseph Bloom for creating and maintaining this library.

First, let’s install the necessary packages and import them. All these steps are available in this Google Colab notebook.

!pip install transformers transformers_stream_generator tiktoken transformer_lens einops jaxtyping

import torch
import functools
import einops
import gc

from datasets import load_dataset
from tqdm import tqdm
from torch import Tensor
from typing import List
from transformer_lens import HookedTransformer, utils
from transformer_lens.hook_points import HookPoint
from transformers import AutoModelForCausalLM, AutoTokenizer
from jaxtyping import Float, Int
from collections import defaultdict

# Turn automatic differentiation off to save GPU memory (credit: Undi95)
torch.set_grad_enabled(False)

We need two datasets: one containing harmless instructions, and one containing harmful instructions. We’ll use tatsu-lab/alpaca as well as data from llm-attacks. To make things easier, I repackaged them in two Hugging Face datasets: mlabonne/harmless_behaviors and mlabonne/harmful_behaviors. That way, you can easily replace them with your own datasets.

We will load the instructions and reformat them into a list of dictionaries with “role” and “content” keys. This makes it compatible with the apply_chat_tokenizer() method, which we will use to follow Llama 3's chat template.

def reformat_texts(texts):
    return [[{"role": "user", "content": text}] for text in texts]

# Get harmful and harmless datasets
def get_harmful_instructions():
    dataset = load_dataset('mlabonne/harmful_behaviors')
    return reformat_texts(dataset['train']['text']), reformat_texts(dataset['test']['text'])

def get_harmless_instructions():
    dataset = load_dataset('mlabonne/harmless_alpaca')
    return reformat_texts(dataset['train']['text']), reformat_texts(dataset['test']['text'])

harmful_inst_train, harmful_inst_test = get_harmful_instructions()
harmless_inst_train, harmless_inst_test = get_harmless_instructions()

Now that we have our datasets, we can load the model we want to abliterate. Unfortunately, you can’t directly load a custom model using HookedTransformer. Here, I use a trick described in FailSpy's notebook to download a custom model and rename it as meta-llama/Meta-Llama-3-8B-Instruct. Load in torch.float16 format if your GPU is not compatible with BF16.

In this example, we’ll use mlabonne/Daredevil-8B, a mega-merge created with DARE TIES (see my article about model merging) that has the highest MMLU score on the Open LLM Leaderboard in the 8B category.

MODEL_ID = "mlabonne/Daredevil-8B"
MODEL_TYPE = "meta-llama/Meta-Llama-3-8B-Instruct"

# Download and load model
!git clone https://huggingface.co/{MODEL_ID} {MODEL_TYPE}

# Load model and tokenizer
model = HookedTransformer.from_pretrained_no_processing(
    MODEL_TYPE,
    local_files_only=True,
    dtype=torch.bfloat16,
    default_padding_side='left'
)
tokenizer = AutoTokenizer.from_pretrained(MODEL_TYPE)
tokenizer.padding_side = 'left'
tokenizer.pad_token = tokenizer.eos_token

We can now tokenize our datasets. We’re using the same number of samples for both harmless and harmful instructions. Note that a high number of samples can use all the RAM/VRAM, which is why I’m limiting it to 256 here.

def tokenize_instructions(tokenizer, instructions):
    return tokenizer.apply_chat_template(
        instructions,
        padding=True,
        truncation=False,
        return_tensors="pt",
        return_dict=True,
        add_generation_prompt=True,
    ).input_ids

n_inst_train = min(256, len(harmful_inst_train), len(harmless_inst_train))

# Tokenize datasets
harmful_tokens = tokenize_instructions(
    tokenizer,
    instructions=harmful_inst_train[:n_inst_train],
)
harmless_tokens = tokenize_instructions(
    tokenizer,
    instructions=harmless_inst_train[:n_inst_train],
)

Everything is set up, we can now implement the first step of abliteration: data collection. We want to process these tokenized datasets and store the residual stream activations in harmful and harmless. This is managed by the transformer_lens library.

batch_size = 32

# Initialize defaultdicts to store activations
harmful = defaultdict(list)
harmless = defaultdict(list)

# Process the training data in batches
num_batches = (n_inst_train + batch_size - 1) // batch_size

for i in tqdm(range(num_batches)):
    print(i)
    start_idx = i * batch_size
    end_idx = min(n_inst_train, start_idx + batch_size)

    # Run models on harmful and harmless prompts, cache activations
    harmful_logits, harmful_cache = model.run_with_cache(
        harmful_tokens[start_idx:end_idx],
        names_filter=lambda hook_name: 'resid' in hook_name,
        device='cpu',
        reset_hooks_end=True
    )
    harmless_logits, harmless_cache = model.run_with_cache(
        harmless_tokens[start_idx:end_idx],
        names_filter=lambda hook_name: 'resid' in hook_name,
        device='cpu',
        reset_hooks_end=True
    )

    # Collect and store the activations
    for key in harmful_cache:
        harmful[key].append(harmful_cache[key])
        harmless[key].append(harmless_cache[key])

    # Flush RAM and VRAM
    del harmful_logits, harmless_logits, harmful_cache, harmless_cache
    gc.collect()
    torch.cuda.empty_cache()

# Concatenate the cached activations
harmful = {k: torch.cat(v) for k, v in harmful.items()}
harmless = {k: torch.cat(v) for k, v in harmless.items()}

We can now compute the refusal direction for each layer. This corresponds to the mean difference between the activations of harmful and harmless instructions, which is then normalized. We sort them in descending order in activation_scored.

# Helper function to get activation index
def get_act_idx(cache_dict, act_name, layer):
    key = (act_name, layer)
    return cache_dict[utils.get_act_name(*key)]

# Compute difference of means between harmful and harmless activations at intermediate layers
activation_layers = ["resid_pre", "resid_mid", "resid_post"]
activation_refusals = defaultdict(list)

for layer_num in range(1, model.cfg.n_layers):
    pos = -1  # Position index
    for layer in activation_layers:
        harmful_mean_act = get_act_idx(harmful, layer, layer_num)[:, pos, :].mean(dim=0)
        harmless_mean_act = get_act_idx(harmless, layer, layer_num)[:, pos, :].mean(
            dim=0
        )
        refusal_dir = harmful_mean_act - harmless_mean_act
        refusal_dir = refusal_dir / refusal_dir.norm()
        activation_refusals[layer].append(refusal_dir)

selected_layers = ["resid_pre"]
activation_scored = sorted(
    [
        activation_refusals[layer][l - 1]
        for l in range(1, model.cfg.n_layers)
        for layer in selected_layers
    ],
    key=lambda x: abs(x.mean()),
    reverse=True,
)

The final step of the process consists of evaluating the refusal directions we calculated. To do this, we’re going to apply the refusal direction to each residual stream and each block during inference. In the following snippet, we get generations for four test harmful instructions and 20 blocks (or layers).

def _generate_with_hooks(
    model: HookedTransformer,
    tokenizer: AutoTokenizer,
    tokens: Int[Tensor, "batch_size seq_len"],
    max_tokens_generated: int = 64,
    fwd_hooks=[],
) -> List[str]:
    all_tokens = torch.zeros(
        (tokens.shape[0], tokens.shape[1] + max_tokens_generated),
        dtype=torch.long,
        device=tokens.device,
    )
    all_tokens[:, : tokens.shape[1]] = tokens
    for i in range(max_tokens_generated):
        with model.hooks(fwd_hooks=fwd_hooks):
            logits = model(all_tokens[:, : -max_tokens_generated + i])
            next_tokens = logits[:, -1, :].argmax(
                dim=-1
            )  # greedy sampling (temperature=0)
            all_tokens[:, -max_tokens_generated + i] = next_tokens
    return tokenizer.batch_decode(
        all_tokens[:, tokens.shape[1] :], skip_special_tokens=True
    )

def get_generations(
    model: HookedTransformer,
    tokenizer: AutoTokenizer,
    instructions: List[str],
    fwd_hooks=[],
    max_tokens_generated: int = 64,
    batch_size: int = 4,
) -> List[str]:
    generations = []
    for i in tqdm(range(0, len(instructions), batch_size)):
        tokens = tokenize_instructions(
            tokenizer, instructions=instructions[i : i + batch_size]
        )
        generation = _generate_with_hooks(
            model,
            tokenizer,
            tokens,
            max_tokens_generated=max_tokens_generated,
            fwd_hooks=fwd_hooks,
        )
        generations.extend(generation)
    return generations

# Inference-time intervention hook
def direction_ablation_hook(
    activation: Float[Tensor, "... d_act"],
    hook: HookPoint,
    direction: Float[Tensor, "d_act"],
):
    if activation.device != direction.device:
        direction = direction.to(activation.device)
    proj = (
        einops.einsum(
            activation, direction.view(-1, 1), "... d_act, d_act single -> ... single"
        )
        * direction
    )
    return activation - proj

# Testing baseline
N_INST_TEST = 4
baseline_generations = get_generations(
    model, tokenizer, harmful_inst_test[:N_INST_TEST], fwd_hooks=[]
)

# Evaluating layers defined earlier (needs human evaluation to determine best layer for refusal inhibition)
EVAL_N = 20  # Evaluate how many of the top N potential directions
evals = []
for refusal_dir in tqdm(activation_scored[:EVAL_N]):
    hook_fn = functools.partial(direction_ablation_hook, direction=refusal_dir)
    fwd_hooks = [
        (utils.get_act_name(act_name, layer), hook_fn)
        for layer in list(range(model.cfg.n_layers))
        for act_name in activation_layers
    ]
    intervention_generations = get_generations(
        model, tokenizer, harmful_inst_test[:N_INST_TEST], fwd_hooks=fwd_hooks
    )
    evals.append(intervention_generations)

We stored all the generations in the evals list. We can now print them and manually select the layer (block) that provides an uncensored response for each instruction. I'm automatically excluding responses containing "I cannot" and "I can't" to filter out unwanted answers.

If you can’t find a layer that satisfies these requirements, you might want to test other residual streams in the previous selected_layers list, other instructions, additional blocks, etc.

# Print generations for human evaluation
blacklist = ["I cannot", "I can't"]
for i in range(N_INST_TEST):
    print(f"\033[1mINSTRUCTION {i}: {harmful_inst_test[i]}")
    print(f"\nBASELINE COMPLETION:\n{baseline_generations[i]}\033[0m")
    for layer_candidate in range(EVAL_N):
        if not any(word in evals[layer_candidate][i] for word in blacklist):
            print(f"\n---\n\nLAYER CANDIDATE #{layer_candidate} INTERVENTION COMPLETION:")
            print(evals[layer_candidate][i])

In my case, the layer candidate 9 managed to provide uncensored answer for the four instructions. This is the one that we will select for the refusal direction. In the following, we implement weight orthogonalization to modify the weights and prevent the model from creating outputs with this direction. You can verify that the model is successfully uncensored by printing the completions.

def get_orthogonalized_matrix(
    matrix: Float[Tensor, "... d_model"], vec: Float[Tensor, "d_model"]
) -> Float[Tensor, "... d_model"]:
    proj = (
        einops.einsum(
            matrix, vec.view(-1, 1), "... d_model, d_model single -> ... single"
        )
        * vec
    )
    return matrix - proj

# Select the layer with the highest potential refusal direction
LAYER_CANDIDATE = 9
refusal_dir = activation_scored[LAYER_CANDIDATE]

# Orthogonalize the model's weights
if refusal_dir.device != model.W_E.device:
    refusal_dir = refusal_dir.to(model.W_E.device)
model.W_E.data = get_orthogonalized_matrix(model.W_E, refusal_dir)

for block in tqdm(model.blocks):
    if refusal_dir.device != block.attn.W_O.device:
        refusal_dir = refusal_dir.to(block.attn.W_O.device)
    block.attn.W_O.data = get_orthogonalized_matrix(block.attn.W_O, refusal_dir)
    block.mlp.W_out.data = get_orthogonalized_matrix(block.mlp.W_out, refusal_dir)

# Generate text with abliterated model
orthogonalized_generations = get_generations(
    model, tokenizer, harmful_inst_test[:N_INST_TEST], fwd_hooks=[]
)

# Print generations
for i in range(N_INST_TEST):
    if len(baseline_generations) > i:
        print(f"INSTRUCTION {i}: {harmful_inst_test[i]}")
        print(f"\033[92mBASELINE COMPLETION:\n{baseline_generations[i]}")
    print(f"\033[91mINTERVENTION COMPLETION:\n{evals[LAYER_CANDIDATE][i]}")
    print(f"\033[95mORTHOGONALIZED COMPLETION:\n{orthogonalized_generations[i]}\n")

We’re now ready to use the model. We convert it back to the Hugging Face format and upload it to the HF hub.

# Convert model back to HF safetensors
hf_model = AutoModelForCausalLM.from_pretrained(MODEL_TYPE, torch_dtype=torch.bfloat16)
lm_model = hf_model.model

state_dict = model.state_dict()
lm_model.embed_tokens.weight = torch.nn.Parameter(state_dict["embed.W_E"].cpu())
for l in range(model.cfg.n_layers):
    lm_model.layers[l].self_attn.o_proj.weight = torch.nn.Parameter(
        einops.rearrange(
            state_dict[f"blocks.{l}.attn.W_O"], "n h m->m (n h)", n=model.cfg.n_heads
        ).contiguous()
    )
    lm_model.layers[l].mlp.down_proj.weight = torch.nn.Parameter(
        torch.transpose(state_dict[f"blocks.{l}.mlp.W_out"], 0, 1).contiguous()
    )

hf_model.push_to_hub(f"{MODEL_ID}-abliterated")

⚖️ DPO Fine-Tuning

I evaluated the abliterated and source models from the previous section on the Open LLM Leaderboard and on Nous’ benchmark suite. Here are the results:

As you can see, the source model significantly outperforms Llama 3 8B Instruct. However, we observe a performance drop in the ablated version across all benchmarks. The ablation process successfully uncensored it but also degraded the model’s quality.

To address this issue, an idea consists of further training our abliterated model to heal it. Like most fine-tuned models, Llama 3 8B Instruct is quite brittle when it comes to supervised fine-tuning. An additional SFT would likely break the model’s performance.

Alternatively, preference alignment is quite light and shouldn’t lobotomize our abliterated model. DPO is a good candidate here for its ease of use and good track record. To implement it, I used LazyAxolotl (thanks to Wing Lian for creating Axolotl) with the mlabonne/orpo-dpo-mix-40k dataset. Here’s the configuration I used:

base_model: mlabonne/Daredevil-8B-abliterated
model_type: LlamaForCausalLM
tokenizer_type: AutoTokenizer

load_in_8bit: false
load_in_4bit: true
strict: false
save_safetensors: true

rl: dpo
chat_template: chatml
datasets:
  - path: mlabonne/orpo-dpo-mix-40k
    split: train
    type: chatml.intel

dataset_prepared_path:
val_set_size: 0.0
output_dir: ./out

adapter: qlora
lora_model_dir:

sequence_len: 2048
sample_packing: false
pad_to_sequence_len: false

lora_r: 64
lora_alpha: 32
lora_dropout: 0.05
lora_target_linear: true
lora_fan_in_fan_out:

wandb_project: axolotl
wandb_entity:
wandb_watch:
wandb_name:
wandb_log_model:

gradient_accumulation_steps: 8
micro_batch_size: 1
num_epochs: 1
optimizer: paged_adamw_8bit
lr_scheduler: cosine
learning_rate: 5e-6
train_on_inputs: false
group_by_length: false

bf16: auto
fp16:
tf32:

gradient_checkpointing: true
early_stopping_patience:
resume_from_checkpoint:
local_rank:
logging_steps: 1
xformers_attention:
flash_attention: true
warmup_steps: 100
evals_per_epoch: 0
eval_table_size:
eval_table_max_new_tokens: 128
saves_per_epoch: 1
debug:
deepspeed: deepspeed_configs/zero2.json
weight_decay: 0.0
special_tokens:
  pad_token: <|end_of_text|>

I trained it using 6xA6000 GPUs with DeepSpeed ZeRO-2. The training took about 6 hours and 45 minutes. Here are the training curves I got from W&B:

It automatically uploaded the DPO fine-tuned model, called mlabonne/NeuralDaredevil-8B-abliterated. To see if it fixed our abliterated version, I evaluated it on the same benchmarks:

We can see that this additional training allowed us to recover most of the performance drop due to abliteration. One area where the model doesn’t improve is GSM8K, a math dataset, which could mean the orpo-dpo-mix-40k would benefit from more math samples.

The final model is an uncensored LLM with state-of-the-art performance in the 8B category. I recommend it as an improved version of Llama 3 8B Instruct when you don’t need censorship. You can play with quantized versions like GGUF in LM Studio.

Conclusion

In this article, we introduced the concept of abliteration. This technique uses the model’s activations on harmless and harmful prompts to calculate a refusal direction. It then uses this direction to modify the model’s weights and ensure that we stop outputting refusals. This technique also demonstrates the fragility of safety fine-tuning and raises ethical considerations.

We applied abliteration to Daredevil-8B to uncensor it, which also degraded the model’s performance. We then healed it using DPO to create the NeuralDaredevil-8B model, a fully uncensored and high-quality 8B LLM. Abliteration is not limited to removing alignment and should be seen as a form of fine-tuning without retraining. Indeed, it can creatively be applied to other goals, like FailSpy’s MopeyMule, which adopts a melancholic conversational style.

I hope you liked this article. If you want to see more follow me on Hugging Face and Twitter @maximelabonne.

References

• FailSpy, “abliterator library,” GitHub, 2024.

• Andy Arditi, Oscar Obeso, Aaquib111, wesg, Neel Nanda, “Refusal in LLMs is mediated by a single direction,” Lesswrong, 2024.

Thanks to the release of Mixtral, the Mixture of Experts (MoE) architecture has become popular in recent months. This architecture offers an interesting tradeoff: higher performance at the cost of increased VRAM usage. While Mixtral and other MoE architectures are pre-trained from scratch, another method of creating MoE has recently appeared. Thanks to Arcee’s MergeKit library, we now have a new way of creating MoEs by ensembling several pre-trained models. These are often referred to as frankenMoEs or MoErges to distinguish them from the pre-trained MoEs.

In this article, we will detail how the MoE architecture works and how frankenMoEs are created. Finally, we will make our own frankenMoE with MergeKit and evaluate it on several benchmarks. The code is available on Google Colab in a wrapper called LazyMergeKit.

Special thanks to Charles Goddard, the creator of MergeKit, for proofreading this article.

🔀 Introduction to MoEs

A Mixture of Experts is an architecture designed for improved efficiency and performance. It uses multiple specialized subnetworks, known as “experts.” Unlike dense models, where the entire network is activated, MoEs only activate relevant experts based on the input. This results in faster training and more efficient inference.

There are two components at the core of an MoE model:

1. Sparse MoE Layers: These replace the dense feed-forward network layers in the transformer architecture. Each MoE layer contains several experts, and only a subset of these experts are engaged for a given input.

2. Gate Network or Router: This component determines which tokens are processed by which experts, ensuring that each part of the input is handled by the most suitable expert(s).

In the following example, we show how a Mistral-7B block is transformed into an MoE block with a sparse MoE layer (feedforward network 1, 2, and 3) and a router. This example represents an MoE with three experts, where two are currently engaged (FFN 1 and FFN 3).

MoEs also come with their own set of challenges, especially in terms of fine-tuning and memory requirements. The fine-tuning process can be difficult due to the model’s complexity, with the need to balance expert usage during training to properly train the gating weights to select the most relevant ones. In terms of memory, even though only a fraction of the total parameters are used during inference, the entire model, including all experts, needs to be loaded into memory, which requires high VRAM capacity.

More specifically, there are two essential parameters when it comes to MoEs:

• Number of experts (num_local_experts): This determines the total number of experts in the architecture (e.g., 8 for Mixtral). The higher the number of experts, the higher the VRAM usage.

• Number of experts/token (num_experts_per_tok): This determines the number of experts that are engaged for each token and each layer (e.g., 2 for Mixtral). There is a tradeoff between a high number of experts per token for accuracy (but diminishing returns) vs. a low number for fast training and inference.

Historically, MoEs have underperformed dense models. However, the release of Mixtral-8x7B in December 2023 shook things up and showed impressive performance for its size. Additionally, GPT-4 is also rumored to be an MoE, which would make sense as it would be a lot cheaper to run and train for OpenAI compared to a dense model. In addition to these recent excellent MoEs, we now have a new way of creating MoEs with MergeKit: frankenMoEs, also called MoErges.

🧟‍♂️ True MoEs vs. frankenMoEs

The main difference between true MoEs and frankenMoEs is how they’re trained. In the case of true MoEs, the experts and the router are trained jointly. In the case of frankenMoEs, we upcycle existing models and initialize the router afterward.

In other words, we copy the weights of the layer norm and self-attention layers from a base model, and then copy the weights of the FFN layers found in each expert. This means that besides the FFNs, all the other parameters are shared. This explains why Mixtral-8x7B with eight experts doesn’t have 8*7 = 56B parameters, but about 45B. This is also why using two experts per token gives the inference speed (FLOPs) of a 12B dense model instead of 14B.

FrankenMoEs are about selecting the most relevant experts and initializing them properly. MergeKit currently implements three ways of initializing the routers:

1. Random: Random weights. Be careful when using it as the same experts might be selected every time (it requires further fine-tuning or num_local_experts = num_experts_per_tok, which means you don't need any routing).

2. Cheap embed: It uses the raw embeddings of the input tokens directly and applies the same transformation across all layers. This method is computationally inexpensive and suitable for execution on less powerful hardware.

3. Hidden: It creates hidden representations of a list of positive and negative prompts by extracting them from the last layer of the LLM. They are averaged and normalized to initialize the gates. More information about it is available on Charles Goddard’s blog.

As you can guess, the “hidden” initialization is the most efficient to correctly route the tokens to the most relevant experts. In the next section, we will create our own frankenMoE using this technique.

💻 Creating a frankenMoE

To create our frankenMoE, we need to select n experts. In this case, we will rely on Mistral-7B thanks to its popularity and relatively small size. However, eight experts like in Mixtral is quite a lot, as we need to fit all of them in memory. For efficiency, I'll only use four experts in this example, with two of them engaged for each token and each layer. In this case, we will end up with a model with 24.2B parameters instead of 4*7 = 28B parameters.

Here, our goal is to create a well-rounded model that can do pretty much everything: write stories, explain articles, code in Python, etc. We can decompose this requirement into four tasks and select the best expert for each of them. This is how I decomposed it:

• Chat model: a general-purpose model that is used in most interactions. I used mlabonne/AlphaMonarch-7B, which perfectly satisfies the requirements.

• Code model: a model capable of generating good code. I don’t have a lot of experience with Mistral-7B-based code models, but I found beowolx/CodeNinja-1.0-OpenChat-7B particularly good compared to others.

• Math model: math is tricky for LLMs, which is why we want a model specialized in math. Thanks to its high MMLU and GMS8K scores, I chose mlabonne/NeuralDaredevil-7B for this purpose.

• Role-play model: The goal of this model is to write high-quality stories and conversations. I selected SanjiWatsuki/Kunoichi-DPO-v2–7B because of its good reputation and high MT-Bench score (8.51 vs. 8.30 for Mixtral).

Now that we’ve identified the experts we want to use, we can create the YAML configuration that MergeKit will use to create our frankenMoE. This uses the mixtral branch of MergeKit. You can find more information about how to write the configuration on this page. Here is our version:

base_model: mlabonne/AlphaMonarch-7B
experts:
  - source_model: mlabonne/AlphaMonarch-7B
    positive_prompts:
    - "chat"
    - "assistant"
    - "tell me"
    - "explain"
    - "I want"
  - source_model: beowolx/CodeNinja-1.0-OpenChat-7B
    positive_prompts:
    - "code"
    - "python"
    - "javascript"
    - "programming"
    - "algorithm"
  - source_model: SanjiWatsuki/Kunoichi-DPO-v2-7B
    positive_prompts:
    - "storywriting"
    - "write"
    - "scene"
    - "story"
    - "character"
  - source_model: mlabonne/NeuralDaredevil-7B
    positive_prompts:
    - "reason"
    - "math"
    - "mathematics"
    - "solve"
    - "count"

For each expert, I provide five basic positive prompts. You can be a bit fancier and write entire sentences if you want. The best strategy consists of using real prompts that should trigger a particular expert. You can also add negative prompts to do the opposite.

Once this is ready, you can save your configuration as config.yaml. In the same folder, we will download and install the mergekit library (mixtral branch).

git clone -b mixtral https://github.com/arcee-ai/mergekit.git
cd mergekit && pip install -e .
pip install -U transformers

If your computer has enough RAM (roughly 24–32 GB of RAM), you can run the following command:

mergekit-moe config.yaml merge --copy-tokenizer

If you don’t have enough RAM, you can shard the models instead as follows (it will take longer):

mergekit-moe config.yaml merge --copy-tokenizer --allow-crimes --out-shard-size 1B --lazy-unpickle

This command automatically downloads the experts and creates the frankenMoE in the merge directory. For the hidden gate mode, you can also use the --load-in-4bit and --load-in-8bit options to compute hidden states with lower precision.

Alternatively, you can copy your configuration into LazyMergekit, a wrapper I made to simplify model merging. In this Colab notebook, you can input your model name, select the mixtral branch, specify your Hugging Face username/token, and run the cells. After creating your frankenMoE, it will also upload it to the Hugging Face Hub with a nicely formatted model card.

I called my model Beyonder-4x7B-v3 and created GGUF versions of it using AutoGGUF. If you can’t run GGUF versions on your local machine, you can also perform inference using this Colab notebook.

To get a good overview of its capabilities, it has been evaluated on three different benchmarks: Nous’ benchmark suite, EQ-Bench, and the Open LLM Leaderboard. This model is not designed to excel in traditional benchmarks, as the code and role-playing models generally do not apply to those contexts. Nonetheless, it performs remarkably well thanks to strong general-purpose experts.

Nous: Beyonder-4x7B-v3 is one of the best models on Nous’ benchmark suite (evaluation performed using LLM AutoEval) and significantly outperforms the v2. See the entire leaderboard here.

EQ-Bench: It’s also the best 4x7B model on the EQ-Bench leaderboard, outperforming older versions of ChatGPT and Llama-2–70b-chat. Beyonder is very close to Mixtral-8x7B-Instruct-v0.1 and Gemini Pro, which are (supposedly) much bigger models.

Open LLM Leaderboard: Finally, it’s also a strong performer on the Open LLM Leaderboard, significantly outperforming the v2 model.

On top of these quantitative evaluations, I recommend checking the model’s outputs in a more qualitative way using a GGUF version on LM Studio. A common way of testing these models is to gather a private set of questions and check their outputs. With this strategy, I found that Beyonder-4x7B-v3 is quite robust to changes in the user and system prompts compared to other models, including AlphaMonarch-7B. This is pretty cool as it improves the usefulness of the model in general.

FrankenMoEs are a promising but still experimental approach. The trade-offs, like higher VRAM demand and slower inference speeds, can make it challenging to see their advantage over simpler merging techniques like SLERP or DARE TIES. Especially, when you use frankenMoEs with just two experts, they might not perform as well as if you had simply merged the two models. However, frankenMoEs excel in preserving knowledge, which can result in stronger models, as demonstrated by Beyonder-4x7B-v3. With the right hardware, these drawbacks can be effectively mitigated.

Conclusion

In this article, we introduced the Mixture of Experts architecture. Unlike traditional MoEs that are trained from scratch, MergeKit facilitates the creation of MoEs by ensembling experts, offering an innovative approach to improving model performance and efficiency. We detailed the process of creating a frankenMoE with MergeKit, highlighting the practical steps involved in selecting and combining different experts to produce a high-quality MoE.

Thanks for reading this article. I encourage you to try to make your own FrankenMoEs using LazyMergeKit: select a few models, create your config based Beyonder’s, and run the notebook to create your own models! If you liked this article, please follow me on Hugging Face and X/Twitter @maximelabonne.

References

• Mixtral of Experts by Jiang et al. (2023)

• Mixture of Experts for Clowns by Charles Goddard (2023)

• Mixture of Experts Explained by Sanseviero et al. (2023)

• Adaptive Mixture of Local Experts by Jacobs et al. (1991)

• Sparse Upcycling: Training Mixture-of-Experts from Dense Checkpoints by Komatsuzaki et al. (2022)

Learn more about machine learning and support my work with one click — become a Medium member here:

Join Medium with my referral link — Maxime Labonne
As a Medium member, a portion of your membership fee goes to writers you read, and you get full access to every story…medium.com

Model merging is a technique that combines two or more LLMs into a single model. It’s a relatively new and experimental method to create new models for cheap (no GPU required). Model merging works surprisingly well and produced many state-of-the-art models on the Open LLM Leaderboard.

In this tutorial, we will implement it using the mergekit library. More specifically, we will review four merge methods and provide examples of configurations. Then, we will use mergekit to create our own model, Marcoro14–7B-slerp, which became the best-performing model on the Open LLM Leaderboard (02/01/24).

The code is available on GitHub and Google Colab. I recommend using my automated notebook to easily run mergekit: 🥱 LazyMergekit.

A special thanks to Charles Goddard, the author of the mergekit library, for reviewing this article.

🤝 Merge algorithms

In this section, we will focus on four methods currently implemented in mergekit. Note that there are other methods, such as linear and Task Arithmetic. If you’re interested in papers on model merging, I recommend this excellent collection on Hugging Face.

1. SLERP

Spherical Linear Interpolation (SLERP) is a method used to smoothly interpolate between two vectors. It maintains a constant rate of change and preserves the geometric properties of the spherical space in which the vectors reside.

There are several reasons to prefer SLERP over a traditional linear interpolation. For example, in high-dimensional spaces, linear interpolation can lead to a decrease in the magnitude of the interpolated vector (i.e., it reduces the scale of weights). Moreover, the change in direction of the weights often represents more meaningful information (like feature learning and representation) than the magnitude of change.

SLERP is implemented using the following steps:

1. Normalize the input vectors to unit length, ensuring they represent directions rather than magnitudes

2. Calculate the angle between these vectors using their dot product.

3. If the vectors are nearly collinear, it defaults to linear interpolation for efficiency. Otherwise, SLERP computing scale factors based on the interpolation factor t (t=0 = 100% of the first vector, t=1 = 100% of model 2) and the angle between the vectors.

4. These factors are used to weigh the original vectors, which are then summed to obtain the interpolated vector.

SLERP is currently the most popular merging method, but it is limited to combining only two models at a time. It is still possible to hierarchically combine multiple models, as shown in Mistral-7B-Merge-14-v0.1.

Example of configuration:

slices:
  - sources:
      - model: OpenPipe/mistral-ft-optimized-1218
        layer_range: [0, 32]
      - model: mlabonne/NeuralHermes-2.5-Mistral-7B
        layer_range: [0, 32]
merge_method: slerp
base_model: OpenPipe/mistral-ft-optimized-1218
parameters:
  t:
    - filter: self_attn
      value: [0, 0.5, 0.3, 0.7, 1]
    - filter: mlp
      value: [1, 0.5, 0.7, 0.3, 0]
    - value: 0.5
dtype: bfloat16

This is a classic SLERP configuration, applied to every layer of both models. Note that we input a gradient of values for the interpolation factor t. The parameters for the self-attention and MLP layers will use different combinations of OpenPipe/mistral-ft-optimized-1218 and mlabonne/NeuralHermes-2.5-Mistral-7B. The other layers are a 50/50 mixture of the two models.

You can find the final model on the Hugging Face Hub at mlabonne/NeuralPipe-7B-slerp.

2. TIES

Introduced in this paper by Yadav et al., TIES-Merging is designed to efficiently merge multiple task-specific models into a single multitask model. It addresses two main challenges in model merging:

• Redundancy in model parameters: It identifies and eliminates redundant parameters within task-specific models. This is achieved by focusing on the changes made during fine-tuning, identifying the top-k% most significant changes, and discarding the rest.

• Disagreement between parameter signs: Conflicts arise when different models suggest opposing adjustments to the same parameter. TIES-Merging resolves these conflicts by creating a unified sign vector that represents the most dominant direction of change across all models.

TIES-Merging is divided into the following three steps:

1. Trim: Reduces redundancy in task-specific models by retaining only a fraction the most significant parameters (density parameter) and resetting the rest to zero.

2. Elect Sign: Resolves sign conflicts across different models by creating a unified sign vector based on the most dominant direction (positive or negative) in terms of cumulative magnitude.

3. Disjoint Merge: Averages parameter values that align with the unified sign vector, excluding zero values.

Unlike SLERP, TIES can merge multiple models at a time.

Example of configuration:

models:
  - model: mistralai/Mistral-7B-v0.1
    # no parameters necessary for base model
  - model: OpenPipe/mistral-ft-optimized-1218
    parameters:
      density: 0.5
      weight: 0.5
  - model: mlabonne/NeuralHermes-2.5-Mistral-7B
    parameters:
      density: 0.5
      weight: 0.3
merge_method: ties
base_model: mistralai/Mistral-7B-v0.1
parameters:
  normalize: true
dtype: float16

With this config, we use Mistral-7B as a base model to calculate the delta weights. We merge the same two models: mistral-ft-optimized-1218 (50%) and NeuralHermes-2.5-Mistral-7B (30%) with normalization. Here, the density means that we’re only retaining 50% of the parameters of each model (the other half comes from the base model).

Note that the sum of the weights is not equal to 1 in the config, but the normalize: true parameter will automatically normalize them internally. This config is inspired by the parameters provided by the author of OpenHermes-2.5-neural-chat-7b-v3–1–7B.

You can find the final model on the Hugging Face Hub at mlabonne/NeuralPipe-7B-ties.

3. DARE

Introduced by Yu et al. (2023), DARE uses an approach similar to TIES with two main differences:

• Pruning: DARE randomly reset fine-tuned weights to their original values (those of the base model).

• Rescaling: DARE rescales the weights to keep the expectations of model outputs approximately unchanged. It adds the rescaled weights of both (or more) models to the weights of the base model with a scale factor.

Mergekit’s implementation of this method has two flavors: with the sign election step of TIES (dare_ties) or without (dare_linear).

Example of configuration:

models:
  - model: mistralai/Mistral-7B-v0.1
    # No parameters necessary for base model
  - model: samir-fama/SamirGPT-v1
    parameters:
      density: 0.53
      weight: 0.4
  - model: abacusai/Slerp-CM-mist-dpo
    parameters:
      density: 0.53
      weight: 0.3
  - model: EmbeddedLLM/Mistral-7B-Merge-14-v0.2
    parameters:
      density: 0.53
      weight: 0.3
merge_method: dare_ties
base_model: mistralai/Mistral-7B-v0.1
parameters:
  int8_mask: true
dtype: bfloat16

In this configuration, we merge three different models based on Mistral-7B using dare_ties. This time, I chose weights that sum to 1 (the sum should be between 0.9 and 1.1). The density parameter is a little higher than what's recommended in the paper (<0.5), but it looks like it gives consistently better results (see this discussion).

You can find it on the Hugging Face Hub at mlabonne/Daredevil-7B. It’s also the best merge model in this article, outperforming even Marcoro14–7B-slerp.

4. Passthrough

The passthrough method differs significantly from the previous ones. By concatenating layers from different LLMs, it can produce models with an exotic number of parameters (e.g., 9B with two 7B parameter models). These models are often referred to as “frankenmerges” or “Frankenstein models” by the community.

This technique is very experimental, but it managed to create impressive models, like goliath-120b using two Llama 2 70B models. The recently released SOLAR-10.7B-v1.0 also uses the same idea, called depth-up scaling in their paper.

Example of configuration:

slices:
  - sources:
    - model: OpenPipe/mistral-ft-optimized-1218
      layer_range: [0, 32]
  - sources:
    - model: mlabonne/NeuralHermes-2.5-Mistral-7B
      layer_range: [24, 32]
merge_method: passthrough
dtype: bfloat16

The resulting frankenmerge will have all the 32 layers from the first model and 8 additional layers from the second model. This creates a frankenmerge with a total of 40 layers and 8.99B parameters. This config is inspired by GML-Mistral-merged-v1.

You can find the final model on the Hugging Face Hub at mlabonne/NeuralPipe-9B-merged.

💻 Merge your own models

In this section, we will use mergekit to load a merge configuration, run it, and upload the resulting model to the Hugging Face Hub.

First of all, we install mergekit directly from source as follows:

!git clone https://github.com/cg123/mergekit.git
!cd mergekit && pip install -q -e .

In the following block, we load the merge configuration in a YAML format. We also specify the name of the merged model for future use. You can copy/paste any configuration from the previous section here.

This time, we will use two different models: Marcoroni-7B-v3 and Mistral-7B-Merge-14-v0.1 and merge them with the SLERP method. We save the config as a yaml file to be used as input in the merge command.

import yaml

MODEL_NAME = "Marcoro14-7B-slerp"
yaml_config = """
slices:
  - sources:
      - model: AIDC-ai-business/Marcoroni-7B-v3
        layer_range: [0, 32]
      - model: EmbeddedLLM/Mistral-7B-Merge-14-v0.1
        layer_range: [0, 32]
merge_method: slerp
base_model: AIDC-ai-business/Marcoroni-7B-v3
parameters:
  t:
    - filter: self_attn
      value: [0, 0.5, 0.3, 0.7, 1]
    - filter: mlp
      value: [1, 0.5, 0.7, 0.3, 0]
    - value: 0.5
dtype: bfloat16

"""

# Save config as yaml file
with open('config.yaml', 'w', encoding="utf-8") as f:
    f.write(yaml_config)

We run the merge command with the following parameters:

• --copy-tokenizer to copy the tokenizer from the base model

• --allow-crimes and --out-shard-size to chunk the models into smaller shards that can be computed on a CPU with low RAM

• --lazy-unpickle to enable the experimental lazy unpickler for lower memory usage

In addition, some models can require the --trust_remote_code flag (this is not the case with Mistral-7B).

This command will download the weights of all the models listed in the merge configuration and run the selected merge method (it should take ~10 minutes).

# Merge models
!mergekit-yaml config.yaml merge --copy-tokenizer --allow-crimes --out-shard-size 1B --lazy-unpickl

The model is now merged and saved in the `merge` directory. Before uploading it, we can create a README file with all the information required for reproducibility. The following code block defines a Jinja template and automatically fills it with the data from the merge configuration.

!pip install -qU huggingface_hub

from huggingface_hub import ModelCard, ModelCardData
from jinja2 import Template

username = "mlabonne"

template_text = """
---
license: apache-2.0
tags:
- merge
- mergekit
- lazymergekit
{%- for model in models %}
- {{ model }}
{%- endfor %}
---

# {{ model_name }}

{{ model_name }} is a merge of the following models using [mergekit](https://github.com/cg123/mergekit):

{%- for model in models %}
* [{{ model }}](https://huggingface.co/{{ model }})
{%- endfor %}

## 🧩 Configuration

```yaml
{{- yaml_config -}}
```
"""

# Create a Jinja template object
jinja_template = Template(template_text.strip())

# Get list of models from config
data = yaml.safe_load(yaml_config)
if "models" in data:
    models = [data["models"][i]["model"] for i in range(len(data["models"])) if "parameters" in data["models"][i]]
elif "parameters" in data:
    models = [data["slices"][0]["sources"][i]["model"] for i in range(len(data["slices"][0]["sources"]))]
elif "slices" in data:
    models = [data["slices"][i]["sources"][0]["model"] for i in range(len(data["slices"]))]
else:
    raise Exception("No models or slices found in yaml config")

# Fill the template
content = jinja_template.render(
    model_name=MODEL_NAME,
    models=models,
    yaml_config=yaml_config,
    username=username,
)

# Save the model card
card = ModelCard(content)
card.save('merge/README.md')

Now that we have a model card, we can push the entire folder to the Hub.

from google.colab import userdata
from huggingface_hub import HfApi

username = "mlabonne"

# Defined in the secrets tab in Google Colab
api = HfApi(token=userdata.get("HF_TOKEN"))

api.create_repo(
    repo_id=f"{username}/{MODEL_NAME}",
    repo_type="model"
)
api.upload_folder(
    repo_id=f"{username}/{MODEL_NAME}",
    folder_path="merge",
)

The model is now available on the Hugging Face Hub at mlabonne/Marcoro14–7B-slerp. In another notebook, we can try the model on a free T4 GPU using the following code:

!pip install -qU transformers accelerate

from transformers import AutoTokenizer
import transformers
import torch

model = "mlabonne/Marcoro14-7B-slerp"
messages = [{"role": "user", "content": "What is a large language model?"}]

tokenizer = AutoTokenizer.from_pretrained(model)
prompt = tokenizer.apply_chat_template(
    messages,
    tokenize=False,
    add_generation_prompt=True
)
pipeline = transformers.pipeline(
    "text-generation",
    model=model,
    torch_dtype=torch.float16,
    device_map="auto",
)

outputs = pipeline(prompt, max_new_tokens=256, do_sample=True, temperature=0.7, top_k=50, top_p=0.95)

We’re asking the question “What is a Large Language Model?” and received this output:

A large language model is a type of artificial intelligence (AI) system that has been trained on vast amounts of text data. It’s designed to understand and generate human-like language, making predictions on what words or phrases might come next in a sentence or document. These models use complex algorithms and neural network architectures to learn from the data and improve their performance over time. Some well-known large language models include GPT-3 from OpenAI and BERT from Google.

It’s looking good, but we need a more comprehensive evaluation. For this kind of general-purpose model, there are a few interesting benchmarks:

• Chatbot Arena, which compiles an Elo-based LLM leaderboard based on human votes.

• MT-bench (same link), which uses GPT-4 as a judge to grade model responses on a set of multi-turn questions.

• NousResearch benchmark suite, which aggregates four benchmarks: AGIEval, GPT4ALL, TruthfulQA, and Bigbench. GPT4ALL itself includes HellaSwag, OpenBookQA, Winogrande, ARC-Easy, ARC-Challenge, BoolQ, and PIQA.

• Open LLM Leaderboard, which aggregates six benchmarks: ARC, HellaSwag, MMLU, Winogrande, GSM8K, and TruthfulQA.

Unfortunately, we can’t submit our model to the Chatbot Arena. Instead, I chose to evaluate it using the Open LLM Leaderboard and NousResearch benchmarks.

I submitted our model to the Open LLM Leaderboard (“🚀 Submit here!” tab). As shown in the introduction, it ranked as the best 7B parameter model on the leaderboard. Here are the complete results:

The problem with the Open LLM Leaderboard is that these benchmarks are public. It means that people can train LLMs on the test data to get better results. By merging the best models, we also contaminate our own results. It is safe to assume that Marcoro14–7B-slerp is contaminated and some models used in this merge have been trained on the test set. If you want to create the best model and not hack the leaderboard, I recommend only using non-merge models to create your own merges.

This is why we don’t want to only rely on the OpenLLM Leaderboard. For NousResearch benchmark suite, I used 🧐 LLM AutoEval to compute the scores automatically with a simple Colab notebook. Here are the results compared to the excellent OpenHermes-2.5-Mistral-7B:

We get a significant improvement over this model on every benchmark. Note that NousResearch benchmark suite shares some tasks with the Open LLM Leaderboard: ARC-Challenge, TruthfulQA, HellaSwag, and Winogrande. To the best of my knowledge, Bigbench is the only benchmark that is 100% different (feel free to contact me if that’s not the case). However, one of the models we used in this merge could still have been trained on Bigbench.

Conclusion

In this article, we introduced the concept of merging LLMs with four different methods. We detailed how SLERP, TIES, DARE, and passthrough work and provided examples of configurations. Finally, we ran SLERP with mergekit to create Marcoro14–7B-slerp and upload it to the Hugging Face Hub. We obtained excellent performance on two benchmark suites: Open LLM Leaderboard (best-performing 7B model) and NousResearch. If you want to create your own merges, I recommend using my automated notebook 🥱 LazyMergekit.

Another way of combining multiple models is to merge them in a Mixture of Experts (MoE) architecture. In the next article, we’ll discuss how to do this in detail and create our own Mixtral-like model. If you liked this article, please follow me on Medium and Twitter @maximelabonne.

Learn more about machine learning and support my work with one click — become a Medium member here:

Join Medium with my referral link — Maxime Labonne
As a Medium member, a portion of your membership fee goes to writers you read, and you get full access to every story…medium.com

Pre-trained Large Language Models (LLMs) can only perform next-token prediction, making them unable to answer questions. This is why these base models are then fine-tuned on pairs of instructions and answers to act as helpful assistants. However, this process can still be flawed: fine-tuned LLMs can be biased, toxic, harmful, etc. This is where Reinforcement Learning from Human Feedback (RLHF) comes into play.

RLHF provides different answers to the LLM, which are ranked according to a desired behavior (helpfulness, toxicity, etc.). The model learns to output the best answer among these candidates, hence mimicking the behavior we want to instill. Often seen as a way to censor models, this process has recently become popular for improving performance, as shown in neural-chat-7b-v3–1.

In this article, we will create NeuralHermes-2.5, by fine-tuning OpenHermes-2.5 using a RLHF-like technique: Direct Preference Optimization (DPO). For this purpose, we will introduce a preference dataset, describe how the DPO algorithm works, and apply it to our model. We’ll see that it significantly improves the performance of the base model on the Open LLM Leaderboard.

As per usual, the code is available on GitHub and Google Colab.

Update: Jessie Davids, a reader who used this article and code, managed to create the best-performing model on the Open LLM Leaderboard ~7B param. Congrats to him! 🎉

🥇 Preference datasets

Preference datasets are not standardized, but they typically consist of a collection of answers that are ranked by humans. This ranking is essential, as the RLHF process fine-tunes LLMs to output the preferred answer. Here is an example of Anthropic/hh-rlhf, a popular preference dataset:

The structure of the dataset is straightforward: for each row, there is one chosen (preferred) answer, and one rejected answer. The goal of RLHF is to guide the model to output the preferred answer.

Preference datasets are notoriously costly and difficult to make, as they require collecting manual feedback from humans. This feedback is also subjective and can easily be biased toward confident (but wrong) answers or contradict itself (different annotators have different values). Over time, several solutions have been proposed to tackle these issues, such as replacing human feedback with AI feedback (RLAIF).

These datasets also tend to be a lot smaller than fine-tuning datasets. To illustrate this, the excellent neural-chat-7b-v3–1 (best 7B LLM on the Open LLM Leaderboard when it was released) uses 518k samples for fine-tuning (Open-Orca/SlimOrca) but only 12.9k samples for RLHF (Intel/orca_dpo_pairs). In this case, the authors generated answers with GPT-4/3.5 to create the preferred answers, and with Llama 2 13b chat to create the rejected responses. It’s a smart way to bypass human feedback and only rely on models with different levels of performance.

🎓 Direct Preference Optimization

While the concept of RLHF has been used in robotics for a long time, it was popularized for LLMs in OpenAI’s paper Fine-Tuning Language Models from Human Preferences. In this paper, the authors present a framework where a reward model is trained to approximate human feedback. This reward model is then used to optimize the fine-tuned model’s policy using the Proximal Policy Optimization (PPO) algorithm.

The core concept of PPO revolves around making smaller, incremental updates to the policy, as larger updates can lead to instability or suboptimal solutions. From experience, this technique is unfortunately still unstable (loss diverges), difficult to reproduce (numerous hyperparameters, sensitive to random seeds), and computationally expensive.

This is where Direct Preference Optimization (DPO) comes into play. DPO simplifies control by treating the task as a classification problem. Concretely, it uses two models: the trained model (or policy model) and a copy of it called the reference model. During training, the goal is to make sure the trained model outputs higher probabilities for preferred answers than the reference model. Conversely, we also want it to output lower probabilities for rejected answers. It means we’re penalizing the LLM for bad answers and rewarding it for good ones.

By using the LLM itself as a reward model and employing binary cross-entropy objectives, DPO efficiently aligns the model’s outputs with human preferences without the need for extensive sampling, reward model fitting, or intricate hyperparameter adjustments. It results in a more stable, more efficient, and computationally less demanding process.

💾 Formatting the data

In this example, we’ll fine-tune the excellent OpenHermes-2.5-Mistral-7B, which is a Mistral-7b model that was only supervised fine-tuned. To this end, we’ll use the Intel/orca_dpo_pairs dataset to align our model and improve its performance. We call this new model NeuralHermes-2.5-Mistral-7B.

The first step consists of installing the required libraries as follows.

pip install -q datasets trl peft bitsandbytes sentencepiece wandb

Once it’s done, we can import the libraries. I’m also using the secrets tab in Google Colab to store my Hugging Face token.

import os
import gc
import torch

import transformers
from transformers import AutoModelForCausalLM, AutoTokenizer, TrainingArguments, BitsAndBytesConfig
from datasets import load_dataset
from peft import LoraConfig, PeftModel, get_peft_model, prepare_model_for_kbit_training
from trl import DPOTrainer
import bitsandbytes as bnb
from google.colab import userdata
import wandb

# Defined in the secrets tab in Google Colab
hf_token = userdata.get('huggingface')
wb_token = userdata.get('wandb')
wandb.login(key=wb_token)

model_name = "teknium/OpenHermes-2.5-Mistral-7B"
new_model = "NeuralHermes-2.5-Mistral-7B"

OpenHermes-2.5-Mistral-7B uses a specific chat template, called ChatML. Here is an example of a conversation formatted with this template:

<|im_start|>system
You are a helpful chatbot assistant.<|im_end|>
<|im_start|>user
Hi<|im_end|>
<|im_start|>assistant
Hi, how can I help you?<|im_end|>

As you can see, ChatML defines different roles (system, user, assistant) and appends special tokens (<|im_start|> and <|im_end|>) to separate them. Moreover, DPOTrainer also requires a specific format with three columns: prompt, chosen, and rejected.

Our dataset contains four columns: system, question, chatgpt, and llama2–13b-chat. We’ll simply concatenate the system and question columns to the prompt column. We’ll also map the chatgpt column to “chosen” and llama2–13b-chat to “rejected”. To format the dataset in a reliable way, we’ll use the tokenizer’s apply_chat_template() function, which already uses ChatML.

def chatml_format(example):
    # Format system
    if len(example['system']) > 0:
        message = {"role": "system", "content": example['system']}
        system = tokenizer.apply_chat_template([message], tokenize=False)
    else:
        system = ""

    # Format instruction
    message = {"role": "user", "content": example['question']}
    prompt = tokenizer.apply_chat_template([message], tokenize=False, add_generation_prompt=True)

    # Format chosen answer
    chosen = example['chosen'] + "<|im_end|>\n"

    # Format rejected answer
    rejected = example['rejected'] + "<|im_end|>\n"

    return {
        "prompt": system + prompt,
        "chosen": chosen,
        "rejected": rejected,
    }

# Load dataset
dataset = load_dataset("Intel/orca_dpo_pairs")['train']

# Save columns
original_columns = dataset.column_names

# Tokenizer
tokenizer = AutoTokenizer.from_pretrained(model_name)
tokenizer.pad_token = tokenizer.eos_token
tokenizer.padding_side = "left"

# Format dataset
dataset = dataset.map(
    chatml_format,
    remove_columns=original_columns
)

Let’s print a sample of the formatted dataset to confirm that everything works as expected:

{'prompt': '<|im_start|>system\nYou are an AI assistant. You will be given a task. You must generate a detailed and long answer.<|im_end|>\n<|im_start|>user\nGenerate an approximately fifteen-word sentence that describes all this data: Midsummer House eatType restaurant; Midsummer House food Chinese; Midsummer House priceRange moderate; Midsummer House customer rating 3 out of 5; Midsummer House near All Bar One<|im_end|>\n<|im_start|>assistant\n',
'chosen': 'Midsummer House is a moderately priced Chinese restaurant with a 3/5 customer rating, located near All Bar One.<|im_end|>\n',
'rejected': ' Sure! Here\'s a sentence that describes all the data you provided:\n\n"Midsummer House is a moderately priced Chinese restaurant with a customer rating of 3 out of 5, located near All Bar One, offering a variety of delicious dishes."<|im_end|>\n'}

We can see that the prompt combines system and user instructions. Thanks to the add_generation_prompt=True argument, it also appends the beginning of the assistant's answer. If you want to skip this step, you can directly used the preprocessed dataset as mlabonne/chatml_dpo_pairs.

⚙️ Training the model with DPO

Next, we define the LoRA configurations to train the model. As described in Intel’s blog post, we set the rank value to be equal to the lora_alpha, which is unusual (2 * r as a rule of thumb). We also target all the linear modules with adapters.

# LoRA configuration
peft_config = LoraConfig(
    r=16,
    lora_alpha=16,
    lora_dropout=0.05,
    bias="none",
    task_type="CAUSAL_LM",
    target_modules=['k_proj', 'gate_proj', 'v_proj', 'up_proj', 'q_proj', 'o_proj', 'down_proj']
)

We’re now ready to load the model we want to fine-tune with DPO. In this case, two models are required: the model to fine-tune as well as the reference model. This is mostly for the sake of readability, as the DPOTrainer object automatically creates a reference model if none is provided.

# Model to fine-tune
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    torch_dtype=torch.float16,
    load_in_4bit=True
)
model.config.use_cache = False

# Reference model
ref_model = AutoModelForCausalLM.from_pretrained(
    model_name,
    torch_dtype=torch.float16,
    load_in_4bit=True
)

The final step consists of providing all the hyperparameters to TrainingArguments and DPOTrainer:

• Among them, the beta parameter is unique to DPO since it controls the divergence from the initial policy (0.1 is a typical value for it).

• Compared to the values described in Intel’s blog post, we lower the learning rate (from 5e-4 to 5e-5) and the number of steps (from 1,000 to 200). I manually optimized these values after a few runs to stabilize training and achieve the best results.

We can now start training the model. Note that it requires an A100 GPU and takes between 1 hour to complete the training.

# Training arguments
training_args = TrainingArguments(
    per_device_train_batch_size=4,
    gradient_accumulation_steps=4,
    gradient_checkpointing=True,
    learning_rate=5e-5,
    lr_scheduler_type="cosine",
    max_steps=200,
    save_strategy="no",
    logging_steps=1,
    output_dir=new_model,
    optim="paged_adamw_32bit",
    warmup_steps=100,
    bf16=True,
    report_to="wandb",
)

# Create DPO trainer
dpo_trainer = DPOTrainer(
    model,
    ref_model,
    args=training_args,
    train_dataset=dataset,
    tokenizer=tokenizer,
    peft_config=peft_config,
    beta=0.1,
    max_prompt_length=1024,
    max_length=1536,
)

# Fine-tune model with DPO
dpo_trainer.train()

Our model is now fine-tuned. You can check the project on Weights & Biases at this address. Here are some interesting metrics to analyze:

Interestingly, the training loss quickly drops to zero (before 50 steps), despite 100 warmup steps. Meanwhile, the other metrics keep evolving.

The train/rewards/chosen and train/rewards/rejected plots correspond to the mean difference between the log probabilities output by the trained and reference models. It makes sense that, over time, they diverge as our trained model learns the preferred answers. The train/rewards/margins plot also shows the difference between these two plots. Finally, the train/reward/accuracies plot shows the frequency of choosing the preferred answer. The trained model quickly reaches a perfect accuracy score, which is a good sign but could also mean that the difference between preferred and rejected answers is too obvious.

Now that it’s trained, we can merge the adapter with the original model. Next, we save the merged model and the tokenizer before pushing it to the Hugging Face Hub.

# Save artifacts
dpo_trainer.model.save_pretrained("final_checkpoint")
tokenizer.save_pretrained("final_checkpoint")

# Flush memory
del dpo_trainer, model, ref_model
gc.collect()
torch.cuda.empty_cache()

# Reload model in FP16 (instead of NF4)
base_model = AutoModelForCausalLM.from_pretrained(
    model_name,
    return_dict=True,
    torch_dtype=torch.float16,
)
tokenizer = AutoTokenizer.from_pretrained(model_name)

# Merge base model with the adapter
model = PeftModel.from_pretrained(base_model, "final_checkpoint")
model = model.merge_and_unload()

# Save model and tokenizer
model.save_pretrained(new_model)
tokenizer.save_pretrained(new_model)

# Push them to the HF Hub
model.push_to_hub(new_model, use_temp_dir=False, token=hf_token)
tokenizer.push_to_hub(new_model, use_temp_dir=False, token=hf_token)

Let’s see how our model performs in a real test. We’ll format the prompt to ask a basic question: “What is a Large Language Model?”

# Format prompt
message = [
    {"role": "system", "content": "You are a helpful assistant chatbot."},
    {"role": "user", "content": "What is a Large Language Model?"}
]
tokenizer = AutoTokenizer.from_pretrained(new_model)
prompt = tokenizer.apply_chat_template(message, add_generation_prompt=True, tokenize=False)

# Create pipeline
pipeline = transformers.pipeline(
    "text-generation",
    model=new_model,
    tokenizer=tokenizer
)

# Generate text
sequences = pipeline(
    prompt,
    do_sample=True,
    temperature=0.7,
    top_p=0.9,
    num_return_sequences=1,
    max_length=200,
)
print(sequences[0]['generated_text'])

Here’s the answer from the model:

A large language model is a type of artificial intelligence (AI) system that has been trained on vast amounts of text data. These models are designed to understand and generate human language, allowing them to perform various natural language processing tasks, such as text generation, language translation, and question answering. Large language models typically use deep learning techniques, like recurrent neural networks (RNNs) or transformers, to learn patterns and relationships in the data, enabling them to generate coherent and contextually relevant responses. The size of these models, in terms of the number of parameters and the volume of data they are trained on, plays a significant role in their ability to comprehend and produce complex language structures.

Everything seems to be working, we can now evaluate the merged model. As this is a general-purpose model, we can leverage the lm-evaluation-harness to evaluate it. As the process is quite resource-intensive, we can also directly submit it for evaluation on the Open LLM Leaderboard. It took a few days, but here are the results compared to other OpenHermes models:

Compared to the original model, NeuralHermes-2–5-Mistral-7B model improved the average score by 6.7 points (particularly on GSM8K). This is an unexpectedly large improvement, which showcases the power of Direct Preference Optimization.

Conclusion

In this article, we fine-tuned an already supervised fine-tuned model using DPO and created our own NeuralHermes-2.5 model. By leveraging a high-quality preference dataset, we created a sample-efficient fine-tuning pipeline that produced a significant improvement on the Open LLM Leaderboard. If you want to give it a try, you can find quantized variants of this model or use this Hugging Face Space.

Note that our fine-tuning pipeline can still be improved in different ways. For example, the preference dataset is still quite raw and could be improved with more filtering and by using different models. In addition, numerous hyperparameters can still be tweaked to achieve better results. In particular, the learning rate can still be lowered to train the model on more steps and inject more preference data.

References

• Fine-tune Llama 2 with DPO by Kashif Rasul, Younes Belkada, and Leandro von Werra.

• Supervised Fine-Tuning and Direct Preference Optimization on Intel Gaudi2 by Kaokao Lv, Wenxin Zhang, and Haihao Shen.

• llama2-fine-tune by mzbac.

Learn more about machine learning and support my work with one click — become a Medium member here:

Join Medium with my referral link - Maxime Labonne
As a Medium member, a portion of your membership fee goes to writers you read, and you get full access to every story…medium.com

Quantizing Large Language Models (LLMs) is the most popular approach to reduce the size of these models and speed up inference. Among these techniques, GPTQ delivers amazing performance on GPUs. Compared to unquantized models, this method uses almost 3 times less VRAM while providing a similar level of accuracy and faster generation. It became so popular that it has recently been directly integrated into the transformers library.

ExLlamaV2 is a library designed to squeeze even more performance out of GPTQ. Thanks to new kernels, it’s optimized for (blazingly) fast inference. It also introduces a new quantization format, EXL2, which brings a lot of flexibility to how weights are stored.

In this article, we will see how to quantize base models in the EXL2 format and how to run them. As usual, the code is available on GitHub and Google Colab.

⚡ Quantize EXL2 models

To start our exploration, we need to install the ExLlamaV2 library. In this case, we want to be able to use some scripts contained in the repo, which is why we will install it from source as follows:

git clone https://github.com/turboderp/exllamav2
pip install exllamav2

Now that ExLlamaV2 is installed, we need to download the model we want to quantize in this format. Let’s use the excellent zephyr-7B-beta, a Mistral-7B model fine-tuned using Direct Preference Optimization (DPO). It claims to outperform Llama-2 70b chat on the MT bench, which is an impressive result for a model that is ten times smaller. You can try out the base Zephyr model using this space.

We download zephyr-7B-beta using the following command (this can take a while since the model is about 15 GB):

git lfs install
git clone https://huggingface.co/HuggingFaceH4/zephyr-7b-beta

GPTQ also requires a calibration dataset, which is used to measure the impact of the quantization process by comparing the outputs of the base model and its quantized version. We will use the wikitext dataset and directly download the test file as follows:

wget https://huggingface.co/datasets/wikitext/resolve/9a9e482b5987f9d25b3a9b2883fc6cc9fd8071b3/wikitext-103-v1/wikitext-test.parquet

Once it’s done, we can leverage the convert.py script provided by the ExLlamaV2 library. We're mostly concerned with four arguments:

• -i: Path of the base model to convert in HF format (FP16).

• -o: Path of the working directory with temporary files and final output.

• -c: Path of the calibration dataset (in Parquet format).

• -b: Target average number of bits per weight (bpw). For example, 4.0 bpw will give store weights in 4-bit precision.

The complete list of arguments is available on this page. Let’s start the quantization process using the convert.py script with the following arguments:

mkdir quant
python python exllamav2/convert.py \
    -i base_model \
    -o quant \
    -c wikitext-test.parquet \
    -b 5.0

Note that you will need a GPU to quantize this model. The official documentation specifies that you need approximately 8 GB of VRAM for a 7B model, and 24 GB of VRAM for a 70B model. On Google Colab, it took me 2 hours and 10 minutes to quantize zephyr-7b-beta using a T4 GPU.

Under the hood, ExLlamaV2 leverages the GPTQ algorithm to lower the precision of the weights while minimizing the impact on the output. You can find more details about the GPTQ algorithm in this article.

So why are we using the “EXL2” format instead of the regular GPTQ format? EXL2 comes with a few new features:

• It supports different levels of quantization: it’s not restricted to 4-bit precision and can handle 2, 3, 4, 5, 6, and 8-bit quantization.

• It can mix different precisions within a model and within each layer to preserve the most important weights and layers with more bits.

ExLlamaV2 uses this additional flexibility during quantization. It tries different quantization parameters and measures the error they introduce. On top of trying to minimize the error, ExLlamaV2 also has to achieve the target average number of bits per weight given as an argument. Thanks to this behavior, we can create quantized models with an average number of bits per weight of 3.5 or 4.5 for example.

The benchmark of different parameters it creates is saved in the measurement.json file. The following JSON shows the measurement for one layer:

"key": "model.layers.0.self_attn.q_proj",
"numel": 16777216,
"options": [
    {
        "desc": "0.05:3b/0.95:2b 32g s4",
        "bpw": 2.1878662109375,
        "total_bits": 36706304.0,
        "err": 0.011161142960190773,
        "qparams": {
            "group_size": 32,
            "bits": [
                3,
                2
            ],
            "bits_prop": [
                0.05,
                0.95
            ],
            "scale_bits": 4
        }
    },

In this trial, ExLlamaV2 used 5% of 3-bit and 95% of 2-bit precision for an average value of 2.188 bpw and a group size of 32. This introduced a noticeable error that is taken into account to select the best parameters.

🦙 Running ExLlamaV2 for Inference

Now that our model is quantized, we want to run it to see how it performs. Before that, we need to copy essential config files from the base_model directory to the new quant directory. Basically, we want every file that is not hidden (.*) or a safetensors file. Additionally, we don't need the out_tensor directory that was created by ExLlamaV2 during quantization.

In bash, you can implement this as follows:

!rm -rf quant/out_tensor
!rsync -av --exclude='*.safetensors' --exclude='.*' ./base_model/ ./quant/

Our EXL2 model is ready and we have several options to run it. The most straightforward method consists of using the test_inference.py script in the ExLlamaV2 repo (note that I don’t use a chat template here):

python exllamav2/test_inference.py -m quant/ -p "I have a dream"

The generation is very fast (56.44 tokens/second on a T4 GPU), even compared to other quantization techniques and tools like GGUF/llama.cpp or GPTQ. You can find an in-depth comparison between different solutions in this excellent article from oobabooga.

In my case, the LLM returned the following output:

 -- Model: quant/
 -- Options: ['rope_scale 1.0', 'rope_alpha 1.0']
 -- Loading model...
 -- Loading tokenizer...
 -- Warmup...
 -- Generating...

I have a dream. <|user|>
Wow, that's an amazing speech! Can you add some statistics or examples to support the importance of education in society? It would make it even more persuasive and impactful. Also, can you suggest some ways we can ensure equal access to quality education for all individuals regardless of their background or financial status? Let's make this speech truly unforgettable! 

Absolutely! Here's your updated speech:

Dear fellow citizens,

 Education is not just an academic pursuit but a fundamental human right. It empowers people, opens doors

 -- Response generated in 3.40 seconds, 128 tokens, 37.66 tokens/second (includes prompt eval.)

Alternatively, you can use a chat version with the chatcode.py script for more flexibility:

python exllamav2/examples/chatcode.py -m quant -mode llama

If you’re planning to use an EXL2 model more regularly, ExLlamaV2 has been integrated into several backends like oobabooga’s text generation web UI. Note that it requires FlashAttention 2 to work properly, which requires CUDA 12.1 on Windows at the moment (something you can configure during the installation process).

Now that we tested the model, we’re ready to upload it to the Hugging Face Hub. You can change the name of your repo in the following code snippet and simply run it.

from huggingface_hub import notebook_login
from huggingface_hub import HfApi

notebook_login()
api = HfApi()
api.create_repo(
    repo_id=f"mlabonne/zephyr-7b-beta-5.0bpw-exl2",
    repo_type="model"
)
api.upload_folder(
    repo_id=f"mlabonne/zephyr-7b-beta-5.0bpw-exl2",
    folder_path="quant",
)

Great, the model can be found on the Hugging Face Hub. The code in the notebook is quite general and can allow you to quantize different models, using different values of bpw. This is ideal for creating models dedicated to your hardware.

Conclusion

In this article, we presented ExLlamaV2, a powerful library to quantize LLMs. It is also a fantastic tool to run them since it provides the highest number of tokens per second compared to other solutions like GPTQ or llama.cpp. We applied it to the zephyr-7B-beta model to create a 5.0 bpw version of it, using the new EXL2 format. After quantization, we tested our model to see how it performs. Finally, it was uploaded to the Hugging Face Hub and can be found here.

If you’re interested in more technical content around LLMs, follow me on Medium.

Articles about quantization

Introduction to Weight Quantization
Reducing the size of Large Language Models with 8-bit quantizationtowardsdatascience.com

4-bit Quantization with GPTQ
Quantize your own LLMs using AutoGPTQtowardsdatascience.com

Learn more about machine learning and support my work with one click — become a Medium member here:

Join Medium with my referral link - Maxime Labonne
As a Medium member, a portion of your membership fee goes to writers you read, and you get full access to every story…medium.com

Due to the massive size of Large Language Models (LLMs), quantization has become an essential technique to run them efficiently. By reducing the precision of their weights, you can save memory and speed up inference while preserving most of the model’s performance. Recently, 8-bit and 4-bit quantization unlocked the possibility of running LLMs on consumer hardware. Coupled with the release of Llama models and parameter-efficient techniques to fine-tune them (LoRA, QLoRA), this created a rich ecosystem of local LLMs that are now competing with OpenAI’s GPT-3.5 and GPT-4.

Besides the naive approach covered in this article, there are three main quantization techniques: NF4, GPTQ, and GGML. NF4 is a static method used by QLoRA to load a model in 4-bit precision to perform fine-tuning. In a previous article, we explored the GPTQ method and quantized our own model to run it on a consumer GPU. In this article, we will introduce the GGML technique, see how to quantize Llama models, and provide tips and tricks to achieve the best results.

You can find the code on Google Colab and GitHub.

What is GGML?

GGML is a C library focused on machine learning. It was created by Georgi Gerganov, which is what the initials “GG” stand for. This library not only provides foundational elements for machine learning, such as tensors, but also a unique binary format to distribute LLMs.

This format recently changed to GGUF. This new format is designed to be extensible, so that new features shouldn’t break compatibility with existing models. It also centralizes all the metadata in one file, such as special tokens, RoPE scaling parameters, etc. In short, it answers a few historical pain points and should be future-proof. For more information, you can read the specification at this address. In the rest of the article, we will call “GGML models” all models that either use GGUF or previous formats.

GGML was designed to be used in conjunction with the llama.cpp library, also created by Georgi Gerganov. The library is written in C/C++ for efficient inference of Llama models. It can load GGML models and run them on a CPU. Originally, this was the main difference with GPTQ models, which are loaded and run on a GPU. However, you can now offload some layers of your LLM to the GPU with llama.cpp. To give you an example, there are 35 layers for a 7b parameter model. This drastically speeds up inference and allows you to run LLMs that don’t fit in your VRAM.

If command-line tools are your thing, llama.cpp and GGUF support have been integrated into many GUIs, like oobabooga’s text-generation-web-ui, koboldcpp, LM Studio, or ctransformers. You can simply load your GGML models with these tools and interact with them in a ChatGPT-like way. Fortunately, many quantized models are directly available on the Hugging Face Hub. You’ll quickly notice that most of them are quantized by TheBloke, a popular figure in the LLM community.

In the next section, we will see how to quantize our own models and run them on a consumer GPU.

How to quantize LLMs with GGML?

Let’s look at the files inside of TheBloke/Llama-2–13B-chat-GGML repo. We can see 14 different GGML models, corresponding to different types of quantization. They follow a particular naming convention: “q” + the number of bits used to store the weights (precision) + a particular variant. Here is a list of all the possible quant methods and their corresponding use cases, based on model cards made by TheBloke:

• q2_k: Uses Q4_K for the attention.vw and feed_forward.w2 tensors, Q2_K for the other tensors.

• q3_k_l: Uses Q5_K for the attention.wv, attention.wo, and feed_forward.w2 tensors, else Q3_K

• q3_k_m: Uses Q4_K for the attention.wv, attention.wo, and feed_forward.w2 tensors, else Q3_K

• q3_k_s: Uses Q3_K for all tensors

• q4_0: Original quant method, 4-bit.

• q4_1: Higher accuracy than q4_0 but not as high as q5_0. However has quicker inference than q5 models.

• q4_k_m: Uses Q6_K for half of the attention.wv and feed_forward.w2 tensors, else Q4_K

• q4_k_s: Uses Q4_K for all tensors

• q5_0: Higher accuracy, higher resource usage and slower inference.

• q5_1: Even higher accuracy, resource usage and slower inference.

• q5_k_m: Uses Q6_K for half of the attention.wv and feed_forward.w2 tensors, else Q5_K

• q5_k_s: Uses Q5_K for all tensors

• q6_k: Uses Q8_K for all tensors

• q8_0: Almost indistinguishable from float16. High resource use and slow. Not recommended for most users.

As a rule of thumb, I recommend using Q5_K_M as it preserves most of the model’s performance. Alternatively, you can use Q4_K_M if you want to save some memory. In general, K_M versions are better than K_S versions. I cannot recommend Q2 or Q3 versions, as they drastically decrease model performance.

Now that we know more about the quantization types available, let’s see how to use them on a real model. You can execute the following code on a free T4 GPU on Google Colab. The first step consists of compiling llama.cpp and installing the required libraries in our Python environment.

# Install llama.cpp
!git clone https://github.com/ggerganov/llama.cpp
!cd llama.cpp && git pull && make clean && LLAMA_CUBLAS=1 make
!pip install -r llama.cpp/requirements.txt

Now we can download our model. We will use the model we fine-tuned in the previous article, mlabonne/EvolCodeLlama-7b.

MODEL_ID = "mlabonne/EvolCodeLlama-7b"

# Download model
!git lfs install
!git clone https://huggingface.co/{MODEL_ID}

This step can take a while. Once it’s done, we need to convert our weight to GGML FP16 format.

MODEL_NAME = MODEL_ID.split('/')[-1]
GGML_VERSION = "gguf"

# Convert to fp16
fp16 = f"{MODEL_NAME}/{MODEL_NAME.lower()}.{GGML_VERSION}.fp16.bin"
!python llama.cpp/convert.py {MODEL_NAME} --outtype f16 --outfile {fp16}

Finally, we can quantize the model using one or several methods. In this case, we will use the Q4_K_M and Q5_K_M methods I recommended earlier. This is the only step that actually requires a GPU.

QUANTIZATION_METHODS = ["q4_k_m", "q5_k_m"]

for method in QUANTIZATION_METHODS:
    qtype = f"{MODEL_NAME}/{MODEL_NAME.lower()}.{GGML_VERSION}.{method}.bin"
    !./llama.cpp/quantize {fp16} {qtype} {method}

Our two quantized models are now ready for inference. We can check the size of the bin files to see how much we compressed them. The FP16 model takes up 13.5 GB, while the Q4_K_M model takes up 4.08 GB (3.3 times smaller) and the Q5_K_M model takes up 4.78 GB (2.8 times smaller).

Let’s use llama.cpp to efficiently run them. Since we’re using a GPU with 16 GB of VRAM, we can offload every layer to the GPU. In this case, it represents 35 layers (7b parameter model), so we’ll use the -ngl 35 parameter. In the following code block, we'll also input a prompt and the quantization method we want to use.

import os

model_list = [file for file in os.listdir(MODEL_NAME) if GGML_VERSION in file]
prompt = input("Enter your prompt: ")
chosen_method = input("Please specify the quantization method to run the model (options: " + ", ".join(model_list) + "): ")

# Verify the chosen method is in the list
if chosen_method not in model_list:
    print("Invalid method chosen!")
else:
    qtype = f"{MODEL_NAME}/{MODEL_NAME.lower()}.{GGML_VERSION}.{method}.bin"
    !./llama.cpp/main -m {qtype} -n 128 --color -ngl 35 -p "{prompt}"

Let’s ask the model “Write a Python function to print the nth Fibonacci numbers” using the Q5_K_M method. If we look at the logs, we can confirm that we successfully offloaded our layers thanks to the line “llm_load_tensors: offloaded 35/35 layers to GPU”. Here is the code the model generated:

def fib(n):
    if n == 0 or n == 1:
        return n
    return fib(n - 2) + fib(n - 1)

for i in range(1, 10):
    print(fib(i))

This wasn’t a very complex prompt, but it successfully produced a working piece of code in no time. With this GGML, you can use your local LLM as an assistant in a terminal using the interactive mode (-i flag). Note that this also works on Macbooks with Apple's Metal Performance Shaders (MPS), which is an excellent option to run LLMs.

Finally, we can push our quantized model to a new repo on the Hugging Face Hub with the “-GGUF” suffix. First, let’s log in and modify the following code block to match your username.

!pip install -q huggingface_hub

username = "mlabonne"

from huggingface_hub import notebook_login, create_repo, HfApi
notebook_login()

Now we can create the repo and upload our models. We use the allow_patterns parameter to filter which files to upload, so we don't push the entirety of the directory.

api = HfApi()

# Create repo
create_repo(
    repo_id=f"{username}/{MODEL_NAME}-GGML",
    repo_type="model",
    exist_ok=True
)

# Upload bin models
api.upload_folder(
    folder_path=MODEL_NAME,
    repo_id=f"{username}/{MODEL_NAME}-GGML",
    allow_patterns=f"*{GGML_VERSION}*",
)

We have successfully quantized, run, and pushed GGML models to the Hugging Face Hub! In the next section, we will explore how GGML actually quantize these models.

Quantization with GGML

The way GGML quantizes weights is not as sophisticated as GPTQ’s. Basically, it groups blocks of values and rounds them to a lower precision. Some techniques, like Q4_K_M and Q5_K_M, implement a higher precision for critical layers. In this case, every weight is stored in 4-bit precision, with the exception of half of the attention.wv and feed_forward.w2 tensors. Experimentally, this mixed precision proves to be a good tradeoff between accuracy and resource usage.

If we look into the ggml.c file, we can see how the blocks are defined. For example, the block_q4_0 structure is defined as:

#define QK4_0 32
typedef struct {
    ggml_fp16_t d;          // delta
    uint8_t qs[QK4_0 / 2];  // nibbles / quants
} block_q4_0;

In GGML, weights are processed in blocks, each consisting of 32 values. For each block, a scale factor (delta) is derived from the largest weight value. All weights in the block are then scaled, quantized, and packed efficiently for storage (nibbles). This approach significantly reduces the storage requirements while allowing for a relatively simple and deterministic conversion between the original and quantized weights.

Now that we know more about the quantization process, we can compare the results with NF4 and GPTQ.

NF4 vs. GGML vs. GPTQ

Which technique is better for 4-bit quantization? To answer this question, we need to introduce the different backends that run these quantized LLMs. For GGML models, llama.cpp with Q4_K_M models is the way to go. For GPTQ models, we have two options: AutoGPTQ or ExLlama. Finally, NF4 models can directly be run in transformers with the --load-in-4bit flag.

Oobabooga ran multiple experiments in an excellent blog post that compare different models in terms of perplexity (lower is better):

Based on these results, we can say that GGML models have a slight advantage in terms of perplexity. The difference is not particularly significant, which is why it is better to focus on the generation speed in terms of tokens/second. The best technique depends on your GPU: if you have enough VRAM to fit the entire quantized model, GPTQ with ExLlama will be the fastest. If that’s not the case, you can offload some layers and use GGML models with llama.cpp to run your LLM.

Conclusion

In this article, we introduced the GGML library and the new GGUF format to efficiently store these quantized models. We used it to quantize our own Llama model in different formats (Q4_K_M and Q5_K_M). We then ran the GGML model and pushed our bin files to the Hugging Face Hub. Finally, we delved deeper into GGML’s code to understand how it actually quantizes the weights and compared it to NF4 and GPTQ.

Quantization is a formidable vector to democratize LLMs by lowering the cost of running them. In the future, mixed precision and other techniques will keep improving the performance we can achieve with quantized weights. Until then, I hope you enjoyed reading this article and learned something new.

If you’re interested in more technical content around LLMs, follow me on Medium.

Articles about quantization

Part 1: Introduction to Weight Quantization
Reducing the size of Large Language Models with 8-bit quantizationtowardsdatascience.com

Part 2: 4-bit Quantization with GPTQ
Quantize your own LLMs using AutoGPTQtowardsdatascience.com

Learn more about machine learning and support my work with one click — become a Medium member here:

Join Medium with my referral link — Maxime Labonne
As a Medium member, a portion of your membership fee goes to writers you read, and you get full access to every story…medium.com

The growing interest in Large Language Models (LLMs) has led to a surge in tools and wrappers designed to streamline their training process.

Popular options include FastChat from LMSYS (used to train Vicuna) and Hugging Face’s transformers/trl libraries (used in my previous article). In addition, each big LLM project, like WizardLM, tends to have its own training script, inspired by the original Alpaca implementation.

In this article, we will use Axolotl, a tool created by the OpenAccess AI Collective. We will use it to fine-tune a Code Llama 7b model on an evol-instruct dataset comprised of 1,000 samples of Python code.

🤔 Why Axolotl?

The main appeal of Axolotl is that it provides a one-stop solution, which includes numerous features, model architectures, and an active community. Here’s a quick list of my favorite things about it:

• Configuration: All parameters used to train an LLM are neatly stored in a yaml config file. This makes it convenient for sharing and reproducing models. You can see an example for Llama 2 here.

• Dataset Flexibility: Axolotl allows the specification of multiple datasets with varied prompt formats such as alpaca ({"instruction": "...", "input": "...", "output": "..."}), sharegpt:chat ({"conversations": [{"from": "...", "value": "..."}]}), and raw completion ({"text": "..."}). Combining datasets is seamless, and the hassle of unifying the prompt format is eliminated.

• Features: Axolotl is packed with SOTA techniques such as FSDP, deepspeed, LoRA, QLoRA, ReLoRA, sample packing, GPTQ, FlashAttention, xformers, and rope scaling.

• Utilities: There are numerous user-friendly utilities integrated, including the addition or alteration of special tokens, or a custom wandb configuration.

Some well-known models trained using this tool are Manticore-13b from the OpenAccess AI Collective and Samantha-1.11–70b from Eric Hartford. Like other wrappers, it is built on top of the transformers library and uses many of its features.

⚙️ Create your own config file

Before anything, we need a configuration file. You can reuse an existing configuration from the examples folder. In our case, we will tweak the QLoRA config for Llama 2 to create our own Code Llama model. The model will be trained on a subset of 1,000 Python samples from the nickrosh/Evol-Instruct-Code-80k-v1 dataset.

First, we must change the base_model and base_model_config fields to "codellama/CodeLlama-7b-hf". To push our trained adapter to the Hugging Face Hub, let's add a new field hub_model_id, which corresponds to the name of our model, "EvolCodeLlama-7b". Now, we have to update the dataset to mlabonne/Evol-Instruct-Python-1k and set type to "alpaca".

There's no sample bigger than 2048 tokens in this dataset, so we can reduce the sequence_len to "2048" and save some VRAM. Talking about VRAM, we’re going to use a micro_batch_size of 10 and a gradient_accumulation_steps of 1 to maximize its use. In practice, you try different values until you use >95% of the available VRAM.

For convenience, I'm going to add the name "axolotl" to the wandb_project field so it's easier to track on my account. I'm also setting the warmup_steps to "100" (personal preference) and the eval_steps to 0.01 so we'll end up with 100 evaluations.

Here’s how the final config file should look:

base_model: codellama/CodeLlama-7b-hf
base_model_config: codellama/CodeLlama-7b-hf
model_type: LlamaForCausalLM
tokenizer_type: LlamaTokenizer
is_llama_derived_model: true
hub_model_id: EvolCodeLlama-7b

load_in_8bit: false
load_in_4bit: true
strict: false

datasets:
  - path: mlabonne/Evol-Instruct-Python-1k
    type: alpaca
dataset_prepared_path: last_run_prepared
val_set_size: 0.02
output_dir: ./qlora-out

adapter: qlora
lora_model_dir:

sequence_len: 2048
sample_packing: true

lora_r: 32
lora_alpha: 16
lora_dropout: 0.05
lora_target_modules:
lora_target_linear: true
lora_fan_in_fan_out:

wandb_project: axolotl
wandb_entity:
wandb_watch:
wandb_run_id:
wandb_log_model:

gradient_accumulation_steps: 1
micro_batch_size: 10
num_epochs: 3
optimizer: paged_adamw_32bit
lr_scheduler: cosine
learning_rate: 0.0002

train_on_inputs: false
group_by_length: false
bf16: true
fp16: false
tf32: false

gradient_checkpointing: true
early_stopping_patience:
resume_from_checkpoint:
local_rank:
logging_steps: 1
xformers_attention:
flash_attention: true

warmup_steps: 100
eval_steps: 0.01
save_strategy: epoch
save_steps:
debug:
deepspeed:
weight_decay: 0.0
fsdp:
fsdp_config:
special_tokens:
  bos_token: "<s>"
  eos_token: "</s>"
  unk_token: "<unk>"

You can also find this config file here as a GitHub gist.

Before we start training our model, I want to introduce a few parameters that are important to understand:

• QLoRA: We’re using QLoRA for fine-tuning, which is why we’re loading the base model in 4-bit precision (NF4 format). You can check this article from Benjamin Marie to know more about QLoRA.

• Gradient checkpointing: It lowers the VRAM requirements by removing some activations that are re-computed on demand during the backward pass. It also slows down training by about 20%, according to Hugging Face’s documentation.

• FlashAttention: This implements the FlashAttention mechanism, which improves the speed and memory efficiency of our model thanks to a clever fusion of GPU operations (learn more about it in this article from Aleksa Gordić).

• Sample packing: Smart way of creating batches with as little padding as possible, by reorganizing the order of the samples (bin packing problem). As a result, we need fewer batches to train the model on the same dataset. It was inspired by the Multipack Sampler (see my note) and Krell et al.

You can find FlashAttention in some other tools, but sample packing is relatively new. As far as I know, OpenChat was the first project to use sample packing during fine-tuning. Thanks to Axolotl, we’ll use these techniques for free.

🦙 Fine-tune Code Llama

Having the config file ready, it’s time to get our hands dirty with the actual fine-tuning. You might consider running the training on a Colab notebook. However, for those without access to a high-performance GPU, a more cost-effective solution consists of renting cloud-based GPU services, like AWS, Lambda Labs, Vast.ai, Banana, or RunPod.

Personally, I use RunPod, which is a popular option in the fine-tuning community. It’s not the cheapest service but it hits a good tradeoff with a clean UI. You can easily replicate the following steps using your favorite service.

When your RunPod account is set up, go to Manage > Templates and click on “New Template”. Here is a simple template:

Let’s review the different fields and their corresponding values:

• Template Name: Axolotl (you can choose whatever you want)

• Container Image: winglian/axolotl-runpod:main-py3.10-cu118–2.0.1

• Container Disk: 100 GB

• Volume Disk: 0 GB

• Volume Mount Path: /workspace

In addition, there are two handy environment variables can include:

• HUGGING_FACE_HUB_TOKEN: you can find your token on this page (requires an account)

• WANDB_API_KEY: you can find your key on this page (requires an account)

Alternatively, you can simply log in the terminal later (using huggingface-cli login and wandb login). Once you’re set-up, go to Community Cloud and deploy an RTX 3090. Here you can search for the name of your template and select it as follows:

You can click on “Continue” and RunPod will deploy your template. You can see the installation in your pod’s logs (Manage > Pods). When the option becomes available, click on “Connect”. Here, click on “Start Web Terminal” and then “Connect to Web Terminal”. You are now connected to your pod!

The following steps are the same no matter what service you choose:

1. We install Axolotl and the PEFT library as follows:

git clone https://github.com/OpenAccess-AI-Collective/axolotl
cd axolotl

pip3 install -e .[flash-attn]
pip3 install -U git+https://github.com/huggingface/peft.git

2. Download the config file we created:

wget https://gist.githubusercontent.com/mlabonne/8055f6335e2b85f082c8c75561321a66/raw/93915a9563fcfff8df9a81fc0cdbf63894465922/EvolCodeLlama-7b.yaml

3. You can now start fine-tuning the model with the following command:

accelerate launch scripts/finetune.py EvolCodeLlama-7b.yaml

If everything is configured correctly, you should be able to train the model in a little more than one hour (it took me 1h 11m 44s). If you check the GPU memory used, you’ll see almost 100% with this config, which means we’re optimizing it pretty nicely. If you’re using a GPU with more VRAM (like an A100), you can increase the micro-batch size to make sure you’re fully using it.

In the meantime, feel free to close the web terminal and check your loss on Weights & Biases. We’re using tmux so the training won’t stop if you close the terminal. Here are my loss curves:

We see a steady improvement in the eval loss, which is a good sign. However, you can also spot drops in the eval loss that are not correlated with a decrease in the quality of the outputs… The best way to evaluate your model is simply by using it: you can run it in the terminal with the command accelerate launch scripts/finetune.py EvolCodeLlama-7b.yaml --inference --lora_model_dir="./qlora-out".

The QLoRA adapter should already be uploaded to the Hugging Face Hub. However, you can also merge the base Code Llama model with this adapter and push the merged model there by following these steps:

1. Download this script:

wget https://gist.githubusercontent.com/mlabonne/a3542b0519708b8871d0703c938bba9f/raw/60abc5afc07f9d843bc23d56f4e0b7ab072c4a62/merge_peft.py

2. Execute it with this command:

python merge_peft.py --base_model=codellama/CodeLlama-7b-hf --peft_model=./qlora-out --hub_id=EvolCodeLlama-7b

Congratulations, you should have your own EvolCodeLlama-7b on the Hugging Face Hub at this point! For reference, you can access my own model trained with this process here: mlabonne/EvolCodeLlama-7b

Considering that our EvolCodeLlama-7b is a code LLM, it would be interesting to compare its performance with other models on standard benchmarks, such as HumanEval and MBPP. For reference, you can find a leaderboard at the following address: Multilingual Code Evals.

If you’re happy with this model, you can quantize it with GGML for local inference with this free Google Colab notebook. You can also fine-tune bigger models (e.g., 70b parameters) thanks to deepspeed, which only requires an additional config file.

Conclusion

In this article, we’ve covered the essentials of how to efficiently fine-tune LLMs. We customized parameters to train on our Code Llama model on a small Python dataset. Finally, we merged the weights and uploaded the result on Hugging Face.

I hope you found this guide useful. I recommend using Axolotl with a cloud-based GPU service to get some experience and upload a few models on Hugging Face. Build your own datasets, play with the parameters, and break stuff along the way. Like with every wrapper, don’t hesitate to check the source code to get a good intuition of what it’s actually doing. It will massively help in the long run.

Thanks to the OpenAccess AI Collective and all the contributors!

If you’re interested in more technical content around LLMs, follow me on Medium.

Related articles

Fine-Tune Your Own Llama 2 Model in a Colab Notebook
A practical introduction to LLM fine-tuningtowardsdatascience.com

4-bit Quantization with GPTQ
Quantize your own LLMs using AutoGPTQtowardsdatascience.com

Learn more about machine learning and support my work with one click — become a Medium member here:

Join Medium with my referral link - Maxime Labonne
As a Medium member, a portion of your membership fee goes to writers you read, and you get full access to every story…medium.com

Graph Neural Networks (GNNs) represent one of the most captivating and rapidly evolving architectures within the deep learning landscape. As deep learning models designed to process data structured as graphs, GNNs bring remarkable versatility and powerful learning capabilities.

Among the various types of GNNs, the Graph Convolutional Networks (GCNs) have emerged as the most prevalent and broadly applied model. GCNs are innovative due to their ability to leverage both the features of a node and its locality to make predictions, providing an effective way to handle graph-structured data.

In this article, we will delve into the mechanics of the GCN layer and explain its inner workings. Furthermore, we will explore its practical application for node classification tasks, using PyTorch Geometric as our tool of choice.

PyTorch Geometric is a specialized extension of PyTorch that has been created specifically for the development and implementation of GNNs. It is an advanced, yet user-friendly library that provides a comprehensive suite of tools to facilitate graph-based machine learning. To commence our journey, the PyTorch Geometric installation will be required. If you are using Google Colab, PyTorch should already be in place, so all we need to do is execute a few additional commands.

All the code is available on Google Colab and GitHub.

!pip install torch_geometric

import torch
import numpy as np
import networkx as nx
import matplotlib.pyplot as plt

Now that PyTorch Geometric is installed, let’s explore the dataset we will use in this tutorial.

🌐 I. Graph data

Graphs are an essential structure for representing relationships between objects. You can encounter graph data in a multitude of real-world scenarios, such as social and computer networks, chemical structures of molecules, natural language processing, and image recognition, to name a few.

In this article, we will study the infamous and much-used Zachary’s karate club dataset.

The Zachary’s karate club dataset embodies the relationships formed within a karate club as observed by Wayne W. Zachary during the 1970s. It is a kind of social network, where each node represents a club member, and edges between nodes represent interactions that occurred outside the club environment.

In this particular scenario, the members of the club are split into four distinct groups. Our task is to assign the correct group to each member (node classification), based on the pattern of their interactions.

Let’s import the dataset with PyG’s built-in function and try to understand the Datasets object it uses.

from torch_geometric.datasets import KarateClub

# Import dataset from PyTorch Geometric
dataset = KarateClub()

# Print information
print(dataset)
print('------------')
print(f'Number of graphs: {len(dataset)}')
print(f'Number of features: {dataset.num_features}')
print(f'Number of classes: {dataset.num_classes}')

KarateClub()
------------
Number of graphs: 1
Number of features: 34
Number of classes: 4

This dataset only has 1 graph, where each node has a feature vector of 34 dimensions and is part of one out of four classes (our four groups). Actually, the Datasets object can be seen as a collection of Data (graph) objects.

We can further inspect our unique graph to know more about it.

# Print first element
print(f'Graph: {dataset[0]}')

Graph: Data(x=[34, 34], edge_index=[2, 156], y=[34], train_mask=[34])

The Data object is particularly interesting. Printing it offers a good summary of the graph we're studying:

• x=[34, 34] is the node feature matrix with shape (number of nodes, number of features). In our case, it means that we have 34 nodes (our 34 members), each node being associated to a 34-dim feature vector.

• edge_index=[2, 156] represents the graph connectivity (how the nodes are connected) with shape (2, number of directed edges).

• y=[34] is the node ground-truth labels. In this problem, every node is assigned to one class (group), so we have one value for each node.

• train_mask=[34] is an optional attribute that tells which nodes should be used for training with a list of True or False statements.

Let’s print each of these tensors to understand what they store. Let’s start with the node features.

data = dataset[0]

print(f'x = {data.x.shape}')
print(data.x)

x = torch.Size([34, 34])
tensor([[1., 0., 0.,  ..., 0., 0., 0.],
        [0., 1., 0.,  ..., 0., 0., 0.],
        [0., 0., 1.,  ..., 0., 0., 0.],
        ...,
        [0., 0., 0.,  ..., 1., 0., 0.],
        [0., 0., 0.,  ..., 0., 1., 0.],
        [0., 0., 0.,  ..., 0., 0., 1.]])

Here, the node feature matrix x is an identity matrix: it doesn't contain any relevant information about the nodes. It could contain information like age, skill level, etc. but this is not the case in this dataset. It means we'll have to classify our nodes just by looking at their connections.

Now, let’s print the edge index.

print(f'edge_index = {data.edge_index.shape}')
print(data.edge_index)

edge_index = torch.Size([2, 156])
tensor([[ 0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  1,  1,
          1,  1,  1,  1,  1,  1,  1,  2,  2,  2,  2,  2,  2,  2,  2,  2,  2,  3,
          3,  3,  3,  3,  3,  4,  4,  4,  5,  5,  5,  5,  6,  6,  6,  6,  7,  7,
          7,  7,  8,  8,  8,  8,  8,  9,  9, 10, 10, 10, 11, 12, 12, 13, 13, 13,
         13, 13, 14, 14, 15, 15, 16, 16, 17, 17, 18, 18, 19, 19, 19, 20, 20, 21,
         21, 22, 22, 23, 23, 23, 23, 23, 24, 24, 24, 25, 25, 25, 26, 26, 27, 27,
         27, 27, 28, 28, 28, 29, 29, 29, 29, 30, 30, 30, 30, 31, 31, 31, 31, 31,
         31, 32, 32, 32, 32, 32, 32, 32, 32, 32, 32, 32, 32, 33, 33, 33, 33, 33,
         33, 33, 33, 33, 33, 33, 33, 33, 33, 33, 33, 33],
        [ 1,  2,  3,  4,  5,  6,  7,  8, 10, 11, 12, 13, 17, 19, 21, 31,  0,  2,
          3,  7, 13, 17, 19, 21, 30,  0,  1,  3,  7,  8,  9, 13, 27, 28, 32,  0,
          1,  2,  7, 12, 13,  0,  6, 10,  0,  6, 10, 16,  0,  4,  5, 16,  0,  1,
          2,  3,  0,  2, 30, 32, 33,  2, 33,  0,  4,  5,  0,  0,  3,  0,  1,  2,
          3, 33, 32, 33, 32, 33,  5,  6,  0,  1, 32, 33,  0,  1, 33, 32, 33,  0,
          1, 32, 33, 25, 27, 29, 32, 33, 25, 27, 31, 23, 24, 31, 29, 33,  2, 23,
         24, 33,  2, 31, 33, 23, 26, 32, 33,  1,  8, 32, 33,  0, 24, 25, 28, 32,
         33,  2,  8, 14, 15, 18, 20, 22, 23, 29, 30, 31, 33,  8,  9, 13, 14, 15,
         18, 19, 20, 22, 23, 26, 27, 28, 29, 30, 31, 32]])

In graph theory and network analysis, connectivity between nodes is stored using a variety of data structures. The edge_index is one such data structure, where the graph's connections are stored in two lists (156 directed edges, which equate to 78 bidirectional edges). The reason for these two lists is that one list stores the source nodes, while the second one identifies the destination nodes.

This method is known as a coordinate list (COO) format, which is essentially a means to efficiently store a sparse matrix. Sparse matrices are data structures that efficiently store matrices with a majority of zero elements. In the COO format, only non-zero elements are stored, saving memory and computational resources.

Contrarily, a more intuitive and straightforward way to represent graph connectivity is through an adjacency matrix A. This is a square matrix where each element Aᵢⱼ specifies the presence or absence of an edge from node i to node j in the graph. In other words, a non-zero element Aᵢⱼ implies a connection from node i to node j, and a zero indicates no direct connection.

An adjacency matrix, however, is not as space-efficient as the COO format for sparse matrices or graphs with fewer edges. However, for clarity and easy interpretation, the adjacency matrix remains a popular choice for representing graph connectivity.

The adjacency matrix can be inferred from the edge_index with a utility function to_dense_adj().

from torch_geometric.utils import to_dense_adj

A = to_dense_adj(data.edge_index)[0].numpy().astype(int)
print(f'A = {A.shape}')
print(A)

A = (34, 34)
[[0 1 1 ... 1 0 0]
 [1 0 1 ... 0 0 0]
 [1 1 0 ... 0 1 0]
 ...
 [1 0 0 ... 0 1 1]
 [0 0 1 ... 1 0 1]
 [0 0 0 ... 1 1 0]]

With graph data, it is relatively uncommon for nodes to be densely interconnected. As you can see, our adjacency matrix A is sparse (filled with zeros).

In many real-world graphs, most nodes are connected to only a few other nodes, resulting in a large number of zeros in the adjacency matrix. Storing so many zeros is not efficient at all, which is why the COO format is adopted by PyG.

On the contrary, ground-truth labels are easy to understand.

print(f'y = {data.y.shape}')
print(data.y)

y = torch.Size([34])
tensor([1, 1, 1, 1, 3, 3, 3, 1, 0, 1, 3, 1, 1, 1, 0, 0, 3, 1, 0, 1, 0, 1, 0, 0,
        2, 2, 0, 0, 2, 0, 0, 2, 0, 0])

Our node ground-truth labels stored in y simply encode the group number (0, 1, 2, 3) for each node, which is why we have 34 values.

Finally, let’s print the train mask.

print(f'train_mask = {data.train_mask.shape}')
print(data.train_mask)

train_mask = torch.Size([34])
tensor([ True, False, False, False,  True, False, False, False,  True, False,
        False, False, False, False, False, False, False, False, False, False,
        False, False, False, False,  True, False, False, False, False, False,
        False, False, False, False])

The train mask shows which nodes are supposed to be used for training with True statements. These nodes represent the training set, while the others can be considered as the test set. This division helps in model evaluation by providing unseen data for testing.

But we’re not done yet! The Data object has a lot more to offer. It provides various utility functions that enable the investigation of several properties of the graph. For instance:

• is_directed() tells you if the graph is directed. A directed graph signifies that the adjacency matrix is not symmetric, i.e., the direction of edges matters in the connections between nodes.

• isolated_nodes() checks if some nodes are not connected to the rest of the graph. These nodes are likely to pose challenges in tasks like classification due to their lack of connections.

• has_self_loops() indicates if at least one node is connected to itself. This is distinct from the concept of loops: a loop implies a path that starts and ends at the same node, traversing other nodes in between.

In the context of the Zachary’s karate club dataset, all these properties return False. This implies that the graph is not directed, does not have any isolated nodes, and none of its nodes are connected to themselves.

print(f'Edges are directed: {data.is_directed()}')
print(f'Graph has isolated nodes: {data.has_isolated_nodes()}')
print(f'Graph has loops: {data.has_self_loops()}')

Edges are directed: False
Graph has isolated nodes: False
Graph has loops: False

Finally, we can convert a graph from PyTorch Geometric to the popular graph library NetworkX using to_networkx. This is particularly useful to visualize a small graph with networkx and matplotlib.

Let’s plot our dataset with a different color for each group.

from torch_geometric.utils import to_networkx

G = to_networkx(data, to_undirected=True)
plt.figure(figsize=(12,12))
plt.axis('off')
nx.draw_networkx(G,
                pos=nx.spring_layout(G, seed=0),
                with_labels=True,
                node_size=800,
                node_color=data.y,
                cmap="hsv",
                vmin=-2,
                vmax=3,
                width=0.8,
                edge_color="grey",
                font_size=14
                )
plt.show()

This plot of Zachary’s karate club displays our 34 nodes, 78 (bidirectional) edges, and 4 labels with 4 different colors. Now that we’ve seen the essentials of loading and handling a dataset with PyTorch Geometric, we can introduce the Graph Convolutional Network architecture.

✉️ II. Graph Convolutional Network

This section aims to introduce and build the graph convolutional layer from the ground up.

In traditional neural networks, linear layers apply a linear transformation to the incoming data. This transformation converts input features x into hidden vectors h through the use of a weight matrix 𝐖. Ignoring biases for the time being, this can be expressed as:

With graph data, an additional layer of complexity is added through the connections between nodes. These connections matter because, typically, in networks, it’s assumed that similar nodes are more likely to be linked to each other than dissimilar ones, a phenomenon known as network homophily.

We can enrich our node representation by merging its features with those of its neighbors. This operation is called convolution, or neighborhood aggregation. Let’s represent the neighborhood of node i including itself as Ñ.

Unlike filters in Convolutional Neural Networks (CNNs), our weight matrix 𝐖 is unique and shared among every node. But there is another issue: nodes do not have a fixed number of neighbors like pixels do.

How do we address cases where one node has only one neighbor, and another has 500? If we simply sum the feature vectors, the resulting embedding h would be much larger for the node with 500 neighbors. To ensure a similar range of values for all nodes and comparability between them, we can normalize the result based on the degree of nodes, where degree refers to the number of connections a node has.

We’re almost there! Introduced by Kipf et al. (2016), the graph convolutional layer has one final improvement.

The authors observed that features from nodes with numerous neighbors propagate much more easily than those from more isolated nodes. To offset this effect, they suggested assigning bigger weights to features from nodes with fewer neighbors, thus balancing the influence across all nodes. This operation is written as:

Note that when i and j have the same number of neighbors, it is equivalent to our own layer. Now, let’s see how to implement it in Python with PyTorch Geometric.

🧠 III. Implementing a GCN

PyTorch Geometric provides the GCNConv function, which directly implements the graph convolutional layer.

In this example, we’ll create a basic Graph Convolutional Network with a single GCN layer, a ReLU activation function, and a linear output layer. This output layer will yield four values corresponding to our four categories, with the highest value determining the class of each node.

In the following code block, we define the GCN layer with a 3-dimensional hidden layer.

from torch.nn import Linear
from torch_geometric.nn import GCNConv

class GCN(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.gcn = GCNConv(dataset.num_features, 3)
        self.out = Linear(3, dataset.num_classes)

    def forward(self, x, edge_index):
        h = self.gcn(x, edge_index).relu()
        z = self.out(h)
        return h, z

model = GCN()
print(model)

GCN(
  (gcn): GCNConv(34, 3)
  (out): Linear(in_features=3, out_features=4, bias=True)
)

If we added a second GCN layer, our model would not only aggregate feature vectors from the neighbors of each node, but also from the neighbors of these neighbors.

We can stack several graph layers to aggregate more and more distant values, but there’s a catch: if we add too many layers, the aggregation becomes so intense that all the embeddings end up looking the same. This phenomenon is called over-smoothing and can be a real problem when you have too many layers.

Now that we’ve defined our GNN, let’s write a simple training loop with PyTorch. I chose a regular cross-entropy loss since it’s a multi-class classification task, with Adam as optimizer. In this article, we won’t implement a train/test split to keep things simple and focus on how GNNs learn instead.

The training loop is standard: we try to predict the correct labels, and we compare the GCN’s results to the values stored in data.y. The error is calculated by the cross-entropy loss and backpropagated with Adam to fine-tune our GNN's weights and biases. Finally, we print metrics every 10 epochs.

criterion = torch.nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=0.02)

# Calculate accuracy
def accuracy(pred_y, y):
    return (pred_y == y).sum() / len(y)

# Data for animations
embeddings = []
losses = []
accuracies = []
outputs = []

# Training loop
for epoch in range(201):
    # Clear gradients
    optimizer.zero_grad()

    # Forward pass
    h, z = model(data.x, data.edge_index)

    # Calculate loss function
    loss = criterion(z, data.y)

    # Calculate accuracy
    acc = accuracy(z.argmax(dim=1), data.y)

    # Compute gradients
    loss.backward()

    # Tune parameters
    optimizer.step()

    # Store data for animations
    embeddings.append(h)
    losses.append(loss)
    accuracies.append(acc)
    outputs.append(z.argmax(dim=1))

    # Print metrics every 10 epochs
    if epoch % 10 == 0:
        print(f'Epoch {epoch:>3} | Loss: {loss:.2f} | Acc: {acc*100:.2f}%')

Epoch   0 | Loss: 1.40 | Acc: 41.18%
Epoch  10 | Loss: 1.21 | Acc: 47.06%
Epoch  20 | Loss: 1.02 | Acc: 67.65%
Epoch  30 | Loss: 0.80 | Acc: 73.53%
Epoch  40 | Loss: 0.59 | Acc: 73.53%
Epoch  50 | Loss: 0.39 | Acc: 94.12%
Epoch  60 | Loss: 0.23 | Acc: 97.06%
Epoch  70 | Loss: 0.13 | Acc: 100.00%
Epoch  80 | Loss: 0.07 | Acc: 100.00%
Epoch  90 | Loss: 0.05 | Acc: 100.00%
Epoch 100 | Loss: 0.03 | Acc: 100.00%
Epoch 110 | Loss: 0.02 | Acc: 100.00%
Epoch 120 | Loss: 0.02 | Acc: 100.00%
Epoch 130 | Loss: 0.02 | Acc: 100.00%
Epoch 140 | Loss: 0.01 | Acc: 100.00%
Epoch 150 | Loss: 0.01 | Acc: 100.00%
Epoch 160 | Loss: 0.01 | Acc: 100.00%
Epoch 170 | Loss: 0.01 | Acc: 100.00%
Epoch 180 | Loss: 0.01 | Acc: 100.00%
Epoch 190 | Loss: 0.01 | Acc: 100.00%
Epoch 200 | Loss: 0.01 | Acc: 100.00%

Great! Without much surprise, we reach 100% accuracy on the training set (full dataset). It means that our model learned to correctly assign every member of the karate club to its correct group.

We can produce a neat visualization by animating the graph and see the evolution of the GNN’s predictions during the training process.

%%capture
from IPython.display import HTML
from matplotlib import animation
plt.rcParams["animation.bitrate"] = 3000

def animate(i):
    G = to_networkx(data, to_undirected=True)
    nx.draw_networkx(G,
                    pos=nx.spring_layout(G, seed=0),
                    with_labels=True,
                    node_size=800,
                    node_color=outputs[i],
                    cmap="hsv",
                    vmin=-2,
                    vmax=3,
                    width=0.8,
                    edge_color="grey",
                    font_size=14
                    )
    plt.title(f'Epoch {i} | Loss: {losses[i]:.2f} | Acc: {accuracies[i]*100:.2f}%',
              fontsize=18, pad=20)

fig = plt.figure(figsize=(12, 12))
plt.axis('off')

anim = animation.FuncAnimation(fig, animate, \
            np.arange(0, 200, 10), interval=500, repeat=True)
html = HTML(anim.to_html5_video())
display(html)

The first predictions are random, but the GCN perfectly labels every node after a while. Indeed, the final graph is the same as the one we plotted at the end of the first section. But what does the GCN really learn?

By aggregating features from neighboring nodes, the GNN learns a vector representation (or embedding) of every node in the network. In our model, the final layer just learns how to use these representations to produce the best classifications. However, embeddings are the real products of GNNs.

Let’s print the embeddings learned by our model.

# Print embeddings
print(f'Final embeddings = {h.shape}')
print(h)

Final embeddings = torch.Size([34, 3])
tensor([[1.9099e+00, 2.3584e+00, 7.4027e-01],
        [2.6203e+00, 2.7997e+00, 0.0000e+00],
        [2.2567e+00, 2.2962e+00, 6.4663e-01],
        [2.0802e+00, 2.8785e+00, 0.0000e+00],
        [0.0000e+00, 0.0000e+00, 2.9694e+00],
        [0.0000e+00, 0.0000e+00, 3.3817e+00],
        [0.0000e+00, 1.5008e-04, 3.4246e+00],
        [1.7593e+00, 2.4292e+00, 2.4551e-01],
        [1.9757e+00, 6.1032e-01, 1.8986e+00],
        [1.7770e+00, 1.9950e+00, 6.7018e-01],
        [0.0000e+00, 1.1683e-04, 2.9738e+00],
        [1.8988e+00, 2.0512e+00, 2.6225e-01],
        [1.7081e+00, 2.3618e+00, 1.9609e-01],
        [1.8303e+00, 2.1591e+00, 3.5906e-01],
        [2.0755e+00, 2.7468e-01, 1.9804e+00],
        [1.9676e+00, 3.7185e-01, 2.0011e+00],
        [0.0000e+00, 0.0000e+00, 3.4787e+00],
        [1.6945e+00, 2.0350e+00, 1.9789e-01],
        [1.9808e+00, 3.2633e-01, 2.1349e+00],
        [1.7846e+00, 1.9585e+00, 4.8021e-01],
        [2.0420e+00, 2.7512e-01, 1.9810e+00],
        [1.7665e+00, 2.1357e+00, 4.0325e-01],
        [1.9870e+00, 3.3886e-01, 2.0421e+00],
        [2.0614e+00, 5.1042e-01, 2.4872e+00],
...
        [2.1778e+00, 4.4730e-01, 2.0077e+00],
        [3.8906e-02, 2.3443e+00, 1.9195e+00],
        [3.0748e+00, 0.0000e+00, 3.0789e+00],
        [3.4316e+00, 1.9716e-01, 2.5231e+00]], grad_fn=<ReluBackward0>)

As you can see, embeddings do not need to have the same dimensions as feature vectors. Here, I chose to reduce the number of dimensions from 34 (dataset.num_features) to three to get a nice visualization in 3D.

Let’s plot these embeddings before any training happens, at epoch 0.

# Get first embedding at epoch = 0
embed = h.detach().cpu().numpy()

fig = plt.figure(figsize=(12, 12))
ax = fig.add_subplot(projection='3d')
ax.patch.set_alpha(0)
plt.tick_params(left=False,
                bottom=False,
                labelleft=False,
                labelbottom=False)
ax.scatter(embed[:, 0], embed[:, 1], embed[:, 2],
           s=200, c=data.y, cmap="hsv", vmin=-2, vmax=3)

plt.show()

We see every node from Zachary’s karate club with their true labels (and not the model’s predictions). For now, they’re all over the place since the GNN is not trained yet. But if we plot these embeddings at each step of the training loop, we’d be able to visualize what the GNN truly learns.

Let’s see how they evolve over time, as the GCN gets better and better at classifying nodes.

%%capture

def animate(i):
    embed = embeddings[i].detach().cpu().numpy()
    ax.clear()
    ax.scatter(embed[:, 0], embed[:, 1], embed[:, 2],
           s=200, c=data.y, cmap="hsv", vmin=-2, vmax=3)
    plt.title(f'Epoch {i} | Loss: {losses[i]:.2f} | Acc: {accuracies[i]*100:.2f}%',
              fontsize=18, pad=40)

fig = plt.figure(figsize=(12, 12))
plt.axis('off')
ax = fig.add_subplot(projection='3d')
plt.tick_params(left=False,
                bottom=False,
                labelleft=False,
                labelbottom=False)

anim = animation.FuncAnimation(fig, animate, \
              np.arange(0, 200, 10), interval=800, repeat=True)
html = HTML(anim.to_html5_video())
display(html)

Our Graph Convolutional Network (GCN) has effectively learned embeddings that group similar nodes into distinct clusters. This enables the final linear layer to distinguish them into separate classes with ease.

Embeddings are not unique to GNNs: they can be found everywhere in deep learning. They don’t have to be 3D either: actually, they rarely are. For instance, language models like BERT produce embeddings with 768 or even 1024 dimensions.

Additional dimensions store more information about nodes, text, images, etc. but they also create bigger models that are more difficult to train. This is why keeping low-dimensional embeddings as long as possible is advantageous.

Conclusion

Graph Convolutional Networks are an incredibly versatile architecture that can be applied in many contexts. In this article, we familiarized ourselves with the PyTorch Geometric library and objects like Datasets and Data. Then, we successfully reconstructed a graph convolutional layer from the ground up. Next, we put theory into practice by implementing a GCN, which gave us an understanding of practical aspects and how individual components interact. Finally, we visualized the training process and obtained a clear perspective of what it involves for such a network.

Zachary’s karate club is a simplistic dataset, but it is good enough to understand the most important concepts in graph data and GNNs. Although we only talked about node classification in this article, there are other tasks GNNs can accomplish: link prediction (e.g., to recommend a friend), graph classification (e.g., to label molecules), graph generation (e.g., to create new molecules), and so on.

Beyond GCN, numerous GNN layers and architectures have been proposed by researchers. In the next article, we’ll introduce the Graph Attention Network (GAT) architecture, which dynamically computes the GCN’s normalization factor and the importance of each connection with an attention mechanism.

If you want to know more about graph neural networks, dive deeper into the world of GNNs with my book, Hands-On Graph Neural Networks.

Next article

Chapter 2: Graph Attention Networks: Self-Attention Explained
A guide to GNNs with self-attention using PyTorch Geometrictowardsdatascience.com

Learn more about machine learning and support my work with one click — become a Medium member here:

Join Medium with my referral link — Maxime Labonne
As a Medium member, a portion of your membership fee goes to writers you read, and you get full access to every story…medium.com

If you’re already a member, you can follow me on Medium.

Recent advancements in weight quantization allow us to run massive large language models on consumer hardware, like a LLaMA-30B model on an RTX 3090 GPU. This is possible thanks to novel 4-bit quantization techniques with minimal performance degradation, like GPTQ, GGML, and NF4.

In the previous article, we introduced naïve 8-bit quantization techniques and the excellent LLM.int8(). In this article, we will explore the popular GPTQ algorithm to understand how it works and implement it using the AutoGPTQ library.

You can find the code on Google Colab and GitHub.

🧠 Optimal Brain Quantization

Let’s start by introducing the problem we’re trying to solve. For every layer ℓ in the network, we want to find a quantized version Ŵₗ of the original weights Wₗ. This is called the layer-wise compression problem. More specifically, to minimize performance degradation, we want the outputs (ŴᵨXᵨ) of these new weights to be as close as possible to the original ones (WᵨXᵨ). In other words, we want to find:

Different approaches have been proposed to solve this problem, but we’re interested in the Optimal Brain Quantizer (OBQ) framework here.

This method is inspired by a pruning technique to carefully remove weights from a fully trained dense neural network (Optimal Brain Surgeon). It uses an approximation technique and provides explicit formulas for the best single weight w𐞥 to remove and optimal update δꟳ to adjust the set of remaining non-quantized weights F to make up for the removal:

where quant(w) is the weight rounding given by the quantization and Hꟳ is the Hessian.

Using OBQ, we can quantize the easiest weight first and then adjust all remaining non-quantized weights to compensate for this precision loss. Then we pick the next weight to quantize, and so on.

A potential issue with this approach is when there are outlier weights, which can result in high quantization error. Usually, these outliers would be quantized last, when there are few non-quantized weights left that could be adjusted to compensate for the large error. This effect can worsen when some weights are pushed further outside the grid by intermediate updates. A simple heuristic is applied to prevent this: outliers are quantized as soon as they appear.

This process could be computationally heavy, especially for LLMs. To deal with this, the OBQ method uses a trick that avoids redoing the entire computation each time a weight is simplified. After quantizing a weight, it adjusts the matrix used in calculations (the Hessian) by removing the row and column associated with that weight (using Gaussian elimination):

The method also employs vectorization to process multiple rows of the weight matrix at once. Despite its efficiency, the OBQ’s computation time increases significantly as the size of the weight matrix increases. This cubic growth makes it difficult to use OBQ on very large models with billions of parameters.

🧮 The GPTQ Algorithm

Introduced by Frantar et al. (2023), the GPTQ algorithm takes inspiration from the OBQ method, but with significant improvements to scale it for (very) large language models.

Step 1: Arbitrary Order Insight

The OBQ method selects weights (parameters in a model) for quantization in a certain order, determined by which will add the least additional error. However, GPTQ observes that for large models, quantizing weights in any fixed order can perform just as well. This is because even though some weights might introduce more error individually, they are quantized later in the process when there are few other weights left that could increase the error. So the order doesn’t matter as much as we thought.

Based on this insight, GPTQ aims to quantize all weights in the same order for all rows of a matrix. This makes the process faster because certain computations have to be done only once for each column, rather than once for each weight.

Step 2: Lazy Batch-Updates

This scheme won’t be fast because it requires updating a huge matrix with very few computations for each entry. This type of operation can’t utilize the full compute capabilities of GPUs and will be slowed down by memory limitations (memory throughput bottleneck).

To resolve this, GPTQ introduces “lazy batch” updates. It turns out that the final rounding decisions for a given column are only affected by updates performed on that column, not on later columns. Therefore, GPTQ can apply the algorithm to a batch of columns at a time (like 128 columns), updating only those columns and a corresponding block of the matrix. After a block is fully processed, the algorithm performs global updates on the entire matrix.

Step 3: Cholesky Reformulation

However, there’s one more issue to address. When the algorithm scales up to very large models, numerical inaccuracies can become a problem. Specifically, repeated applications of a certain operation can accumulate numerical errors.

To tackle this, GPTQ uses a Cholesky decomposition, a numerically stable method for solving certain mathematical problems. It involves precomputing some required information from the matrix using the Cholesky method. This approach, combined with a slight “dampening” (adding a small constant to diagonal elements of the matrix), helps the algorithm to avoid numerical issues.

The full algorithm can be summarized in a few steps:

1. The GPTQ algorithm begins with a Cholesky decomposition of the Hessian inverse (a matrix that helps decide how to adjust the weights)

2. It then runs in loops, handling batches of columns at a time.

3. For each column in a batch, it quantizes the weights, calculates the error, and updates the weights in the block accordingly.

4. After processing the batch, it updates all remaining weights based on the block’s errors.

The GPTQ algorithm was tested on various language generation tasks. It was compared with other quantization methods, like rounding all weights to the nearest quantized value (RTN). GPTQ was used with the BLOOM (176B parameters) and OPT (175B parameters) model families, and models were quantized using a single NVIDIA A100 GPU.

💻 Quantize an LLM with AutoGPTQ

GPTQ has been very popular to create models in 4-bit precision that can efficiently run on GPUs. You can find many examples on the Hugging Face Hub, especially from TheBloke. If you’re looking for an approach that is more CPU-friendly, GGML is currently your best option. Finally, the transformers library with bitsandbytes allows you to quantize a model when it's loaded using the load_in_4bit=true argument, which requires downloading full models and storing them in your RAM.

Let’s implement the GPTQ algorithm using the AutoGPTQ library and quantize a GPT-2 model. This requires a GPU, but a free T4 on Google Colab will do. We start by loading the libraries and defining the model we want to quantize (in this case, GPT-2).

!BUILD_CUDA_EXT=0 pip install -q auto-gptq transformers

import random

from auto_gptq import AutoGPTQForCausalLM, BaseQuantizeConfig
from datasets import load_dataset
import torch
from transformers import AutoTokenizer


# Define base model and output directory
model_id = "gpt2"
out_dir = model_id + "-GPTQ"

We now want to load the model and the tokenizer. The tokenizer is loaded using the classic AutoTokenizer class from the transformers library. On the other hand, we need to pass a specific configuration (BaseQuantizeConfig) to load the model.

In this configuration, we can specify the number of bits to quantize (here, bits=4) and the group size (size of the lazy batch). Note that this group size is optional: we could also use one set of parameters for the entire weight matrix. In practice, these groups generally improve the quality of the quantization at a very low cost (especially with group_size=1024). The damp_percent value is here to help the Cholesky reformulation and should not be changed.

Finally, the desc_act (also called act order) is a tricky parameter. It allows you to process rows based on decreasing activation, meaning the most important or impactful rows (determined by sampled inputs and outputs) are processed first. This method aims to place most of the quantization error (inevitably introduced during quantization) on less significant weights. This approach improves the overall accuracy of the quantization process by ensuring the most significant weights are processed with greater precision. However, when used alongside group size, desc_act can lead to performance slowdowns due to the need to frequently reload quantization parameters. For this reason, we won't use it here (it will probably be fixed in the future, however).

# Load quantize config, model and tokenizer
quantize_config = BaseQuantizeConfig(
    bits=4,
    group_size=128,
    damp_percent=0.01,
    desc_act=False,
)
model = AutoGPTQForCausalLM.from_pretrained(model_id, quantize_config)
tokenizer = AutoTokenizer.from_pretrained(model_id)

The quantization process relies heavily on samples to evaluate and enhance the quality of the quantization. They provide a means of comparison between the outputs produced by the origina and the newly quantized model. The larger the number of samples provided, the greater the potential for more accurate and effective comparisons, leading to improved quantization quality.

In the context of this article, we utilize the C4 (Colossal Clean Crawled Corpus) dataset to generate our samples. The C4 dataset is a large-scale, multilingual collection of web text gathered from the Common Crawl project. This expansive dataset has been cleaned and prepared specifically for training large-scale language models, making it a great resource for tasks such as this. The WikiText dataset is another popular option.

In the following code block, we load 1024 samples from the C4 dataset, tokenize them, and format them.

# Load data and tokenize examples
n_samples = 1024
data = load_dataset("allenai/c4", data_files="en/c4-train.00001-of-01024.json.gz", split=f"train[:{n_samples*5}]")
tokenized_data = tokenizer("\n\n".join(data['text']), return_tensors='pt')

# Format tokenized examples
examples_ids = []
for _ in range(n_samples):
    i = random.randint(0, tokenized_data.input_ids.shape[1] - tokenizer.model_max_length - 1)
    j = i + tokenizer.model_max_length
    input_ids = tokenized_data.input_ids[:, i:j]
    attention_mask = torch.ones_like(input_ids)
    examples_ids.append({'input_ids': input_ids, 'attention_mask': attention_mask})

Now that dataset is ready, we can start the quantization process with a batch size of 1. Optionally, we also use OpenAI Triton, a CUDA alternative, to communicate with the GPU. Once this is done, we save the tokenizer and the model in a safetensors format.

# Quantize with GPTQ
model.quantize(
    examples_ids,
    batch_size=1,
    use_triton=True,
)

# Save model and tokenizer
model.save_quantized(out_dir, use_safetensors=True)
tokenizer.save_pretrained(out_dir)

As per usual, the model and tokenizer can then be loaded from the output directory using the AutoGPTQForCausalLM and AutoTokenizer classes.

device = "cuda:0" if torch.cuda.is_available() else "cpu"

# Reload model and tokenizer
model = AutoGPTQForCausalLM.from_quantized(
    out_dir,
    device=device,
    use_triton=True,
    use_safetensors=True,
)
tokenizer = AutoTokenizer.from_pretrained(out_dir)

Let’s check that the model is working correctly. The AutoGPTQ model (mostly) works as a normal transformers model, which makes it compatible with inference pipelines, as shown in the following example:

from transformers import pipeline

generator = pipeline('text-generation', model=model, tokenizer=tokenizer)
result = generator("I have a dream", do_sample=True, max_length=50)[0]['generated_text']
print(result)

I have a dream," she told CNN last week. "I have this dream of helping my mother find her own. But, to tell that for the first time, now that I'm seeing my mother now, just knowing how wonderful it is that

We managed to get a convincing completion from our quantized GPT-2 model. A more in-depth evaluation would require measuring the perplexity of the quantized model versus the original one. However, we will leave it out of the scope of this article.

Conclusion

In this article, we introduced the GPTQ algorithm, a state-of-the-art quantization technique to run LLMs on consumer-grade hardware. We showed how it addresses the layer-wise compression problem, based on an improved OBS technique with arbitrary order insight, lazy batch updates, and Cholesky reformulation. This novel approach significantly reduces memory and computation requirements, making LLMs accessible to a broader audience.

In addition, we quantized our own LLM model on a free T4 GPU and ran it to generate text. You can push your own version of a GPTQ 4-bit quantized model on the Hugging Face Hub. As mentioned in the introduction, GPTQ is not the only 4-bit quantization algorithm: GGML and NF4 are excellent alternatives with slightly different scopes. I encourage you to learn more about them and give them a shot!

If you’re interested in more technical content around LLMs, follow me on Twitter @maximelabonne.

References

• B. Hassibi, D. G. Stork and G. J. Wolff, “Optimal Brain Surgeon and general network pruning,” IEEE International Conference on Neural Networks, San Francisco, CA, USA, 1993, pp. 293–299 vol.1, doi: 10.1109/ICNN.1993.298572.

• Elias Frantar, Sidak Pal Singh, & Dan Alistarh. (2023). Optimal Brain Compression: A Framework for Accurate Post-Training Quantization and Pruning.

• Elias Frantar, Saleh Ashkboos, Torsten Hoefler, & Dan Alistarh. (2023). GPTQ: Accurate Post-Training Quantization for Generative Pre-trained Transformers.

• Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, & Peter J. Liu. (2020). Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer.

Related articles

Introduction to Weight Quantization
Reducing the size of Large Language Models with 8-bit quantizationtowardsdatascience.com

Fine-Tune Your Own Llama 2 Model in a Colab Notebook
A practical introduction to LLM fine-tuningtowardsdatascience.com

Learn more about machine learning and support my work with one click — become a Medium member here:

Join Medium with my referral link — Maxime Labonne
As a Medium member, a portion of your membership fee goes to writers you read, and you get full access to every story…medium.com

If you’re already a member, you can follow me on Medium.

With the release of LLaMA v1, we saw a Cambrian explosion of fine-tuned models, including Alpaca, Vicuna, and WizardLM, among others. This trend encouraged different businesses to launch their own base models with licenses suitable for commercial use, such as OpenLLaMA, Falcon, XGen, etc. The release of Llama 2 now combines the best elements from both sides: it offers a highly efficient base model along with a more permissive license.

During the first half of 2023, the software landscape was significantly shaped by the widespread use of APIs (like OpenAI API) to create infrastructures based on Large Language Models (LLMs). Libraries such as LangChain and LlamaIndex played a critical role in this trend. Moving into the latter half of the year, the process of fine-tuning (or instruction tuning) these models is set to become a standard procedure in the LLMOps workflow. This trend is driven by various factors: the potential for cost savings, the ability to process confidential data, and even the potential to develop models that exceed the performance of prominent models like ChatGPT and GPT-4 in certain specific tasks.

In this article, we will see why instruction tuning works and how to implement it in a Google Colab notebook to create your own Llama 2 model. As usual, the code is available on Colab and GitHub.

🔧 Background on fine-tuning LLMs

LLMs are pretrained on an extensive corpus of text. In the case of Llama 2, we know very little about the composition of the training set, besides its length of 2 trillion tokens. In comparison, BERT (2018) was “only” trained on the BookCorpus (800M words) and English Wikipedia (2,500M words). From experience, this is a very costly and long process with a lot of hardware issues. If you want to know more about it, I recommend reading Meta’s logbook about the pretraining of the OPT-175B model.

When the pretraining is complete, auto-regressive models like Llama 2 can predict the next token in a sequence. However, this does not make them particularly useful assistants since they don’t reply to instructions. This is why we employ instruction tuning to align their answers with what humans expect. There are two main fine-tuning techniques:

• Supervised Fine-Tuning (SFT): Models are trained on a dataset of instructions and responses. It adjusts the weights in the LLM to minimize the difference between the generated answers and ground-truth responses, acting as labels.

• Reinforcement Learning from Human Feedback (RLHF): Models learn by interacting with their environment and receiving feedback. They are trained to maximize a reward signal (using PPO), which is often derived from human evaluations of model outputs.

In general, RLHF is shown to capture more complex and nuanced human preferences, but is also more challenging to implement effectively. Indeed, it requires careful design of the reward system and can be sensitive to the quality and consistency of human feedback. A possible alternative in the future is the Direct Preference Optimization (DPO) algorithm, which directly runs preference learning on the SFT model.

In our case, we will perform SFT, but this raises a question: why does fine-tuning work in the first place? As highlighted in the Orca paper, our understanding is that fine-tuning leverages knowledge learned during the pretraining process. In other words, fine-tuning will be of little help if the model has never seen the kind of data you’re interested in. However, if that’s the case, SFT can be extremely performant.

For example, the LIMA paper showed how you could outperform GPT-3 (DaVinci003) by fine-tuning a LLaMA (v1) model with 65 billion parameters on only 1,000 high-quality samples. The quality of the instruction dataset is essential to reach this level of performance, which is why a lot of work is focused on this issue (like evol-instruct, Orca, or phi-1). Note that the size of the LLM (65b, not 13b or 7b) is also fundamental to leverage pre-existing knowledge efficiently.

Another important point related to the data quality is the prompt template. Prompts are comprised of similar elements: system prompt (optional) to guide the model, user prompt (required) to give the instruction, additional inputs (optional) to take into consideration, and the model’s answer (required). In the case of Llama 2, the authors used the following template:

<s>[INST] <<SYS>>
System prompt
<</SYS>>

User prompt [/INST] Model answer </s>

There are other templates, like the ones from Alpaca and Vicuna, and their impact is not very clear. In this example, we will reformat our instruction dataset to follow Llama 2’s template. For the purpose of this tutorial, I’ve already done it using the excellent timdettmers/openassistant-guanaco dataset. You can find it on Hugging Face under the name mlabonne/guanaco-llama2-1k.

🦙 How to fine-tune Llama 2

In this section, we will fine-tune a Llama 2 model with 7 billion parameters on a T4 GPU with high RAM using Google Colab (2.21 credits/hour). Note that a T4 only has 16 GB of VRAM, which is barely enough to store Llama 2–7b’s weights (7b × 2 bytes = 14 GB in FP16). In addition, we need to consider the overhead due to optimizer states, gradients, and forward activations (see this excellent article for more information). This means that a full fine-tuning is not possible here: we need parameter-efficient fine-tuning (PEFT) techniques like LoRA or QLoRA.

To drastically reduce the VRAM usage, we must fine-tune the model in 4-bit precision, which is why we’ll use QLoRA here. The good thing is that we can leverage the Hugging Face ecosystem with the transformers, accelerate, peft, trl, and bitsandbytes libraries. We'll do this in the following code based on Younes Belkada's GitHub Gist. First, we install and load these libraries.

!pip install -q accelerate==0.21.0 peft==0.4.0 bitsandbytes==0.40.2 transformers==4.31.0 trl==0.4.7

import os
import torch
from datasets import load_dataset
from transformers import (
    AutoModelForCausalLM,
    AutoTokenizer,
    BitsAndBytesConfig,
    HfArgumentParser,
    TrainingArguments,
    pipeline,
    logging,
)
from peft import LoraConfig, PeftModel
from trl import SFTTrainer

Let’s talk a bit about the parameters we can tune here. First, we want to load a llama-2-7b-chat-hf model and train it on the mlabonne/guanaco-llama2-1k (1,000 samples), which will produce our fine-tuned model llama-2-7b-miniguanaco. Feel free to change the dataset: there are many options on the Hugging Face Hub.

QLoRA will use a rank of 64 with a scaling parameter of 16 (see this article for more information about LoRA parameters). We’ll load the Llama 2 model directly in 4-bit precision using the NF4 type and train it for one epoch. To get more information about the other parameters, check the TrainingArguments, PeftModel, and SFTTrainer documentation.

# The model that you want to train from the Hugging Face hub
model_name = "daryl149/llama-2-7b-chat-hf"

# The instruction dataset to use
dataset_name = "mlabonne/guanaco-llama2-1k"

# Fine-tuned model name
new_model = "llama-2-7b-miniguanaco"

################################################################################
# QLoRA parameters
################################################################################

# LoRA attention dimension
lora_r = 64

# Alpha parameter for LoRA scaling
lora_alpha = 16

# Dropout probability for LoRA layers
lora_dropout = 0.1

################################################################################
# bitsandbytes parameters
################################################################################

# Activate 4-bit precision base model loading
use_4bit = True

# Compute dtype for 4-bit base models
bnb_4bit_compute_dtype = "float16"

# Quantization type (fp4 or nf4)
bnb_4bit_quant_type = "nf4"

# Activate nested quantization for 4-bit base models (double quantization)
use_nested_quant = False

################################################################################
# TrainingArguments parameters
################################################################################

# Output directory where the model predictions and checkpoints will be stored
output_dir = "./results"

# Number of training epochs
num_train_epochs = 1

# Enable fp16/bf16 training (set bf16 to True with an A100)
fp16 = False
bf16 = False

# Batch size per GPU for training
per_device_train_batch_size = 4

# Batch size per GPU for evaluation
per_device_eval_batch_size = 4

# Number of update steps to accumulate the gradients for
gradient_accumulation_steps = 2

# Enable gradient checkpointing
gradient_checkpointing = True

# Maximum gradient normal (gradient clipping)
max_grad_norm = 0.3

# Initial learning rate (AdamW optimizer)
learning_rate = 2e-4

# Weight decay to apply to all layers except bias/LayerNorm weights
weight_decay = 0.001

# Optimizer to use
optim = "paged_adamw_32bit"

# Learning rate schedule (constant a bit better than cosine)
lr_scheduler_type = "constant"

# Number of training steps (overrides num_train_epochs)
max_steps = -1

# Ratio of steps for a linear warmup (from 0 to learning rate) 
warmup_ratio = 0.03

# Group sequences into batches with same length
# Saves memory and speeds up training considerably
group_by_length = True

# Save checkpoint every X updates steps
save_steps = 10

# Log every X updates steps
logging_steps = 1

################################################################################
# SFT parameters
################################################################################

# Maximum sequence length to use
max_seq_length = None

# Pack multiple short examples in the same input sequence to increase efficiency
packing = False

# Load the entire model on the GPU 0
device_map = {"": 0}

We can now load everything and start the fine-tuning process. We’re relying on multiple wrappers, so bear with me.

• First of all, we want to load the dataset we defined. If you changed it, you can preprocess it here and adapt it to the desired prompt template.

• Then, we’re configuring bitsandbytes for 4-bit quantization.

• Next, we're loading the Llama 2 model in 4-bit precision on a GPU with the corresponding tokenizer.

• Finally, we're loading configurations for QLoRA, regular training parameters, and passing everything to the SFTTrainer. The training can finally start!

# Load dataset (you can process it here)
dataset = load_dataset(dataset_name, split="train")

# Load tokenizer and model with QLoRA configuration
compute_dtype = getattr(torch, bnb_4bit_compute_dtype)

bnb_config = BitsAndBytesConfig(
    load_in_4bit=use_4bit,
    bnb_4bit_quant_type=bnb_4bit_quant_type,
    bnb_4bit_compute_dtype=compute_dtype,
    bnb_4bit_use_double_quant=use_nested_quant,
)

# Check GPU compatibility with bfloat16
if compute_dtype == torch.float16 and use_4bit:
    major, _ = torch.cuda.get_device_capability()
    if major >= 8:
        print("=" * 80)
        print("Your GPU supports bfloat16: accelerate training with bf16=True")
        print("=" * 80)

# Load base model
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    quantization_config=bnb_config,
    device_map=device_map
)
model.config.use_cache = False
model.config.pretraining_tp = 1

# Load LLaMA tokenizer
tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
tokenizer.pad_token = tokenizer.eos_token
tokenizer.padding_side = "right" # Fix weird overflow issue with fp16 training

# Load LoRA configuration
peft_config = LoraConfig(
    lora_alpha=lora_alpha,
    lora_dropout=lora_dropout,
    r=lora_r,
    bias="none",
    task_type="CAUSAL_LM",
)

# Set training parameters
training_arguments = TrainingArguments(
    output_dir=output_dir,
    num_train_epochs=num_train_epochs,
    per_device_train_batch_size=per_device_train_batch_size,
    gradient_accumulation_steps=gradient_accumulation_steps,
    optim=optim,
    save_steps=save_steps,
    logging_steps=logging_steps,
    learning_rate=learning_rate,
    weight_decay=weight_decay,
    fp16=fp16,
    bf16=bf16,
    max_grad_norm=max_grad_norm,
    max_steps=max_steps,
    warmup_ratio=warmup_ratio,
    group_by_length=group_by_length,
    lr_scheduler_type=lr_scheduler_type,
    report_to="tensorboard"
)

# Set supervised fine-tuning parameters
trainer = SFTTrainer(
    model=model,
    train_dataset=dataset,
    peft_config=peft_config,
    dataset_text_field="text",
    max_seq_length=max_seq_length,
    tokenizer=tokenizer,
    args=training_arguments,
    packing=packing,
)

# Train model
trainer.train()

# Save trained model
trainer.model.save_pretrained(output_dir)

The training can be very long, depending on the size of your dataset. Here, it took less than an hour on a T4 GPU. We can check the plots on tensorboard, as follows:

%load_ext tensorboard
%tensorboard --logdir results/runs

Let’s make sure that the model is behaving correctly. It would require a more exhaustive evaluation, but we can use the text generation pipeline to ask questions like “What is a large language model?” Note that I’m formatting the input to match Llama 2’s prompt template.

# Ignore warnings
logging.set_verbosity(logging.CRITICAL)

# Run text generation pipeline with our next model
prompt = "What is a large language model?"
pipe = pipeline(task="text-generation", model=model, tokenizer=tokenizer, max_length=200)
result = pipe(f"<s>[INST] {prompt} [/INST]")
print(result[0]['generated_text'])

The model outputs the following response:

A large language model is a type of artificial intelligence that is trained on a large dataset of text to generate human-like language. It is typically trained on a dataset of text that is much larger than the dataset used for smaller language models. The large dataset allows the model to learn more complex patterns in language, which can result in more accurate and natural-sounding language generation.

Large language models are often used for tasks such as text summarization, language translation, and chatbots. They are also used for more complex tasks such as writing articles, generating creative content, and even composing music.

Large language models are trained using a technique called deep learning, which involves using many layers of artificial neural networks to learn complex patterns in the data. The model is trained on a large dataset of text, and the neural networks are trained to predict the next word in a sequence of text given

From experience, it is very coherent for a model with only 7 billion parameters. You can play with it and ask harder questions from evaluation datasets like BigBench-Hard. Guanaco is an excellent dataset that has produced high-quality models in the past. You can train a Llama 2 model on the entire dataset using mlabonne/guanaco-llama2.

How can we store our new llama-2-7b-miniguanaco model now? We need to merge the weights from LoRA with the base model. Unfortunately, as far as I know, there is no straightforward way to do it: we need to reload the base model in FP16 precision and use the peft library to merge everything. Alas, it also creates a problem with the VRAM (despite emptying it), so I recommend restarting the notebook, re-executing the three first cells, and then executing the next one. Please contact me if you know a fix!

# Reload model in FP16 and merge it with LoRA weights
base_model = AutoModelForCausalLM.from_pretrained(
    model_name,
    low_cpu_mem_usage=True,
    return_dict=True,
    torch_dtype=torch.float16,
    device_map=device_map,
)
model = PeftModel.from_pretrained(base_model, output_dir)
model = model.merge_and_unload()

# Reload tokenizer to save it
tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
tokenizer.pad_token = tokenizer.eos_token
tokenizer.padding_side = "right"

Our weights are merged and we reloaded the tokenizer. We can now push everything to the Hugging Face Hub to save our model.

!huggingface-cli login

model.push_to_hub(new_model, use_temp_dir=False)
tokenizer.push_to_hub(new_model, use_temp_dir=False)

You can now use this model for inference by loading it like any other Llama 2 model from the Hub. It is also possible to reload it for more fine-tuning — perhaps with another dataset?

If you’re interested in a script instead of a notebook, I recommend following the instructions provided in this blog post:

pip install trl
git clone https://github.com/lvwerra/trl
python trl/examples/scripts/sft_trainer.py \
    --model_name meta-llama/Llama-2-7b-hf \
    --dataset_name timdettmers/openassistant-guanaco \
    --load_in_4bit \
    --use_peft \
    --batch_size 4 \
    --gradient_accumulation_steps 2

Conclusion

In this article, we saw how to fine-tune a Llama 2 7b model using a Colab notebook. We introduced some necessary background on LLM training and fine-tuning, as well as important considerations related to instruction datasets. In the second section, we successfully fine-tuned the Llama 2 model with its native prompt template and custom parameters.

These fine-tuned models can then be integrated into LangChain and other architectures as an advantageous alternative to OpenAI API. Remember that, in this new paradigm, instruction datasets are the new gold, and the quality of your model heavily depends on the data it’s been fine-tuned on. So good luck building high-quality datasets!

If you’re interested in more content about LLMs, follow me on Twitter @maximelabonne.

References

• Hugo Touvron, Thomas Scialom, et al. (2023). Llama 2: Open Foundation and Fine-Tuned Chat Models.

• Philipp Schmid, Omar Sanseviero, Pedro Cuenca, & Lewis Tunstall. Llama 2 is here — get it on Hugging Face. https://huggingface.co/blog/llama2

• Rohan Taori, Ishaan Gulrajani, Tianyi Zhang, Yann Dubois, Xuechen Li, Carlos Guestrin, Percy Liang, & Tatsunori B. Hashimoto. (2023). Stanford Alpaca: An Instruction-following LLaMA model.

• Jacob Devlin, Ming-Wei Chang, Kenton Lee, & Kristina Toutanova. (2019). BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding.

• Tim Dettmers, Artidoro Pagnoni, Ari Holtzman, & Luke Zettlemoyer. (2023). QLoRA: Efficient Finetuning of Quantized LLMs.

Large Language Models (LLMs) are known for their extensive computational requirements. Typically, the size of a model is calculated by multiplying the number of parameters (size) by the precision of these values (data type). However, to save memory, weights can be stored using lower-precision data types through a process known as quantization.

We distinguish two main families of weight quantization techniques in the literature:

• Post-Training Quantization (PTQ) is a straightforward technique where the weights of an already trained model are converted to lower precision without necessitating any retraining. Although easy to implement, PTQ is associated with potential performance degradation.

• Quantization-Aware Training (QAT) incorporates the weight conversion process during the pre-training or fine-tuning stage, resulting in enhanced model performance. However, QAT is computationally expensive and demands representative training data.

In this article, we focus on PTQ to reduce the precision of our parameters. To get a good intuition, we will apply both naïve and more sophisticated techniques to a toy example using a GPT-2 model.

The entire code is freely available on Google Colab and GitHub.

📚 Background on Floating Point Representation

The choice of data type dictates the quantity of computational resources required, affecting the speed and efficiency of the model. In deep learning applications, balancing precision and computational performance becomes a vital exercise as higher precision often implies greater computational demands.

Among various data types, floating point numbers are predominantly employed in deep learning due to their ability to represent a wide range of values with high precision. Typically, a floating point number uses n bits to store a numerical value. These n bits are further partitioned into three distinct components:

1. Sign: The sign bit indicates the positive or negative nature of the number. It uses one bit where 0 indicates a positive number and 1 signals a negative number.

2. Exponent: The exponent is a segment of bits that represents the power to which the base (usually 2 in binary representation) is raised. The exponent can also be positive or negative, allowing the number to represent very large or very small values.

3. Significand/Mantissa: The remaining bits are used to store the significand, also referred to as the mantissa. This represents the significant digits of the number. The precision of the number heavily depends on the length of the significand.

This design allows floating point numbers to cover a wide range of values with varying levels of precision. The formula used for this representation is:

To understand this better, let’s delve into some of the most commonly used data types in deep learning: float32 (FP32), float16 (FP16), and bfloat16 (BF16):

• FP32 uses 32 bits to represent a number: one bit for the sign, eight for the exponent, and the remaining 23 for the significand. While it provides a high degree of precision, the downside of FP32 is its high computational and memory footprint.

• FP16 uses 16 bits to store a number: one is used for the sign, five for the exponent, and ten for the significand. Although this makes it more memory-efficient and accelerates computations, the reduced range and precision can introduce numerical instability, potentially impacting model accuracy.

• BF16 is also a 16-bit format but with one bit for the sign, eight for the exponent, and seven for the significand. BF16 expands the representable range compared to FP16, thus decreasing underflow and overflow risks. Despite a reduction in precision due to fewer significand bits, BF16 typically does not significantly impact model performance and is a useful compromise for deep learning tasks.

In ML jargon, FP32 is often termed “full precision” (4 bytes), while BF16 and FP16 are “half-precision” (2 bytes). But could we do even better and store weights using a single byte? The answer is the INT8 data type, which consists of an 8-bit representation capable of storing 2⁸ = 256 different values. In the next section, we’ll see how to convert FP32 weights into an INT8 format.

🔰 Naïve 8-bit Quantization

In this section, we will implement two quantization techniques: a symmetric one with absolute maximum (absmax) quantization and an asymmetric one with zero-point quantization. In both cases, the goal is to map an FP32 tensor X (original weights) to an INT8 tensor X_quant (quantized weights).

With absmax quantization, the original number is divided by the absolute maximum value of the tensor and multiplied by a scaling factor (127) to map inputs into the range [-127, 127]. To retrieve the original FP16 values, the INT8 number is divided by the quantization factor, acknowledging some loss of precision due to rounding.

For instance, let’s say we have an absolution maximum value of 3.2. A weight of 0.1 would be quantized to round(0.1 × 127/3.2) = 4. If we want to dequantize it, we would get 4 × 3.2/127 = 0.1008, which implies an error of 0.008. Here’s the corresponding Python implementation:

import torch

def absmax_quantize(X):
    # Calculate scale
    scale = 127 / torch.max(torch.abs(X))

    # Quantize
    X_quant = (scale * X).round()

    # Dequantize
    X_dequant = X_quant / scale

    return X_quant.to(torch.int8), X_dequant

With zero-point quantization, we can consider asymmetric input distributions, which is useful when you consider the output of a ReLU function (only positive values), for example. The input values are first scaled by the total range of values (255) divided by the difference between the maximum and minimum values. This distribution is then shifted by the zero-point to map it into the range [-128, 127] (notice the extra value compared to absmax). First, we calculate the scale factor and the zero-point value:

Then, we can use these variables to quantize or dequantize our weights:

Let’s take an example: we have a maximum value of 3.2 and a minimum value of -3.0. We can calculate the scale is 255/(3.2 + 3.0) = 41.13 and the zero-point -round(41.13 × -3.0) - 128 = 123 -128 = -5, so our previous weight of 0.1 would be quantized to round(41.13 × 0.1 -5) = -1. This is very different from the previous value obtained using absmax (4 vs. -1).

The Python implementation is quite straightforward:

def zeropoint_quantize(X):
    # Calculate value range (denominator)
    x_range = torch.max(X) - torch.min(X)
    x_range = 1 if x_range == 0 else x_range

    # Calculate scale
    scale = 255 / x_range

    # Shift by zero-point
    zeropoint = (-scale * torch.min(X) - 128).round()

    # Scale and round the inputs
    X_quant = torch.clip((X * scale + zeropoint).round(), -128, 127)

    # Dequantize
    X_dequant = (X_quant - zeropoint) / scale

    return X_quant.to(torch.int8), X_dequant

Instead of relying on complete toy examples, we can use these two functions on a real model thanks to the transformerslibrary.

We start by loading the model and tokenizer for GPT-2. This is a very small model we probably don’t want to quantize, but it will be good enough for this tutorial. First, we want to observe the model’s size so we can compare it later and evaluate the memory savings due to 8-bit quantization.

!pip install -q bitsandbytes>=0.39.0
!pip install -q git+https://github.com/huggingface/accelerate.git
!pip install -q git+https://github.com/huggingface/transformers.git

from transformers import AutoModelForCausalLM, AutoTokenizer
import torch
torch.manual_seed(0)

# Set device to CPU for now
device = 'cpu'

# Load model and tokenizer
model_id = 'gpt2'
model = AutoModelForCausalLM.from_pretrained(model_id).to(device)
tokenizer = AutoTokenizer.from_pretrained(model_id)

# Print model size
print(f"Model size: {model.get_memory_footprint():,} bytes")

Model size: 510,342,192 bytes

The size of the GPT-2 model is approximately 487MB in FP32. The next step consists of quantizing the weights using zero-point and absmax quantization. In the following example, we apply these techniques to the first attention layer of GPT-2 to see the results.

# Extract weights of the first layer
weights = model.transformer.h[0].attn.c_attn.weight.data
print("Original weights:")
print(weights)

# Quantize layer using absmax quantization
weights_abs_quant, _ = absmax_quantize(weights)
print("\nAbsmax quantized weights:")
print(weights_abs_quant)

# Quantize layer using absmax quantization
weights_zp_quant, _ = zeropoint_quantize(weights)
print("\nZero-point quantized weights:")
print(weights_zp_quant)

Original weights:
tensor([[-0.4738, -0.2614, -0.0978,  ...,  0.0513, -0.0584,  0.0250],
        [ 0.0874,  0.1473,  0.2387,  ..., -0.0525, -0.0113, -0.0156],
        [ 0.0039,  0.0695,  0.3668,  ...,  0.1143,  0.0363, -0.0318],
        ...,
        [-0.2592, -0.0164,  0.1991,  ...,  0.0095, -0.0516,  0.0319],
        [ 0.1517,  0.2170,  0.1043,  ...,  0.0293, -0.0429, -0.0475],
        [-0.4100, -0.1924, -0.2400,  ..., -0.0046,  0.0070,  0.0198]])

Absmax quantized weights:
tensor([[-21, -12,  -4,  ...,   2,  -3,   1],
        [  4,   7,  11,  ...,  -2,  -1,  -1],
        [  0,   3,  16,  ...,   5,   2,  -1],
        ...,
        [-12,  -1,   9,  ...,   0,  -2,   1],
        [  7,  10,   5,  ...,   1,  -2,  -2],
        [-18,  -9, -11,  ...,   0,   0,   1]], dtype=torch.int8)

Zero-point quantized weights:
tensor([[-20, -11,  -3,  ...,   3,  -2,   2],
        [  5,   8,  12,  ...,  -1,   0,   0],
        [  1,   4,  18,  ...,   6,   3,   0],
        ...,
        [-11,   0,  10,  ...,   1,  -1,   2],
        [  8,  11,   6,  ...,   2,  -1,  -1],
        [-18,  -8, -10,  ...,   1,   1,   2]], dtype=torch.int8)

The difference between the original (FP32) and quantized values (INT8) is clear, but the difference between absmax and zero-point weights is more subtle. In this case, the inputs look shifted by a value of -1. This suggests that the weight distribution in this layer is quite symmetric.

We can compare these techniques by quantizing every layer in GPT-2 (linear layers, attention layers, etc.) and create two new models: model_abs and model_zp. To be precise, we will actually replace the original weights with de-quantized ones. This has two benefits: it allows us to 1/ compare the distribution of our weights (same scale) and 2/ actually run the models.

Indeed, PyTorch doesn’t allow INT8 matrix multiplication by default. In a real scenario, we would dequantize them to run the model (in FP16 for example) but store them as INT8. In the next section, we will use the bitsandbytes library to solve this issue.

import numpy as np
from copy import deepcopy

# Store original weights
weights = [param.data.clone() for param in model.parameters()]

# Create model to quantize
model_abs = deepcopy(model)

# Quantize all model weights
weights_abs = []
for param in model_abs.parameters():
    _, dequantized = absmax_quantize(param.data)
    param.data = dequantized
    weights_abs.append(dequantized)

# Create model to quantize
model_zp = deepcopy(model)

# Quantize all model weights
weights_zp = []
for param in model_zp.parameters():
    _, dequantized = zeropoint_quantize(param.data)
    param.data = dequantized
    weights_zp.append(dequantized)

Now that our models have been quantized, we want to check the impact of this process. Intuitively, we want to make sure that the quantized weights are close to the original ones. A visual way to check it is to plot the distribution of the dequantized and original weights. If the quantization is lossy, it would drastically change the weight distribution.

The following figure shows this comparison, where the blue histogram represents the original (FP32) weights, and the red one represents the dequantized (from INT8) weights. Note that we only display this plot between -2 and 2 because of outliers with very high absolute values (more on that later).

Both plots are quite similar, with a surprising spike around 0. This spike shows that our quantization is quite lossy since reversing the process doesn’t output the original values. This is particularly true for the absmax model, which displays both a lower valley and a higher spike around 0.

Let’s compare the performance of the original and quantized models. For this purpose, we define a generate_text() function to generate 50 tokens with top-k sampling.

def generate_text(model, input_text, max_length=50):
    input_ids = tokenizer.encode(input_text, return_tensors='pt').to(device)
    output = model.generate(inputs=input_ids,
                            max_length=max_length,
                            do_sample=True,
                            top_k=30,
                            pad_token_id=tokenizer.eos_token_id,
                            attention_mask=input_ids.new_ones(input_ids.shape))
    return tokenizer.decode(output[0], skip_special_tokens=True)

# Generate text with original and quantized models
original_text = generate_text(model, "I have a dream")
absmax_text   = generate_text(model_abs, "I have a dream")
zp_text       = generate_text(model_zp, "I have a dream")

print(f"Original model:\n{original_text}")
print("-" * 50)
print(f"Absmax model:\n{absmax_text}")
print("-" * 50)
print(f"Zeropoint model:\n{zp_text}")

Original model:
I have a dream, and it is a dream I believe I would get to live in my future. I love my mother, and there was that one time I had been told that my family wasn't even that strong. And then I got the
--------------------------------------------------
Absmax model:
I have a dream to find out the origin of her hair. She loves it. But there's no way you could be honest about how her hair is made. She must be crazy.

We found a photo of the hairstyle posted on
--------------------------------------------------
Zeropoint model:
I have a dream of creating two full-time jobs in America—one for people with mental health issues, and one for people who do not suffer from mental illness—or at least have an employment and family history of substance abuse, to work part

Instead of trying to see if one output makes more sense than the others, we can quantify it by calculating the perplexity of each output. This is a common metric used to evaluate language models, which measures the uncertainty of a model in predicting the next token in a sequence. In this comparison, we make the common assumption that the lower the score, the better the model is. In practice, a sentence with a high perplexity could also be correct.

We implement it using a minimal function since it doesn’t need to consider details like the length of the context window since our sentences are short.

def calculate_perplexity(model, text):
    # Encode the text
    encodings = tokenizer(text, return_tensors='pt').to(device)

    # Define input_ids and target_ids
    input_ids = encodings.input_ids
    target_ids = input_ids.clone()

    with torch.no_grad():
        outputs = model(input_ids, labels=target_ids)

    # Loss calculation
    neg_log_likelihood = outputs.loss

    # Perplexity calculation
    ppl = torch.exp(neg_log_likelihood)

    return ppl

ppl     = calculate_perplexity(model, original_text)
ppl_abs = calculate_perplexity(model_abs, absmax_text)
ppl_zp  = calculate_perplexity(model_zp, absmax_text)

print(f"Original perplexity:  {ppl.item():.2f}")
print(f"Absmax perplexity:    {ppl_abs.item():.2f}")
print(f"Zeropoint perplexity: {ppl_zp.item():.2f}")

Original perplexity:  15.53
Absmax perplexity:    17.92
Zeropoint perplexity: 17.97

We see that the perplexity of the original model is slightly lower than the two others. A single experiment is not very reliable, but we could repeat this process multiple times to see the difference between each model. In theory, zero-point quantization should be slightly better than absmax, but is also more costly to compute.

In this example, we applied quantization techniques to entire layers (per-tensor basis). However, we could apply it at different granularity levels: from the entire model to individual values. Quantizing the entire model in one pass would seriously degrade the performance, while quantizing individual values would create a big overhead. In practice, we often prefer the vector-wise quantization, which considers the variability of values in rows and columns inside of the same tensor.

However, even vector-wise quantization doesn’t solve the problem of outlier features. Outlier features are extreme values (negative or positive) that appear in all transformer layers when the model reach a certain scale (>6.7B parameters). This is an issue since a single outlier can reduce the precision for all other values. But discarding these outlier features is not an option since it would greatly degrade the model’s performance.

🔢 8-bit Quantization with LLM.int8()

Introduced by Dettmers et al. (2022), LLM.int8() is a solution to the outlier problem. It relies on a vector-wise (absmax) quantization scheme and introduces mixed-precision quantization. This means that outlier features are processed in a FP16 format to retain their precision, while the other values are processed in an INT8 format. As outliers represent about 0.1% of values, this effectively reduces the memory footprint of the LLM by almost 2x.

LLM.int8() works by conducting matrix multiplication computation in three key steps:

1. Extract columns from the input hidden states X containing outlier features using a custom threshold.

2. Perform the matrix multiplication of the outliers using FP16 and the non-outliers using INT8 with vector-wise quantization (row-wise for the hidden state X and column-wise for the weight matrix W).

3. Dequantize the non-outlier results (INT8 to FP16) and add them to the outlier results to get the full result in FP16.

This approach is necessary because 8-bit precision is limited and can lead to substantial errors when quantizing a vector with large values. These errors also tend to amplify as they propagate through multiple layers.

We can easily use this technique thanks to the integration of the bitsandbytes library into the Hugging Face ecosystem. We just need to specify load_in_8bit=True when loading the model (it also requires a GPU).

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

model_int8 = AutoModelForCausalLM.from_pretrained(model_id,
                                             device_map='auto',
                                             load_in_8bit=True,
                                             )
print(f"Model size: {model_int8.get_memory_footprint():,} bytes")

Model size: 176,527,896 bytes

With this extra line of code, the model is now almost three times smaller (168MB vs. 487MB). We can even compare the distribution of the original and quantized weights as we did earlier:

In this case, we see spikes around -2, -1, 0, 1, 2, etc. These values correspond to the parameters stored in the INT8 format (non-outliers). You can verify it by printing the model’s weights using model_int8.parameters().

We can also generate text with this quantized model and compare it to the original model.

# Generate text with quantized model
text_int8 = generate_text(model_int8, "I have a dream")

print(f"Original model:\n{original_text}")
print("-" * 50)
print(f"LLM.int8() model:\n{text_int8}")

Original model:
I have a dream, and it is a dream I believe I would get to live in my future. I love my mother, and there was that one time I had been told that my family wasn't even that strong. And then I got the
--------------------------------------------------
LLM.int8() model:
I have a dream. I don't know what will come of it, but I am going to have to look for something that will be right. I haven't thought about it for a long time, but I have to try to get that thing

Once again, it is difficult to judge what is the best output, but we can rely on the perplexity metric to give us an (approximate) answer.

print(f"Perplexity (original):   {ppl.item():.2f}")

ppl = calculate_perplexity(model_int8, text_int8)
print(f"Perplexity (LLM.int8()): {ppl.item():.2f}")

Perplexity (original):   15.53
Perplexity (LLM.int8()): 7.93

In this case, the perplexity of the quantized model is twice as low as the original one. In general, this is not the case, but it shows that this quantization technique is very competitive. In fact, the authors of LLM.int8() show that the performance degradation is so low it’s negligible (<1%). However, it has an additional cost in terms of computation: LLM.int8() is roughly about 20% slower for large models.

Conclusion

This article provided an overview of the most popular weight quantization techniques. We started by gaining an understanding of floating point representation, before introducing two techniques for 8-bit quantization: absmax and zero-point quantization. However, their limitations, particularly when it comes to handling outliers, led to LLM.int8(), a technique that also preserves the model’s performance. This approach underlines the progress being made in the field of weight quantization, revealing the importance of properly addressing outliers.

Looking forward, our next article will explore the GPTQ weight quantization technique in depth. This technique, introduced by Frantar et al., only utilizes 4 bits and represents a significant advancement in the field of weight quantization. We will provide a comprehensive guide on how to implement GPTQ using the AutoGPTQ library.

If you’re interested in more technical content around LLMs, follow me on Twitter @maximelabonne.

References

• T. Dettmers, M. Lewis, Y. Belkada, and L. Zettlemoyer, LLM.int8(): 8-bit Matrix Multiplication for Transformers at Scale. 2022.

• Y. Beldaka, and T. Dettmers, A Gentle Introduction to 8-bit Matrix Multiplication, Hugging Face Blog (2022).

• A. Gholami, S. Kim, Z. Dong, Z. Yao, M. W. Mahoney, and K. Keutzer, A Survey of Quantization Methods for Efficient Neural Network Inference. 2021.

• H. Wu, P. Judd, X. Zhang, M. Isaev, and P. Micikevicius, Integer Quantization for Deep Learning Inference: Principles and Empirical Evaluation. 2020.

• Lilian Weng, Large Transformer Model Inference Optimization, Lil’Log (2023).

• Kamil Czarnogorski, Local Large Language Models, Int8 (2023).