File size: 998 Bytes
76b78ee | 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 | """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
|