| """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) |
|
|
|
|
| |
| @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 |
|
|
|
|
| |
| @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 |
|
|
|
|
| @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 |
| model.save_pretrained(tmp_path) |
| |
| 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") |
|
|
