ZAdapter

TPU-native bottleneck FFN adapter for parameter-efficient fine-tuning.

Unlike LoRA (parallel low-rank matrices merged into existing weights), ZAdapter injects a small bottleneck feed-forward block serially after attention and MLP in each transformer layer. This keeps the computation graph static — no merge/unmerge step, no dynamic branching — which plays well with XLA compilation on TPU.

input -> down_proj [d_model -> r] -> activation -> up_proj [r -> d_model] -> output
output = input + adapter(input)

Install

pip install zadapter

Usage

from transformers import AutoModelForCausalLM
from zadapter import inject_adapter, get_trainable_params

model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen2.5-Coder-7B")
model = inject_adapter(model, r=64)

trainable_params = get_trainable_params(model)
optimizer = torch.optim.AdamW(trainable_params, lr=1e-4)

Why not LoRA

ZAdapter trades a slightly larger parameter count and a small inference overhead for a simpler, more XLA-friendly training graph — useful when targeting TPU where static graphs compile more predictably.

Companion library

For TPU sharding (data-parallel, tensor-parallel) and quantization (int8, NF4) to use alongside ZAdapter, see zenbit.

License

MIT