#!/usr/bin/env python3 from __future__ import annotations import argparse import json import random import time from contextlib import nullcontext from pathlib import Path from typing import Any import torch import torch.nn.functional as F from datasets import load_dataset from transformers import AutoModelForCausalLM, AutoTokenizer from tinycenn_lm.smollm2_memory_fusion import ( DEFAULT_SMOLLM2, MemoryFusionLlamaAttention, SmolMemoryFusionConfig, parameter_summary, replace_attention_layers, structural_summary, ) def parse_args() -> argparse.Namespace: p = argparse.ArgumentParser( description=( "Sequential teacher-guided SmolLM2 Memory Fusion conversion with " "per-layer acceptance gates and staged whole-model distillation." ) ) p.add_argument("--base-model", default=DEFAULT_SMOLLM2) p.add_argument("--output-dir", default="checkpoints/smollm2-memory-fusion-sequential-r64") p.add_argument("--dataset", default="HuggingFaceFW/fineweb-edu") p.add_argument("--dataset-config", default="sample-10BT") p.add_argument("--split", default="train") p.add_argument("--text-field", default="text") p.add_argument("--context-length", type=int, default=128) p.add_argument("--batch-size", type=int, default=1) p.add_argument("--feature-dim", type=int, default=32) p.add_argument("--memory-rank", type=int, choices=(32, 48, 64), default=64) p.add_argument("--seed", type=int, default=73) p.add_argument("--shuffle-buffer", type=int, default=2048) p.add_argument("--min-layer-steps", type=int, default=50) p.add_argument("--max-layer-steps", type=int, default=300) p.add_argument("--check-every", type=int, default=25) p.add_argument("--layer-lr", type=float, default=2e-4) p.add_argument("--teacher-alpha-start", type=float, default=0.90) p.add_argument("--teacher-alpha-end", type=float, default=0.00) p.add_argument("--layer-kl-weight", type=float, default=0.10) p.add_argument("--layer-ce-weight", type=float, default=0.05) p.add_argument("--cosine-weight", type=float, default=0.25) p.add_argument("--accept-nmse", type=float, default=0.20) p.add_argument("--accept-cosine", type=float, default=0.90) p.add_argument("--accept-incremental-delta-nll", type=float, default=0.015) p.add_argument("--accept-cumulative-delta-nll", type=float, default=0.05) p.add_argument("--probe-blocks", type=int, default=4) p.add_argument("--probe-context", type=int, default=128) p.add_argument("--strict-acceptance", action=argparse.BooleanOptionalAction, default=True) p.add_argument("--core-o-tokens", type=int, default=50_000) p.add_argument("--core-o-lr", type=float, default=3e-5) p.add_argument("--norm-tokens", type=int, default=50_000) p.add_argument("--norm-lr", type=float, default=8e-6) p.add_argument("--full-tokens", type=int, default=100_000) p.add_argument("--full-lr", type=float, default=3e-6) p.add_argument("--grad-accum", type=int, default=4) p.add_argument("--temperature", type=float, default=2.0) p.add_argument("--ce-weight", type=float, default=0.25) p.add_argument("--kl-weight", type=float, default=1.0) p.add_argument("--hidden-weight", type=float, default=0.5) p.add_argument("--resume", action=argparse.BooleanOptionalAction, default=True) p.add_argument("--backup-every-updates", type=int, default=50) p.add_argument("--max-runtime-minutes", type=float, default=240.0) p.add_argument("--log-every", type=int, default=10) return p.parse_args() def set_seed(seed: int) -> None: random.seed(seed) torch.manual_seed(seed) if torch.cuda.is_available(): torch.cuda.manual_seed_all(seed) def choose_dtype(device: torch.device) -> torch.dtype: if device.type != "cuda": return torch.float32 return torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16 def amp_factory(device: torch.device, dtype: torch.dtype): if device.type == "cuda": return lambda: torch.autocast("cuda", dtype=dtype) return nullcontext def token_blocks(dataset, tokenizer, text_field: str, context_length: int): eos = tokenizer.eos_token_id if eos is None: raise ValueError("tokenizer must define eos_token_id") buffer: list[int] = [] for row in dataset: text = str(row.get(text_field, "")).strip() if not text: continue ids = tokenizer(text, add_special_tokens=False, verbose=False)["input_ids"] if not ids: continue buffer.extend(ids) buffer.append(eos) while len(buffer) >= context_length: yield torch.tensor(buffer[:context_length], dtype=torch.long) del buffer[:context_length] def batches(blocks, batch_size: int): pending = [] for block in blocks: pending.append(block) if len(pending) == batch_size: yield torch.stack(pending) pending.clear() def make_optimizer(params, lr: float, device: torch.device): params = list(params) if not params: raise ValueError("optimizer received no trainable parameters") try: return torch.optim.AdamW(params, lr=lr, weight_decay=0.01, fused=(device.type == "cuda")) except Exception: return torch.optim.AdamW(params, lr=lr, weight_decay=0.01) def distillation_kl(student_logits, teacher_logits, temperature: float) -> torch.Tensor: s = student_logits.float() / temperature t = teacher_logits.float() / temperature per_token = F.kl_div( F.log_softmax(s, dim=-1), F.softmax(t, dim=-1), reduction="none", ).sum(dim=-1) return per_token.mean() * (temperature ** 2) def representation_loss(student_hidden, teacher_hidden) -> torch.Tensor: max_idx = min(len(student_hidden), len(teacher_hidden)) - 1 candidates = (6, 12, 18, 24, 30) indices = [i for i in candidates if i <= max_idx] if not indices: indices = [max_idx] terms = [] for idx in indices: s = student_hidden[idx].float() t = teacher_hidden[idx].float() cosine = 1.0 - F.cosine_similarity(s, t, dim=-1).mean() nmse = (s - t).square().mean() / t.square().mean().clamp_min(1e-5) terms.append(cosine + 0.25 * nmse) return torch.stack(terms).mean() def alignment_metrics(pred: torch.Tensor, target: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: pred32 = pred.float() target32 = target.float() mse = (pred32 - target32).square().mean() nmse = mse / target32.square().mean().clamp_min(1e-5) cosine = F.cosine_similarity(pred32, target32, dim=-1).mean() return nmse, cosine def alignment_loss(pred: torch.Tensor, target: torch.Tensor, cosine_weight: float) -> torch.Tensor: nmse, cosine = alignment_metrics(pred, target) return nmse + cosine_weight * (1.0 - cosine) def _detach_tree(value): if torch.is_tensor(value): return value.detach() if isinstance(value, tuple): return tuple(_detach_tree(v) for v in value) if isinstance(value, list): return [_detach_tree(v) for v in value] return value def capture_attention_input(model, ids: torch.Tensor, layer_idx: int, amp, *, with_output: bool): capture: dict[str, Any] = {} module = model.model.layers[layer_idx].self_attn def pre_hook(mod, args, kwargs): hidden = args[0] if args else kwargs.get("hidden_states") if hidden is None: raise RuntimeError("attention hidden_states not found") capture["hidden"] = hidden for key in ( "position_embeddings", "position_ids", "attention_mask", "cache_position", ): if key in kwargs and kwargs[key] is not None: capture[key] = _detach_tree(kwargs[key]) handle = module.register_forward_pre_hook(pre_hook, with_kwargs=True) try: if with_output: with amp(): out = model( input_ids=ids, labels=ids, use_cache=False, output_hidden_states=True, return_dict=True, ) else: with torch.no_grad(), amp(): out = model( input_ids=ids, use_cache=False, output_hidden_states=False, return_dict=True, ) finally: handle.remove() if "hidden" not in capture: raise RuntimeError(f"failed to capture layer {layer_idx} attention input") return capture, out def attention_kwargs_from_capture(capture: dict[str, Any]) -> dict[str, Any]: kwargs: dict[str, Any] = {"use_cache": False} for key in ( "position_embeddings", "position_ids", "attention_mask", "cache_position", ): if key in capture: kwargs[key] = capture[key] return kwargs def call_attention(module, hidden: torch.Tensor, kwargs: dict[str, Any]) -> torch.Tensor: out = module(hidden, **kwargs) if isinstance(out, (tuple, list)): return out[0] return out def alpha_for_step(step: int, max_steps: int, start: float, end: float) -> float: if max_steps <= 1: return float(end) progress = min(max((step - 1) / (max_steps - 1), 0.0), 1.0) return float(start + progress * (end - start)) def freeze_current_layer_only(student, layer_idx: int) -> list[torch.nn.Parameter]: for p in student.parameters(): p.requires_grad = False module = student.model.layers[layer_idx].self_attn if not isinstance(module, MemoryFusionLlamaAttention): raise TypeError(f"layer {layer_idx} is not MemoryFusionLlamaAttention") trainable = [] for p in module.core.parameters(): p.requires_grad = True trainable.append(p) for p in module.o_proj.parameters(): p.requires_grad = True trainable.append(p) return trainable def select_integrated_stage_params(student, stage: str) -> list[torch.nn.Parameter]: for p in student.parameters(): p.requires_grad = False if stage in {"core_o", "core_o_norm"}: for module in student.modules(): if isinstance(module, MemoryFusionLlamaAttention): for p in module.core.parameters(): p.requires_grad = True for p in module.o_proj.parameters(): p.requires_grad = True if stage == "core_o_norm": for name, p in student.named_parameters(): if "norm" in name.lower(): p.requires_grad = True elif stage == "full": for p in student.parameters(): p.requires_grad = True else: raise ValueError(f"unknown integrated stage {stage!r}") return [p for p in student.parameters() if p.requires_grad] def _layer_prefix(idx: int) -> str: return f"model.layers.{idx}.self_attn." def progressive_attention_state(student, accepted_layers: list[int]) -> dict[str, torch.Tensor]: prefixes = tuple(_layer_prefix(i) for i in accepted_layers) return { k: v.detach().cpu() for k, v in student.state_dict().items() if prefixes and k.startswith(prefixes) } def save_progress( output_dir: Path, student, config: SmolMemoryFusionConfig, accepted_layers: list[int], layer_reports: list[dict], *, stage: str, ) -> None: output_dir.mkdir(parents=True, exist_ok=True) payload = { "format_version": 1, "stage": stage, "accepted_layers": accepted_layers, "config": config.to_dict(), "layer_reports": layer_reports, "attention_state": progressive_attention_state(student, accepted_layers), } tmp = output_dir / "sequential_progress.tmp" final = output_dir / "sequential_progress.pt" torch.save(payload, tmp) tmp.replace(final) (output_dir / "sequential_progress.json").write_text( json.dumps( { "format_version": 1, "stage": stage, "accepted_layers": accepted_layers, "config": config.to_dict(), "layer_reports": layer_reports, }, indent=2, ), encoding="utf-8", ) def load_progress_if_available( output_dir: Path, student, config: SmolMemoryFusionConfig, *, resume: bool, ) -> tuple[list[int], list[dict]]: path = output_dir / "sequential_progress.pt" if not resume or not path.exists(): return [], [] payload = torch.load(path, map_location="cpu", weights_only=False) accepted = [int(x) for x in payload.get("accepted_layers", [])] if payload.get("config", {}).get("memory_rank") != config.memory_rank: raise RuntimeError("resume checkpoint memory rank differs from requested rank") if accepted: replace_attention_layers(student, config, accepted) incompatible = student.load_state_dict(payload["attention_state"], strict=False) expected = { k for k in student.state_dict() if any(k.startswith(_layer_prefix(i)) for i in accepted) } missing = [k for k in incompatible.missing_keys if k in expected] if missing: raise RuntimeError(f"resume checkpoint missing accepted-layer keys: {missing[:8]}") print(f"RESUME: accepted layers={accepted}") return accepted, list(payload.get("layer_reports", [])) def save_full_state( output_dir: Path, student, config: SmolMemoryFusionConfig, metadata: dict, *, filename: str = "smollm2_memory_fusion_sequential_full.pt", ) -> None: output_dir.mkdir(parents=True, exist_ok=True) state_path = output_dir / filename tmp = output_dir / (filename + ".tmp") torch.save({k: v.detach().cpu() for k, v in student.state_dict().items()}, tmp) tmp.replace(state_path) meta = { "format_version": 1, "architecture": "smollm2-memory-fusion-sequential-full", "base_model": metadata["base_model"], "memory_fusion": config.to_dict(), "training": metadata, "state_file": filename, } (output_dir / "smollm2_memory_fusion_sequential_config.json").write_text( json.dumps(meta, indent=2), encoding="utf-8", ) def build_probe_blocks(tokenizer, context: int, count: int) -> list[torch.Tensor]: try: wiki = load_dataset( "Salesforce/wikitext", "wikitext-2-raw-v1", split="validation", ) except Exception: texts = [ "The history of science is shaped by observation, measurement, and careful comparison.", "A language model predicts the next token from the sequence that came before it.", "Vienna is the capital of Austria and has a long history of music, science, and public administration.", "Neural networks can approximate complex functions when their parameters are trained on representative data.", "The experiment should separate training data from the probe used to decide whether a replacement is acceptable.", "Reliable evaluation requires the same inputs, tokenizer, context length, and scoring rule for every model.", "A recurrent memory can carry information through a sequence without constructing a full dense attention matrix.", "When a model component is replaced, small approximation errors can accumulate across many layers.", ] stream = "\n".join(texts) else: stream = "\n".join(str(x) for x in wiki["text"] if str(x).strip()) ids = tokenizer(stream, add_special_tokens=False)["input_ids"] need = context blocks = [] for start in range(0, len(ids) - need + 1, need): blocks.append(torch.tensor(ids[start : start + need], dtype=torch.long)) if len(blocks) >= count: break if not blocks: raise RuntimeError("could not construct probe blocks") return blocks @torch.no_grad() def probe_nll(model, probe_blocks: list[torch.Tensor], device: torch.device, amp) -> float: model.eval() losses = [] for block in probe_blocks: x = block.unsqueeze(0).to(device) with amp(): out = model(input_ids=x, labels=x, use_cache=False, return_dict=True) losses.append(float(out.loss.detach().float())) return sum(losses) / len(losses) @torch.no_grad() def real_hidden_function_metrics( teacher, student, ids: torch.Tensor, layer_idx: int, amp, ) -> tuple[float, float]: t_capture, _ = capture_attention_input(teacher, ids, layer_idx, amp, with_output=False) s_capture, _ = capture_attention_input(student, ids, layer_idx, amp, with_output=False) real_hidden = s_capture["hidden"].detach() kwargs = attention_kwargs_from_capture(t_capture) target = call_attention(teacher.model.layers[layer_idx].self_attn, real_hidden, kwargs) pred = call_attention(student.model.layers[layer_idx].self_attn, real_hidden, kwargs) nmse, cosine = alignment_metrics(pred, target) return float(nmse), float(cosine) def acceptance_passes( *, nmse: float, cosine: float, incremental_delta_nll: float, cumulative_delta_nll: float, args: argparse.Namespace, ) -> bool: return ( nmse <= args.accept_nmse and cosine >= args.accept_cosine and incremental_delta_nll <= args.accept_incremental_delta_nll and cumulative_delta_nll <= args.accept_cumulative_delta_nll ) def train_one_replacement( *, teacher, student, layer_idx: int, batch_iter, probe_blocks, teacher_probe_nll: float, pre_replacement_probe_nll: float, args, device, amp, ) -> dict: trainable = freeze_current_layer_only(student, layer_idx) student.train() optimizer = make_optimizer(trainable, args.layer_lr, device) scaler = torch.amp.GradScaler( "cuda", enabled=(device.type == "cuda" and choose_dtype(device) == torch.float16) ) best: dict[str, float | int | bool] | None = None accepted = False for step in range(1, args.max_layer_steps + 1): ids = next(batch_iter).to(device, non_blocking=True) t_capture, t_out = capture_attention_input( teacher, ids, layer_idx, amp, with_output=True ) s_capture, s_out = capture_attention_input( student, ids, layer_idx, amp, with_output=True ) alpha = alpha_for_step( step, args.max_layer_steps, args.teacher_alpha_start, args.teacher_alpha_end, ) mixed_hidden = ( alpha * t_capture["hidden"].detach() + (1.0 - alpha) * s_capture["hidden"].detach() ) kwargs = attention_kwargs_from_capture(t_capture) with torch.no_grad(), amp(): target = call_attention( teacher.model.layers[layer_idx].self_attn, mixed_hidden, kwargs, ) with amp(): pred = call_attention( student.model.layers[layer_idx].self_attn, mixed_hidden, kwargs, ) functional = alignment_loss(pred, target, args.cosine_weight) kl = distillation_kl(s_out.logits, t_out.logits, args.temperature) total = functional + args.layer_kl_weight * kl + args.layer_ce_weight * s_out.loss if not torch.isfinite(total): raise RuntimeError(f"non-finite loss at layer {layer_idx}, step {step}") optimizer.zero_grad(set_to_none=True) scaler.scale(total).backward() scaler.unscale_(optimizer) grad_norm = float(torch.nn.utils.clip_grad_norm_(trainable, 1.0)) scaler.step(optimizer) scaler.update() if step == 1 or step % args.log_every == 0: nmse_batch, cosine_batch = alignment_metrics(pred.detach(), target.detach()) print( f"layer={layer_idx:02d} step={step:03d}/{args.max_layer_steps} " f"alpha={alpha:.3f} functional={float(functional.detach()):.4f} " f"nmse={float(nmse_batch):.4f} cos={float(cosine_batch):.4f} " f"kl={float(kl.detach()):.4f} ce={float(s_out.loss.detach()):.4f} " f"grad={grad_norm:.3f}" ) should_check = ( step >= args.min_layer_steps and (step % args.check_every == 0 or step == args.max_layer_steps) ) if should_check: nmse, cosine = real_hidden_function_metrics( teacher, student, ids, layer_idx, amp ) current_probe_nll = probe_nll(student, probe_blocks, device, amp) student.train() incremental = current_probe_nll - pre_replacement_probe_nll cumulative = current_probe_nll - teacher_probe_nll passed = acceptance_passes( nmse=nmse, cosine=cosine, incremental_delta_nll=incremental, cumulative_delta_nll=cumulative, args=args, ) candidate = { "step": step, "nmse": nmse, "cosine": cosine, "probe_nll": current_probe_nll, "incremental_delta_nll": incremental, "cumulative_delta_nll": cumulative, "passed": passed, } if best is None or ( candidate["nmse"] < best["nmse"] and candidate["cumulative_delta_nll"] <= best["cumulative_delta_nll"] + 0.01 ): best = candidate print( f" ACCEPTANCE CHECK layer={layer_idx:02d}: " f"NMSE={nmse:.4f} (≤{args.accept_nmse:.4f}) " f"cos={cosine:.4f} (≥{args.accept_cosine:.4f}) " f"ΔNLL_inc={incremental:+.5f} (≤{args.accept_incremental_delta_nll:+.5f}) " f"ΔNLL_total={cumulative:+.5f} (≤{args.accept_cumulative_delta_nll:+.5f}) " f"=> {'PASS' if passed else 'continue'}" ) if passed: accepted = True best = candidate break del t_out, s_out, pred, target, total if best is None: raise RuntimeError("acceptance was never evaluated") return { "layer": layer_idx, "accepted": accepted, "steps": int(best["step"]), "nmse": float(best["nmse"]), "cosine": float(best["cosine"]), "probe_nll": float(best["probe_nll"]), "incremental_delta_nll": float(best["incremental_delta_nll"]), "cumulative_delta_nll": float(best["cumulative_delta_nll"]), } def run_integrated_stage( *, teacher, student, batch_iter, stage_name: str, token_budget: int, lr: float, args, device, amp, output_dir: Path, config: SmolMemoryFusionConfig, report: dict, ) -> dict: if token_budget <= 0: return {"stage": stage_name, "tokens": 0, "updates": 0, "skipped": True} trainable = select_integrated_stage_params(student, stage_name) optimizer = make_optimizer(trainable, lr, device) scaler = torch.amp.GradScaler( "cuda", enabled=(device.type == "cuda" and choose_dtype(device) == torch.float16) ) student.train() optimizer.zero_grad(set_to_none=True) seen_tokens = 0 micro = 0 updates = 0 last = {} start = time.perf_counter() while seen_tokens < token_budget: ids = next(batch_iter).to(device, non_blocking=True) with torch.no_grad(), amp(): t_out = teacher( input_ids=ids, use_cache=False, output_hidden_states=True, return_dict=True, ) with amp(): s_out = student( input_ids=ids, labels=ids, use_cache=False, output_hidden_states=True, return_dict=True, ) kl = distillation_kl(s_out.logits, t_out.logits, args.temperature) hidden = representation_loss(s_out.hidden_states, t_out.hidden_states) loss = ( args.ce_weight * s_out.loss + args.kl_weight * kl + args.hidden_weight * hidden ) scaled = loss / args.grad_accum if not torch.isfinite(scaled): raise RuntimeError(f"non-finite loss during integrated stage {stage_name}") scaler.scale(scaled).backward() micro += 1 seen_tokens += ids.numel() last = { "ce": float(s_out.loss.detach().float()), "kl": float(kl.detach().float()), "hidden": float(hidden.detach().float()), "total": float(loss.detach().float()), } del t_out, s_out if micro % args.grad_accum: continue scaler.unscale_(optimizer) grad_norm = float(torch.nn.utils.clip_grad_norm_(trainable, 1.0)) scaler.step(optimizer) scaler.update() optimizer.zero_grad(set_to_none=True) updates += 1 if updates == 1 or updates % args.log_every == 0: print( f"stage={stage_name} update={updates} tokens={seen_tokens:,}/{token_budget:,} " f"ce={last['ce']:.4f} kl={last['kl']:.4f} hidden={last['hidden']:.4f} " f"grad={grad_norm:.3f}" ) if args.backup_every_updates > 0 and updates % args.backup_every_updates == 0: backup_meta = dict(report) backup_meta["integrated_stage"] = stage_name backup_meta["integrated_stage_tokens"] = seen_tokens save_full_state( output_dir, student, config, backup_meta, filename="live_sequential_full_state.pt", ) print(" persistent full-state backup saved") elapsed = time.perf_counter() - start return { "stage": stage_name, "tokens": seen_tokens, "updates": updates, "lr": lr, "elapsed_minutes": elapsed / 60.0, "last": last, "trainable_parameters": sum(p.numel() for p in trainable), } def main() -> None: args = parse_args() set_seed(args.seed) output_dir = Path(args.output_dir) output_dir.mkdir(parents=True, exist_ok=True) device = torch.device("cuda" if torch.cuda.is_available() else "cpu") dtype = choose_dtype(device) amp = amp_factory(device, dtype) if device.type == "cuda": torch.cuda.reset_peak_memory_stats() torch.backends.cuda.matmul.allow_tf32 = True print(f"device={device} dtype={dtype} memory_rank={args.memory_rank}") print(f"output_dir={output_dir}") tokenizer = AutoTokenizer.from_pretrained(args.base_model, use_fast=True) if tokenizer.pad_token_id is None: tokenizer.pad_token = tokenizer.eos_token teacher = AutoModelForCausalLM.from_pretrained(args.base_model, dtype=dtype).to(device) teacher.eval() teacher.config.use_cache = False teacher.requires_grad_(False) student = AutoModelForCausalLM.from_pretrained(args.base_model, dtype=dtype).to(device) student.config.use_cache = False config = SmolMemoryFusionConfig( feature_dim=args.feature_dim, memory_rank=args.memory_rank, train_output_projection=True, ) accepted_layers, layer_reports = load_progress_if_available( output_dir, student, config, resume=args.resume ) raw = load_dataset( args.dataset, name=args.dataset_config, split=args.split, streaming=True, ).shuffle(seed=args.seed, buffer_size=args.shuffle_buffer) batch_iter = iter( batches( token_blocks(raw, tokenizer, args.text_field, args.context_length), args.batch_size, ) ) probe_blocks = build_probe_blocks( tokenizer, context=args.probe_context, count=args.probe_blocks, ) teacher_probe_nll = probe_nll(teacher, probe_blocks, device, amp) print(f"teacher probe NLL={teacher_probe_nll:.6f}") start_time = time.perf_counter() num_layers = int(student.config.num_hidden_layers) for layer_idx in range(num_layers): if layer_idx in accepted_layers: continue if accepted_layers != list(range(layer_idx)): raise RuntimeError( f"accepted layers must be a contiguous prefix before layer {layer_idx}: " f"{accepted_layers}" ) if (time.perf_counter() - start_time) / 60.0 >= args.max_runtime_minutes * 0.70: raise RuntimeError( "runtime budget reached during sequential replacement; progress was " "saved and the same command can resume from the last accepted layer" ) print("\n" + "=" * 110) print(f"SEQUENTIAL REPLACEMENT: layer {layer_idx}/{num_layers - 1}") print("=" * 110) pre_probe_nll = probe_nll(student, probe_blocks, device, amp) print( f"before replacement: probe NLL={pre_probe_nll:.6f}, " f"Δ vs teacher={pre_probe_nll - teacher_probe_nll:+.6f}" ) replace_attention_layers(student, config, [layer_idx]) layer_report = train_one_replacement( teacher=teacher, student=student, layer_idx=layer_idx, batch_iter=batch_iter, probe_blocks=probe_blocks, teacher_probe_nll=teacher_probe_nll, pre_replacement_probe_nll=pre_probe_nll, args=args, device=device, amp=amp, ) layer_reports.append(layer_report) if not layer_report["accepted"] and args.strict_acceptance: save_progress( output_dir, student, config, accepted_layers, layer_reports, stage=f"layer_{layer_idx}_rejected", ) (output_dir / "sequential_training_report.json").write_text( json.dumps( { "status": "stopped_on_rejected_layer", "base_model": args.base_model, "memory_rank": args.memory_rank, "accepted_layers": accepted_layers, "layer_reports": layer_reports, "thresholds": { "nmse": args.accept_nmse, "cosine": args.accept_cosine, "incremental_delta_nll": args.accept_incremental_delta_nll, "cumulative_delta_nll": args.accept_cumulative_delta_nll, }, }, indent=2, ), encoding="utf-8", ) raise RuntimeError( f"layer {layer_idx} did not satisfy acceptance criteria; " "the next Transformer layer was NOT replaced" ) if not layer_report["accepted"]: print("WARNING: relaxed mode accepts a layer that did not pass all thresholds") accepted_layers.append(layer_idx) save_progress( output_dir, student, config, accepted_layers, layer_reports, stage=f"accepted_layer_{layer_idx}", ) print(f"✅ accepted layer {layer_idx}; persistent progress saved") if accepted_layers != list(range(num_layers)): raise RuntimeError("not all layers were accepted") summary = structural_summary(student) if summary["memory_fusion_layers"] != num_layers or summary["transformer_attention_layers"] != 0: raise RuntimeError(f"unexpected final structure: {summary}") report = { "status": "all_layers_accepted", "architecture": "smollm2-memory-fusion-sequential-full", "base_model": args.base_model, "memory_rank": args.memory_rank, "feature_dim": args.feature_dim, "context_length": args.context_length, "teacher_probe_nll": teacher_probe_nll, "accepted_layers": accepted_layers, "layer_reports": layer_reports, "thresholds": { "nmse": args.accept_nmse, "cosine": args.accept_cosine, "incremental_delta_nll": args.accept_incremental_delta_nll, "cumulative_delta_nll": args.accept_cumulative_delta_nll, }, "integrated_stages": [], } print("\n" + "=" * 110) print("ALL 30 REPLACEMENTS ACCEPTED — STARTING INTEGRATED TRAINING") print("=" * 110) integrated_specs = [ ("core_o", args.core_o_tokens, args.core_o_lr), ("core_o_norm", args.norm_tokens, args.norm_lr), ("full", args.full_tokens, args.full_lr), ] for stage_name, token_budget, lr in integrated_specs: if (time.perf_counter() - start_time) / 60.0 >= args.max_runtime_minutes: print("runtime budget reached before remaining integrated stages") report["status"] = "runtime_budget_after_acceptance" break print(f"\n--- integrated stage: {stage_name} ---") stage_report = run_integrated_stage( teacher=teacher, student=student, batch_iter=batch_iter, stage_name=stage_name, token_budget=token_budget, lr=lr, args=args, device=device, amp=amp, output_dir=output_dir, config=config, report=report, ) report["integrated_stages"].append(stage_report) save_full_state( output_dir, student, config, report, filename=f"stage_{stage_name}_full_state.pt", ) print(f"✅ completed {stage_name}; full persistent checkpoint saved") student.eval() final_probe_nll = probe_nll(student, probe_blocks, device, amp) report["final_probe_nll"] = final_probe_nll report["final_probe_delta_nll"] = final_probe_nll - teacher_probe_nll report["parameters"] = parameter_summary(student) report["structure"] = structural_summary(student) report["elapsed_minutes"] = (time.perf_counter() - start_time) / 60.0 report["peak_vram_gib"] = ( torch.cuda.max_memory_allocated() / (1024 ** 3) if device.type == "cuda" else 0.0 ) save_full_state(output_dir, student, config, report) (output_dir / "sequential_training_report.json").write_text( json.dumps(report, indent=2), encoding="utf-8", ) tokenizer.save_pretrained(output_dir) print("\nFINAL REPORT") print(json.dumps(report, indent=2)) print("saved:", output_dir) if __name__ == "__main__": main()