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

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

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:

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.