| # 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. |
|
|
