lixiaochuan2020

qwen3-4b-rl-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. Load with dtype=torch.float32, or cast to bfloat16 for inference.
Branch main: the policy after RL step 300 of 300, converted from the miles FSDP checkpoint to Hugging Face format.
Conversion check against the starting model: same 4,022,468,096 parameters, all finite, relative L2 change of the weights 0.0214

python

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

Starting model

lixiaochuan2020/qwen3-4b-midtrain-local-kv-per-document: Qwen/Qwen3-4B mid-trained for 1 epoch (4,763 steps, 8 packed rows of 18,432 tokens per step, trajectories up to 16,384 tokens) on 297,720 judged-correct search-agent trajectories written by a Qwen/Qwen3-32B teacher with wiki-18 search. Muon for the decoder-layer matrices and AdamW for the rest, learning rate 2e-5, cosine schedule, warmup 0.03, float32 weights with bf16 compute, FlexAttention with per-document local-KV attention and detached passage K/V. Details are on its card.

RL data

Table with columns: Item, Value
Item	Value
Question pool	Search-R1 `nq_hotpotqa_train`
Training questions	32,000
HotpotQA questions	16,524
NQ questions	15,476

The RL questions are a random sample (seed 0) of the pool after removing questions in the AFM SFT data and questions asked more than once. None is a mid-training question or an eval question. Gold answers are used only by the reward.

Reward

1 when Qwen/Qwen3.5-9B (no thinking, temperature 0) judges the rollout's final answer correct against the gold answers, otherwise 0.

Final answer: in the last assistant turn, after the reasoning, the content of the last <answer> or <final_answer> tag, otherwise the last paragraph.
No answer (reward 0): the reasoning was not closed, the last turn still calls a tool, or the answer is empty or over 50 words.
Judge prompt: the question, the gold answers and the final answer; the judge decides whether the answer refers to the same thing as any gold answer (wording, formatting and aliases do not matter; a different entity, date or number, or a list of several candidates, is incorrect) and replies Correct or Incorrect.

RL config

Table with columns: Setting, Value
Setting	Value
Framework	miles (FSDP training, SGLang rollouts)
Algorithm	GRPO
Starting model and KL reference	`lixiaochuan2020/qwen3-4b-midtrain-local-kv-per-document`
RL steps (rollout steps)	300
Prompts per step	56
Samples per prompt	16
Optimizer steps per RL step	4
Rollout temperature	1.0
New tokens per turn	4,096
Max context (tokens)	16,384
Max searches per rollout	10
Search	`wiki_search` over wiki-18, e5-base-v2, top 5 passages
KL loss coefficient	0.001
KL estimator	low-variance KL (`low_var_kl`)
Clip ε	0.2
Entropy coefficient	0
Optimizer	Adam
Adam β1	0.9
Adam β2	0.999
Learning rate	1e-6
Warmup (optimizer steps)	40
Schedule	constant
Weight decay	0
Max gradient norm	1.0
GPUs for training and rollouts (colocated)	7
GPUs for search and the reward judge	1

Attention

Per-document local-KV attention in both rollouts and training:

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.
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.
Rollouts: SGLang with the local-KV prefill (flex_attention backend), fa3 decode, chunked prefill off, so a tool turn is never split.
Training: FlexAttention with the local-KV mask; passage K/V are detached, as in mid-training, so the loss on later tokens does not train how passages are encoded.

Prompt format

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

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 or RL 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: vLLM 0.24 (kv_midtrain.evaluation.vllm_generate), batched over all open conversations; Local mode runs through the local-KV vLLM patch. Checked against the PyTorch loop used for the mid-trained models on the Qwen3-1.7B local-KV mid-trained model (ID questions, 3 samples each): judged accuracy differs by +0.010 (CI −0.005 to +0.026) in Full mode and +0.002 (CI −0.011 to +0.016) in Local mode.
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.504	0.699	0.672	0.655	0.631	0.678	0.852
Local	0.558	0.694	0.662	0.660	0.637	0.684	0.949

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.648	0.927	0.367	0.593	0.631	0.601	0.660	0.873
Local	0.680	0.933	0.423	0.570	0.647	0.618	0.674	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.937	0.995
MATH-500	500	0.924	0.980

Training run

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

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`	mid-training; starting model of this one	per-document local-KV
`lixiaochuan2020/qwen3-4b-rl-baseline`	GRPO from the mid-trained baseline	full
`lixiaochuan2020/qwen3-4b-rl-local-kv-per-document` (this model)	GRPO from the mid-trained local-KV 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

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

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. Load with dtype=torch.float32, or cast to bfloat16 for inference.
Branch main: the policy after RL step 300 of 300, converted from the miles FSDP checkpoint to Hugging Face format.
Conversion check against the starting model: same 4,022,468,096 parameters, all finite, relative L2 change of the weights 0.0214

python

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

Starting model

RL data

Table with columns: Item, Value
Item	Value
Question pool	Search-R1 `nq_hotpotqa_train`
Training questions	32,000
HotpotQA questions	16,524
NQ questions	15,476

Reward

1 when Qwen/Qwen3.5-9B (no thinking, temperature 0) judges the rollout's final answer correct against the gold answers, otherwise 0.

Final answer: in the last assistant turn, after the reasoning, the content of the last <answer> or <final_answer> tag, otherwise the last paragraph.
No answer (reward 0): the reasoning was not closed, the last turn still calls a tool, or the answer is empty or over 50 words.
Judge prompt: the question, the gold answers and the final answer; the judge decides whether the answer refers to the same thing as any gold answer (wording, formatting and aliases do not matter; a different entity, date or number, or a list of several candidates, is incorrect) and replies Correct or Incorrect.

RL config

Table with columns: Setting, Value
Setting	Value
Framework	miles (FSDP training, SGLang rollouts)
Algorithm	GRPO
Starting model and KL reference	`lixiaochuan2020/qwen3-4b-midtrain-local-kv-per-document`
RL steps (rollout steps)	300
Prompts per step	56
Samples per prompt	16
Optimizer steps per RL step	4
Rollout temperature	1.0
New tokens per turn	4,096
Max context (tokens)	16,384
Max searches per rollout	10
Search	`wiki_search` over wiki-18, e5-base-v2, top 5 passages
KL loss coefficient	0.001
KL estimator	low-variance KL (`low_var_kl`)
Clip ε	0.2
Entropy coefficient	0
Optimizer	Adam
Adam β1	0.9
Adam β2	0.999
Learning rate	1e-6
Warmup (optimizer steps)	40
Schedule	constant
Weight decay	0
Max gradient norm	1.0
GPUs for training and rollouts (colocated)	7
GPUs for search and the reward judge	1

Attention

Per-document local-KV attention in both rollouts and training:

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.
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.
Rollouts: SGLang with the local-KV prefill (flex_attention backend), fa3 decode, chunked prefill off, so a tool turn is never split.
Training: FlexAttention with the local-KV mask; passage K/V are detached, as in mid-training, so the loss on later tokens does not train how passages are encoded.

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 or RL 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: vLLM 0.24 (kv_midtrain.evaluation.vllm_generate), batched over all open conversations; Local mode runs through the local-KV vLLM patch. Checked against the PyTorch loop used for the mid-trained models on the Qwen3-1.7B local-KV mid-trained model (ID questions, 3 samples each): judged accuracy differs by +0.010 (CI −0.005 to +0.026) in Full mode and +0.002 (CI −0.011 to +0.016) in Local mode.
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.504	0.699	0.672	0.655	0.631	0.678	0.852
Local	0.558	0.694	0.662	0.660	0.637	0.684	0.949

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.648	0.927	0.367	0.593	0.631	0.601	0.660	0.873
Local	0.680	0.933	0.423	0.570	0.647	0.618	0.674	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.937	0.995
MATH-500	500	0.924	0.980

Training run

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

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`	mid-training; starting model of this one	per-document local-KV
`lixiaochuan2020/qwen3-4b-rl-baseline`	GRPO from the mid-trained baseline	full
`lixiaochuan2020/qwen3-4b-rl-local-kv-per-document` (this model)	GRPO from the mid-trained local-KV model	per-document local-KV

Collection: kv-midtrain.

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

README

Access

Weights

Starting model

RL data

Reward

RL config

Attention

Prompt format

Evaluation

Training run

Related models

Explore FriendliAI today

README

Access

Weights

Starting model

RL data

Reward

RL config

Attention

Prompt format

Evaluation

Training run

Related models