Grafted Gemma-4-E4B: GitHub PR Task (carlesoctav/grafting-gh_pr_task)

This model is created using Grafting as described in:

"Off-Policy Merging Beats On-Policy Self-Distillation for Continual Learning"
Chen Henry Wu, Thomas Zhang, & Aditi Raghunathan (CMU, 2026)
Paper: arXiv:2610.05872
Project: ar-forum.github.io/grafting-web

---

1. How Grafting Works

Rather than fine-tuning the post-trained reasoning model directly (which causes catastrophic forgetting on prior capabilities due to narrow off-policy gradient routing), Grafting adapts an earlier donor checkpoint—specifically the pretrained base model—on the target reasoning traces, and transplants the task vector into the deployed post-trained model:

$$\theta (\lambda ,\rho )={\theta }_{\text{post}}+\lambda \cdot M_{\rho }\odot ({\theta }_{\text{donor-ft}}-{\theta }_{\text{donor}})\backslash theta(\backslash lambda,\; \backslash rho)\; =\; \backslash theta\_\{\backslash text\{post\}\}\; +\; \backslash lambda\; \backslash cdot\; M\_\backslash rho\; \backslash odot\; (\backslash theta\_\{\backslash text\{donor-ft\}\}\; -\; \backslash theta\_\{\backslash text\{donor\}\})$$

In this model:

• Target Post Model ($\theta_{\text{post}}$): google/gemma-4-E4B-it (instruction/reasoning tuned checkpoint)
• Donor Base Model ($\theta_{\text{donor}}$): google/gemma-4-e4b (pretrained base checkpoint)
• Task Vector Update: Trained with 4-bit QLoRA on GitHub PR reasoning traces (sft_dataset_inkling_f2p.jsonl).
• Grafting Scale ($\lambda$): 1.0
• Protection ($M_\rho$): High-sensitivity embedding coordinates (Layer 0 token embeddings and 42-layer Engram embeddings totaling 3.49B parameters) remain frozen at their post-trained values ($\Delta W = 0$), preserving 100% of the reasoning control token representations (<|think|>, <|turn>, <|channel>, <|tool_call>).
• Linear Projections: All 258 attention and MLP weight matrices transplanted via $\Delta W = \frac{\alpha}{r}(B @ A)$.

---

2. Donor Training Details

Parameter	Value
Base Model	google/gemma-4-e4b
LoRA Rank $r$	16
LoRA Alpha $\alpha$	16 (scaling factor = 1.0)
Target Modules	q_proj, k_proj, v_proj, o_proj, gate_proj, up_proj, down_proj (258 modules)
Target Sequence Length	24,000 tokens (packed with BFD)
Training Steps	12 steps across 3 epochs (4x NVIDIA L4 DDP)
Loss Trajectory	2.926 -> 1.330 (Ep 1) -> 1.159 (Ep 2) -> 0.9242 (Ep 3)
Mean Loss	1.434
Final Gradient Norm	0.4636 (clean sub-1.0 unclipped convergence)

---

3. Usage with Transformers

import torch
from transformers import AutoTokenizer, Gemma4ForConditionalGeneration

model_id = "carlesoctav/grafting-gh_pr_task"

tokenizer = AutoTokenizer.from_pretrained(model_id)
model = Gemma4ForConditionalGeneration.from_pretrained(
    model_id,
    torch_dtype=torch.bfloat16,
    device_map="auto"
)

messages = [
    {"role": "user", "content": "Analyze the following repository diff and write a pull request description."}
]

prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
inputs = tokenizer(prompt, return_tensors="pt").to(model.device)

outputs = model.generate(
    **inputs,
    max_new_tokens=2048,
    temperature=0.7,
    top_p=0.95
)

print(tokenizer.decode(outputs[0][inputs.input_ids.shape[1]:], skip_special_tokens=True))

---

4. Citation

@article{wu2026grafting,
  title   = {Off-Policy Merging Beats On-Policy Self-Distillation for Continual Learning},
  author  = {Chen Henry Wu and Thomas Zhang and Aditi Raghunathan},
  journal = {arXiv preprint arXiv:2610.05872},
  year    = {2026}
}