---
license: other
license_name: research-only
base_model: Qwen/Qwen3-4B-Base
datasets:
- jepetolee/AMQ3-Math-ShortCoT-0k7k
language:
- en
pipeline_tag: text-generation
tags:
- math
- reasoning
- chain-of-thought
- qwen3
- sft
library_name: transformers
---
# Qwen3-4B AMQ3 Math Short-CoT SFT
Supervised fine-tune of **Qwen/Qwen3-4B-Base** on ~292K short chain-of-thought math
solutions distilled from **Qwen3-235B-A22B** (via `a-m-team/AM-Qwen3-Distilled`).
Intended as a clean math-reasoning cold-start checkpoint (e.g. before RLVR).
## Training
| | |
|---|---|
| Base model | `Qwen/Qwen3-4B-Base` |
| Data | [`jepetolee/AMQ3-Math-ShortCoT-0k7k`](https://huggingface.co/datasets/jepetolee/AMQ3-Math-ShortCoT-0k7k) — 292,375 examples, problem+CoT ≤ ~7K tokens |
| Format | official Qwen3 chat template, `…` reasoning + `\boxed{}` answer |
| Epochs | 1 (full dataset) |
| Effective batch | 32 · lr 1e-5 · warmup 0.03 · max_len 9216 |
| Final loss | 0.48 (token-weighted, full dataset) |
| Tokens seen | ~0.87B (96.9% on the target span) |
## Prompt format
The model is trained to open reasoning with `\n` right after the assistant
header. Use the chat template and let it generate the `` block:
```
<|im_start|>system
Please reason step by step, and put your final answer within \boxed{}.<|im_end|>
<|im_start|>user
{question}<|im_end|>
<|im_start|>assistant
```
## Usage (vLLM)
```python
from vllm import LLM, SamplingParams
llm = LLM(model="jepetolee/Qwen3-4B-AMQ3-Math-SFT", max_model_len=9216)
sp = SamplingParams(temperature=0.7, top_p=0.9, max_tokens=8192)
prompt = (
"<|im_start|>system\n"
"Please reason step by step, and put your final answer within \\boxed{}.<|im_end|>\n"
"<|im_start|>user\nWhat is the sum of the first 10 primes?<|im_end|>\n"
"<|im_start|>assistant\n\n"
)
print(llm.generate([prompt], sp)[0].outputs[0].text)
```
The `generation_config.json` sets `eos_token_id = [151645, 151643]` so generation
stops on `<|im_end|>` out of the box — no manual `stop_token_ids` needed.
## Recommended sampling
Hygiene degrades sharply above temperature 0.75 (holdout sweep, no logit processors):
| temperature | fully-clean generations | ends without an answer |
|---|---|---|
| 0.5 | 80% | 20% |
| 0.75 | 75% | 17% |
| 1.0 | 33% | 33% |
**Use temperature ≤ 0.75**, or use the generation recipe below, which keeps
higher temperatures usable by construction.
## Generation recipe: think/answer budget split (logits processors)
The failure mode behind the table above: on hard prompts the model keeps thinking
until the token cap and never closes ``, so the run truncates with no
`\boxed{}` answer. Instead of only lowering temperature, we split the generation
budget — **total 10240 tokens = up to 6144 think + ~4096 answer** — and enforce it
with two vLLM V1 logits processors, shipped in this repo:
| file | role |
|---|---|
| [`vllm_think_format.py`](./vllm_think_format.py) | `` tag grammar + think-budget cut + forced seal |
| [`vllm_repetition_abort.py`](./vllm_repetition_abort.py) | early EOS for n-gram repetition runaways |
How it works:
1. **Prefill `\n`** after the assistant header (see Prompt format) — the
think block opens exactly once, by construction.
2. **Tag grammar (token-id state machine, no decoding)**: while think is open,
`` is banned; after the first `` both tags are banned forever;
optionally `<|im_start|>` is banned (blocks fake new-turn hallucinations).
3. **Think-budget cut with forced seal**: when the think span reaches
`max_think_tokens`, the processor force-prefills `\n\n` and constrains
the *first* answer token to a whitelist of answer-opening tokens
(`To / We / Let / The / Given / ### / ( / First / In` — ≥96% coverage of answer
openers measured on the 292K SFT set). The model then writes a normal answer
with the remaining budget, so a `\boxed{}` answer still appears even when
thinking was cut.
4. **Repetition abort**: if a rollout's 7-gram repetition ratio exceeds 0.9
(checked every 512 tokens, after the first 2048), logits are masked to EOS-only
for that request. Rollouts that never trigger are **bit-identical** to running
without the processor.
> **vLLM caveat (important)**: pass `async_scheduling=False` to the engine.
> vLLM V1's async scheduling fills `output_tok_ids` with `-1` placeholders, which
> silently disables any logits processor that reads output tokens.
```python
from vllm import LLM, SamplingParams
from transformers import AutoTokenizer
# Download vllm_think_format.py / vllm_repetition_abort.py from this repo
# and put them on your PYTHONPATH.
from vllm_think_format import build_think_format_extra_args
from vllm_repetition_abort import build_repetition_abort_extra_args
model_id = "jepetolee/Qwen3-4B-AMQ3-Math-SFT"
tok = AutoTokenizer.from_pretrained(model_id)
llm = LLM(
model=model_id,
max_model_len=32768,
async_scheduling=False, # REQUIRED for the custom processors
logits_processors=[
"vllm_think_format:ThinkFormatLogitsProcessor",
"vllm_repetition_abort:RepetitionEosLogitsProcessor",
],
)
extra_args = {}
extra_args.update(build_think_format_extra_args(
{"think_format": {
"enabled": True,
"prefilled_open": True, # prompt ends with "\n"
"ban_im_start": True,
"max_think_tokens": 6144, # think budget
"force_close_prefill": True, # seal "\n\n" + whitelist on cut
}},
tok, prefilled_open=True) or {})
extra_args.update(build_repetition_abort_extra_args(
{"repetition_abort": {
"enabled": True, "ngram": 7, "threshold": 0.9,
"min_tokens": 2048, "check_interval": 512,
}},
eos_token_id=tok.convert_tokens_to_ids("<|im_end|>")) or {})
sp = SamplingParams(
temperature=0.7, top_p=0.95,
max_tokens=10240, # total budget: think 6144 + answer ~4096
extra_args=extra_args,
)
prompt = (
"<|im_start|>system\n"
"Please reason step by step, and put your final answer within \\boxed{}.<|im_end|>\n"
"<|im_start|>user\n{question}<|im_end|>\n"
"<|im_start|>assistant\n\n"
)
print(llm.generate([prompt], sp)[0].outputs[0].text)
```
Notes:
- The processor code is research code from our RL training stack (docstrings are in
Korean); requests whose `extra_args` omit the config blocks are ignored entirely,
so the processors are safe to register globally.
- Budget scaling: with 10240 total on this model, per-problem worst-case decode cost
scales roughly with the square of the total length — 12288 costs ~2× and 16384
~3.5× of an 8192 budget. 6144/4096 was chosen as the stability/cost sweet spot.
- With the recipe active, temperature 1.0 remains usable: unclosed-think truncations
are eliminated by construction (thinking is force-sealed and the answer budget is
reserved).
## Limitations
- Math only (English). MCQ items were filtered out of the training data.
- Answers are `\boxed{}`; grading assumes boxed-answer extraction.
- Distilled from a single teacher (Qwen3-235B-A22B); inherits its style and blind spots.
## License
Base model `Qwen/Qwen3-4B-Base` is Apache-2.0, but training data derives from
`a-m-team/AM-Qwen3-Distilled`, which restricts use to **research purposes only**.
This checkpoint therefore carries the same research-only restriction: no commercial
use, no potentially harmful application. The bundled logits-processor files are
released under the same research-only terms.