Berlm/hyb16-gdn2-s1

hyb16-gdn2-s1: GDN2 + 6 full-attention layers

Seed 1 of the 16K hybrid GDN2 model, trained for approximately 15B tokens on FineWeb-Edu. The s0/s1/s2 runs use the same architecture and training recipe with different random seeds.

Related runs: gdn2 s0, gdn2 s2, edm s0, edm s1, edm s2.

Architecture

• —18 Gated DeltaNet-2 (GDN2) layers + 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: 365,281,644 total, 332,513,644 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.
• —14,998,831,104 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 1. 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

Final checkpoint (step 9536). Accuracy/recall scores are percentages; loss and perplexity are unscaled. HellaSwag and OpenBookQA use normalized accuracy; recall uses contains scoring. Raw results are included in eval_lm.json, eval_recall.json, and eval_siqa.json.

Single-needle retrieval

Protocol niah-v2, 500 examples per task/context length with random key locations. Values are percent correct. Raw results: eval_niah.json and eval_niah_long.json.

Key location × context length

Protocol niah-grid-v1, 200 examples per cell. Columns show nominal key location within the haystack (0% beginning, 100% end); rows show the configured context budget. Heatmap cells are fractions correct (0–1). Needle-1 uses repeated text/number, needle-2 essay/number, and needle-3 essay/UUID. Raw cells: `eval_niah_grid.json`.

[image]

These are single-seed results. Across seeds 0–2, EDM has consistently lower validation loss, while needle-2/3 retrieval varies substantially with seed.

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-gdn2-s1"
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.

The bundled shared RA/EDM backward code preserves saved activations for repeated backward calls on a retained graph.