lixiaochuan2020

qwen3-4b-midtrain-local-kv-per-document

Deploy Dedicated

README

License: apache-2.0

Access

Public: anyone can download it, no access request or login needed.

Weights

Saved dtype: float32 (as written by training). Load with dtype=torch.float32, or cast to bfloat16 for inference.
Branch main: the final checkpoint, step 4,763 of 4,763. Intermediate checkpoints, if added later, go on branches named step-<N>.

python

import torchfrom transformers import AutoModelForCausalLM, AutoTokenizer repo = "lixiaochuan2020/qwen3-4b-midtrain-local-kv-per-document"tokenizer = AutoTokenizer.from_pretrained(repo)model = AutoModelForCausalLM.from_pretrained(repo, dtype=torch.float32)  # full attention; see the inference note

Mid-training data

A Qwen/Qwen3-32B teacher (thinking mode) answered NQ and HotpotQA training questions as a search agent, 4 times per question, using a wiki_search tool over wiki-18 (e5-base-v2, 5 passages per search, at most 10 searches, at most 16,384 tokens per trajectory). Trajectories were kept if they called at least one search and the Qwen3-32B answer judge marked the final answer correct. Gold answers were used only for that check, never as training text.

Table with columns: Item, Value
Item	Value
Questions answered	128,875
Trajectories kept	297,720
Supervised (assistant) tokens	243.78M
Tool-result (passage) tokens	314.24M
Total tokens	684.86M

The loss covers assistant tokens only.

Mid-training config

Table with columns: Setting, Value
Setting	Value
Base model	`Qwen/Qwen3-4B`
Optimizer for decoder-layer weight matrices	Muon
Optimizer for embeddings and norms	AdamW
Learning rate	2e-5
Schedule	cosine
Warmup (fraction of steps)	0.03
Epochs	1
Steps	4,763
Packed row length (tokens)	18,432 (padding-free; each trajectory attends only to itself)
Max trajectory length (tokens)	16,384
Rows per step	8
GPUs	8
Weight decay	0
Max gradient norm	1.0
Weights	float32
Compute	bf16 autocast
Attention kernel	FlexAttention
Gradient checkpointing	on
Seed	42

Attention

Each retrieved passage, from its Doc k(Title: token through its final newline, is its own document. Its tokens attend only to earlier tokens of the same passage: not to the conversation before the search, and not to the other passages.
Passage K/V are detached: the loss on later tokens does not train how passages are encoded.
The tool-turn wrapper tokens (<|im_start|>user\n<tool_response>\n before the passages and </tool_response><|im_end|>\n after them) attend normally.
Every token after the tool turn attends to all passages.
Passages keep their absolute positions in the sequence.

Prompt format

Qwen3 chat template with one tool, wiki_search, called in Qwen3's <tool_call> format. Each tool result lists the top 5 wiki-18 passages as Doc k(Title: …) text lines. The system prompt and tool schema used in training are SYSTEM_PROMPT and TOOLS in src/kv_midtrain/data_loading/afm_mhqa.py of ashmitkx/kv-midtrain (private repository).

Evaluation

ID search: 1,200 held-out questions (HotpotQA dev 720, NQ test 240, 2WikiMultihopQA dev 240), 3 samples each. None is a mid-training question.
OOD search: 1,025 questions (Bamboogle 125, FEVER 300, MuSiQue 300, PopQA 300), 1 sample each.
Search settings: the training system prompt, a Question: … user message and the wiki_search tool; temperature 0.7, top-p 0.95, 4,096 new tokens per turn, up to 10 wiki-18 searches, top-5 passages.
Generation: a PyTorch loop over the Hugging Face model (kv_midtrain.evaluation.evaluate), one token per forward pass; Local mode uses the local-KV forward of src/kv_midtrain/training/local_kv.py.
Modes: Full is standard attention. Local is per-document local-KV attention, as in training (see Attention).
Math: GSM8K test and MATH-500, one sample per problem with thinking, no tools; temperature 0.6, top-p 0.95, top-k 20, up to 16,384 tokens.
Judge: Qwen/Qwen3.5-9B (no thinking, temperature 0) gets the question, the dataset's gold answers and the extracted final answer, and replies Correct or Incorrect. Search: the content of the last <answer> or <final_answer> tag after the reasoning, otherwise the last paragraph; there is no answer when the reasoning was not closed or the last turn calls a tool. Math: the last \boxed{} after the reasoning. A missing answer, or one over 50 words, counts as incorrect.
Columns: source columns and Accuracy are the share of rollouts judged correct (Accuracy over all rollouts of the group); CI low and CI high come from 2,000 bootstrap resamples of questions; Answer rate is the share of rollouts with a final answer.

ID search:

Table with columns: Mode, 2WikiMultihopQA, HotpotQA, NQ, Accuracy, CI low, CI high, Answer rate
Mode	2WikiMultihopQA	HotpotQA	NQ	Accuracy	CI low	CI high	Answer rate
Full	0.535	0.610	0.597	0.593	0.569	0.616	0.946
Local	0.485	0.556	0.593	0.549	0.525	0.574	0.966

OOD search:

Table with columns: Mode, Bamboogle, FEVER, MuSiQue, PopQA, Accuracy, CI low, CI high, Answer rate
Mode	Bamboogle	FEVER	MuSiQue	PopQA	Accuracy	CI low	CI high	Answer rate
Full	0.576	0.907	0.303	0.487	0.567	0.537	0.597	0.963
Local	0.592	0.910	0.263	0.463	0.551	0.520	0.582	0.978

Math (no passages, so Full and Local are the same):

Table with columns: Dataset, Problems, Accuracy, Boxed answer rate
Dataset	Problems	Accuracy	Boxed answer rate
GSM8K test	1,319	0.901	0.954
MATH-500	500	0.864	0.932

Training run

W&B: local-KV-per-document (project cxcscmu/kv-midtrain-qwen3-4b, run id qwen3-4b-midtrain-muon-local_kv_per_document_detached).

Table with columns: Model, Training, Attention in training
Model	Training	Attention in training
`lixiaochuan2020/qwen3-4b-midtrain-baseline`	mid-training	full
`lixiaochuan2020/qwen3-4b-midtrain-local-kv-per-document` (this model)	mid-training	per-document local-KV
`lixiaochuan2020/qwen3-4b-rl-baseline`	GRPO from the mid-trained baseline	full
`lixiaochuan2020/qwen3-4b-rl-local-kv-per-document`	GRPO from this model	per-document local-KV

Collection: kv-midtrain.

Available on FriendliAI

Dedicated Endpoints

Run this model inference on single tenant GPU with unmatched speed and reliability at scale.

Learn more

Model Details

Model Provider

lixiaochuan2020

Model Tree

Base

Qwen/Qwen3-4B

Fine-tuned

this model

Input Modalities

Text

Output Modalities

Text

Supported Functionality

Dedicated Endpoints

Explore FriendliAI today

Get started Talk to an engineer

README

License: apache-2.0

Access

Public: anyone can download it, no access request or login needed.

Weights

Saved dtype: float32 (as written by training). Load with dtype=torch.float32, or cast to bfloat16 for inference.
Branch main: the final checkpoint, step 4,763 of 4,763. Intermediate checkpoints, if added later, go on branches named step-<N>.

python

import torchfrom transformers import AutoModelForCausalLM, AutoTokenizer repo = "lixiaochuan2020/qwen3-4b-midtrain-local-kv-per-document"tokenizer = AutoTokenizer.from_pretrained(repo)model = AutoModelForCausalLM.from_pretrained(repo, dtype=torch.float32)  # full attention; see the inference note

import torchfrom transformers import AutoModelForCausalLM, AutoTokenizer repo = "lixiaochuan2020/qwen3-4b-midtrain-local-kv-per-document"tokenizer = AutoTokenizer.from_pretrained(repo)model = AutoModelForCausalLM.from_pretrained(repo, dtype=torch.float32)  # full attention; see the inference note

Mid-training data

Table with columns: Item, Value
Item	Value
Questions answered	128,875
Trajectories kept	297,720
Supervised (assistant) tokens	243.78M
Tool-result (passage) tokens	314.24M
Total tokens	684.86M

The loss covers assistant tokens only.

Mid-training config

Table with columns: Setting, Value
Setting	Value
Base model	`Qwen/Qwen3-4B`
Optimizer for decoder-layer weight matrices	Muon
Optimizer for embeddings and norms	AdamW
Learning rate	2e-5
Schedule	cosine
Warmup (fraction of steps)	0.03
Epochs	1
Steps	4,763
Packed row length (tokens)	18,432 (padding-free; each trajectory attends only to itself)
Max trajectory length (tokens)	16,384
Rows per step	8
GPUs	8
Weight decay	0
Max gradient norm	1.0
Weights	float32
Compute	bf16 autocast
Attention kernel	FlexAttention
Gradient checkpointing	on
Seed	42

Attention

Each retrieved passage, from its Doc k(Title: token through its final newline, is its own document. Its tokens attend only to earlier tokens of the same passage: not to the conversation before the search, and not to the other passages.
Passage K/V are detached: the loss on later tokens does not train how passages are encoded.
The tool-turn wrapper tokens (<|im_start|>user\n<tool_response>\n before the passages and </tool_response><|im_end|>\n after them) attend normally.
Every token after the tool turn attends to all passages.
Passages keep their absolute positions in the sequence.

Prompt format

Evaluation

ID search: 1,200 held-out questions (HotpotQA dev 720, NQ test 240, 2WikiMultihopQA dev 240), 3 samples each. None is a mid-training question.
OOD search: 1,025 questions (Bamboogle 125, FEVER 300, MuSiQue 300, PopQA 300), 1 sample each.
Search settings: the training system prompt, a Question: … user message and the wiki_search tool; temperature 0.7, top-p 0.95, 4,096 new tokens per turn, up to 10 wiki-18 searches, top-5 passages.
Generation: a PyTorch loop over the Hugging Face model (kv_midtrain.evaluation.evaluate), one token per forward pass; Local mode uses the local-KV forward of src/kv_midtrain/training/local_kv.py.
Modes: Full is standard attention. Local is per-document local-KV attention, as in training (see Attention).
Math: GSM8K test and MATH-500, one sample per problem with thinking, no tools; temperature 0.6, top-p 0.95, top-k 20, up to 16,384 tokens.
Judge: Qwen/Qwen3.5-9B (no thinking, temperature 0) gets the question, the dataset's gold answers and the extracted final answer, and replies Correct or Incorrect. Search: the content of the last <answer> or <final_answer> tag after the reasoning, otherwise the last paragraph; there is no answer when the reasoning was not closed or the last turn calls a tool. Math: the last \boxed{} after the reasoning. A missing answer, or one over 50 words, counts as incorrect.
Columns: source columns and Accuracy are the share of rollouts judged correct (Accuracy over all rollouts of the group); CI low and CI high come from 2,000 bootstrap resamples of questions; Answer rate is the share of rollouts with a final answer.

ID search:

Table with columns: Mode, 2WikiMultihopQA, HotpotQA, NQ, Accuracy, CI low, CI high, Answer rate
Mode	2WikiMultihopQA	HotpotQA	NQ	Accuracy	CI low	CI high	Answer rate
Full	0.535	0.610	0.597	0.593	0.569	0.616	0.946
Local	0.485	0.556	0.593	0.549	0.525	0.574	0.966

OOD search:

Table with columns: Mode, Bamboogle, FEVER, MuSiQue, PopQA, Accuracy, CI low, CI high, Answer rate
Mode	Bamboogle	FEVER	MuSiQue	PopQA	Accuracy	CI low	CI high	Answer rate
Full	0.576	0.907	0.303	0.487	0.567	0.537	0.597	0.963
Local	0.592	0.910	0.263	0.463	0.551	0.520	0.582	0.978

Math (no passages, so Full and Local are the same):

Table with columns: Dataset, Problems, Accuracy, Boxed answer rate
Dataset	Problems	Accuracy	Boxed answer rate
GSM8K test	1,319	0.901	0.954
MATH-500	500	0.864	0.932

Training run

W&B: local-KV-per-document (project cxcscmu/kv-midtrain-qwen3-4b, run id qwen3-4b-midtrain-muon-local_kv_per_document_detached).

Table with columns: Model, Training, Attention in training
Model	Training	Attention in training
`lixiaochuan2020/qwen3-4b-midtrain-baseline`	mid-training	full
`lixiaochuan2020/qwen3-4b-midtrain-local-kv-per-document` (this model)	mid-training	per-document local-KV
`lixiaochuan2020/qwen3-4b-rl-baseline`	GRPO from the mid-trained baseline	full
`lixiaochuan2020/qwen3-4b-rl-local-kv-per-document`	GRPO from this model	per-document local-KV

Collection: kv-midtrain.

qwen3-4b-midtrain-local-kv-per-document

README

Access

Weights

Mid-training data

Mid-training config

Attention

Prompt format

Evaluation

Training run

Related models

Explore FriendliAI today

README

Access

Weights

Mid-training data

Mid-training config

Attention

Prompt format

Evaluation

Training run

Related models