memory-lora-gemma4 / README.md
moncefem's picture
Rewrite the model card
ef8744f
|
Raw
History Blame Contribute Delete
8.93 kB
---
license: apache-2.0
base_model: google/gemma-4-E2B
library_name: peft
tags:
- hypernetwork
- lora
- code
- repository-adaptation
- gemma
pipeline_tag: text-generation
---
# Memory-LoRA β€” a hypernetwork that writes repo-specific LoRA adapters for Gemma-4-E2B
Give it a code repository; it returns a **LoRA adapter for that repository in one
forward pass**. No fine-tuning, no retrieval, and **zero repository tokens at
inference time** β€” the knowledge lives in the weights.
```
repo ──► 6-view embedding (12288-d) ──► hypernetwork ──► LoRA adapter ──► frozen Gemma-4-E2B
frozen Qwen3-Embedding-0.6B 750M params rank 16, Ξ± 32
```
On repositories **absent from the training corpus**, the generated adapter makes
correct answers roughly **200Γ— more likely** than the frozen model (βˆ’5.31 nats)
and wins **9 of 9** benchmark family/repo combinations.
| question (`pallets/click`, never trained on) | frozen Gemma-4-E2B | + generated adapter |
|---|---|---|
| What testing framework does this use? | "Jest" βœ— | **pytest** βœ“ |
| How is this project built and packaged? | "a Dockerfile" βœ— | **setuptools, setup.py** βœ“ |
| What documentation tool? | "a Gantt chart" βœ— | **Sphinx** βœ“ |
Code, demo app, and full engineering write-up:
**https://github.com/moncifem/memory-lora-gemma4** Β· deep dive in
[`docs/DEEP_DIVE.md`](docs/DEEP_DIVE.md).
---
## What's in this repository
| path | what it is |
|---|---|
| `runs/all_lora/all_lora_best_cpt.pth` | **the model.** Merged corpus (2,146 repos / 55,700 QA). `cr_val` 2.6811, `cr_test` 2.6266 |
| `runs/h200_run/head.best.pt` | prose-QA-only variant, kept as a fallback |
| `data/embeddings/*.parquet` | precomputed 12288-d repo embeddings |
| `data/qna/*.jsonl` | training QA corpora |
| `data/real_code2lora/` | RepoPeftBench (Code2LoRA's benchmark) |
| `memory_lora/` `scripts/` `app/` `deploy/` | code, mirrored from GitHub |
Earlier `sixview_v1` / `sixview_v2` weights were **removed**: they predate the
input-standardisation fix, and `sixview_v2` measurably *degrades* the base model
(see below). Their `metrics.jsonl` remain for provenance.
---
## Using the hypernetwork
### 1. Get the pieces
```bash
git lfs install
git clone https://huggingface.co/moncefem/memory-lora-gemma4
cd memory-lora-gemma4
git lfs pull --include="runs/all_lora/all_lora_best_cpt.pth"
pip install torch transformers peft safetensors pyarrow numpy
python app/engine/fetch_base_model.py # google/gemma-4-E2B, 10.25 GB
```
### 2. Embed a repository (6 views β†’ 12288-d)
The head is conditioned on a specific representation: six views of the repo β€”
call graph, architecture, git history, contracts/tests, conventions, ops β€” each
embedded by a frozen `Qwen3-Embedding-0.6B` and concatenated. **A different
embedding will not work.**
```bash
git clone --depth 80 https://github.com/pallets/click /tmp/click
python app/engine/build_embedding.py --repo /tmp/click --out /tmp/click.npy
```
### 3. Generate the adapter
```bash
python app/engine/generate_and_merge.py \
--embedding /tmp/click.npy \
--checkpoint runs/all_lora/all_lora_best_cpt.pth \
--adapter-out /tmp/click-adapter \
--no-merge
```
That writes a **standard PEFT adapter**. Use `--merged-out DIR` instead of
`--no-merge` for a self-contained merged model (for vLLM).
### 4. Use it
```python
from peft import PeftModel
from transformers import AutoModelForImageTextToText, AutoTokenizer
tok = AutoTokenizer.from_pretrained("models/gemma-4-E2B")
model = AutoModelForImageTextToText.from_pretrained("models/gemma-4-E2B")
model = PeftModel.from_pretrained(model, "/tmp/click-adapter")
prompt = "Q: What testing framework does this repository use?\nA:"
out = model.generate(**tok(prompt, return_tensors="pt"), max_new_tokens=32)
print(tok.decode(out[0], skip_special_tokens=True))
```
`with model.disable_adapter():` gives the frozen baseline for an A/B β€” the
comparison that actually matters.
### Calling the head directly
```python
import numpy as np, torch, sys
sys.path.insert(0, "app/engine")
from generate_and_merge import load_head
head, cfg, alpha = load_head("runs/all_lora/all_lora_best_cpt.pth")
emb = np.load("/tmp/click.npy").astype("float32") # 12288-d, RAW
out = head(torch.from_numpy(emb).unsqueeze(0))
# out["A"][type] -> [1, r, in_features]
# out["B"][type] -> [1, out_features, r]
```
The head standardises its input internally using statistics stored in the
checkpoint β€” **pass raw embeddings**, do not normalise them yourself.
Its update is `Ξ” = (Ξ±/r)Β·(xΒ·Aα΅€)Β·Bα΅€` with `A:[r,in]`, `B:[out,r]` β€” identically
PEFT's LoRA convention, so the output drops straight into a standard adapter.
One `(A, B)` pair per shape-qualified module type, shared across the transformer
layers of that shape (205 target modules, 14 types).
---
## Results
Cross-repo held-out loss versus the **frozen base on identical data**:
| checkpoint | `cr_val` | `cr_test` | vs base | verdict |
|---|---|---|---|---|
| sixview_v2 *(removed)* | 2.606 | β€” | **+0.198** | worse than no adapter |
| `h200_run/head.best.pt` | 2.7168 | 2.7155 | βˆ’5.191 | good |
| **`all_lora_best_cpt.pth`** | **2.6811** | **2.6266** | **βˆ’5.246 / βˆ’5.310** | **best** |
Absolute losses are *not* comparable across rows β€” the eval sets differ. The
delta against the frozen model is.
Benchmark on three unseen repositories, three task families each (FACT = the
trained Q&A format, CODE = real source-line completion, TEXT = repo prose):
| repo | FACT | CODE | TEXT | keyword accuracy |
|---|---|---|---|---|
| `psf/requests` | βˆ’10.31 | βˆ’5.38 | βˆ’2.83 | 0% β†’ 83% |
| `pallets/click` | βˆ’11.74 | βˆ’7.66 | βˆ’3.26 | 25% β†’ 75% |
| `OpenLLM-France/AudioBench` | βˆ’5.08 | βˆ’3.81 | βˆ’1.73 | 0% β†’ 0% |
**100% win rate on all nine.**
### Limits, stated plainly
- It learns a repo's **stack and conventions**, not what the project *does*. On
`requests` it answers "XML library" instead of HTTP. That is exact factual
recall, which a rank-16 LoRA structurally cannot hold β€” retrieval covers it.
- `AudioBench` keyword accuracy stayed at **0%**: projects whose identity is not
inferable from structure transfer poorly.
- Base losses of 12–16 on short gold targets inflate the deltas. The keyword
accuracy and the generations are the trustworthy evidence.
---
## The failure worth knowing about
An earlier checkpoint reached a healthy-looking eval loss of **2.606** while
being **worse than applying no adapter at all** β€” and worse than random noise of
matched magnitude.
**Cause:** 64% of every repo embedding is a constant vector shared by all
repositories (the frozen encoder's mean response to "source code"). It dominated
the trunk, which collapsed to emitting essentially the same adapter for every
repo.
| stage | mean pairwise cosine across repos |
|---|---|
| input embedding | 0.73 |
| **after trunk** | **0.978** ← discriminability destroyed |
| emitted adapter | 0.96 |
| input, centered | **0.00** ← the signal was there all along |
**Fix:** `MemoryLoRAHead.fit_input_stats()` standardises the conditioning input
with training-set statistics, stored in the checkpoint so training and inference
apply the same transform. Emitted-adapter cosine went **0.96 β†’ 0.21 in 40 steps**;
the shipped model sits at **0.32**.
**Why it went unnoticed:** training logged only the adapted loss. A number like
2.606 says nothing without the frozen-model baseline beside it. Every eval now
reports `delta_vs_baseline` and `diag/adapter_cosine`.
Before trusting any checkpoint:
```bash
python app/scripts/diagnose_head.py --job <id> --checkpoints runs/<run>/head.best.pt
```
It scores against `none`, `random` noise of matched scale, and a `zero-B` control
that must reproduce the baseline *exactly*. A head that cannot beat `random` has
not learned the mapping.
---
## Training
Reproduce or continue on a single GPU:
```bash
python3 deploy/h200/preflight.py # validates before spending GPU time
GATE=1 bash deploy/h200/train_h200.sh # ~20-min go/no-go
bash deploy/h200/train_h200.sh # full run
```
`preflight.py` fails fast on git-LFS pointers masquerading as data,
pre-standardised embeddings, a head without `fit_input_stats`, and missing eval
splits β€” then auto-tunes the largest micro-batch that fits.
---
## Credits
Reimplements and extends **Code2LoRA** (arXiv 2606.06492) β€” a static hypernetwork
mapping a repository embedding to a LoRA adapter β€” retargeted to
`google/gemma-4-E2B`, extended from single-view code completion to a six-view
representation, and wrapped in a serving stack that speaks the OpenAI and
Anthropic APIs. Training data includes RepoPeftBench from that work.
Base model: [`google/gemma-4-E2B`](https://huggingface.co/google/gemma-4-E2B) Β·
Encoder: [`Qwen/Qwen3-Embedding-0.6B`](https://huggingface.co/Qwen/Qwen3-Embedding-0.6B)