Buckets:

Rishik001
/

nmt-loop

Files

xet

Rishik001/nmt-loop / looping-1709 /main.py

Rishik001

about 2 months ago

download

raw

17.5 kB

"""
pip install pyyaml datasets wandb transformers muon torch
"""

import os
import sys
import yaml
import shutil
import time
from datetime import datetime

import torch
import torch.nn as nn
import torch.optim as optim

from model import LoopedTransformerAutoEncoder
from data import (
    setup_tokenizer,
    create_dataloader,
    load_reconstruction_dataset,
    apply_token_dropout,
)


def load_config(path: str) -> dict:
    with open(path) as f:
        return yaml.safe_load(f)


class TrapezoidLR:
    """Trapezoidal LR schedule: linear warmup -> hold at peak -> linear decay."""

    def __init__(self, optimizer, peak_lr, warmup_steps, hold_steps, total_steps, min_lr=1e-7):
        self.optimizer = optimizer
        self.peak_lr = peak_lr
        self.min_lr = min_lr
        self.warmup_steps = warmup_steps
        self.hold_steps = hold_steps
        self.total_steps = total_steps
        self.decay_steps = max(1, total_steps - warmup_steps - hold_steps)
        self._step = 0
        self._set_lr(min_lr)

    def step(self):
        self._step += 1
        lr = self._compute_lr(self._step)
        self._set_lr(lr)

    def _compute_lr(self, step):
        if step < self.warmup_steps:
            return self.min_lr + (self.peak_lr - self.min_lr) * (step / max(1, self.warmup_steps))
        elif step < self.warmup_steps + self.hold_steps:
            return self.peak_lr
        else:
            decay_progress = (step - self.warmup_steps - self.hold_steps) / self.decay_steps
            decay_progress = min(1.0, decay_progress)
            return self.peak_lr - (self.peak_lr - self.min_lr) * decay_progress

    def _set_lr(self, lr):
        for pg in self.optimizer.param_groups:
            pg["lr"] = lr

    def state_dict(self):
        return {"step": self._step}

    def load_state_dict(self, d):
        self._step = d["step"]


def compute_component_grad_norms(model):
    """Compute grad norms for each major component of the model."""
    components = {
        "embedding": [],
        "encoder_layer": [],
        "decoder_layer": [],
        "enc_proj": [],
        "lm_head": [],
    }

    for name, param in model.named_parameters():
        if param.grad is None:
            continue
        norm = param.grad.detach().norm().item()
        if "embedding" in name or "pos_encoding" in name:
            components["embedding"].append(norm)
        elif "encoder_layer" in name:
            components["encoder_layer"].append(norm)
        elif "decoder_layer" in name:
            components["decoder_layer"].append(norm)
        elif "enc_proj" in name or "enc_to_token" in name:
            components["enc_proj"].append(norm)
        elif "lm_head" in name:
            components["lm_head"].append(norm)

    result = {}
    for key, norms in components.items():
        if norms:
            result[key] = (sum(n ** 2 for n in norms)) ** 0.5
    return result


def validate_config(cfg: dict):
    for section in ["model", "training", "data", "checkpoint"]:
        assert section in cfg, f"Missing config section: {section}"

    assert cfg["training"]["optimizer"] in ("adam", "adamw", "muon")

    cfg.setdefault("dry_run", False)
    cfg.setdefault("dry_run_samples", 10)
    cfg["checkpoint"].setdefault("save_intermediate", True)
    cfg["checkpoint"].setdefault("save_every_n_steps", 1000)
    cfg.setdefault("wandb", {"project": "loop-reconstruction", "enabled": False})
    cfg["training"].setdefault("token_dropout", 0.15)
    cfg["training"].setdefault("validation_steps", 1000)
    cfg["training"].setdefault("lr_schedule", "trapezoid")
    cfg["training"].setdefault("warmup_steps", 1000)
    cfg["training"].setdefault("hold_steps", 5000)
    cfg["training"].setdefault("log_every_n_steps", 1)
    cfg["training"].setdefault("print_every_n_steps", 50)


def create_optimizer(model: nn.Module, cfg: dict):
    opt_name = cfg["training"]["optimizer"]
    lr = cfg["training"]["learning_rate"]
    wd = cfg["training"].get("weight_decay", 0.01)

    if opt_name == "adam":
        return optim.Adam(model.parameters(), lr=lr)
    elif opt_name == "adamw":
        return optim.AdamW(model.parameters(), lr=lr, weight_decay=wd)
    elif opt_name == "muon":
        try:
            from muon import Muon
        except ImportError:
            raise ImportError("Muon not found. Install with: pip install muon")

        muon_params = [p for _, p in model.named_parameters() if p.ndim >= 2]
        adamw_params = [p for _, p in model.named_parameters() if p.ndim < 2]
        return Muon(
            muon_params=muon_params,
            lr=lr * 10,
            adamw_params=adamw_params,
            adamw_lr=lr,
            adamw_wd=wd,
        )


def create_run_dir(cfg: dict) -> str:
    base = cfg["checkpoint"]["save_dir"]
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    loops = cfg["model"]["n_encoder_loops"]
    run_name = f"loop{loops}_{cfg['training']['optimizer']}_{timestamp}"
    run_dir = os.path.join(base, run_name)
    os.makedirs(run_dir, exist_ok=True)
    return run_dir


def save_checkpoint(model, optimizer, scheduler, step, epoch, loss, run_dir, tag="step"):
    ckpt = {
        "step": step,
        "epoch": epoch,
        "loss": loss,
        "model_state_dict": model.state_dict(),
        "optimizer_state_dict": optimizer.state_dict(),
        "scheduler_state_dict": scheduler.state_dict() if scheduler else None,
    }
    path = os.path.join(run_dir, f"checkpoint_{tag}_{step}.pt")
    torch.save(ckpt, path)
    print(f"Saved checkpoint: {path}")
    return path


def compute_loss(logits, target_ids, target_mask, pad_id):
    min_len = min(logits.size(1), target_ids.size(1))
    pred = logits[:, :min_len, :]
    labels = target_ids[:, :min_len].clone()
    mask = target_mask[:, :min_len]
    labels[mask == 0] = -100

    return nn.CrossEntropyLoss(ignore_index=-100)(
        pred.reshape(-1, pred.size(-1)), labels.reshape(-1)
    )


@torch.no_grad()
def compute_accuracy(logits, target_ids, target_mask):
    min_len = min(logits.size(1), target_ids.size(1))
    preds = logits[:, :min_len].argmax(dim=-1)
    labels = target_ids[:, :min_len]
    mask = target_mask[:, :min_len]

    correct = ((preds == labels).float() * mask).sum()
    total = mask.sum()
    return (correct / (total + 1e-8)).item()


@torch.no_grad()
def validate(model, val_loader, device, cfg, tokenizer):
    model.eval()
    total_loss = 0.0
    total_correct = 0
    total_tokens = 0
    pad_id = tokenizer.pad_token_id
    num_mem = cfg["model"]["num_memory_tokens"]

    for batch in val_loader:
        ids = batch["input_ids"].to(device)
        mask = batch["attention_mask"].to(device)
        target_ids = batch["original_input_ids"].to(device)
        target_mask = batch["original_attention_mask"].to(device)

        with torch.amp.autocast("cuda", enabled=device.type == "cuda"):
            logits, _ = model(ids, mask, target_ids, target_mask)
            loss = compute_loss(logits, target_ids, target_mask, pad_id)

        min_len = min(logits.size(1), target_ids.size(1))
        valid_mask = target_mask[:, :min_len].bool()
        token_count = valid_mask.sum().item()
        total_loss += loss.item() * token_count
        total_correct += (
            (logits[:, :min_len].argmax(dim=-1) == target_ids[:, :min_len]) & valid_mask
        ).sum().item()
        total_tokens += token_count

    model.train()
    return total_loss / max(1, total_tokens), total_correct / max(1, total_tokens)


@torch.no_grad()
def validate_free_running(model, val_loader, device):
    """Autoregressively reconstruct every sequence in the validation set."""
    model.eval()
    total_correct = 0
    total_tokens = 0
    exact_sequences = 0
    total_sequences = 0

    for batch in val_loader:
        ids = batch["input_ids"].to(device)
        mask = batch["attention_mask"].to(device)
        targets = batch["original_input_ids"].to(device)
        target_mask = batch["original_attention_mask"].to(device).bool()
        with torch.amp.autocast("cuda", enabled=device.type == "cuda"):
            generated = model.generate(ids, mask, max_length=targets.size(1))
        matches = (generated == targets) & target_mask
        total_correct += matches.sum().item()
        total_tokens += target_mask.sum().item()
        exact_sequences += (matches | ~target_mask).all(dim=1).sum().item()
        total_sequences += targets.size(0)

    model.train()
    return (
        total_correct / max(1, total_tokens),
        exact_sequences / max(1, total_sequences),
        total_sequences,
    )


def main(config_path: str):
    cfg = load_config(config_path)
    validate_config(cfg)
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"Device: {device}")

    run_dir = create_run_dir(cfg)
    shutil.copy2(config_path, os.path.join(run_dir, "config.yaml"))
    print(f"Run directory: {run_dir}")

    # Data
    mcfg = cfg["model"]
    num_mem = mcfg["num_memory_tokens"]
    tokenizer = setup_tokenizer(num_memory_tokens=num_mem)

    train_ds, val_ds = load_reconstruction_dataset(cfg, tokenizer)

    train_loader = create_dataloader(
        cfg["training"]["batch_size"], train_ds, tokenizer,
        shuffle=True, num_memory_tokens=num_mem,
        max_length=cfg["data"].get("max_length", 32),
        slot_source=mcfg.get("slot_source", "appended"),
    )
    val_loader = create_dataloader(
        cfg["training"].get("val_batch_size", cfg["training"]["batch_size"]),
        val_ds, tokenizer,
        shuffle=False, num_memory_tokens=num_mem,
        max_length=cfg["data"].get("max_length", 32),
        slot_source=mcfg.get("slot_source", "appended"),
    )

    # Model
    vocab_size = len(tokenizer)
    model = LoopedTransformerAutoEncoder(
        vocab_size=vocab_size,
        d_model=mcfg["d_model"],
        n_heads=mcfg["n_heads"],
        d_ff=mcfg.get("d_ff", mcfg["d_model"] * 4),
        n_encoder_loops=mcfg["n_encoder_loops"],
        n_decoder_loops=mcfg["n_decoder_loops"],
        num_memory_tokens=num_mem,
        dropout=mcfg.get("dropout", 0.0),
        max_seq_len=mcfg.get("max_seq_len", 512),
        pad_id=tokenizer.pad_token_id,
        bos_id=tokenizer.bos_token_id,
        memory_conditioning=mcfg.get("memory_conditioning", "prefix"),
        slot_source=mcfg.get("slot_source", "appended"),
    ).to(device)

    num_params = sum(p.numel() for p in model.parameters())
    print(f"Model parameters: {num_params:,}")
    print(f"Encoder loops: {mcfg['n_encoder_loops']}, Decoder loops: {mcfg['n_decoder_loops']}")

    # Optimizer and scheduler
    optimizer = create_optimizer(model, cfg)
    tcfg = cfg["training"]
    total_steps = len(train_loader) * tcfg["num_epochs"]

    lr_schedule = tcfg.get("lr_schedule", "trapezoid")
    if lr_schedule == "trapezoid":
        scheduler = TrapezoidLR(
            optimizer,
            peak_lr=tcfg["learning_rate"],
            warmup_steps=tcfg.get("warmup_steps", 1000),
            hold_steps=tcfg.get("hold_steps", 5000),
            total_steps=total_steps,
            min_lr=tcfg.get("eta_min", 1e-7),
        )
        print(f"LR schedule: trapezoid (warmup={tcfg.get('warmup_steps', 1000)}, hold={tcfg.get('hold_steps', 5000)}, total={total_steps})")
    else:
        scheduler = optim.lr_scheduler.CosineAnnealingLR(
            optimizer, T_max=total_steps, eta_min=tcfg.get("eta_min", 1e-7)
        )
        print(f"LR schedule: cosine (T_max={total_steps})")

    scaler = torch.amp.GradScaler("cuda", enabled=device.type == "cuda")
    pad_id = tokenizer.pad_token_id

    # Wandb
    if cfg["wandb"].get("enabled"):
        import wandb
        wandb.init(
            project=cfg["wandb"]["project"],
            entity=cfg["wandb"].get("entity"),
            config=cfg,
            name=os.path.basename(run_dir),
        )
        wandb.define_metric("global_step")
        wandb.define_metric("train/*", step_metric="global_step")
        wandb.define_metric("grad_norm/*", step_metric="global_step")
        wandb.define_metric("val/*", step_metric="global_step")

    # Training
    global_step = 0
    best_val_loss = float("inf")
    max_grad_norm = tcfg.get("max_grad_norm", 1.0)
    save_every = cfg["checkpoint"]["save_every_n_steps"]
    val_every = tcfg.get("validation_steps", 1000)
    log_every = tcfg.get("log_every_n_steps", 1)
    print_every = tcfg.get("print_every_n_steps", 50)

    print(f"\nTraining for {tcfg['num_epochs']} epochs, ~{total_steps} steps")
    print("=" * 80)

    for epoch in range(tcfg["num_epochs"]):
        model.train()
        epoch_loss = 0.0
        t0 = time.time()

        for step, batch in enumerate(train_loader):
            ids = batch["input_ids"].to(device)
            mask = batch["attention_mask"].to(device)
            target_ids = batch["original_input_ids"].to(device)
            target_mask = batch["original_attention_mask"].to(device)

            protected_memory_tokens = (
                num_mem if mcfg.get("slot_source", "appended") == "appended" else 0
            )
            corrupted, mask = apply_token_dropout(
                ids, mask, tokenizer,
                dropout_prob=tcfg["token_dropout"],
                num_memory_tokens=protected_memory_tokens,
            )

            with torch.amp.autocast("cuda", enabled=device.type == "cuda"):
                logits, enc_vec = model(corrupted, mask, target_ids, target_mask)
                loss = compute_loss(logits, target_ids, target_mask, pad_id)

            scaler.scale(loss).backward()
            scaler.unscale_(optimizer)

            # Per-component grad norms (before clipping)
            grad_norm_total = torch.nn.utils.clip_grad_norm_(model.parameters(), float("inf")).item()
            component_norms = compute_component_grad_norms(model)

            torch.nn.utils.clip_grad_norm_(model.parameters(), max_grad_norm)
            scaler.step(optimizer)
            scaler.update()
            optimizer.zero_grad()
            scheduler.step()

            epoch_loss += loss.item()
            global_step += 1

            if global_step % log_every == 0:
                lr = optimizer.param_groups[0]["lr"]
                avg = epoch_loss / (step + 1)
                acc = compute_accuracy(logits, target_ids, target_mask)

                if cfg["wandb"].get("enabled"):
                    import wandb
                    log_dict = {
                        "train/loss": loss.item(),
                        "train/accuracy": acc,
                        "train/lr": lr,
                        "train/grad_norm": grad_norm_total,
                        "global_step": global_step,
                    }
                    for comp, norm in component_norms.items():
                        log_dict[f"grad_norm/{comp}"] = norm
                    wandb.log(log_dict)

            if global_step % print_every == 0:
                lr = optimizer.param_groups[0]["lr"]
                avg = epoch_loss / (step + 1)
                acc = compute_accuracy(logits, target_ids, target_mask)
                elapsed = time.time() - t0
                print(
                    f"Epoch {epoch+1} | Step {global_step} | "
                    f"Loss: {loss.item():.4f} | Avg: {avg:.4f} | "
                    f"Acc: {acc:.4f} | LR: {lr:.2e} | "
                    f"GradNorm: {grad_norm_total:.3f} | Time: {elapsed:.0f}s"
                )

            if cfg["checkpoint"]["save_intermediate"] and save_every > 0 and global_step % save_every == 0:
                save_checkpoint(model, optimizer, scheduler, global_step, epoch, loss.item(), run_dir)

            if val_every > 0 and global_step % val_every == 0:
                val_loss, val_acc = validate(model, val_loader, device, cfg, tokenizer)
                print(f"  Validation | Loss: {val_loss:.4f} | Acc: {val_acc:.4f}")
                if val_loss < best_val_loss:
                    best_val_loss = val_loss
                    save_checkpoint(model, optimizer, scheduler, global_step, epoch, val_loss, run_dir, tag="best")

                if cfg["wandb"].get("enabled"):
                    import wandb
                    wandb.log({"val/loss": val_loss, "val/accuracy": val_acc, "global_step": global_step})

        # End of epoch
        val_loss, val_acc = validate(model, val_loader, device, cfg, tokenizer)
        print(
            f"\nEpoch {epoch+1} | Train Avg: {epoch_loss / len(train_loader):.4f} | "
            f"Val Loss: {val_loss:.4f} | Val Acc: {val_acc:.4f}"
        )
        print("=" * 80)

    save_checkpoint(model, optimizer, scheduler, global_step, epoch, loss.item(), run_dir, tag="final")
    free_acc, exact_match, evaluated_sequences = validate_free_running(model, val_loader, device)
    print(
        f"Full free-running validation | Acc: {free_acc:.4f} | "
        f"Exact: {exact_match:.4f} | Sequences: {evaluated_sequences}"
    )
    if cfg["wandb"].get("enabled"):
        import wandb
        wandb.log({
            "val/free_running_accuracy": free_acc,
            "val/exact_sequence_accuracy": exact_match,
            "val/evaluated_sequences": evaluated_sequences,
            "global_step": global_step,
        })
    print(f"\nTraining complete. Results in: {run_dir}")

    if cfg["wandb"].get("enabled"):
        import wandb
        wandb.finish()


if __name__ == "__main__":
    config_path = sys.argv[1] if len(sys.argv) > 1 else "config.yaml"
    main(config_path)

Xet Storage Details

Size:: 17.5 kB
Xet hash:: b5245a827c56026b02a72874f295476bf290b23e690b81bcf3512d20aa7f6219

Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.

	"""
	pip install pyyaml datasets wandb transformers muon torch
	"""

	import os
	import sys
	import yaml
	import shutil
	import time
	from datetime import datetime

	import torch
	import torch.nn as nn
	import torch.optim as optim

	from model import LoopedTransformerAutoEncoder
	from data import (
	setup_tokenizer,
	create_dataloader,
	load_reconstruction_dataset,
	apply_token_dropout,
	)


	def load_config(path: str) -> dict:
	with open(path) as f:
	return yaml.safe_load(f)


	class TrapezoidLR:
	"""Trapezoidal LR schedule: linear warmup -> hold at peak -> linear decay."""

	def __init__(self, optimizer, peak_lr, warmup_steps, hold_steps, total_steps, min_lr=1e-7):
	self.optimizer = optimizer
	self.peak_lr = peak_lr
	self.min_lr = min_lr
	self.warmup_steps = warmup_steps
	self.hold_steps = hold_steps
	self.total_steps = total_steps
	self.decay_steps = max(1, total_steps - warmup_steps - hold_steps)
	self._step = 0
	self._set_lr(min_lr)

	def step(self):
	self._step += 1
	lr = self._compute_lr(self._step)
	self._set_lr(lr)

	def _compute_lr(self, step):
	if step < self.warmup_steps:
	return self.min_lr + (self.peak_lr - self.min_lr) * (step / max(1, self.warmup_steps))
	elif step < self.warmup_steps + self.hold_steps:
	return self.peak_lr
	else:
	decay_progress = (step - self.warmup_steps - self.hold_steps) / self.decay_steps
	decay_progress = min(1.0, decay_progress)
	return self.peak_lr - (self.peak_lr - self.min_lr) * decay_progress

	def _set_lr(self, lr):
	for pg in self.optimizer.param_groups:
	pg["lr"] = lr

	def state_dict(self):
	return {"step": self._step}

	def load_state_dict(self, d):
	self._step = d["step"]


	def compute_component_grad_norms(model):
	"""Compute grad norms for each major component of the model."""
	components = {
	"embedding": [],
	"encoder_layer": [],
	"decoder_layer": [],
	"enc_proj": [],
	"lm_head": [],
	}

	for name, param in model.named_parameters():
	if param.grad is None:
	continue
	norm = param.grad.detach().norm().item()
	if "embedding" in name or "pos_encoding" in name:
	components["embedding"].append(norm)
	elif "encoder_layer" in name:
	components["encoder_layer"].append(norm)
	elif "decoder_layer" in name:
	components["decoder_layer"].append(norm)
	elif "enc_proj" in name or "enc_to_token" in name:
	components["enc_proj"].append(norm)
	elif "lm_head" in name:
	components["lm_head"].append(norm)

	result = {}
	for key, norms in components.items():
	if norms:
	result[key] = (sum(n 2 for n in norms)) 0.5
	return result


	def validate_config(cfg: dict):
	for section in ["model", "training", "data", "checkpoint"]:
	assert section in cfg, f"Missing config section: {section}"

	assert cfg["training"]["optimizer"] in ("adam", "adamw", "muon")

	cfg.setdefault("dry_run", False)
	cfg.setdefault("dry_run_samples", 10)
	cfg["checkpoint"].setdefault("save_intermediate", True)
	cfg["checkpoint"].setdefault("save_every_n_steps", 1000)
	cfg.setdefault("wandb", {"project": "loop-reconstruction", "enabled": False})
	cfg["training"].setdefault("token_dropout", 0.15)
	cfg["training"].setdefault("validation_steps", 1000)
	cfg["training"].setdefault("lr_schedule", "trapezoid")
	cfg["training"].setdefault("warmup_steps", 1000)
	cfg["training"].setdefault("hold_steps", 5000)
	cfg["training"].setdefault("log_every_n_steps", 1)
	cfg["training"].setdefault("print_every_n_steps", 50)


	def create_optimizer(model: nn.Module, cfg: dict):
	opt_name = cfg["training"]["optimizer"]
	lr = cfg["training"]["learning_rate"]
	wd = cfg["training"].get("weight_decay", 0.01)

	if opt_name == "adam":
	return optim.Adam(model.parameters(), lr=lr)
	elif opt_name == "adamw":
	return optim.AdamW(model.parameters(), lr=lr, weight_decay=wd)
	elif opt_name == "muon":
	try:
	from muon import Muon
	except ImportError:
	raise ImportError("Muon not found. Install with: pip install muon")

	muon_params = [p for _, p in model.named_parameters() if p.ndim >= 2]
	adamw_params = [p for _, p in model.named_parameters() if p.ndim < 2]
	return Muon(
	muon_params=muon_params,
	lr=lr * 10,
	adamw_params=adamw_params,
	adamw_lr=lr,
	adamw_wd=wd,
	)


	def create_run_dir(cfg: dict) -> str:
	base = cfg["checkpoint"]["save_dir"]
	timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
	loops = cfg["model"]["n_encoder_loops"]
	run_name = f"loop{loops}_{cfg['training']['optimizer']}_{timestamp}"
	run_dir = os.path.join(base, run_name)
	os.makedirs(run_dir, exist_ok=True)
	return run_dir


	def save_checkpoint(model, optimizer, scheduler, step, epoch, loss, run_dir, tag="step"):
	ckpt = {
	"step": step,
	"epoch": epoch,
	"loss": loss,
	"model_state_dict": model.state_dict(),
	"optimizer_state_dict": optimizer.state_dict(),
	"scheduler_state_dict": scheduler.state_dict() if scheduler else None,
	}
	path = os.path.join(run_dir, f"checkpoint_{tag}_{step}.pt")
	torch.save(ckpt, path)
	print(f"Saved checkpoint: {path}")
	return path


	def compute_loss(logits, target_ids, target_mask, pad_id):
	min_len = min(logits.size(1), target_ids.size(1))
	pred = logits[:, :min_len, :]
	labels = target_ids[:, :min_len].clone()
	mask = target_mask[:, :min_len]
	labels[mask == 0] = -100

	return nn.CrossEntropyLoss(ignore_index=-100)(
	pred.reshape(-1, pred.size(-1)), labels.reshape(-1)
	)


	@torch.no_grad()
	def compute_accuracy(logits, target_ids, target_mask):
	min_len = min(logits.size(1), target_ids.size(1))
	preds = logits[:, :min_len].argmax(dim=-1)
	labels = target_ids[:, :min_len]
	mask = target_mask[:, :min_len]

	correct = ((preds == labels).float() * mask).sum()
	total = mask.sum()
	return (correct / (total + 1e-8)).item()


	@torch.no_grad()
	def validate(model, val_loader, device, cfg, tokenizer):
	model.eval()
	total_loss = 0.0
	total_correct = 0
	total_tokens = 0
	pad_id = tokenizer.pad_token_id
	num_mem = cfg["model"]["num_memory_tokens"]

	for batch in val_loader:
	ids = batch["input_ids"].to(device)
	mask = batch["attention_mask"].to(device)
	target_ids = batch["original_input_ids"].to(device)
	target_mask = batch["original_attention_mask"].to(device)

	with torch.amp.autocast("cuda", enabled=device.type == "cuda"):
	logits, _ = model(ids, mask, target_ids, target_mask)
	loss = compute_loss(logits, target_ids, target_mask, pad_id)

	min_len = min(logits.size(1), target_ids.size(1))
	valid_mask = target_mask[:, :min_len].bool()
	token_count = valid_mask.sum().item()
	total_loss += loss.item() * token_count
	total_correct += (
	(logits[:, :min_len].argmax(dim=-1) == target_ids[:, :min_len]) & valid_mask
	).sum().item()
	total_tokens += token_count

	model.train()
	return total_loss / max(1, total_tokens), total_correct / max(1, total_tokens)


	@torch.no_grad()
	def validate_free_running(model, val_loader, device):
	"""Autoregressively reconstruct every sequence in the validation set."""
	model.eval()
	total_correct = 0
	total_tokens = 0
	exact_sequences = 0
	total_sequences = 0

	for batch in val_loader:
	ids = batch["input_ids"].to(device)
	mask = batch["attention_mask"].to(device)
	targets = batch["original_input_ids"].to(device)
	target_mask = batch["original_attention_mask"].to(device).bool()
	with torch.amp.autocast("cuda", enabled=device.type == "cuda"):
	generated = model.generate(ids, mask, max_length=targets.size(1))
	matches = (generated == targets) & target_mask
	total_correct += matches.sum().item()
	total_tokens += target_mask.sum().item()
	exact_sequences += (matches \| ~target_mask).all(dim=1).sum().item()
	total_sequences += targets.size(0)

	model.train()
	return (
	total_correct / max(1, total_tokens),
	exact_sequences / max(1, total_sequences),
	total_sequences,
	)


	def main(config_path: str):
	cfg = load_config(config_path)
	validate_config(cfg)
	device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
	print(f"Device: {device}")

	run_dir = create_run_dir(cfg)
	shutil.copy2(config_path, os.path.join(run_dir, "config.yaml"))
	print(f"Run directory: {run_dir}")

	# Data
	mcfg = cfg["model"]
	num_mem = mcfg["num_memory_tokens"]
	tokenizer = setup_tokenizer(num_memory_tokens=num_mem)

	train_ds, val_ds = load_reconstruction_dataset(cfg, tokenizer)

	train_loader = create_dataloader(
	cfg["training"]["batch_size"], train_ds, tokenizer,
	shuffle=True, num_memory_tokens=num_mem,
	max_length=cfg["data"].get("max_length", 32),
	slot_source=mcfg.get("slot_source", "appended"),
	)
	val_loader = create_dataloader(
	cfg["training"].get("val_batch_size", cfg["training"]["batch_size"]),
	val_ds, tokenizer,
	shuffle=False, num_memory_tokens=num_mem,
	max_length=cfg["data"].get("max_length", 32),
	slot_source=mcfg.get("slot_source", "appended"),
	)

	# Model
	vocab_size = len(tokenizer)
	model = LoopedTransformerAutoEncoder(
	vocab_size=vocab_size,
	d_model=mcfg["d_model"],
	n_heads=mcfg["n_heads"],
	d_ff=mcfg.get("d_ff", mcfg["d_model"] * 4),
	n_encoder_loops=mcfg["n_encoder_loops"],
	n_decoder_loops=mcfg["n_decoder_loops"],
	num_memory_tokens=num_mem,
	dropout=mcfg.get("dropout", 0.0),
	max_seq_len=mcfg.get("max_seq_len", 512),
	pad_id=tokenizer.pad_token_id,
	bos_id=tokenizer.bos_token_id,
	memory_conditioning=mcfg.get("memory_conditioning", "prefix"),
	slot_source=mcfg.get("slot_source", "appended"),
	).to(device)

	num_params = sum(p.numel() for p in model.parameters())
	print(f"Model parameters: {num_params:,}")
	print(f"Encoder loops: {mcfg['n_encoder_loops']}, Decoder loops: {mcfg['n_decoder_loops']}")

	# Optimizer and scheduler
	optimizer = create_optimizer(model, cfg)
	tcfg = cfg["training"]
	total_steps = len(train_loader) * tcfg["num_epochs"]

	lr_schedule = tcfg.get("lr_schedule", "trapezoid")
	if lr_schedule == "trapezoid":
	scheduler = TrapezoidLR(
	optimizer,
	peak_lr=tcfg["learning_rate"],
	warmup_steps=tcfg.get("warmup_steps", 1000),
	hold_steps=tcfg.get("hold_steps", 5000),
	total_steps=total_steps,
	min_lr=tcfg.get("eta_min", 1e-7),
	)
	print(f"LR schedule: trapezoid (warmup={tcfg.get('warmup_steps', 1000)}, hold={tcfg.get('hold_steps', 5000)}, total={total_steps})")
	else:
	scheduler = optim.lr_scheduler.CosineAnnealingLR(
	optimizer, T_max=total_steps, eta_min=tcfg.get("eta_min", 1e-7)
	)
	print(f"LR schedule: cosine (T_max={total_steps})")

	scaler = torch.amp.GradScaler("cuda", enabled=device.type == "cuda")
	pad_id = tokenizer.pad_token_id

	# Wandb
	if cfg["wandb"].get("enabled"):
	import wandb
	wandb.init(
	project=cfg["wandb"]["project"],
	entity=cfg["wandb"].get("entity"),
	config=cfg,
	name=os.path.basename(run_dir),
	)
	wandb.define_metric("global_step")
	wandb.define_metric("train/*", step_metric="global_step")
	wandb.define_metric("grad_norm/*", step_metric="global_step")
	wandb.define_metric("val/*", step_metric="global_step")

	# Training
	global_step = 0
	best_val_loss = float("inf")
	max_grad_norm = tcfg.get("max_grad_norm", 1.0)
	save_every = cfg["checkpoint"]["save_every_n_steps"]
	val_every = tcfg.get("validation_steps", 1000)
	log_every = tcfg.get("log_every_n_steps", 1)
	print_every = tcfg.get("print_every_n_steps", 50)

	print(f"\nTraining for {tcfg['num_epochs']} epochs, ~{total_steps} steps")
	print("=" * 80)

	for epoch in range(tcfg["num_epochs"]):
	model.train()
	epoch_loss = 0.0
	t0 = time.time()

	for step, batch in enumerate(train_loader):
	ids = batch["input_ids"].to(device)
	mask = batch["attention_mask"].to(device)
	target_ids = batch["original_input_ids"].to(device)
	target_mask = batch["original_attention_mask"].to(device)

	protected_memory_tokens = (
	num_mem if mcfg.get("slot_source", "appended") == "appended" else 0
	)
	corrupted, mask = apply_token_dropout(
	ids, mask, tokenizer,
	dropout_prob=tcfg["token_dropout"],
	num_memory_tokens=protected_memory_tokens,
	)

	with torch.amp.autocast("cuda", enabled=device.type == "cuda"):
	logits, enc_vec = model(corrupted, mask, target_ids, target_mask)
	loss = compute_loss(logits, target_ids, target_mask, pad_id)

	scaler.scale(loss).backward()
	scaler.unscale_(optimizer)

	# Per-component grad norms (before clipping)
	grad_norm_total = torch.nn.utils.clip_grad_norm_(model.parameters(), float("inf")).item()
	component_norms = compute_component_grad_norms(model)

	torch.nn.utils.clip_grad_norm_(model.parameters(), max_grad_norm)
	scaler.step(optimizer)
	scaler.update()
	optimizer.zero_grad()
	scheduler.step()

	epoch_loss += loss.item()
	global_step += 1

	if global_step % log_every == 0:
	lr = optimizer.param_groups[0]["lr"]
	avg = epoch_loss / (step + 1)
	acc = compute_accuracy(logits, target_ids, target_mask)

	if cfg["wandb"].get("enabled"):
	import wandb
	log_dict = {
	"train/loss": loss.item(),
	"train/accuracy": acc,
	"train/lr": lr,
	"train/grad_norm": grad_norm_total,
	"global_step": global_step,
	}
	for comp, norm in component_norms.items():
	log_dict[f"grad_norm/{comp}"] = norm
	wandb.log(log_dict)

	if global_step % print_every == 0:
	lr = optimizer.param_groups[0]["lr"]
	avg = epoch_loss / (step + 1)
	acc = compute_accuracy(logits, target_ids, target_mask)
	elapsed = time.time() - t0
	print(
	f"Epoch {epoch+1} \| Step {global_step} \| "
	f"Loss: {loss.item():.4f} \| Avg: {avg:.4f} \| "
	f"Acc: {acc:.4f} \| LR: {lr:.2e} \| "
	f"GradNorm: {grad_norm_total:.3f} \| Time: {elapsed:.0f}s"
	)

	if cfg["checkpoint"]["save_intermediate"] and save_every > 0 and global_step % save_every == 0:
	save_checkpoint(model, optimizer, scheduler, global_step, epoch, loss.item(), run_dir)

	if val_every > 0 and global_step % val_every == 0:
	val_loss, val_acc = validate(model, val_loader, device, cfg, tokenizer)
	print(f" Validation \| Loss: {val_loss:.4f} \| Acc: {val_acc:.4f}")
	if val_loss < best_val_loss:
	best_val_loss = val_loss
	save_checkpoint(model, optimizer, scheduler, global_step, epoch, val_loss, run_dir, tag="best")

	if cfg["wandb"].get("enabled"):
	import wandb
	wandb.log({"val/loss": val_loss, "val/accuracy": val_acc, "global_step": global_step})

	# End of epoch
	val_loss, val_acc = validate(model, val_loader, device, cfg, tokenizer)
	print(
	f"\nEpoch {epoch+1} \| Train Avg: {epoch_loss / len(train_loader):.4f} \| "
	f"Val Loss: {val_loss:.4f} \| Val Acc: {val_acc:.4f}"
	)
	print("=" * 80)

	save_checkpoint(model, optimizer, scheduler, global_step, epoch, loss.item(), run_dir, tag="final")
	free_acc, exact_match, evaluated_sequences = validate_free_running(model, val_loader, device)
	print(
	f"Full free-running validation \| Acc: {free_acc:.4f} \| "
	f"Exact: {exact_match:.4f} \| Sequences: {evaluated_sequences}"
	)
	if cfg["wandb"].get("enabled"):
	import wandb
	wandb.log({
	"val/free_running_accuracy": free_acc,
	"val/exact_sequence_accuracy": exact_match,
	"val/evaluated_sequences": evaluated_sequences,
	"global_step": global_step,
	})
	print(f"\nTraining complete. Results in: {run_dir}")

	if cfg["wandb"].get("enabled"):
	import wandb
	wandb.finish()


	if __name__ == "__main__":
	config_path = sys.argv[1] if len(sys.argv) > 1 else "config.yaml"
	main(config_path)