Spaces:
Running on Zero
Running on Zero
File size: 3,226 Bytes
5f653cc | 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 | import math
import torch
import numpy as np
from scipy.interpolate import make_interp_spline
from safetensors.torch import load_file
def freeze_parameters(module: torch.nn.Module):
for p in module.parameters():
p.requires_grad_(False)
module.eval()
def summarize_model(model):
# Calculate parameters
trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
total_params = sum(p.numel() for p in model.parameters())
# Calculate size on disk/memory (assuming float32 = 4 bytes)
param_size = 0
for param in model.parameters():
param_size += param.nelement() * param.element_size()
buffer_size = 0
for buffer in model.buffers():
buffer_size += buffer.nelement() * buffer.element_size()
size_all_mb = (param_size + buffer_size) / 1024**2
return {
"total_params": total_params,
"trainable_params": trainable_params,
"trainable_ratio(%)": trainable_params / total_params,
"model_size_mb(MB)": size_all_mb,
}
def lora_scale_states(
anchor_states: list[str],
query_point: float,
device: torch.device,
dtype: torch.dtype,
):
state = load_file(anchor_states[0])
new_state = {}
for k, p in state.items():
if "projector.weight" in k:
new_state[k] = torch.full_like(p, fill_value=2 * query_point, device=device, dtype=dtype)
elif "projector.bias" in k:
new_state[k] = torch.zeros_like(p, device=device, dtype=dtype)
return new_state
def neighbor_interpolate_states(
anchor_states: list[str],
query_point: float,
device: torch.device,
dtype: torch.dtype,
):
num_levels = len(anchor_states)
level = math.floor(query_point * num_levels)
interp_state = {}
if level == 0:
state = load_file(anchor_states[0])
w = query_point * num_levels - level
for k, p in state.items():
interp_state[k] = (w * p).to(device=device, dtype=dtype)
elif level == num_levels:
interp_state = load_file(anchor_states[-1])
for k, p in interp_state.items():
interp_state[k] = p.to(device=device, dtype=dtype)
else:
state1 = load_file(anchor_states[level - 1])
state2 = load_file(anchor_states[level])
w = query_point * num_levels - level
for k, p1 in state1.items():
p2 = state2[k]
interp_state[k] = ((1 - w) * p1 + w * p2).to(device=device, dtype=dtype)
return interp_state
def bspline_interpolate_states(
anchor_states: list[str],
query_point: float,
X,
device: torch.device,
dtype: torch.dtype,
order=3,
):
states = load_file(anchor_states[0])
states = {k: [p] for k, p in states.items()}
for sp in anchor_states[1:]:
s = load_file(sp)
for k, p in s.items():
states[k].append(p)
interp_state = {}
# X = [0.2, 0.4, 0.6, 0.8, 1.0]
for k, pl in states.items():
Y = torch.stack(pl, dim=0).to(torch.float32).numpy()
bsp = make_interp_spline(X, Y, k=order)
y = torch.from_numpy(bsp(np.array(query_point))).to(device=device, dtype=dtype)
interp_state[k] = y
return interp_state
|