wangyue114514/small-agent-1b-wave4-zh-math

small-agent-1b — wave4 (zh + math focused)

A 1B-parameter hybrid-attention language model (Qwen3-Next architecture: Gated DeltaNet linear attention interleaved with gated full attention).

This repository holds the final checkpoint of the wave4 training phase, which deliberately focused on Chinese and mathematics.

Model details

Parameters	1,005,213,696
Architecture	Qwen3NextForCausalLM (32 layers)
Layer pattern	LLLA x 8 — 24 linear-attention + 8 full-attention layers
Hidden size	1536
Attention heads	16 (4 KV heads), head_dim 96
Linear-attention heads	16 key / 16 value, head_dim 96, conv kernel 4
Intermediate size	4096
Vocab size	32768 (byte-level BPE)
Max position	8192
Experts	0 (no MoE)
Tied embeddings	yes
Weight dtype	float32 master weights (trained under bfloat16 autocast)

Training

Phase	wave4 (resumed from the wave3 checkpoint)
Tokens in this phase	5,432,427,745
Optimizer steps	82,960
World size	8 (data parallel)
Sequence length	4096
Micro batch	2 x 1 accumulation
Peak LR	1e-4 (cosine to 2e-5)
Hardware	8 x NVIDIA L40
Wall time	38.4 h (96.4% compute efficiency)

Data mixture in this phase (unique-token share):

Domain	Tokens	Share
general_en	2,562,713,662	47.2%
math	1,927,099,783	35.5%
general_zh	883,156,629	16.3%
code	59,227,381	1.1%
technical	230,290	0.004%

Evaluation (held-out NLL, one fixed panel across all milestones)

Measured from the 250M-token boundary of this phase to the final step.

Domain	start	final	change
Chinese (in-distribution)	5.4286	3.2024	-41.0%
Math (in-distribution)	1.6385	1.0023	-38.8%
English (in-distribution)	2.5850	2.3715	-8.3%
Chinese (retention panel)	4.2924	2.7316	-36.4%
English (retention panel)	2.4349	2.3428	-3.8%
Code (retention panel)	1.0825	1.4155	+30.8%
Technical (retention panel)	1.9321	2.1631	+12.0%

PIQA (0-shot, 128 samples): acc 0.6328 -> 0.6797, acc_norm 0.6797 -> 0.7109 (within the standard error of this small sample size).

Caveats. Code and technical-document performance regressed, consistent with their very low share of this phase's data mixture (1.1% and 0.004%) combined with a no-replay data policy. The code retention curve had already begun recovering before the phase ended, but did not return to its starting level. This checkpoint is best suited to Chinese and mathematical workloads, not code generation.

Usage

from transformers import AutoModelForCausalLM, AutoTokenizer

model_id = "wangyue114514/small-agent-1b-wave4-zh-math"
tok = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(model_id, dtype="auto")

ids = tok("def add(a, b):", return_tensors="pt").input_ids
print(tok.decode(model.generate(ids, max_new_tokens=64)[0]))

Requires transformers >= 4.57 (Qwen3NextForCausalLM).

Chat template is not bundled; the tokenizer defines these special tokens: <|system|>, <|user|>, <|assistant|>, <|tool|>, <|tool_result|>, <|end|>, <|fim_prefix|>, <|fim_middle|>, <|fim_suffix|>, <|memory|>, <|memory_end|>.

Notes

• The checkpoint stores the tied input/output embedding matrix under model.embed_tokens.weight (standard Hugging Face tied layout); lm_head.weight is inferred on load.
• dtype in config.json reflects the training configuration, not the on-disk tensor dtype. Tensors are stored in float32; load with dtype="auto" or dtype=torch.float32 to reproduce the recorded numbers exactly.