hyb16-swa1k-s0: SWA-1K + 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-edm-s0, hyb16-swa2k-s0, hyb16-mix-s0.

Architecture

• 18 gated sliding-window attention layers (window 1024, RoPE theta 10000) + 6 gated full-attention layers (no RoPE) 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: 366,265,344 total, 333,497,344 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.

metric	this model	gdn2	edm	swa1k	swa2k	mix
final val loss (nats)	2.3245	2.3055	2.2977	2.3245	2.3289	2.3156
arc_easy / arc_challenge	57.0 / 25.0	59.4 / 26.1	59.7 / 25.2	57.0 / 25.0	55.0 / 25.0	57.1 / 25.1
hellaswag (acc_norm)	40.7	41.5	42.0	40.7	40.8	41.5
piqa / winogrande / boolq	66.5 / 50.4 / 61.0	66.5 / 52.5 / 56.5	67.3 / 52.2 / 59.5	66.5 / 50.4 / 61.0	64.7 / 52.5 / 58.7	67.1 / 53.0 / 58.8
openbookqa (acc_norm)	32.0	33.4	31.8	32.0	31.2	31.6
lambada acc / ppl	33.4 / 37.4	35.5 / 29.9	35.9 / 28.2	33.4 / 37.4	37.4 / 29.3	34.6 / 31.5
wikitext word ppl	25.1	24.8	24.3	25.1	25.3	25.1
swde / fda / squad_completion	47.9 / 51.5 / 21.8	56.0 / 66.2 / 36.7	55.5 / 62.1 / 38.4	47.9 / 51.5 / 21.8	50.9 / 13.2 / 21.5	53.8 / 13.6 / 35.4
social_iqa (prompt-qualified)	37.8	38.8	39.0	37.8	37.5	37.9

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

metric	this model	gdn2	edm	swa1k	swa2k	mix
needle-1 @1K/2K/4K/8K	1.00/0.98/0.96/0.94	1.00/1.00/1.00/1.00	1.00/1.00/1.00/1.00	1.00/0.98/0.96/0.94	1.00/1.00/0.52/0.28	1.00/1.00/0.53/0.28
needle-1 @16K/32K	0.95/0.94	1.00/1.00	1.00/1.00	0.95/0.94	0.15/0.05	0.15/0.05
needle-2 @1K/2K/4K/8K	1.00/0.99/0.98/0.81	1.00/1.00/1.00/1.00	1.00/1.00/1.00/0.97	1.00/0.99/0.98/0.81	0.94/0.71/0.42/0.25	1.00/1.00/0.48/0.20
needle-2 @16K/32K	0.88/0.63	0.93/0.37	0.94/0.65	0.88/0.63	0.15/0.08	0.13/0.08
needle-3 @1K/2K/4K/8K	0.99/0.94/0.76/0.37	0.98/0.98/0.91/0.63	0.96/0.95/0.84/0.54	0.99/0.94/0.76/0.37	0.18/0.30/0.29/0.25	0.98/0.98/0.30/0.01
needle-3 @16K/32K	0.16/0.09	0.17/0.04	0.45/0.20	0.16/0.09	0.16/0.06	0.00/0.00

Reading for this model: Retrieval degrades gently beyond the 1K window (needle-1 0.94 at 8K, 0.94 at 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-swa1k-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.