open-lm-3b-202407-stage3-docuse (UPDATED → diverse GRPO checkpoint-300, best all-around)

3B OpenLM (knowledge cutoff 2024-07, context-extended to 8k) for faithful retrieval-augmented document use: read retrieved documents, reason in <think>...</think>, commit to a short answer. Pipeline = diversity-curated RAFT SFT (variable doc count 1–43 + short/long docs) → GRPO (binary correctness reward, β=0.03 KL). This revision is robust across the whole distribution — it handles 1 doc or 40 docs, short intros or long passages, and (unlike the earlier ckpt-150) does not fail to answer on long stacked contexts.

⚠️ READ BEFORE EVAL — three things that will silently break your numbers

1. DO NOT use vLLM. This custom arch (OpenLM, LayerNorm + QK-norm) is unfaithful under vLLM (~15% vs ~33% on the same data — vLLM mishandles the qk-norm). Use HF transformers.generate.
2. trust_remote_code=True required (custom OpenLMForCausalLM in modeling_open_lm_hf.py); use attn_implementation="sdpa".
3. EOS = BOS = <|endoftext|> = id 0. Set eos_token_id=0, pad_token_id=0, tokenizer.padding_side="left".

Exact prompt format (must match)

User content:

Question: {question}

Retrieved documents:
[1] {title_1}: {text_1}
[2] {title_2}: {text_2}
...

Read the documents and reason step by step inside <think>...</think>, then give a SHORT final answer on a new line as "Answer: <answer>". If the documents are insufficient or only partially relevant, combine them with your own knowledge and reasoning to give your best answer. Always commit to a concrete answer — never refuse or say there is no information.

Then wrap with the chat template (add_generation_prompt=True) → Human: {user_content}<|endoftext|>\nAssistant:. Output: <think> reasoning citing the docs </think>\nAnswer: <short answer>. Extract = text after Answer:.

Decoding (to reproduce our numbers)

• single-answer EM: do_sample=True, temperature=0.6, top_p=0.95, max_new_tokens>=512, eos_token_id=0.
• pass@k: temperature=0.8.

Minimal eval snippet

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
MP = "mattwang123/open-lm-3b-202407-stage3-docuse"
tok = AutoTokenizer.from_pretrained(MP, trust_remote_code=True); tok.padding_side="left"
if tok.pad_token is None: tok.pad_token = tok.eos_token
m = AutoModelForCausalLM.from_pretrained(MP, trust_remote_code=True, dtype=torch.bfloat16,
        attn_implementation="sdpa").cuda().eval()

def build_user_content(docs, question):  # docs: list of {"title","text"}
    lines = [f"Question: {question}", "", "Retrieved documents:"]
    for i, d in enumerate(docs):
        lines.append(f"[{i+1}] {d['title']}: {d['text']}")
    lines += ["", ('Read the documents and reason step by step inside <think>...</think>, then give a '
        'SHORT final answer on a new line as "Answer: <answer>". If the documents are insufficient or '
        'only partially relevant, combine them with your own knowledge and reasoning to give your best '
        'answer. Always commit to a concrete answer — never refuse or say there is no information.')]
    return "\n".join(lines)

prompt = tok.apply_chat_template([{"role":"user","content":build_user_content(docs, question)}],
                                 tokenize=False, add_generation_prompt=True)
enc = tok(prompt, return_tensors="pt").to(m.device)
g = m.generate(**enc, do_sample=True, temperature=0.6, top_p=0.95, max_new_tokens=512,
               eos_token_id=0, pad_token_id=0)
print(tok.decode(g[0, enc.input_ids.shape[1]:], skip_special_tokens=True))

Our internal results (curated_eval, N=120/slice, with-docs EM = exact+judge)

slice	EM	exact_EM	format_miss	notes
core_multi (~10 docs)	0.567	0.38	0.00	realistic multi-hop
withhold_hop (1 gold dropped → use prior)	0.392	0.28	0.00	doc+prior fusion
long_stack (~43 docs, ~6k tok)	0.358	0.25	0.00	now works (prev ckpt-150 was 0.10 / format_miss 0.53)
closed-book core_multi (no docs)	0.108	—	—	with-docs 0.567 vs 0.108 ⇒ genuinely uses docs

pass@8: core_multi 0.82, long_stack 0.66 (T=0.8, N=50). All training data content-filtered leakage-free (pre-2013-equivalent) for the 2013–2024 timestamp suite.