Skip to content

Latest commit

Β 

History

6 Commits

Folders and files

NameName
Last commit message
Last commit date
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 

Repository files navigation

SEAL: Self-Adapting Language Models on an 8GB GPU

Python PyTorch PEFT License Single file Read on Medium

SEAL: a language model that writes its own study notes

A minimal, single-file implementation of SEAL (Self-Adapting Language Models, arXiv:2506.10943): the knowledge-incorporation variant trained with ReST-EM (rejection sampling), not PPO. Built to run on one RTX 4060 (8 GB) and to degrade gracefully to CPU.

A normal LLM is frozen after training. SEAL teaches a model to write its own fine-tuning data (synthetic "study notes"), test whether those notes measurably improve its answers, and reinforce the habit of writing good ones. This repo makes that loop correct, readable, and runnable end to end on low-end hardware.

πŸ“– Full write-up: How I Implemented Self-Adapting Language Models with LoRA and QLoRA on an RTX 4060 (Medium).


Highlights

  • Fits in 8 GB. One frozen base loaded once in 4-bit (QLoRA), never reloaded.
  • Two LoRA adapters on one shared base. A persistent generator and a transient inner scratch adapter, no second copy of the model.
  • Strict separation of the two update types. requires_grad gating guarantees the generator improves only through the filtered outer SFT.
  • Hand-rolled inner loop. Bare forward / backward / step, no Trainer overhead for 3 to 5 step updates.
  • Graceful degradation. Auto-disables 4-bit without CUDA/bitsandbytes and falls back to an ungated model if a gated one is not yet approved.
  • ~500 lines, heavily commented, so the file doubles as reference notes.

Requirements

  • Python 3.10+
  • An NVIDIA GPU with ~8 GB VRAM for the 4-bit path (a 4060 is the reference). CPU works too, just slower.
  • Dependencies in requirements.txt: torch, transformers, peft, accelerate, bitsandbytes (GPU 4-bit), datasets (SQuAD), python-dotenv.

For the GPU path, install a CUDA build of PyTorch (e.g. pip install torch --index-url https://download.pytorch.org/whl/cu124); a plain pip install torch gives a CPU-only build and the pipeline will report 4bit=False.


Quickstart

python -m venv .venv
.venv\Scripts\activate            # Windows  (macOS/Linux: source .venv/bin/activate)
pip install -r requirements.txt

# 1) Prove the loop in seconds on a tiny model (CPU is fine)
python seal_pipeline.py --smoke

# 2) A real run on real data (SQuAD), on an 8GB GPU
python seal_pipeline.py --squad --passages 10 --outer-iters 10 --n-edits 8 \
    --inner-steps 5 --edit-tokens 200 --max-seq-len 768 --save seal_generator

Authentication and gated models

  • .env for your token. Copy .env.example to .env and set HF_TOKEN=hf_.... The script auto-loads it at startup via python-dotenv, so the secret stays out of the code and shell history. .env is gitignored.
  • Gated models need approval. meta-llama/Llama-3.2-1B-Instruct requires both a token and manual access approval on its model page. A valid token alone still returns 403 until you are on the authorized list.
  • Automatic fallback. If the configured model is inaccessible (403 / access pending), resolve_model_name() probes access up front and transparently switches to Config.FALLBACK_MODEL (Qwen/Qwen2.5-0.5B-Instruct, ungated), so the pipeline never crashes. Once approved, the same command uses Llama again with no code change. Set Config.FALLBACK_MODEL = None to disable.

Ungated drop-ins with the same LoRA targets (no approval needed):

python seal_pipeline.py --model Qwen/Qwen2.5-0.5B-Instruct     # CPU-friendly
python seal_pipeline.py --model Qwen/Qwen2.5-1.5B-Instruct
python seal_pipeline.py --model HuggingFaceTB/SmolLM2-1.7B-Instruct

How SEAL works

The single most important invariant: the generator improves ONLY through the filtered outer SFT. The inner adaptation is transient and never leaks back. The code enforces this with two separate LoRA adapters plus requires_grad gating.

The two nested loops of SEAL

Text version of the same flow:

                    β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€ OUTER loop (persistent, ReST-EM) ───────────────────────┐
                    β”‚                                                                                β”‚
  passage ─────────┼──▢ [generator adapter] ──sample N──▢  self-edit₁ … self-editβ‚™                   β”‚
                    β”‚                                          β”‚                                     β”‚
                    β”‚            β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€ INNER loop (transient, thrown away) ──────────────┐  β”‚
                    β”‚            β”‚  reset inner β–Ά SFT inner on editα΅’ β–Ά eval QA β–Ά reward, DISCARD  β”‚  β”‚
                    β”‚            β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜  β”‚
                    β”‚                                          β”‚                                     β”‚
                    β”‚                        reward = adapted_acc βˆ’ base_acc                         β”‚
                    β”‚                                          β”‚                                     β”‚
                    β”‚                     keep edits with reward > 0   (rejection sampling)          β”‚
                    β”‚                                          β”‚                                     β”‚
                    β”‚             SFT the GENERATOR on the winning (passage β†’ edit) pairs β—€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜
                    β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜

Inner loop, the reward signal (transient)

For one self-edit, take a few gradient steps of the ordinary causal-LM loss over the edit's tokens on a scratch adapter:

$$\theta' = \theta - \eta , \nabla_\theta , L(\text{self-edit})$$

where $L$ is the causal-LM loss $L = -\frac{1}{T}\sum_{t} \log p_\theta(x_t \mid x_{<t})$.

Then score how useful that transient $\theta'$ was, and throw it away:

$$\text{reward} = \text{acc}(\theta') - \text{acc}(\theta)$$

Outer loop, ReST-EM (persistent)

ReST-EM approximately maximizes expected reward without PPO, using filtered behavior cloning:

  • E-step (rejection sampling): keep only samples with reward > 0.
  • M-step (maximization): ordinary maximum-likelihood SFT of the generator on the kept edits.

$$\theta_{\text{gen}} \leftarrow \arg\max_{\theta}\ \sum_{(\text{passage},,\text{edit})\ :\ \text{reward}>0} \log p_\theta(\text{edit}\mid \text{passage})$$

No value function, no importance weights, no clipping. That simplicity is exactly why SEAL uses ReST-EM here.


Why the architecture is efficient (the whole point)

One frozen 4-bit base shared by two small LoRA adapters

Trick What it saves Where
Load base once, in 4-bit NF4 base weights at ~ΒΌ the memory, never reloaded build_model()
Two LoRA adapters on one frozen base no second copy of the 1B weights per phase get_peft_model(...) + add_adapter("inner")
Reset inner in place by re-zeroing LoRA B no delete/rebuild of modules per passage reset_inner_adapter()
Hand-rolled inner train loop removes Trainer per-call setup that dominates a 3-step update inner_adapt()
requires_grad gating guarantees the two updates stay separate set_only_trainable()
Gradient checkpointing + short seq len + small max_new_tokens activation VRAM build_model(), Config
No mid-run merge you cannot merge cleanly into 4-bit, only optional merge at the very end run_seal()

Why zeroing B resets the adapter. A LoRA layer adds Ξ”W = (Ξ±/r)Β·BΒ·A, with B initialized to 0. So B = 0 ⟹ Ξ”W = 0 ⟹ the adapter is the identity and the model behaves exactly like the frozen base. Zeroing B between passages restarts adaptation cleanly, with no rebuild.


CLI reference

Flag Default Meaning
--smoke off Use the tiny smoke-test model (sshleifer/tiny-gpt2), CPU-ok
--model NAME Config.MODEL_NAME Override the base model
--squad off Use the SQuAD loader instead of the toy passages
--passages N 10 How many SQuAD passages to train on (with --squad)
--outer-iters N Config.OUTER_ITERS ReST-EM iterations
--n-edits N Config.N_SELF_EDITS Self-edits sampled per passage
--inner-steps N Config.INNER_STEPS Transient inner SFT steps
--outer-steps N Config.OUTER_STEPS Generator SFT epochs over winners
--edit-tokens N Config.EDIT_MAX_NEW_TOKENS Max tokens per self-edit (raise for SQuAD)
--max-seq-len N Config.MAX_SEQ_LEN Max training/eval sequence length (raise for long passages)
--save DIR none Save the trained generator adapter
--merge off Also merge + save a standalone model (non-4bit only)

All Config knobs live at the top of seal_pipeline.py and can be edited directly.


Reading the output

--- Passage 1/10 ---
  base_acc = 0.000
  edit 0: adapted_acc=0.333 reward=+0.333 KEEP  | '1. Early Cretaceous fossils of ctenophores were discovered…'
  edit 1: adapted_acc=0.667 reward=+0.667 KEEP  | '* Origin of ctenophores: Ctenophores are believed to have…'
  edit 6: adapted_acc=0.000 reward=+0.000 drop  | 'These are some short, self-contained factual statements…'

  [outer] iter 1 summary: kept 19/80 edits (kept_frac=0.24), mean_reward=+0.083
  [outer] generator updated via filtered SFT (persistent).
  • Per-edit rewards and KEEP/drop show the filter working.
  • The training curve printed at the end tracks kept_frac per iteration. If SEAL is learning, it should trend upward.
  • With the tiny smoke model, winners may be empty, which is fine and expected. To force the persistent-update branch, set Config.REWARD_THRESHOLD = -1e-9 (keep reward >= 0).

Example results (Llama-3.2-1B, 4-bit, RTX 4060)

Closed-book QA on SQuAD passages (the model answers with no passage in the prompt, so base_acc starts near 0 by design). The generator visibly improves across iterations:

Iteration kept_frac mean_reward
1 0.24 +0.083
2 ~0.5 (kept-rate roughly doubled) higher

Kept-fraction trending upward across training iterations

Concretely, passage 1 (ctenophores) went from 6/8 kept in iteration 1 to 8/8 in iteration 2, and passage 3 (which kept nothing in iteration 1) started producing winners once the generator had been updated. Absolute accuracy stays modest by design; the point is that the reward and kept-fraction climb, i.e. the model teaches itself to write more useful notes. Trainable parameters: 1.7M of 1.24B (0.14%).


Repository layout

seal_pipeline.py     # the whole pipeline, heavily commented
requirements.txt     # dependencies with GPU/CPU install notes
.env.example         # copy to .env and add HF_TOKEN
images/              # diagrams and figures used in this README
LICENSE              # MIT
README.md            # this file

Citation

If you use this, please cite the original paper:

@article{seal2025,
  title  = {Self-Adapting Language Models},
  year   = {2025},
  eprint = {2506.10943},
  archivePrefix = {arXiv}
}

This repository is an independent, educational reimplementation and is not affiliated with the paper's authors.


License

Released under the MIT License. Add a LICENSE file if you have not already (the badge above assumes MIT).

About

No description, website, or topics provided.

Resources

Stars

1 star

Watchers

0 watching

Forks

Releases

Packages

Contributors

Languages