| """V4-only initialization, task identity, normalized updates and stop policy.""" |
| import copy |
| import math |
| from pathlib import Path |
|
|
| import torch |
|
|
| from vimeml.posttraining.data import TOKENIZER_SHA |
| from vimeml.posttraining.infer import SCORE_CONTRACT, READING_CONTRACT |
| from vimeml.posttraining.streams import microbatches |
| from vimeml.training.data import file_sha |
| from vimeml.training.model_v4 import GPTV4Config |
| from vimeml.training.train import amp_context |
|
|
| STAGE = "v4-c1" |
| TASK_CONTRACT = { |
| "format": "BOS_CTX_context_READ_reading_OUT_surface_EOS_v1", |
| "labels": "already_shifted_OUT_and_EOS_only_PAD_minus100", |
| "lm": "ordinary_unmarked_LM_all_valid_targets", |
| "normalization": "independent_task_total_supervised_tokens_per_update", |
| "score": SCORE_CONTRACT, "reading": READING_CONTRACT, |
| "context_length": 128, "tokenizer_sha256": TOKENIZER_SHA, |
| "marker_ids": {"<CTX>": 4, "<READ>": 5, "<OUT>": 6}, |
| } |
|
|
|
|
| def validate_task_checkpoint(saved): |
| if (saved.get("format") != "vimeml_tiny_gpt_v4" or saved.get("architecture") != "tiny_gpt_v4" |
| or saved.get("model_config") != GPTV4Config().__dict__ |
| or saved.get("posttraining", {}).get("stage") not in (STAGE, "v4.1") |
| or saved["posttraining"].get("task_contract") != TASK_CONTRACT): |
| raise ValueError("Explicit compatible V4-C1 checkpoint required") |
| expected = {"tokenizer_model_sha256": TOKENIZER_SHA, "vocab_size": 8192, |
| "special_ids": {"pad": 0, "unk": 1, "bos": 2, "eos": 3}} |
| if saved.get("tokenizer_identity") != expected: |
| raise ValueError("Conditional checkpoint tokenizer identity mismatch") |
|
|
|
|
| def initialize_v4(model, source, expected_sha, expected_step, tokenizer_identity, token_signature, window_signature): |
| """Cross-stage initialization loads weights only, never AdamW/RNG/cursors.""" |
| if file_sha(source) != expected_sha: |
| raise ValueError("Source checkpoint SHA256 mismatch") |
| saved = torch.load(source, map_location="cpu", weights_only=True) |
| if (saved.get("format") != "vimeml_tiny_gpt_v4" or saved.get("architecture") != "tiny_gpt_v4" |
| or saved.get("posttraining") or saved.get("model_config") != model.configuration() |
| or saved.get("step") != expected_step or saved.get("tokenizer_identity") != tokenizer_identity |
| or saved.get("signatures", {}).get("tokens") != token_signature |
| or saved.get("signatures", {}).get("windows") != window_signature): |
| raise ValueError("Source must be the frozen matching V4-B LM checkpoint") |
| model.load_state_dict(saved["model"], strict=True) |
| return {"checkpoint": str(Path(source)), "checkpoint_sha256": expected_sha, |
| "source_step": saved["step"], "source_total_tokens": saved["total_tokens"], |
| "source_epoch": saved["epoch"], "optimizer": "fresh_AdamW_no_source_state_restored"} |
|
|
|
|
| def update(model, optimizer, conversion, lm, settings, device, precision): |
| """One atomic logical update; never average unequal microbatch means.""" |
| denominators = {name: sum(int((b["labels"] != -100).sum()) for b in batches) |
| for name, batches in (("conversion", conversion), ("lm", lm))} |
| if min(denominators.values()) <= 0: |
| raise ValueError("Each task needs supervised targets") |
| optimizer.zero_grad(set_to_none=True) |
| losses, inputs, examples, positions = {}, {}, {}, {} |
| for name, batches in (("conversion", conversion), ("lm", lm)): |
| pieces = [] |
| inputs[name] = sum(int(b["lengths"].sum()) for b in batches) |
| examples[name] = sum(len(b["lengths"]) for b in batches) |
| positions[name] = 0 |
| for batch in microbatches(batches, settings[f"{name}_microbatch"]): |
| ids = batch["input_ids"].to(device, non_blocking=True) |
| labels = batch["labels"].to(device, non_blocking=True) |
| with amp_context(device, precision): |
| result = model(ids, labels) |
| loss = result["loss_sum"] * (settings[f"{name}_weight"] / denominators[name]) |
| pieces.append(result["loss_sum"].detach()) |
| |
| loss.backward() |
| positions[name] += ids.numel() |
| value = float(torch.stack(pieces).sum()) |
| if not math.isfinite(value): |
| raise FloatingPointError(f"Non-finite {name} training loss") |
| losses[name] = value / denominators[name] |
| grad_norm = float(torch.nn.utils.clip_grad_norm_(model.parameters(), settings["grad_clip"], error_if_nonfinite=True)) |
| optimizer.step() |
| return {"losses": losses, "supervised_tokens": denominators, "input_tokens": inputs, |
| "examples": examples, "padded_positions": positions, "grad_norm": grad_norm, |
| "weighted_loss": sum(losses[n] * settings[f"{n}_weight"] for n in losses)} |
|
|
|
|
| def initial_state(): |
| return {"step": 0, "epoch": 0, "conversion_cursor": 0, |
| "lm_cursor": {"epoch": 0, "batch": 0, "offset": 0}, |
| "epoch_results": [], "conversion_reference": None, "plateau_count": 0, |
| "forgetting_count": 0, "best_conversion": None, "best_lm": None, |
| "best_rerank_key": None, "stop_reason": None, |
| "totals": {name: {"examples": 0, "input_tokens": 0, "supervised_tokens": 0, "padded_positions": 0} |
| for name in ("conversion", "lm")}, "training_seconds": 0.0} |
|
|
|
|
| def apply_epoch_result(state, conversion_loss, lm_bpc, ime_correct, settings): |
| """Returns checkpoint roles; caller applies exactly once per completed epoch.""" |
| if not all(math.isfinite(x) and x > 0 for x in (conversion_loss, lm_bpc)): |
| raise ValueError("Invalid full validation metrics") |
| roles = [] |
| if state["best_conversion"] is None or conversion_loss < state["best_conversion"]: |
| state["best_conversion"] = conversion_loss |
| roles.append("best-conversion") |
| if state["best_lm"] is None or lm_bpc < state["best_lm"]: |
| state["best_lm"] = lm_bpc |
| roles.append("best-lm") |
| reference = state["conversion_reference"] |
| if reference is None or conversion_loss < reference * (1 - settings["early_stopping_relative_delta"]): |
| state["conversion_reference"], state["plateau_count"] = conversion_loss, 0 |
| else: |
| state["plateau_count"] += 1 |
| allowed = lm_bpc <= settings["baseline_lm_bpc"] * (1 + settings["lm_relative_degradation"]) |
| state["forgetting_count"] = 0 if allowed else state["forgetting_count"] + 1 |
| key = [int(ime_correct), -conversion_loss, -state["step"]] |
| enforce_lm = settings.get("lm_retention_policy", "enforce") == "enforce" |
| if (allowed or not enforce_lm) and (state["best_rerank_key"] is None or key > state["best_rerank_key"]): |
| state["best_rerank_key"] = key |
| roles += ["best-rerank", "best"] |
| if enforce_lm and state["forgetting_count"] >= settings["lm_degradation_patience"]: |
| state["stop_reason"] = "lm_retention_guard" |
| elif state["plateau_count"] >= settings["early_stopping_patience"]: |
| state["stop_reason"] = "conversion_plateau" |
| return roles, allowed |
|
|
|
|
| def restore_checkpoint(saved, model, optimizer, config, signatures, precision): |
| validate_task_checkpoint(saved) |
| if saved["config"] != config or saved["signatures"] != signatures or saved["precision"] != precision: |
| raise ValueError("C1 resume requires identical config, code, data, task and precision") |
| if saved["step"] != saved["training_state"]["step"] or saved["epoch"] != saved["training_state"]["epoch"]: |
| raise ValueError("Checkpoint header and committed state disagree") |
| model.load_state_dict(saved["model"], strict=True) |
| optimizer.load_state_dict(saved["optimizer"]) |
| return copy.deepcopy(saved["training_state"]), saved["rng"], saved["initialization"] |
|
|
