Hugging Face's logo

mokshjain
/
pico

Text Generation
JAX
Safetensors
English
agentic
tool-use
reinforcement-learning
grpo
nanochat
from-scratch
pico
Model card Files
xet
 Community

Pico: A 48M-Parameter Agentic Language Model Trained from Scratch in JAX

Pico is a 48.2M-parameter decoder-only transformer trained from scratch in pure JAX (no Flax, no optax), following Karpathy's nanochat recipe and extending it to build a complete tool-calling agent from pretraining through on-policy reinforcement learningβ€”without any distillation or synthetic LLM-generated data.

Highlights

  1. Agentic Pretraining (10B tokens, final_10b_mila):
    • Trained on 80% ClimbMix web text + 20% code (shell, python, sysadmin_qa), with the final 15% annealed on high-quality reference text (finepdfs_edu, fineweb_edu_hq, finemath, wikipedia) + 9% agentic tool-use trajectories (terminal_nemotron, terminal_openthoughts, exec_traces, nl2cmd_chat).
    • Reaches 0.1273 full CORE (beating GPT-2 124M's 0.114 with 40% of the parameters) and 0.304 TOOL (72% 4-way NLβ†’bash accuracy, 100% well-formed tool calls on terminal episodes).
  2. Free Long-Context Extrapolation via Global NoPE (final_10b_mila_16k_lr01):
    • Uses an SSSL attention pattern: 7 local layers with a fixed 512-token sliding window + RoPE (base 10k), and 3 global causal layers (layers 3, 7, 9) with no positional encoding (NoPE).
    • Extrapolates zero-shot to 32k passkey retrieval (100% vs 20% for all-RoPE). A short 315M-token extension stage at 16,384 context lowers 8–16k FinePDFs-Edu loss from 0.918 to 0.847–0.853 bpb and lifts 4-needle retrieval at 32k from 0.18 to 0.75–0.76 (full CORE 0.1321–0.1343).
  3. Post-Training Without Distillation (mt_tool20c β†’ rl_web5 β†’ rl_web7):
    • Mid-training (0.9B tokens at 4k): Teaches multi-turn chat, math-as-python, and web/function-calling tool formats (bash, python, Wikipedia search/fetch, weather, declared JSON functions, and irrelevant-function refusals).
    • On-Policy GRPO RL (rl_web7): Synchronous GRPO in fp32 over 8 real execution environments with exact trainer–engine logprob parity ($\max |\Delta \log p| \le 2 \times 10^{-4}$), std-normalized advantages, and reference-guided exploration (--guide, injecting ~230 programmatic reference solutions only into zero-pass groups to break the zero-variance plateau).

Available Checkpoints

Every pipeline stage is included in this repository under runs/<run_name>/ckpt.npz (weights + config metadata, ~193 MB each), alongside the root model.safetensors for rl_web7:

Run Name Stage Context Training Key Metrics Path
rl_web7 Final Multi-Tool RL (Stage 2) 4k (16k base) +300 GRPO steps (8 envs incl. NL2SH) Terminal 1.00 (clean unseen 0.333), NL2SH 0.303, GSM8K 0.145, QA w/ search 0.210, Weather 0.635, BFCL 0.33, CORE 0.131 runs/rl_web7/ckpt.npz & model.safetensors
rl_web5 Multi-Tool RL (Stage 1) 4k (16k base) 600 GRPO steps (7 envs, --guide) Terminal 1.00 (clean unseen 0.320), QA w/ search 0.201, Weather 0.560, BFCL 0.33, CORE 0.125 runs/rl_web5/ckpt.npz
mt_tool20c Chat + Tool Mid-Training 4,096 0.9B tokens (LR Γ—0.5, anneal mask) Terminal 0.138, GSM8K 0.092, BFCL 0.32 (simple 0.41, irrelevance 0.58), CORE 0.123 runs/mt_tool20c/ckpt.npz
rl_mt05b Terminal + GSM8K RL 4k (16k base) 300 GRPO steps (terminal:1,gsm8k:1) Held-out terminal 0.558, GSM8K 0.165, NL2SH 0.267, CORE 0.120 runs/rl_mt05b/ckpt.npz
mt_lr05 Chat Mid-Training 4,096 0.9B tokens (LR Γ—0.5, anneal mask) Held-out terminal 0.167, GSM8K 0.093, NL2SH 0.283, CORE 0.121 runs/mt_lr05/ckpt.npz
final_10b_mila_16k_lr01 16k Context Extension (LR Γ—0.1) 16,384 315M tokens (--long_doc_rows=True) Full CORE 0.1343, TOOL 0.297, 8–16k bpb 0.853, 4-needle@32k 0.76 runs/final_10b_mila_16k_lr01/ckpt.npz
final_10b_mila_16k 16k Context Extension (LR Γ—0.3) 16,384 315M tokens (--long_doc_rows=True) Full CORE 0.1321, TOOL 0.293, 8–16k bpb 0.847, 4-needle@32k 0.75 runs/final_10b_mila_16k/ckpt.npz
final_10b_mila 10B Pretrained Base 2,048 10B tokens (38,146 steps, WSD + anneal) Full CORE 0.1273, TOOL 0.304, ClimbMix val bpb 0.9090, shell_cmd bpb 0.707 runs/final_10b_mila/ckpt.npz

End-to-End Evaluation Summary

All agent evaluations run greedily with real tool execution in isolated sandboxes (agent_eval.py). Open-domain QA uses the local 20.8M-passage Wikipedia BM25+rerank index (search_v2), and weather uses the strict single-temperature scorer:

Metric final_10b_mila_16k_lr01 (Base) mt_lr05 rl_mt05b mt_tool20c rl_web5 rl_web7 (Final)
Full CORE (22 tasks, centered) 0.134 0.121 0.120 0.123 0.125 0.131
Terminal (312 held-out instances) 0.074 0.167 0.558 0.138 1.000 1.000
Terminal_new (6 clean unseen task types) β€” 0.000 0.250 β€” 0.320 0.333
NL2SH Execution Match (300) 0.280 0.283 0.267 β€” 0.223 0.303
GSM8K with Python Tool (1,319) 0.002 0.093 0.165 0.092 0.137 0.145
QA with Wikipedia Search (2,000) β€” 0.056 0.063 0.062 0.201 0.210
QA Closed-Book (2,000) β€” 0.048 0.028 0.051 0.035 0.044
Weather Value (strict) / No-Tool β€” 0 / 1.0 0 / 1.0 0.10 / 0.15 0.56 / 0.98 0.635 / 0.99
BFCL Accuracy (Simple / Irrelevance) β€” 0.20 (0 / 1.0) 0.19 (0 / 0.96) 0.32 (0.41 / 0.58) 0.33 (0.48 / 0.44) 0.33 (0.46 / 0.48)
Chat No-End Rate (CLI Sampling) β€” 0.02 0.02 0.04 0.02 0.01

Quickstart

1. Interactive Terminal CLI with Tools

Clone the repository and run cli.py (automatically downloads rl_web7 and tokenizer v2 from mokshjain/pico on first run; works on CPU or GPU):

uv sync
uv run python cli.py --run_name=rl_web7
# Use live Wikipedia & Open-Meteo APIs instead of the local index / simulator:
uv run python cli.py --run_name=rl_web7 --live

2. Loading Any Stage in Python (via Repo) 

import checkpoint
from tokenizer import Tokenizer
from engine import Engine
from agent import Agent, Episode
from cli import CLI_SAMPLING
from tools import Sandbox
import chat

# Automatically downloads from mokshjain/pico if not cached locally
params, cfg, meta = checkpoint.load_model(checkpoint.run_path("rl_web7"), dtype="float32")
tok = Tokenizer.load(cfg.tokenizer)

engine = Engine(params, cfg, n_slots=1, max_len=4096)
agent = Agent(engine, tok, CLI_SAMPLING)

with Sandbox() as sandbox:
    seq_tokens = agent.prompt([
        chat.system(chat.tools_prompt(["bash", "python", "search", "fetch", "weather"])),
        chat.user("What is 2**32 - 1, and is it prime? Check with python."),
    ])
    from engine import Sequence
    ep = Episode(Sequence(seq_tokens), sandbox)
    agent.run([ep])
    print(tok.decode(ep.seq.tokens))

3. Loading model.safetensors Directly 

import json, pickle
from huggingface_hub import hf_hub_download
from safetensors.numpy import load_file

tensors = load_file(hf_hub_download("mokshjain/pico", "model.safetensors"))
config = json.load(open(hf_hub_download("mokshjain/pico", "config.json")))
enc = pickle.load(open(hf_hub_download("mokshjain/pico", "tokenizer.pkl"), "rb"))

print("Loaded tensors:", len(tensors), "total params:", sum(x.size for x in tensors.values()))

Architecture Details

Component Shape Parameters
Token embedding wte (untied) 16,384 Γ— 512 8,388,608
Per layer: wq, wk, wv, wo 4 Γ— 512 Γ— 512 1,048,576
Per layer: MLP w1, w2 (ReLUΒ², 4Γ— expansion) 512 Γ— 2048 + 2048 Γ— 512 2,097,152
10 transformer blocks (SSSL window pattern) 31,457,280
Output head lm_head (untied) 512 Γ— 16,384 8,388,608
Per-layer residual / input-mixing scalars (resid_lambdas, x0_lambdas) 2 Γ— 10 20
Total 48,234,516

Forward pass (model.py): x0=RMSNorm(wte[t]),x←x0for each layer i:x←λiresid x+Ξ»ix0 x0x←x+Attni(RMSNorm(x)),x←x+W2 ReLU(W1 RMSNorm(x))2logits=15tanh⁑ ⁣(RMSNorm(x) lm_head/15) \begin{aligned} x_0 &= \mathrm{RMSNorm}(\mathrm{wte}[t]), \qquad x \leftarrow x_0\\ \text{for each layer } i:\quad x &\leftarrow \lambda^{\text{resid}}_i\, x + \lambda^{x_0}_i\, x_0\\ x &\leftarrow x + \mathrm{Attn}_i(\mathrm{RMSNorm}(x)), \qquad x \leftarrow x + W_2\,\mathrm{ReLU}(W_1\,\mathrm{RMSNorm}(x))^2\\ \text{logits} &= 15 \tanh\!\big(\mathrm{RMSNorm}(x)\, \mathrm{lm\_head} / 15\big) \end{aligned}

Chat & Tool Special Tokens (chat.py) 

<|bos|>[<|system_start|>...<|system_end|>]<|user_start|>task<|user_end|>
<|assistant_start|>optional text<|tool_start|>bash
<command><|tool_end|><|output_start|><output>
[exit code N]<|output_end|> ... final answer<|assistant_end|>
Downloads last month
235
Safetensors
Model size
48.2M params
Tensor type
F32
Β·
Inference Providers NEW
Text Generation
This model isn't deployed by any Inference Provider. πŸ™‹ Ask for provider support