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 — 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)

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 <think> tag grammar + think-budget cut + forced seal
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.

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.

Downloads last month
974
Safetensors
Model size
1B params
Tensor type
BF16
·
Inference Providers NEW
This model isn't deployed by any Inference Provider. 🙋 Ask for provider support

Model tree for jepetolee/Qwen3-4B-AMQ3-Math-SFT

Finetuned
(386)
this model

Dataset used to train jepetolee/Qwen3-4B-AMQ3-Math-SFT