File size: 5,697 Bytes
29f25be | 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 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 | """Optional training compilation and frozen-weight initialization for V2."""
from types import MethodType
from pathlib import Path
import torch
from torch.nn import functional as F
from vimeml.training.model_v2 import RMSNorm
def rmsnorm_operations(hidden, weight, eps):
normalized = hidden.float()
variance = normalized.pow(2).mean(-1, keepdim=True)
normalized = normalized * torch.rsqrt(variance + eps)
return normalized.to(hidden.dtype) * weight
def fuse_rmsnorm(model):
"""Compile only pure norm operations, leaving variable-length logits eager."""
compiled = torch.compile(rmsnorm_operations, fullgraph=True, dynamic=True)
def forward(module, hidden):
return compiled(hidden, module.weight, module.eps)
count = 0
for module in model.modules():
if isinstance(module, RMSNorm):
module.forward = MethodType(forward, module)
count += 1
if not count:
raise ValueError("RMSNorm fusion requires a V2 model.")
return {
"implementation": "torch.compile RMSNorm only",
"norm_modules": count,
"dynamic_shapes": True,
"state_dict_unchanged": True,
"checkpoint_inference": "Plain primitive RMSNorm; training fusion is optional.",
}
def compile_backbone(model):
"""Compile the full hidden-state stack, keeping variable-size token loss eager."""
def hidden_stack(input_ids):
positions = torch.arange(input_ids.shape[1], device=input_ids.device)
hidden = model.dropout(
model.token_embedding(input_ids) + model.position_embedding(positions)
)
for block in model.blocks:
hidden = block(hidden)
return model.final_norm(hidden)
compiled = torch.compile(hidden_stack, fullgraph=True, dynamic=True)
def forward(module, input_ids, labels=None):
if (
input_ids.ndim != 2
or not 1 <= input_ids.shape[1] <= module.config.context_length
):
raise ValueError("Expected [batch, time] within the model context.")
hidden = compiled(input_ids)
if labels is None:
return module.lm_head(hidden)
if labels.shape != input_ids.shape:
raise ValueError("Labels must match input shape.")
valid = labels != -100
logits = module.lm_head(hidden[valid])
return {
"loss_sum": F.cross_entropy(logits, labels[valid], reduction="sum"),
"token_count": valid.sum(),
}
model.forward = MethodType(forward, model)
return {
"implementation": "torch.compile full hidden stack; eager valid-token head and CE",
"dynamic_shapes": True,
"state_dict_unchanged": True,
"checkpoint_inference": "Plain TinyGPTV2",
}
def initialize_weights(model, checkpoint, signatures, optimizer=None, precision=None):
"""Start a separate run; optionally carry compatible AdamW moments forward."""
path = Path(checkpoint)
saved = torch.load(path, map_location="cpu", weights_only=True)
if (
saved.get("format") != "vimeml_tiny_gpt_v2"
or saved["model_config"] != model.configuration()
):
raise ValueError("Initialization requires the same V2 model configuration.")
if any(
saved["signatures"][name] != signatures[name] for name in ("tokens", "windows")
):
raise ValueError(
"Initialization requires the same tokenizer/token store and window index."
)
if optimizer is not None:
if saved.get("precision") != precision or not saved.get("optimizer", {}).get(
"state"
):
raise ValueError(
"Optimizer continuation requires matching precision and nonempty state."
)
source_groups = saved["optimizer"]["param_groups"]
if len(source_groups) != len(optimizer.param_groups):
raise ValueError("Optimizer parameter groups differ.")
for source, target in zip(source_groups, optimizer.param_groups):
if (
len(source["params"]) != len(target["params"])
or source["betas"] != target["betas"]
or source["weight_decay"] != target["weight_decay"]
or source["eps"] != target["eps"]
or source.get("amsgrad", False) != target.get("amsgrad", False)
):
raise ValueError(
"Optimizer continuation requires compatible AdamW groups."
)
model.load_state_dict(saved["model"])
if optimizer is not None:
learning_rates = [group["lr"] for group in optimizer.param_groups]
optimizer.load_state_dict(saved["optimizer"])
for group, rate in zip(optimizer.param_groups, learning_rates):
group["lr"] = rate
details = {
"checkpoint": str(path.resolve()),
"source_step": saved["step"],
"source_trained_windows": saved["total_windows"],
"source_epoch": saved["epoch"],
"source_batch_cursor": saved["batch_cursor"],
"source_signatures": saved["signatures"],
"file_bytes": path.stat().st_size,
"optimizer": (
"restored AdamW moments and parameter step counters; new local scheduler and data cursors"
if optimizer is not None
else "fresh AdamW; source optimizer, scheduler and cursors are not restored"
),
"optimizer_state_count": len(saved["optimizer"]["state"])
if optimizer is not None
else 0,
"verification": "Model configuration and small dataset manifest signatures; no repeated checkpoint SHA256",
}
del saved
return details
|