hyb16-edm-s0

hyb16-edm-s0: GDN2-EDM + 6 full-attention layers

One of five 366M-parameter hybrid language models trained under an identical recipe (15B tokens at sequence length 16,384) to compare token mixers at long context. Only the 18 non-full-attention layers differ between the five; the six gated full-attention layers, the MLPs, the tokenizer, the data, the optimizer, and the schedule are shared. This card reports all five so the numbers can be read side by side.

Siblings: hyb16-gdn2-s0, hyb16-swa1k-s0, hyb16-swa2k-s0, hyb16-mix-s0.

Architecture

• 18 GDN2 layers with the EDM variant (decoupled read/write addresses, factored gate rank 64) + 6 gated full-attention layers at 3, 7, 11, 15, 19, 23.
• 24 layers, hidden size 1024, 6 heads of dimension 128 in the linear/EDM layers (8 heads of 128 in attention layers), SwiGLU MLP with intermediate size 2816, RMSNorm, tied input/output embeddings.
• Parameters: 369,410,412 total, 336,642,412 non-embedding (embedding 32,768,000).
• Vocabulary: Llama-2 tokenizer (32,000), EOS id 2.
• Full-attention layers use no positional encoding (NoPE); sliding-window layers use RoPE with theta 10,000; every attention layer has a sigmoid output gate.

Training

• Data: FineWeb-Edu, Llama-2 tokens, documents packed to 16,384 without document masking. Median document 686 tokens; half of all tokens sit in documents longer than 2K.
• 15,000,000,000 tokens (9,536 optimizer steps), global batch 1,572,864 tokens (96 sequences of 16,384).
• AdamW, peak LR 8e-4, warmup 1.5B tokens, warmup-stable-decay schedule with the decay starting at 13.5B tokens. Seed 0. bf16.
• Precision note: the linear layers run flash-linear-attention's GDN2 chunk kernels; attention layers run flash-attn 2 (fa2) with the window as a left-only sliding window.

Evaluation

Zero-shot suite (lm-eval-harness 0.4.12, batch 8 for the LM suite, batch 16 for the recall trio), accuracy unless noted. The "this model" column repeats the model's own values in bold.

Single-needle retrieval (S-NIAH, RULER-style, 500 samples per cell, fraction correct):

Reading for this model: Best validation loss of the wave and the best needle-3 retention at 16K/32K.

All numbers are one seed. Generation-task scores move by a few tenths of a point between hardware and batch sizes; treat differences under about half a point as noise.

Usage

The model class is not in transformers; the modeling code ships with the repo and is loaded with trust_remote_code=True. It needs a CUDA GPU and these packages, which are not bundled:

torch==2.10.*            # what the release was verified with
transformers==5.14.1
flash-linear-attention==0.5.2
flash-attn==2.8.3        # prebuilt wheel; the fa2 attention backend is the default
triton>=3.6
einops

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

repo = "Berlm/hyb16-edm-s0"
tok = AutoTokenizer.from_pretrained(repo, trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(repo, dtype=torch.bfloat16, trust_remote_code=True).cuda().eval()

ids = tok("The capital of France is", return_tensors="pt").to("cuda")
out = model.generate(**ids, max_new_tokens=16, do_sample=False)
print(tok.decode(out[0]))

Set HYBRIDLM_ATTN_BACKEND=flex to use a PyTorch flex-attention backend instead of flash-attn (slower, bitwise-reproducible), or fla for flash-linear-attention's parallel attention kernel.

The weights in model.safetensors are the bf16 export of the final training checkpoint (step 9,536); train_state.json records the step, token count, and the tokenizer fingerprint the evaluations used.

Limitations

These are small research models trained on 15B tokens of web text. They are not instruction-tuned, not safety-tuned, and will produce incorrect or offensive text. They are released to support research on long-context token mixers.