PCSS-Granite-4.1-3B

GitHub Repository | Blog Post | Reproduction Notebook

PCSS-Granite-4.1-3B is an experimental checkpoint demonstrating that fine-tuning small base models on pure logic deduction traces can elicit substantial out-of-domain reasoning generalization. Starting from ibm-granite/granite-4.1-3b-base, this model was fine-tuned on 500 5x5 zebra puzzles from the tamewild/instruct5 dataset, which contains zero mathematical data.

For fine-tuning on these long-form logic traces, the model was trained using PCSS (Per-Example Calibrated Sigmoid Scaler), an adaptive loss scaler derived from KTO. The training process leveraged MiSS (Matrix Shard Sharing) for parameter-efficient fine-tuning (PEFT), completing in approximately 24 minutes on a single NVIDIA H200 NVL GPU. (Note: The accompanying reproduction notebook was run on an H100 SXM, producing comparable results but with minor variance due to training non-determinism).

Experimental Checkpoint Notice: PCSS-Granite-4.1-3B is an exploratory research artifact developed solely to study out-of-domain reasoning generalization. Unlike PCSS-Qwen3.5-9B, this model has a significantly narrower capability profile and lacks general conversational robustness. It is not intended for chat, multi-turn dialogue, or general assistant tasks.

---

Benchmark Results

Our fine-tuned checkpoint was evaluated via vLLM (0.19.1) using greedy decoding with 16,384 max completion tokens. All reported scores are Pass@1 estimates. To account for runtime batching non-determinism in vLLM, scores for MATH-500 and 5x5 Zebra were averaged over 3 evaluation runs, and AIME 2025 was averaged over 6 runs.

For comparison, the official post-trained model (ibm-granite/granite-4.1-3b) was evaluated on the same zero-shot Chain-of-Thought prompts using its official chat template.

Benchmark	Granite 4.1 3B Base (PCSS + MiSS)	Granite 4.1 3B Official Instruct
MATH-500	77.73% (+11.13% delta)	66.60%
AIME 2025	19.44% (+12.77% delta)	6.67%
5x5 Zebra	20.00%	—

---

Training Configuration & Hyperparameters

For full training details and mathematical derivations, please refer to the Reproduction Notebook and Blog Post.

• Base Model: ibm-granite/granite-4.1-3b-base
• Dataset: tamewild/instruct5 (500 5x5 zebra puzzles)
• PEFT Method: MiSS (Matrix Shard Sharing), rank r=512
• Loss Function: PCSS (beta=0.65, peak_scale=5.0)
• Optimizer: 8-bit AdamW (LR: 5e-6, beta2: 0.99994, Weight Decay: 0.01)
• LR Scheduler: Constant learning rate with 1-epoch linear warmup
• Training Duration: 5 epochs at batch size 1 (2,500 steps total)

---

Usage & Inference

Serving with vLLM

The checkpoint bundles the pre-configured chat_template.jinja. To serve the model via vLLM:

vllm serve tamewild/PCSS-Granite-4.1-3B \
  --max-model-len 32000 \
  --generation-config vllm \
  --host 127.0.0.1 \
  --port 18000 \
  --gpu-memory-utilization 0.90

Prompt Formats

Evaluations were conducted using the following zero-shot Chain-of-Thought templates:

1. Mathematical Benchmarks (MATH-500, AIME 2025)

{problem}

Please reason step by step, and put your final answer within \boxed{}

2. Logic Grid Benchmarks (Zebra Puzzles)

{problem}

Provide the solution grid as the final answer.

Please reason step by step, and put your final answer within \boxed{}