Day 66: Adapter Lifecycles: Serialization, Merging & Zero-Latency Serving

def merge(self):
# Fold the adapter into the base weight: W_merged = W0 + (alpha / r) * B @ A.
# B @ A has shape (out, r) @ (r, in) = (out, in), the same shape as linear.weight.
if not self.merged: # guard: merging twice would add the delta twice
delta = (self.alpha / self.r) * torch.matmul(self.lora_B, self.lora_A)
self.linear.weight.data.add_(delta)
self.merged = True # forward() now skips the adapter path

def check_merge(layer, x, atol=1e-5):
# Worked check. Before: y = W0 x + b + s * B (A x). After: y = (W0 + s B A) x + b.
# Matrix products are linear, so the two are equal in exact arithmetic; fp32 rounds the
# sums in a different order, so compare with a tolerance (about 1e-6 for small layers).
y_before = layer(x)
layer.merge()
y_after = layer(x)
return (y_after - y_before).abs().max().item() < atol

# Latency: the merged layer is one GEMM, exactly like the base model, so serving cost is
# unchanged. Example: in=out=4096, r=16 adds 2 * 16 * (4096 + 4096) = 262,144 FLOPs per token
# per layer when unmerged; merged, the extra cost is zero.

def unmerge(self):
# Inverse of merge(): W0 = W_merged - (alpha / r) * B @ A, using the SAME A, B and scale.
if self.merged: # guard: unmerging an unmerged layer would corrupt W0
delta = (self.alpha / self.r) * torch.matmul(self.lora_B, self.lora_A)
self.linear.weight.data.sub_(delta)
self.merged = False # forward() adds the adapter path again

def swap_cycle(layer, adapters, x, cycles=3, atol=1e-6):
# adapters: list of (A, B) pairs, e.g. Adapter_Medical then Adapter_Legal.
# After every unmerge the base weight must be back to W0; record the worst drift seen.
w0 = layer.linear.weight.detach().clone()
worst = 0.0
for _ in range(cycles):
for A, B in adapters:
layer.lora_A.data.copy_(A)
layer.lora_B.data.copy_(B)
layer.merge()
layer(x) # serve one request with this adapter
layer.unmerge()
worst = max(worst, (layer.linear.weight - w0).abs().max().item())
return worst < atol, worst

# Why the drift is not exactly zero: (w0 + d) - d is not always w0 in floating point. In fp32
# the worst drift on a 16 x 8 test layer was 4.2e-7; the same cycle in bf16 drifted by 3.6e-2.
# So never merge into bf16 weights you plan to unmerge; keep an fp32 or original copy of W0.

import json
import os
import torch
from safetensors.torch import save_file, load_file

def save_adapter(model, path, r=8, alpha=16, target_modules=('linear',), base_model='toy-base'):
os.makedirs(path, exist_ok=True)
cfg = {'r': r, 'lora_alpha': alpha, 'target_modules': list(target_modules),
'lora_dropout': 0.0, 'base_model_name_or_path': base_model}
with open(os.path.join(path, 'adapter_config.json'), 'w') as f:
json.dump(cfg, f, indent=2)
# Only adapter tensors; the frozen base weights stay out of the artifact.
weights = {k: v.detach().contiguous() for k, v in model.state_dict().items() if 'lora_' in k}
save_file(weights, os.path.join(path, 'adapter_model.safetensors'))
return sorted(weights)

def load_adapter(model, path):
with open(os.path.join(path, 'adapter_config.json')) as f:
cfg = json.load(f)
weights = load_file(os.path.join(path, 'adapter_model.safetensors'))
# strict=False because the file deliberately lacks the base weights; check what was missed.
missing, unexpected = model.load_state_dict(weights, strict=False)
assert not unexpected, unexpected
assert all('lora_' not in k for k in missing), missing # every adapter tensor was loaded
return cfg

# Size check for a LoRALinear(512, 512, r=8) (Day 64): A is 8 x 512 and B is 512 x 8,
# 8,192 fp32 values = 32,768 bytes, plus a small JSON header. The base 512 x 512 weight alone
# is 1 MiB, which is why the adapter directory stays far below 1 MB.
# Older exports used adapter_model.bin (torch.save, a pickle); prefer .safetensors (Day 60).

import os
import tempfile
import torch

# Uses LoRALinear (Day 64 Activity 4), merge/unmerge (Activities 1-2) and
# save_adapter/load_adapter (Activity 3).
torch.manual_seed(0)
model = LoRALinear(16, 16, r=4, lora_alpha=8)
LoRALinear.merge, LoRALinear.unmerge = merge, unmerge
x_train, y_train = torch.randn(64, 16), torch.randn(64, 16)
opt = torch.optim.AdamW([model.lora_A, model.lora_B], lr=1e-2)

loss_history = []
for _ in range(15):
loss = torch.nn.functional.mse_loss(model(x_train), y_train)
opt.zero_grad()
loss.backward()
opt.step()
loss_history.append(loss.item())

out_dir = os.path.join(tempfile.mkdtemp(), 'exported_adapter')
save_adapter(model, out_dir, r=4, alpha=8)

# A fresh model gets the SAME base weights (in practice: the same base checkpoint) plus the adapter.
fresh = LoRALinear(16, 16, r=4, lora_alpha=8)
fresh.linear.load_state_dict(model.linear.state_dict())
load_adapter(fresh, out_dir)

x = torch.randn(8, 16)
y_unmerged = fresh(x)
W0_before_merge = fresh.linear.weight.detach().clone()
fresh.merge()
y_merged = fresh(x)

assert loss_history[-1] < loss_history[0]
assert torch.equal(fresh(x), fresh.linear(x)) # merged forward is one GEMM
assert not torch.equal(fresh.linear.weight, W0_before_merge)
assert (y_unmerged - y_merged).abs().max().item() < 1e-5
assert torch.allclose(y_unmerged, model(x), atol=1e-6) # reload reproduced the trained model
print('LoRA end-to-end lifecycle verified.')