Qwen3-4B-AMQ3-Math-SFT/README.md

---
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, `<think>…</think>` 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 `<think>\n` right after the assistant
header. Use the chat template and let it generate the `<think>` 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
<think>
```

## 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<think>\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 `</think>`, 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) | `<think>` 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 `<think>\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,
   `<think>` is banned; after the first `</think>` 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 `</think>\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 "<think>\n"
        "ban_im_start": True,
        "max_think_tokens": 6144,      # think budget
        "force_close_prefill": True,   # seal "</think>\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<think>\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.