tests/test_lora_fold.py

"""Regression: fold_state_dict must preserve bias (base.bias -> name.bias),
not overwrite name.weight (v23 resume crash root cause, 2026-08-13)."""

import sys
from pathlib import Path
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))

import torch
import torch.nn as nn

from train.train_lora import fold_state_dict, wrap_lora


class TinyNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.lin = nn.Linear(8, 4, bias=True)

    def forward(self, x):
        return self.lin(x)


def test_fold_preserves_bias():
    net = TinyNet()
    w_before = net.lin.weight.clone()
    b_before = net.lin.bias.clone()
    wrapped = wrap_lora(net, r=2, alpha=4.0, dropout=0.0)
    folded = fold_state_dict(net.state_dict(), wrapped)
    assert "lin.weight" in folded and "lin.bias" in folded
    assert folded["lin.bias"].shape == b_before.shape
    # weight must still be a 2-D weight (bias must NOT overwrite it)
    assert folded["lin.weight"].shape == w_before.shape