mirror of
https://github.com/wassname/evil_MoE.git
synced 2026-09-11 14:52:47 +08:00
spec
This commit is contained in:
wassname
@@ -0,0 +1,153 @@
|
||||
"""AntiPaSTO full-rank adapter for projected-GRPO.
|
||||
|
||||
Per spec.md: wrap nn.Linear with frozen U, S, Vh (full rank = min(d_in, d_out)).
|
||||
Trainable: delta_S only, shape [r]. No rotation (would break v_hack basis invariance).
|
||||
|
||||
Forward:
|
||||
y = ((x @ Vh.T) * (S + delta_S)) @ U.T + b
|
||||
|
||||
At delta_S=0, output == original linear up to fp32 SVD round-trip precision.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import hashlib
|
||||
from pathlib import Path
|
||||
|
||||
import torch
|
||||
from jaxtyping import Float
|
||||
from loguru import logger
|
||||
from torch import Tensor, nn
|
||||
|
||||
|
||||
class AntiPaSTOLinear(nn.Module):
|
||||
"""Drop-in replacement for nn.Linear with full-rank SVD + learnable delta_S.
|
||||
|
||||
Buffers (frozen): U[d_out, r], S[r], Vh[r, d_in], optional bias[d_out].
|
||||
Trainable: delta_S[r].
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
U: Float[Tensor, "d_out r"],
|
||||
S: Float[Tensor, "r"],
|
||||
Vh: Float[Tensor, "r d_in"],
|
||||
bias: Float[Tensor, "d_out"] | None,
|
||||
dtype: torch.dtype = torch.float32,
|
||||
):
|
||||
super().__init__()
|
||||
r = S.shape[0]
|
||||
self.register_buffer("U", U.to(dtype).contiguous())
|
||||
self.register_buffer("S", S.to(dtype).contiguous())
|
||||
self.register_buffer("Vh", Vh.to(dtype).contiguous())
|
||||
if bias is not None:
|
||||
self.register_buffer("bias", bias.to(dtype).contiguous())
|
||||
else:
|
||||
self.bias = None
|
||||
self.delta_S = nn.Parameter(torch.zeros(r, dtype=dtype))
|
||||
|
||||
@property
|
||||
def r(self) -> int:
|
||||
return self.S.shape[0]
|
||||
|
||||
def forward(self, x: Float[Tensor, "... d_in"]) -> Float[Tensor, "... d_out"]:
|
||||
# x @ Vh.T : [..., r]; * (S+dS) : elementwise; @ U.T : [..., d_out]
|
||||
h = x @ self.Vh.transpose(-1, -2)
|
||||
h = h * (self.S + self.delta_S)
|
||||
y = h @ self.U.transpose(-1, -2)
|
||||
if self.bias is not None:
|
||||
y = y + self.bias
|
||||
return y
|
||||
|
||||
|
||||
def _model_svd_dir(model_name: str, cache_root: Path) -> Path:
|
||||
safe = model_name.replace("/", "__")
|
||||
return cache_root / safe
|
||||
|
||||
|
||||
def svd_cached(
|
||||
W: Float[Tensor, "d_out d_in"],
|
||||
cache_path: Path,
|
||||
device: torch.device,
|
||||
) -> tuple[Tensor, Tensor, Tensor]:
|
||||
"""SVD with disk cache. Compute on `device` in fp32, save as fp32 cpu tensors.
|
||||
|
||||
Cache key = sha256(W.cpu fp32 bytes)[:16] in filename suffix, so weight change
|
||||
invalidates the cache automatically (fail-loud, no silent stale).
|
||||
"""
|
||||
cache_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
W_fp32 = W.detach().to(torch.float32).cpu().contiguous()
|
||||
sha = hashlib.sha256(W_fp32.numpy().tobytes()).hexdigest()[:16]
|
||||
final = cache_path.with_suffix(f".{sha}.pt")
|
||||
if final.exists():
|
||||
d = torch.load(final, map_location="cpu", weights_only=True)
|
||||
return d["U"], d["S"], d["Vh"]
|
||||
W_gpu = W_fp32.to(device)
|
||||
U, S, Vh = torch.linalg.svd(W_gpu, full_matrices=False)
|
||||
U, S, Vh = U.cpu(), S.cpu(), Vh.cpu()
|
||||
torch.save({"U": U, "S": S, "Vh": Vh}, final)
|
||||
logger.info(f"SVD cached: {final.name} shape U={tuple(U.shape)} S0={S[0]:.3f} S-1={S[-1]:.3e}")
|
||||
return U, S, Vh
|
||||
|
||||
|
||||
TARGET_SUFFIXES = (
|
||||
# full attention (Qwen3.5 has 6 full-attn layers)
|
||||
"q_proj",
|
||||
"k_proj",
|
||||
"v_proj",
|
||||
"o_proj",
|
||||
# linear-attention / GatedDeltaNet (Qwen3.5 has 18 linear-attn layers)
|
||||
"in_proj_qkv",
|
||||
"in_proj_z",
|
||||
"in_proj_a",
|
||||
"in_proj_b",
|
||||
"out_proj",
|
||||
# MLP (24 layers)
|
||||
"up_proj",
|
||||
"gate_proj",
|
||||
"down_proj",
|
||||
)
|
||||
|
||||
|
||||
def is_target(name: str) -> bool:
|
||||
return name.split(".")[-1] in TARGET_SUFFIXES
|
||||
|
||||
|
||||
def wrap_model_with_antipasto(
|
||||
model: nn.Module,
|
||||
model_name: str,
|
||||
cache_root: Path = Path("svd_cache"),
|
||||
svd_device: torch.device | str = "cuda",
|
||||
adapter_dtype: torch.dtype = torch.float32,
|
||||
) -> dict[str, AntiPaSTOLinear]:
|
||||
"""Replace every target nn.Linear in `model` (in place) with AntiPaSTOLinear.
|
||||
|
||||
SVD is computed on `svd_device` per layer, cached to disk by weight hash.
|
||||
Returns dict[module_qualified_name -> wrapper] for downstream v_hack code.
|
||||
"""
|
||||
svd_device_t = torch.device(svd_device) if isinstance(svd_device, str) else svd_device
|
||||
svd_dir = _model_svd_dir(model_name, cache_root)
|
||||
wrappers: dict[str, AntiPaSTOLinear] = {}
|
||||
|
||||
# Collect first to avoid mutating during iteration.
|
||||
targets: list[tuple[str, nn.Linear, nn.Module, str]] = []
|
||||
for name, m in model.named_modules():
|
||||
if isinstance(m, nn.Linear) and is_target(name):
|
||||
parent_name = name.rsplit(".", 1)[0]
|
||||
child_name = name.rsplit(".", 1)[1]
|
||||
parent = model.get_submodule(parent_name)
|
||||
targets.append((name, m, parent, child_name))
|
||||
|
||||
logger.info(f"AntiPaSTO wrap: {len(targets)} target Linear modules in {model_name}")
|
||||
for i, (name, linear, parent, child_name) in enumerate(targets):
|
||||
W = linear.weight.data
|
||||
bias = linear.bias.data if linear.bias is not None else None
|
||||
cache_path = svd_dir / f"{name}.pt"
|
||||
U, S, Vh = svd_cached(W, cache_path, device=svd_device_t)
|
||||
# Place wrapper on the same device as the original module's weight.
|
||||
target_device = W.device
|
||||
wrap = AntiPaSTOLinear(U, S, Vh, bias, dtype=adapter_dtype).to(target_device)
|
||||
setattr(parent, child_name, wrap)
|
||||
wrappers[name] = wrap
|
||||
if (i + 1) % 20 == 0 or i == len(targets) - 1:
|
||||
logger.info(f" wrapped {i+1}/{len(targets)} last={name}")
|
||||
return wrappers
|
||||
@@ -0,0 +1,102 @@
|
||||
"""Step-1 verify: wrapped Qwen3.5-0.8B output == base output at delta_S=0.
|
||||
|
||||
SHOULD: max abs diff < 1e-3 over 3 prompts of different lengths.
|
||||
ELSE: SVD round-trip is bad (numerical, dtype, or shape bug).
|
||||
|
||||
Run: uv run python -m projected_grpo.verify_antipasto_identity
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import copy
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
import torch
|
||||
from loguru import logger
|
||||
from tabulate import tabulate
|
||||
from transformers import AutoModelForCausalLM, AutoTokenizer
|
||||
|
||||
from .antipasto import wrap_model_with_antipasto
|
||||
|
||||
|
||||
MODEL = "Qwen/Qwen3.5-0.8B"
|
||||
PROMPTS = [
|
||||
"Hello",
|
||||
"Write a Python function that returns the sum of two integers.",
|
||||
(
|
||||
"You are an expert programmer. Solve the following LeetCode problem:\n"
|
||||
"Given an integer array nums, find the contiguous subarray with the largest sum.\n"
|
||||
"Return the sum."
|
||||
),
|
||||
]
|
||||
CACHE_ROOT = Path("svd_cache")
|
||||
|
||||
|
||||
def main() -> int:
|
||||
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
||||
logger.info(f"device={device} model={MODEL}")
|
||||
|
||||
tokenizer = AutoTokenizer.from_pretrained(MODEL)
|
||||
base = AutoModelForCausalLM.from_pretrained(
|
||||
MODEL, dtype=torch.float32, attn_implementation="sdpa"
|
||||
).to(device)
|
||||
base.eval()
|
||||
|
||||
wrapped = copy.deepcopy(base)
|
||||
wrappers = wrap_model_with_antipasto(
|
||||
wrapped,
|
||||
model_name=MODEL,
|
||||
cache_root=CACHE_ROOT,
|
||||
svd_device=device,
|
||||
adapter_dtype=torch.float32,
|
||||
)
|
||||
wrapped.eval()
|
||||
|
||||
n_wrapped = len(wrappers)
|
||||
n_params_trainable = sum(p.numel() for w in wrappers.values() for p in w.parameters() if p.requires_grad)
|
||||
n_params_base = sum(p.numel() for p in base.parameters())
|
||||
logger.info(
|
||||
f"wrapped={n_wrapped} modules "
|
||||
f"delta_S params={n_params_trainable:,} "
|
||||
f"base params={n_params_base:,} "
|
||||
f"ratio={n_params_trainable / n_params_base:.4%}"
|
||||
)
|
||||
|
||||
rows = []
|
||||
all_ok = True
|
||||
for i, prompt in enumerate(PROMPTS):
|
||||
ids = tokenizer(prompt, return_tensors="pt").input_ids.to(device)
|
||||
with torch.no_grad():
|
||||
y_base = base(ids).logits
|
||||
y_wrap = wrapped(ids).logits
|
||||
diff = (y_base - y_wrap).abs()
|
||||
max_diff = diff.max().item()
|
||||
mean_diff = diff.mean().item()
|
||||
scale = y_base.abs().mean().item()
|
||||
ok = max_diff < 1e-3
|
||||
all_ok = all_ok and ok
|
||||
rows.append(
|
||||
dict(
|
||||
idx=i,
|
||||
seq_len=ids.shape[1],
|
||||
logit_scale=f"{scale:.3f}",
|
||||
max_abs_diff=f"{max_diff:.2e}",
|
||||
mean_abs_diff=f"{mean_diff:.2e}",
|
||||
ok=("PASS" if ok else "FAIL"),
|
||||
)
|
||||
)
|
||||
|
||||
print(tabulate(rows, headers="keys", tablefmt="pipe"))
|
||||
logger.info(
|
||||
"SHOULD: max_abs_diff < 1e-3 on all rows. "
|
||||
"ELSE: SVD round-trip broken (dtype downcast, shape bug, or wrong forward)."
|
||||
)
|
||||
if not all_ok:
|
||||
logger.error("IDENTITY CHECK FAILED")
|
||||
return 1
|
||||
logger.info(f"IDENTITY CHECK PASSED ({n_wrapped} modules, {n_params_trainable:,} delta_S scalars)")
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(main())
|
||||
Reference in New Issue
Block a user