File size: 9,113 Bytes
49dddf5 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 | """pytest tests for TensorPEFT (CPU, tiny random models)."""
import math
import pytest
import torch
import torch.nn.functional as F
import transformers as T
from tensorpeft import TensorPEFTConfig, TensorPEFTModel, get_tensorpeft_model
from tensorpeft.cp import cp_value, contract_except, outer, rank1_cp
V = 97
def tiny(arch):
common = dict(vocab_size=V, max_position_embeddings=64)
if arch == "gpt2":
return T.GPT2LMHeadModel(T.GPT2Config(n_layer=3, n_embd=32, n_head=4, vocab_size=V, n_positions=64))
if arch == "llama":
return T.LlamaForCausalLM(T.LlamaConfig(num_hidden_layers=3, hidden_size=32, intermediate_size=48,
num_attention_heads=4, num_key_value_heads=2, **common))
if arch == "qwen2":
return T.Qwen2ForCausalLM(T.Qwen2Config(num_hidden_layers=3, hidden_size=32, intermediate_size=48,
num_attention_heads=4, num_key_value_heads=2, **common))
if arch == "mistral":
return T.MistralForCausalLM(T.MistralConfig(num_hidden_layers=3, hidden_size=32, intermediate_size=48,
num_attention_heads=4, num_key_value_heads=2, **common))
if arch == "phi":
return T.PhiForCausalLM(T.PhiConfig(num_hidden_layers=3, hidden_size=32, intermediate_size=48,
num_attention_heads=4, **common))
if arch == "gemma":
return T.GemmaForCausalLM(T.GemmaConfig(num_hidden_layers=3, hidden_size=32, intermediate_size=48,
num_attention_heads=4, num_key_value_heads=2, head_dim=8, **common))
raise ValueError(arch)
ARCHS = ["gpt2", "llama", "qwen2", "mistral", "phi", "gemma"]
def batches(n=2, b=3, t=17, seed=0):
g = torch.Generator().manual_seed(seed)
return [torch.randint(0, V, (b, t), generator=g) for _ in range(n)]
def lm_loss(model, x):
logits = model(x[:, :-1]).logits
return F.cross_entropy(logits.float().flatten(0, 1), x[:, 1:].flatten())
@pytest.fixture(autouse=True)
def _seed():
torch.manual_seed(0)
# ---------------------------------------------------------------- CP / HOPM
@pytest.mark.parametrize("shape", [(5, 30, 40), (6, 4, 30, 20)])
def test_hopm_recovers_planted_rank1(shape):
us = [F.normalize(torch.randn(n, dtype=torch.float64), dim=0) for n in shape]
T_ = 5.0 * outer(us) + 0.05 * torch.randn(shape, dtype=torch.float64)
est, sig = rank1_cp(T_, iters=50)
assert min(abs(float(a @ b)) for a, b in zip(us, est)) > 0.99
assert abs(sig - 5.0) < 0.3
def test_hopm_fixed_point():
T_ = torch.randn(4, 3, 20, 25, dtype=torch.float64)
est, sig = rank1_cp(T_, iters=500, tol=1e-14)
res = max(float((contract_except(T_, est, n) / sig - est[n]).abs().max()) for n in range(4))
assert res < 1e-6
# ---------------------------------------------------------------- model wrapping
@pytest.mark.parametrize("arch", ARCHS)
def test_unchanged_before_and_after_init(arch):
base = tiny(arch).eval()
x = batches(1)[0]
ref = base(x).logits.detach().clone()
model = get_tensorpeft_model(base, TensorPEFTConfig(atoms_per_group=2))
assert torch.equal(model(x).logits, ref)
model.init_atoms(batches(2), verbose=False)
assert model.is_initialized
assert float((model(x).logits - ref).abs().max()) < 1e-6 # a = 0 at init
@pytest.mark.parametrize("arch", ARCHS)
def test_initial_gradient_is_steepest_cp_direction(arch):
"""dL/d alpha_raw_i at init = -alpha_scale * sigma_i (atoms are descent directions)."""
torch.set_default_dtype(torch.float64)
try:
base = tiny(arch).eval()
model = get_tensorpeft_model(base, TensorPEFTConfig(atoms_per_group=3))
xs = batches(2)
sig = model.init_atoms(xs, verbose=False)
loss = sum(lm_loss(model, x) for x in xs) / len(xs)
grads = torch.autograd.grad(loss, [g.alpha_raw for g in model.groups])
for g, gr, s in zip(model.groups, grads, sig):
expect = -g.alpha_scale * torch.tensor(s, dtype=gr.dtype)
assert torch.allclose(gr, expect, rtol=1e-6, atol=1e-10), (arch, gr, expect)
finally:
torch.set_default_dtype(torch.float32)
@pytest.mark.parametrize("arch", ARCHS)
def test_merge_equals_hooks_and_save_load_roundtrip(arch, tmp_path):
torch.set_default_dtype(torch.float64)
try:
base = tiny(arch).eval()
state = {k: v.clone() for k, v in base.state_dict().items()}
x = batches(1, seed=5)[0]
ref = base(x).logits.detach().clone()
model = get_tensorpeft_model(base, TensorPEFTConfig(atoms_per_group=2))
model.init_atoms(batches(2), verbose=False)
with torch.no_grad():
for p in model.parameters():
if p.requires_grad:
p.add_(0.3 * torch.randn_like(p))
hooked = model(x).logits.detach().clone()
assert float((hooked - ref).abs().max()) > 1e-3 # the adapter does change the model
model.save_pretrained(tmp_path)
# reload onto a fresh copy of the same base
base2 = tiny(arch).eval()
base2.load_state_dict(state)
re = TensorPEFTModel.from_pretrained(base2, str(tmp_path))
assert float((re(x).logits - hooked).abs().max()) < 1e-10
merged = model.merge_and_unload()
assert float((merged(x).logits - hooked).abs().max()) < 1e-10
finally:
torch.set_default_dtype(torch.float32)
def test_param_count_and_grouping_gpt2():
base = tiny("gpt2")
m = get_tensorpeft_model(base, TensorPEFTConfig(atoms_per_group=4))
names = [g.name for g in m.groups]
assert names[0].startswith("attn.c_attn.k+attn.c_attn.q+attn.c_attn.v+attn.c_proj [3, 4, 32, 32]")
expect = 4 * ((1 + 3 + 4 + 32 + 32) + (1 + 3 + 32 + 128) + (1 + 3 + 128 + 32))
assert m.get_nb_trainable_parameters()[0] == expect
assert all(not p.requires_grad for p in base.parameters())
def test_target_modules_subset():
base = tiny("llama")
m = get_tensorpeft_model(base, TensorPEFTConfig(target_modules=["q_proj", "v_proj"], atoms_per_group=1))
roles = sorted({s.role for s in m.slots})
assert roles == ["self_attn.q_proj", "self_attn.v_proj"]
def test_wrong_architecture_rejected(tmp_path):
m = get_tensorpeft_model(tiny("llama"), TensorPEFTConfig(atoms_per_group=1))
m.init_atoms(batches(1), verbose=False)
m.save_pretrained(tmp_path)
with pytest.raises(ValueError):
TensorPEFTModel.from_pretrained(tiny("gpt2"), str(tmp_path))
def test_dict_batches_with_attention_mask_and_labels():
base = tiny("llama").eval()
m = get_tensorpeft_model(base, TensorPEFTConfig(atoms_per_group=1))
x = batches(1)[0]
m.init_atoms([{"input_ids": x, "attention_mask": torch.ones_like(x), "labels": x}], verbose=False)
assert m.is_initialized
def test_bf16_base():
base = tiny("llama").to(torch.bfloat16).eval()
m = get_tensorpeft_model(base, TensorPEFTConfig(atoms_per_group=2))
m.init_atoms(batches(2), verbose=False)
with torch.no_grad():
for g in m.groups:
g.alpha_raw.add_(0.01)
out = m(batches(1)[0]).logits
assert out.dtype == torch.bfloat16 and torch.isfinite(out).all()
assert all(p.dtype == torch.float32 for p in m.parameters() if p.requires_grad)
def test_one_adam_lr_trains():
base = tiny("gpt2").eval()
m = get_tensorpeft_model(base, TensorPEFTConfig(atoms_per_group=2))
xs = batches(2, b=4, t=33)
m.init_atoms(xs, verbose=False)
opt = torch.optim.AdamW([p for p in m.parameters() if p.requires_grad], lr=3e-3, weight_decay=0.0)
l0 = float(lm_loss(m, xs[0]))
for _ in range(10):
opt.zero_grad()
lm_loss(m, xs[0]).backward()
opt.step()
assert float(lm_loss(m, xs[0])) < l0
def test_hf_trainer_integration(tmp_path):
from tensorpeft import TensorPEFTTrainer
base = tiny("llama")
m = get_tensorpeft_model(base, TensorPEFTConfig(atoms_per_group=2))
data = [{"input_ids": x, "attention_mask": torch.ones_like(x), "labels": x.clone()}
for x in torch.randint(0, V, (16, 24)).unbind(0)]
m.init_atoms(torch.utils.data.DataLoader(data, batch_size=4), verbose=False)
class DS(torch.utils.data.Dataset):
def __len__(self):
return len(data)
def __getitem__(self, i):
return data[i]
args = T.TrainingArguments(output_dir=str(tmp_path / "out"), per_device_train_batch_size=4, max_steps=6,
learning_rate=3e-3, warmup_steps=1, logging_steps=1, save_steps=3,
report_to=[], use_cpu=True)
tr = TensorPEFTTrainer(model=m, args=args, train_dataset=DS())
hist = tr.train()
assert math.isfinite(hist.training_loss)
tr.save_model(str(tmp_path / "final"))
re = TensorPEFTModel.from_pretrained(tiny("llama"), str(tmp_path / "final"))
assert re.is_initialized
import os
assert os.path.exists(tmp_path / "out" / "checkpoint-3" / "adapter_model.safetensors")
|