# TensorPEFT

   ---
   license: apache-2.0
   library_name: tensorpeft
   tags:
   - peft
   - fine-tuning
   - tensor-decomposition
   - lora-alternative
   ---
     
Parameter-efficient fine-tuning with **tensor gradient-pursuit adapters**. TensorPEFT is a drop-in
alternative to LoRA for Hugging Face models, designed for very small trainable-parameter budgets.

LoRA gives every weight matrix its own low-rank update. TensorPEFT instead stacks all
equal-shaped matrices of the network into one tensor, for example
`layers x {q,k,v,o} x d_in x d_out`, and learns a few **rank-one CP atoms** of that tensor:

    dW = sum_i  a_i  u_i(layer) o u_i(role) o u_i(in) o u_i(out)

Each atom couples every layer and projection while costing only
`L + P + d_in + d_out + 1` parameters. The atoms are not random. They are the leading rank-one
CP components of the negative gradient tensor (higher-order power iteration with deflation),
computed from **one gradient probe** before training. The coefficients start at zero, so the
model is unchanged at step 0, and every atom starts on a steepest-descent direction. After that,
training is ordinary: any optimiser and the HF `Trainer` both work.

## Install

    pip install tensorpeft          # once published; until then: pip install .

It requires `torch>=2.1`, `transformers>=4.40`, `safetensors` and `huggingface_hub`.

## Quick start

```python
from transformers import AutoModelForCausalLM
from tensorpeft import TensorPEFTConfig, get_tensorpeft_model, TensorPEFTModel

base = AutoModelForCausalLM.from_pretrained("gpt2")
model = get_tensorpeft_model(base, TensorPEFTConfig(atoms_per_group=4))

model.init_atoms(train_dataloader, n_batches=4)   # the one step LoRA does not have
model.print_trainable_parameters()                # trainable params: 37,036 || ... || 0.0298%

# train as usual: AdamW, lr ~3e-3, short warm-up, weight_decay 0  (or TensorPEFTTrainer below)

model.save_pretrained("my-adapter")               # adapter_config.json + adapter_model.safetensors
model.push_to_hub("your-name/my-adapter")

base = AutoModelForCausalLM.from_pretrained("gpt2")
model = TensorPEFTModel.from_pretrained(base, "your-name/my-adapter")   # local dir or Hub id
merged = model.merge_and_unload()                 # plain HF model, zero inference overhead
```

`train_dataloader` can yield dicts (`input_ids`, optional `attention_mask` / `labels`) or plain
input-id tensors. For causal LMs, labels default to the inputs.

### With the Hugging Face Trainer

```python
from tensorpeft import TensorPEFTTrainer
trainer = TensorPEFTTrainer(model=model, args=training_args, train_dataset=ds, data_collator=collator)
trainer.train()             # checkpoints contain only the adapter
```

`TensorPEFTTrainer` is the stock `Trainer` except that it saves adapters instead of full model
weights. A complete script is in `examples/train_with_trainer.py`.

## Configuration

| `TensorPEFTConfig` field | default | meaning |
|---|---|---|
| `atoms_per_group` | 4 | rank-one atoms per tensor group; this sets the budget |
| `target_modules` | `"all-linear"` | or a list of name suffixes, e.g. `["q_proj", "v_proj"]` (`"c_attn.q"` for GPT-2) |
| `exclude_modules` | `lm_head`, embeddings | never adapted |
| `split_fused_qkv` | True | split GPT-2's fused `c_attn` into q, k and v |
| `init_batches` | 4 | batches used by `init_atoms` |
| `hopm_iters` | 30 | power-iteration sweeps per atom |
| `alpha_scale` | 33.3 | coefficient reparametrisation so that one learning rate (~3e-3) fits coefficients and factors |

**Budget.** Groups are formed automatically by matrix shape. GPT-2 small has three groups:
`{q,k,v,o}` (12x4x768x768), MLP-in (12x768x3072) and MLP-out (12x3072x768). One atom per group
is 9,259 parameters, four per group is 37,036, and six per group is 55,554.
`print_trainable_parameters()` shows the exact count for any model.

**Recipe.** Use AdamW at lr 3e-3 with about 10 warm-up steps, then a constant schedule and
weight decay 0. Use all atoms from the start: adding atoms gradually during training was clearly
worse in our tests.

## Results

Setup: GPT-2 124M (English base) fine-tuned on Bengali news (3M tokens), 300 AdamW steps,
batch 8 x 128 tokens. Validation cross-entropy is in nats/token, lower is better; the base model
scores 1.708. All methods saw identical data, and baseline learning rates follow their papers.

Three seeds, best baseline per budget:

| budget | TensorPEFT | best baseline | difference | seeds won | p (paired t) |
|---|---:|---|---:|:---:|---:|
| ~9k params | **1.3245 ± 0.004** | prompt tuning: 1.3684 ± 0.010 | **-0.044** | 3/3 | 0.008 |
| ~37k | 1.2558 ± 0.010 | LayerNorm tuning (38.4k): 1.2504 ± 0.003 | +0.005 | 1/3 | 0.37 (tie) |
| ~55k | **1.2368 ± 0.017** | (IA)³: 1.2880 ± 0.003 | **-0.051** | 3/3 | 0.026 |

All baselines at each budget (seed 0):

| budget | TensorPEFT | LoRA | QLoRA | VeRA | (IA)³ | prompt tuning | LayerNorm tuning | BitFit |
|---|---:|---:|---:|---:|---:|---:|---:|---:|
| ~1.5k | **1.461** | - | - | - | - | 1.493 (768 params) | 1.600 | - |
| ~9k | **1.327** | 1.430 | - | - | 1.432 | 1.370 | 1.407 | - |
| ~37k | 1.255 | 1.318 | 1.335 | 1.393 (24.6k) | - | 1.327 | **1.253** | - |
| ~55k | **1.231** | - | - | - | 1.289 | - | - | - |
| ~83k | - | - | - | - | - | - | - | 1.326 |

TensorPEFT budgets map to configs as follows: ~1.5k is one atom on the attention tensor only
(`target_modules=["c_attn.q", "c_attn.k", "c_attn.v", "attn.c_proj"]`, `atoms_per_group=1`),
~9k is `atoms_per_group=1`, ~37k is 4 and ~55k is 6. A dash means the method cannot be
configured at that budget on GPT-2, or was not run there.
LoRA's smallest all-layer configuration is r=1 on q,v (37k parameters). At 9k it is LoRA r=1 on
q,v of the last 3 layers.

**What this does and does not show.** TensorPEFT beats every baseline from about 1.5k to
about 9k parameters and at about 55k, and ties LayerNorm tuning at about 37k. With larger
budgets (≥ 115k) LoRA was better in our tests. The evidence covers one base model, one task and
short training. Results on other models and downstream tasks are not yet available.

`examples/reproduce_gpt2_bengali.py` reruns these benchmarks with this package
(`--atoms-per-group 1 | 4 | 6`). The
data comes from the Bengali news corpus at github.com/soham96/Bengali_news_classifier,
tokenised with the GPT-2 tokenizer.

## Supported models

Any model whose transformer blocks use `nn.Linear` or GPT-2-style `Conv1D` layers inside
numbered blocks (`...layers.N...` / `...h.N...`). Tested on GPT-2, Llama (GQA), Qwen2, Mistral,
Phi and Gemma configurations. Fused projections other than GPT-2's `c_attn` are treated as a
single matrix.

## Limitations (v0.1)

- **GPU support is untested.** The code is device- and dtype-agnostic (adapter parameters in
  fp32, bf16/fp16 bases supported), but all tests and benchmarks so far ran on CPU.
- **Probe memory.** `init_atoms` needs one backward pass with gradients on the targeted weights,
  plus an fp32 copy of each group's gradient tensor. That is roughly 2x the size of the targeted
  weights, temporarily. Use `target_modules` to limit it for large models.
- **Quantised bases** (bitsandbytes 4/8-bit) are not supported yet, because the probe needs
  weight gradients.
- **No checkpoint resume.** `Trainer` resume-from-checkpoint is not implemented; reload with
  `TensorPEFTModel.from_pretrained(..., is_trainable=True)`.

## Tests

    pip install .[test]
    pytest tests

## License

Apache-2.0. Copyright 2026 Angshul Majumdar.