Cola2 on 16 GB: it fits with offload, scores 56% on ARC-Easy, but the paper's harness won't run

On November 6, 2026 ByteDance Seed uploaded Cola2 to Hugging Face: a four-billion-parameter language model that does not generate text word by word. Instead, it compresses text into continuous vectors, predicts them in blocks with a diffusion transformer, and only at the end turns them back into tokens.

The paper reports a 60.10 average over eight academic benchmarks trained on 346 billion tokens. The repository itself ships the scorer, the JSONL files for the eight benchmarks, and a fifteen-line inference example. It is the first time in a long while that a lab publishes not only the weights but the exact rules it measured with.

I pulled it onto an RTX 4070 Ti SUPER with 16 GB. What follows are the three things I could measure there, and the one I could not.

What changes compared to a normal language model

An autoregressive model writes one token, reads it, writes the next, reads it, and so on. Cola2 reorganizes the work into three stages.

The VAE encoder splits text into dynamic-length semantic spans, one span per “idea” it detects, and projects each span to a continuous 64-dimension vector. The same vector thought of as an idea rather than as a token sequence.

The DiT (thirty-four layers, 2,560 hidden dimensions, 10,240 in the intermediate layer) receives the prefix of already-committed vectors and produces the next block of eight vectors through Flow Matching. It does not sample from a discrete distribution: it integrates a vector field over a fixed number of steps.

The decoder takes the block of vectors, conditioned on the text prefix, and decompresses it to discrete tokens. Then the VAE encoder reads those tokens again and re-compresses them under the same layout, so the continuous and the discrete do not drift apart in long conversations.

This has an immediate practical consequence: every output block is generated in one shot, not in sequence. The promise is that for the same real output speed, the model can be larger than its autoregressive equivalent.

What happened when I pulled it

The snapshot_download of microsoft/Mage-VL-Cowork was a walk; this one less so. The download brings 31 files and 8.77 GB in float32 across eight safetensors shards. Two accidental parallel downloads stepped on each other and left me with four of the eight; just not calling the same snapshot_download twice fixed it.

With the weights on disk, the repository example is literally a single page:

from transformers import AutoModel, AutoTokenizer
import torch

model = AutoModel.from_pretrained(
    "/data/lab/cola2",
    trust_remote_code=True,
    torch_dtype=torch.bfloat16,
    low_cpu_mem_usage=True,
).to("cuda").eval()

It did not load. CUDA out of memory. The model weighs 4.4 GB in bfloat16, but the DiT instantiated on GPU, plus the VAE, plus the activations of a first forward, pushes past 16 GB. At least the first signal.

On top of that, the model’s own GenerationConfig has a bug: it passes is_init=is_init to super().validate() and in transformers 4.56 and 4.57 that argument does not exist. Three sed lines fix the configuration_cola2.py that Hugging Face cached. I left it documented in the blog repository.

With device_map="auto" and max_memory={0: "14GiB", "cpu": "60GiB"}, the model loads in twenty-six and a half seconds. Accelerate splits the pieces like this:

• GPU: the whole VAE, segmenter, pool, latent encoder, latent decoder, duration head, expander, decoder cond stack, and decoder blocks 0 to 20.
• CPU: decoder blocks 21 to 33, cond_in, and the entire DiT, its thirty-four blocks.

Once loaded, the first inference worked: "Question: What is the capital of Chile? Answer:" returned Santiago, correct. By sixty tokens it started to degrade: "largest largest of South and Central and? Answer: Brazil". The next Flow Matching block, repeated many times through the offload, loses coherence.

The paper’s official harness won’t run on 16 GB

The repository includes scripts/run_benchmark_paper.sh, which runs the eight benchmarks (ARC-Easy, HellaSwag, LAMBADA, MMLU, NarrativeQA, RACE, Story Cloze, TruthfulQA) three times with three different seeds and aggregates the results. What it reports at the end is the mean of the eight tasks, averaged over the three seeds: that is the 60.10 from the paper.

To try it on 16 GB I took the quick path first: one seed, one benchmark (arc_easy), fifty samples. The orchestrator does not accept device_map via command line, so I hand-patched benchmarks/inference.py so that the Cola2 load path uses the same device_map="auto" + max_memory that worked for me in the example.

I ran it with GPU max_memory at 14 GiB. OOM. Dropped to 9 GiB. OOM. Dropped to 6 GiB. OOM. Dropped to 4 GiB. OOM, always with the same message: “CUDA out of memory. Allocated 15.16 GiB”. low_cpu_mem_usage and device_map split the initial load well (I can print the hf_device_map after and the heavy blocks are on CPU), but the harness’s first generation step has accelerate bring the whole DiT back to GPU to run it, and that is where it blows.

PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True changes nothing because the problem is not fragmentation: it is a 34-layer DiT with hidden 2,560 processing 32 tokens at a time.

The minibench that did run

Since the orchestrator would not let me, I wrote a thirty-line homegrown benchmark that does not go through orchestrate: it loads the model once with the offload that works, opens eval_data/arc_easy.jsonl, and for each question builds the prompt, generates 32 tokens, and marks a hit if any of the options appears literally in the answer or if the answer matches the ground truth. That is the scoring rule the repository’s own prompt_template.py uses, at least for multichoice.

The code is in a gist and in lab/cola2-16gb/results.json in the blog repository. Fifty questions, one seed, one run.

Cola2: 28 out of 50, 56.0%. The “always first option” baseline: 11 out of 50, 22.0%.

The thirty-four-point margin is crystal clear: the model does much more than guess. On the other hand, 56% on ARC-Easy is not spectacular; the paper does not publish ARC-Easy separately, but the paper’s mean over the eight benchmarks is 60.10, and ARC-Easy tends to be one of the easier tasks in the set. My 56% is consistent with a model that loses some performance to the offload and the max_prompt_length=512 window, but does not collapse.

The hits are textual and sensible. Question 1: “Which piece of safety equipment is used to keep mold spores from entering the respiratory system?” The model answers breathing mask, the ground truth is breathing mask. Question 8: “Plants use sunlight to make”, answers food, correct. Question 18, “Using a softball bat to hit a softball is an example of using which simple machine?”, answers lever, correct.

The errors are text that wants to be the answer but comes out garbled. Question 0, ground truth Sunlight is the source of energy for nearly all ecosystems., Cola2 answers Sun is the source of energy for nearly the ecosystems. It loses words. Question 17: fish that live in coral reree in in shallow wates. The doubled repetitions and missing letters are the pattern.

The visible loop

Each question: load the prompt, call generate, wait for the DiT to generate the latents of the first block (one Flow Matching step per layer × 34 layers, with parameters coming up from RAM), decode to tokens, repeat the block if more tokens are still missing. 10.54 seconds per answer. For 32 tokens, that is 3 tokens per second. On a GPU without offload the same model should be at 30–40 tokens per second, an order of magnitude above.

What can be concluded

Cola2 fits in 16 GB for free-form conversation, with an order-of-magnitude penalty on latency. It loads in 26.6 s, takes 14.5 GB with offload, and answers short questions at 3 tokens per second. Quality holds for the first thirty tokens and starts to degrade after that; for a home chatbot it does not serve, for experimenting with the architecture it does.

It beats the ARC-Easy baseline by a clear margin. 56.0% against 22.0%. The model is not broken; the continuous-latent architecture works.

The full paper harness, on this GPU, cannot run. To reproduce the 60.10 the paper publishes you need a bigger GPU, or a hand-rolled scorer that, like mine, skips the accelerate path and accepts that the DiT cannot be moved whole to GPU and back for each question.

What I did not measure

HellaSwag, LAMBADA, MMLU, NarrativeQA, RACE, Story Cloze, TruthfulQA. The other seven benchmarks. The ficha talked about all eight. The time each 50-question run took (nine minutes) and the orchestrate block made it impractical to run the eight today. What does stand as a data point: the offload configuration I tried is the same for all eight, so the numbers would consistently be “16 GB with offload” mode, not paper mode.

Qwen3-4B as a reference. An autoregressive model of similar size without offload would give the latency number to contrast with the 10.54 s. I have not measured it yet.

The seed. One single seed (0 for the random.shuffle client seed and sample_seed=42 for the orchestrator). The paper reports three different seeds. A diffusion-latent model is stochastic; the README itself warns that the numbers dance. Mine can dance by several points.

Sources

• Model repository on Hugging Face: ByteDance-Seed/Cola2
• Inference code and scorer repository: ByteDance-Seed/Cola2 on GitHub
• Paper: Cola2: Co-Adaptive Hierarchical Language Modeling with Continuous Semantic Compression and Balanced Discrete Generation (arXiv:2610.31305)
• Project page with visualizations: Cola2 Project Page
• Raw data and scripts: lab/cola2-16gb/ in the blog repository

November 6, 2026 · RTX 4070 Ti SUPER 16 GB on NixOS 26.05 · transformers 4.57.0 · torch 2.6.0+cu124 · accelerate 1.15.0 · is_init patch applied to configuration_cola2.py · min-ARC-Easy 50 samples, one seed, max_new_tokens 32, bfloat16, GPU max_memory 14 GiB.