File size: 7,242 Bytes
c83c617 | 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 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 | """Optional training compilation and frozen-weight initialization for V2."""
import json
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,
allow_new_data=False, tokenizer_identity=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.")
data_changed = any(
saved["signatures"][name] != signatures[name] for name in ("tokens", "windows")
)
tokenizer_check = None
if data_changed:
if not allow_new_data or optimizer is not None or tokenizer_identity is None:
raise ValueError("New-data initialization requires explicit permission, matching tokenizer and a fresh optimizer.")
source_identity = saved.get("tokenizer_identity")
if source_identity is None:
root = Path(__file__).resolve().parents[3]
source_manifest = root / saved["config"]["token_dir"] / "manifest.json"
source_tokens = json.loads(source_manifest.read_text(encoding="utf-8"))
source_hash = source_tokens.get("tokenizer_model_sha256")
if source_hash is None:
candidates = [value for name, value in source_tokens.get("input_sha256", {}).items()
if name.replace("\\", "/").endswith("/tokenizer.model")]
if len(candidates) != 1:
raise ValueError("Source tokenizer identity is unavailable in the frozen manifest.")
source_hash = candidates[0]
source_identity = {"tokenizer_model_sha256": source_hash,
"vocab_size": source_tokens["vocab_size"], "special_ids": source_tokens["special_ids"]}
if source_identity != tokenizer_identity:
raise ValueError("Source and target tokenizer identities differ.")
tokenizer_check = "Frozen tokenizer identity matched; new corpus allowed; optimizer reset"
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"],
"new_training_data": data_changed,
"tokenizer_compatibility": tokenizer_check,
"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
|