Cola2 en 16 GB: cabe con offload, acierta 56 % en ARC-Easy, pero el banco del paper no corre

El 6 de noviembre de 2026 ByteDance Seed subió a Hugging Face Cola2: un modelo de lenguaje de cuatro mil millones de parámetros que no genera texto palabra por palabra. En vez de eso, comprime el texto a vectores continuos, los predice en bloques con un transformer de difusión, y recién al final los traduce de vuelta a tokens.

El paper reporta 60,10 de media en ocho bancos académicos entrenando sobre 346 mil millones de tokens. El propio repositorio trae el scorer, los JSONL de los ocho bancos, y un ejemplo de inferencia de quince líneas. Es la primera vez en mucho tiempo que un laboratorio publica no solo los pesos, sino las reglas exactas con las que midió.

Lo bajé a una RTX 4070 Ti SUPER de 16 GB. Lo que sigue son las tres cosas que pude medir ahí, y la que no.

Qué cambia respecto a un modelo de lenguaje normal

Un modelo autoregresivo escribe un token, lo lee, escribe el siguiente, lo lee, y así. Cola2 reorganiza el trabajo en tres etapas.

El VAE encoder parte el texto en tramos semánticos de longitud dinámica, un tramo por cada “idea” que detecta, y proyecta cada tramo a un vector continuo de 64 dimensiones. Es el mismo vector pensado como una idea en vez de como una secuencia de tokens.

El DiT (treinta y cuatro capas, 2 560 dimensiones ocultas, 10 240 en la capa intermedia) recibe el prefijo de vectores ya comprometidos y produce el siguiente bloque de ocho vectores por Flow Matching. No muestrea de una distribución discreta: integra un campo vectorial en un número fijo de pasos.

El decoder toma el bloque de vectores, condicionado por el prefijo de texto, y lo descomprime a tokens discretos. Después el VAE encoder vuelve a leer esos tokens y los re-comprime bajo el mismo layout, para que lo continuo y lo discreto no se desalineen en largas conversaciones.

Esto tiene una consecuencia práctica inmediata: cada bloque de salida se genera de una vez, no en serie. La promesa es que para una misma velocidad real de salida, el modelo puede ser más grande que su equivalente autoregresivo.

Qué pasó al bajarlo

El snapshot_download del microsoft/Mage-VL-Cowork fue un paseo; este no tanto. La descarga trae 31 archivos y 8,77 GB en float32 repartidos en ocho shards de safetensors. Dos descargas en paralelo accidentales pisaron archivos y me dejaron cuatro de los ocho; solo rebajar no llamando al mismo snapshot_download dos veces resolvió el problema.

Con los pesos en disco, el ejemplo del repositorio es literalmente una página:

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()

No cargó. CUDA out of memory. El modelo pesa 4,4 GB en bfloat16, pero el DiT instanciado en GPU, más el VAE, más las activaciones de un primer forward, supera los 16 GB. Al menos la primera señal.

Aparte, el GenerationConfig del propio modelo tiene un bug: pasa is_init=is_init a super().validate() y en transformers 4.56 y 4.57 ese argumento no existe. Tres líneas de sed arreglan el configuration_cola2.py cacheado por Hugging Face. Lo dejé documentado en el repositorio del blog.

Con device_map="auto" y max_memory={0: "14GiB", "cpu": "60GiB"}, el modelo carga en veintiséis segundos y medio. Accelerate reparte las piezas así:

• GPU: VAE entero, segmenter, pool, latent encoder, latent decoder, duration head, expander, decoder cond stack, y los bloques 0 a 20 del decoder.
• CPU: bloques 21 a 33 del decoder, cond_in, y el DiT completo, sus treinta y cuatro bloques.

Una vez cargado, la primera inferencia funcionó: "Question: What is the capital of Chile? Answer:" devolvió Santiago, correcto. A los sesenta tokens empezó a degradarse: "largest largest of South and Central and? Answer: Brazil". El bloque siguiente de Flow Matching, al repetirse muchas veces por el offload, pierde coherencia.

El banco oficial del paper no corre en 16 GB

El repositorio incluye scripts/run_benchmark_paper.sh, que corre los ocho bancos (ARC-Easy, HellaSwag, LAMBADA, MMLU, NarrativeQA, RACE, Story Cloze, TruthfulQA) tres veces con tres semillas distintas, y agrega los resultados. Lo que reporta al final es la media de las ocho tareas, promediada sobre las tres semillas: eso es el 60,10 del paper.

Para probarlo en 16 GB probé el camino rápido primero: un solo seed, un solo banco (arc_easy), cincuenta muestras. El orchestrador no acepta device_map por línea de comandos, así que parché a mano benchmarks/inference.py para que la ruta de carga de Cola2 use el mismo device_map="auto" + max_memory que me funcionó en el ejemplo.

Lo corrí con max_memory de GPU en 14 GiB. OOM. Bajé a 9 GiB. OOM. Bajé a 6 GiB. OOM. Bajé a 4 GiB. OOM, siempre con el mismo mensaje: “CUDA out of memory. Allocated 15,16 GiB”. El low_cpu_mem_usage y el device_map reparten bien la carga inicial (puedo imprimir el hf_device_map después y los bloques pesados están en CPU), pero el primer paso de generación del harness hace que accelerate traiga de vuelta todo el DiT a GPU para correrlo, y ahí revienta.

La variable PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True no cambia nada porque el problema no es fragmentación: son activaciones de un DiT de 34 capas con hidden 2 560 procesando 32 tokens a la vez.

El minibench que sí corrió

Como el orchestrador no me dejaba, escribí un banco propio de treinta líneas que no pasa por orchestrate: carga el modelo una sola vez con el offload que funciona, abre eval_data/arc_easy.jsonl, y por cada pregunta arma el prompt, genera 32 tokens, y marca un acierto si alguna de las opciones aparece literal en la respuesta o si la respuesta coincide con la ground truth. Esa es la regla de scoring que el propio prompt_template.py del repositorio usa, al menos para multichoice.

El código está en un gist y en lab/cola2-16gb/results.json del repositorio del blog. Cincuenta preguntas, un solo seed, una sola corrida.

Cola2: 28 de 50, 56,0 %. La línea base “siempre primera opción”: 11 de 50, 22,0 %.

El margen de treinta y cuatro puntos es clarísimo: el modelo hace mucho más que adivinar. Por otro lado, 56 % en ARC-Easy no es espectacular; el paper no publica ARC-Easy por separado, pero la media del paper sobre los ocho bancos es 60,10, y ARC-Easy suele ser una de las tareas más fáciles del lote. Mi 56 % es consistente con un modelo que pierde algo de desempeño por el offload y la ventana de max_prompt_length=512, pero no se desploma.

Los aciertos son textuales y sensatos. Pregunta 1: “Which piece of safety equipment is used to keep mold spores from entering the respiratory system?” El modelo responde breathing mask, la ground truth es breathing mask. Pregunta 8: “Plants use sunlight to make”, responde food, correcto. Pregunta 18, “Using a softball bat to hit a softball is an example of using which simple machine?”, responde lever, correcto.

Los errores son texto que quiere ser la respuesta pero sale corrido. Pregunta 0, ground truth Sunlight is the source of energy for nearly all ecosystems., Cola2 responde Sun is the source of energy for nearly the ecosystems. Pierde palabras. Pregunta 17: fish that live in coral reree in in shallow wates. Las repeticiones dobles y las letras que faltan son el patrón.

El bucle visible

Cada pregunta: cargar el prompt, llamar a generate, esperar a que el DiT genere los latentes del primer bloque (un paso de Flow Matching por capa × 34 capas, con los parámetros bajando de la RAM), decodificar a tokens, repetir el bloque si siguen faltando tokens. 10,54 segundos por respuesta. Para 32 tokens, eso es 3 tokens por segundo. En una GPU sin offload el mismo modelo debería estar en 30-40 tokens por segundo, un orden de magnitud por encima.

Qué sí se puede concluir

Cola2 entra en 16 GB para conversación libre, con una penalización de un orden de magnitud en latencia. Carga en 26,6 s, ocupa 14,5 GB con offload, y responde preguntas cortas a 3 tokens por segundo. La calidad aguanta los primeros treinta tokens y empieza a degradarse después; para un chatbot casero no sirve, para experimentar con la arquitectura sí.

Acierta ARC-Easy por encima de la línea base por un margen claro. 56,0 % contra 22,0 %. El modelo no está roto; la arquitectura de latentes continuos funciona.

El banco entero del paper, en esta GPU, no se puede correr. Para reproducir los 60,10 que el paper publica hace falta una GPU más grande, o un scorer manual que, como el mío, salte la ruta de accelerate y acepte que el DiT no se puede mover completo a GPU y vuelta cada pregunta.

Qué no medí

HellaSwag, LAMBADA, MMLU, NarrativeQA, RACE, Story Cloze, TruthfulQA. Los otros siete bancos. La ficha hablaba de los ocho. El tiempo que llevó cada corrida de 50 preguntas (nueve minutos) y el bloqueo con orchestrate hicieron inviable correr los ocho hoy. Lo que sí vale como dato: la configuración de offload que probé es la misma para los ocho, así que los números serían consistentemente de modo “16 GB con offload”, no de modo paper.

Qwen3-4B como referencia. Un modelo autoregresivo de similar tamaño sin offload daría el número de latencia justo para contrastar los 10,54 s. No lo medí todavía.

La semilla. Una sola semilla (0 por el seed del random.shuffle del cliente y el sample_seed=42 del orchestrador). El paper reporta tres semillas distintas. Un modelo diffusion-latent es estocástico; el propio README advierte que los números bailan. Los míos pueden bailar varios puntos.

Fuentes

• Repositorio del modelo en Hugging Face: ByteDance-Seed/Cola2
• Repositorio del código de inferencia y scorer: ByteDance-Seed/Cola2 en GitHub
• Paper: Cola2: Co-Adaptive Hierarchical Language Modeling with Continuous Semantic Compression and Balanced Discrete Generation (arXiv:2610.31305)
• Página del proyecto con visualizaciones: Cola2 Project Page
• Datos crudos y scripts: lab/cola2-16gb/ en el repositorio del blog

6 de noviembre de 2026 · RTX 4070 Ti SUPER 16 GB sobre NixOS 26.05 · transformers 4.57.0 · torch 2.6.0+cu124 · accelerate 1.15.0 · patch is_init aplicado a configuration_cola2.py · min-ARC-Easy 50 muestras, un seed, max_new_tokens 32, bfloat16, max_memory GPU 14 GiB.