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.weightis inferred on load. dtypeinconfig.jsonreflects the training configuration, not the on-disk tensor dtype. Tensors are stored in float32; load withdtype="auto"ordtype=torch.float32to reproduce the recorded numbers exactly.