scripts/03_training/utils.py

"""
Training utilities: FastText init, FORGETTER, checkpointing, optimizers, schedulers.
"""

import math
import copy
from pathlib import Path
from typing import Optional

import torch
import torch.nn as nn
from torch.optim import Adam, AdamW
from torch.optim.lr_scheduler import LambdaLR


ROOT = Path(__file__).resolve().parent.parent.parent


# ═══════════════════════════════════════════════════════
# FastText Embedding Initialization
# ═══════════════════════════════════════════════════════

def init_embeddings_fasttext(model, tokenizer, fasttext_path: str):
    """
    Initialize model's word embeddings from pre-trained FastText vectors.
    +2-3pp BLiMP in 2025 papers.

    Args:
        model: model with .word_embedding or .embedding.word_embedding
        tokenizer: HuggingFace tokenizer
        fasttext_path: path to FastText .bin file
    """
    try:
        from gensim.models import FastText
    except ImportError:
        print("WARNING: gensim not installed, skipping FastText init. pip install gensim")
        return

    print(f"Loading FastText model from {fasttext_path}...")
    ft_model = FastText.load(fasttext_path)

    # Find the embedding layer
    if hasattr(model, 'word_embedding'):
        embed = model.word_embedding
    elif hasattr(model, 'embedding') and hasattr(model.embedding, 'word_embedding'):
        embed = model.embedding.word_embedding
    else:
        print("WARNING: Could not find word_embedding layer, skipping FastText init")
        return

    vocab = tokenizer.get_vocab()
    hidden_size = embed.weight.shape[1]
    ft_dim = ft_model.wv.vector_size

    # Project if dimensions don't match
    if ft_dim != hidden_size:
        print(f"  FastText dim={ft_dim} != hidden_size={hidden_size}, using linear projection")
        proj = torch.randn(ft_dim, hidden_size) * (1.0 / math.sqrt(ft_dim))
    else:
        proj = None

    initialized = 0
    with torch.no_grad():
        for token, idx in vocab.items():
            # Clean BPE token
            clean = token.replace('Ġ', ' ').replace('▁', ' ').strip()
            if not clean:
                continue
            try:
                vec = torch.tensor(ft_model.wv[clean], dtype=torch.float32)
                if proj is not None:
                    vec = vec @ proj
                embed.weight[idx] = vec
                initialized += 1
            except KeyError:
                pass

    print(f"  Initialized {initialized}/{len(vocab)} token embeddings from FastText")


# ═══════════════════════════════════════════════════════
# FORGETTER: Reset optimizer state each epoch
# ═══════════════════════════════════════════════════════

def reset_optimizer_state(optimizer):
    """
    FORGETTER Sleep 1 (Light Sleep): Reset optimizer state (momentum, variance).
    Model weights are kept. From Yamamoto & Miura 2025.
    """
    for group in optimizer.param_groups:
        for p in group['params']:
            if p in optimizer.state:
                optimizer.state[p] = {}
    print("  FORGETTER Sleep 1: optimizer state reset")


def deep_sleep(model):
    """
    FORGETTER Sleep 2 (Deep Sleep): Reinitialize all model state EXCEPT weights.
    Saves weights, constructs fresh model, loads weights back.
    From Yamamoto & Miura 2025 — used for the final epoch.

    Returns a new model instance with same weights but fresh buffers/state.
    """
    state_dict = {k: v.clone() for k, v in model.state_dict().items()
                  if 'weight' in k or 'bias' in k}
    # Re-initialize model
    device = next(model.parameters()).device
    model_class = type(model)
    # This requires the model to store its init args — caller should handle
    print("  FORGETTER Sleep 2: deep sleep — reinitialize model, keep weights")
    return state_dict


# ═══════════════════════════════════════════════════════
# Checkpoint Management
# ═══════════════════════════════════════════════════════

def save_checkpoint(model, optimizer, epoch: int, step: int, loss: float,
                    save_dir: str, cfg=None):
    """Save training checkpoint."""
    save_path = Path(save_dir)
    save_path.mkdir(parents=True, exist_ok=True)

    ckpt = {
        'epoch': epoch,
        'step': step,
        'loss': loss,
        'model_state_dict': model.state_dict(),
        'optimizer_state_dict': optimizer.state_dict(),
    }
    if cfg is not None:
        from .config import config_to_dict
        ckpt['config'] = config_to_dict(cfg)

    path = save_path / f"checkpoint_epoch{epoch}.pt"
    torch.save(ckpt, path)
    print(f"  Saved checkpoint: {path}")
    return path


def load_checkpoint(model, optimizer, checkpoint_path: str):
    """Load training checkpoint."""
    ckpt = torch.load(checkpoint_path, map_location='cpu')
    model.load_state_dict(ckpt['model_state_dict'])
    if optimizer is not None:
        optimizer.load_state_dict(ckpt['optimizer_state_dict'])
    print(f"  Loaded checkpoint from epoch {ckpt['epoch']}, step {ckpt['step']}")
    return ckpt['epoch'], ckpt['step']


def average_checkpoints(checkpoint_paths: list, model) -> dict:
    """
    Average multiple checkpoint model weights.
    Used as final model for evaluation.

    Args:
        checkpoint_paths: list of checkpoint file paths
        model: model instance (for structure reference)

    Returns:
        averaged state_dict
    """
    print(f"Averaging {len(checkpoint_paths)} checkpoints...")

    avg_state = None
    for i, path in enumerate(checkpoint_paths):
        ckpt = torch.load(path, map_location='cpu')
        state = ckpt['model_state_dict']

        if avg_state is None:
            avg_state = {k: v.clone().float() for k, v in state.items()}
        else:
            for k in avg_state:
                avg_state[k] += state[k].float()

    # Divide by number of checkpoints
    n = len(checkpoint_paths)
    for k in avg_state:
        avg_state[k] /= n

    model.load_state_dict(avg_state)
    print(f"  Checkpoint averaging complete")
    return avg_state


# ═══════════════════════════════════════════════════════
# Learning Rate Scheduler
# ═══════════════════════════════════════════════════════

def get_cosine_schedule_with_warmup(optimizer, num_warmup_steps: int,
                                      num_training_steps: int,
                                      min_lr_ratio: float = 0.1,
                                      num_cooldown_steps: int = 0):
    """
    Cosine annealing with linear warmup and optional linear cooldown.
    Matches official GPT-BERT scheduler (warmup 1.6%, cooldown 1.6%, min_lr=0.1).
    """

    def lr_lambda(current_step):
        # Warmup phase
        if current_step < num_warmup_steps:
            return float(current_step) / float(max(1, num_warmup_steps))
        # Cooldown phase
        if num_cooldown_steps > 0 and current_step > num_training_steps - num_cooldown_steps:
            remaining = num_training_steps - current_step
            return max(0.0, min_lr_ratio * float(remaining) / float(max(1, num_cooldown_steps)))
        # Cosine phase
        progress = float(current_step - num_warmup_steps) / float(
            max(1, num_training_steps - num_warmup_steps - num_cooldown_steps))
        return max(min_lr_ratio, 0.5 * (1.0 + math.cos(math.pi * progress)))

    return LambdaLR(optimizer, lr_lambda)


# ═══════════════════════════════════════════════════════
# Optimizer Factory
# ═══════════════════════════════════════════════════════

class _CombinedOptimizer(torch.optim.Optimizer):
    """Wraps two optimizers (e.g., Muon for 2D + AdamW for 1D) into one.

    Inherits from Optimizer so LR schedulers accept it via isinstance() check.
    We initialize the base class with a dummy param to satisfy the non-empty check,
    then replace param_groups with the real ones from both sub-optimizers.
    """

    def __init__(self, opt1, opt2):
        # Create a dummy param to satisfy Optimizer.__init__'s non-empty check
        dummy = torch.nn.Parameter(torch.zeros(1), requires_grad=False)
        super().__init__([dummy], defaults={"lr": 0.0})
        # Now replace with real param_groups from sub-optimizers
        self.param_groups = list(opt1.param_groups)
        self.opt1 = opt1
        self.opt2 = opt2

    def step(self, closure=None):
        # Muon.step() doesn't accept closure argument
        self.opt1.step()
        self.opt2.step()

    def zero_grad(self, set_to_none=False):
        self.opt1.zero_grad(set_to_none=set_to_none)
        self.opt2.zero_grad(set_to_none=set_to_none)

    def state_dict(self):
        return {"opt1": self.opt1.state_dict(), "opt2": self.opt2.state_dict()}

    def load_state_dict(self, state_dict):
        self.opt1.load_state_dict(state_dict["opt1"])
        self.opt2.load_state_dict(state_dict["opt2"])


def build_optimizer(model, opt_cfg, training_cfg):
    """Build optimizer from config."""

    # Normalize optimizer type to lowercase
    opt_cfg.type = opt_cfg.type.lower()

    # Separate weight decay groups (matching official GPT-BERT)
    no_decay = {"bias", "layer_norm"}
    params = [
        {
            "params": [p for n, p in model.named_parameters()
                      if not any(nd in n for nd in no_decay) and p.requires_grad],
            "weight_decay": training_cfg.weight_decay,
        },
        {
            "params": [p for n, p in model.named_parameters()
                      if any(nd in n for nd in no_decay) and p.requires_grad],
            "weight_decay": 0.0,
        },
    ]

    # ── Adam (Kingma & Ba, 2014) ──
    # 最经典的自适应学习率优化器。为每个参数维护一阶动量(梯度均值 m)和二阶动量
    # (梯度平方均值 v)，用 m/√v 自适应调整每个参数的更新步长。
    # 优点：收敛快、对超参数不敏感；缺点：L2正则化与自适应学习率耦合，
    # 导致正则化效果被削弱（这正是 AdamW 要解决的问题）。
    if opt_cfg.type == "adam":
        optimizer = Adam(
            params,
            lr=training_cfg.learning_rate,
            betas=opt_cfg.betas,
            eps=opt_cfg.eps,
        )

    # ── AdamW (Loshchilov & Hutter, 2017/2019) ──
    # Adam 的修正版本，将权重衰减(weight decay)从梯度更新中解耦出来。
    # 原始 Adam 把 L2 正则加在梯度上再做自适应缩放，导致大梯度参数的正则化
    # 被削弱；AdamW 改为在参数更新后直接乘以衰减系数 (θ *= 1 - λ·lr)，
    # 使正则化强度与自适应学习率无关。几乎所有现代 Transformer 训练的标配。
    elif opt_cfg.type == "adamw":
        optimizer = AdamW(
            params,
            lr=training_cfg.learning_rate,
            betas=opt_cfg.betas,
            eps=opt_cfg.eps,
        )

    # ── LAMB (You et al., 2019) ──
    # Layer-wise Adaptive Moments optimizer，专为大 batch 训练设计。
    # 在 AdamW 基础上增加了逐层的信赖域归一化(trust ratio)：
    # 每层的更新步长 = ‖θ_layer‖ / ‖update_layer‖，使得不同层的参数更新
    # 幅度与该层权重范数成正比，避免大 batch 下某些层更新过大/过小。
    # GPT-BERT 官方使用 LAMB + lr=1.41e-2，在 BabyLM 规模下效果优于 AdamW。
    elif opt_cfg.type == "lamb":
        try:
            from torch_optimizer import Lamb
            optimizer = Lamb(
                params,
                lr=training_cfg.learning_rate,
                betas=opt_cfg.betas,
                eps=opt_cfg.eps,
            )
        except ImportError:
            print("WARNING: torch-optimizer not installed, falling back to AdamW")
            optimizer = AdamW(params, lr=training_cfg.learning_rate,
                            betas=opt_cfg.betas, eps=opt_cfg.eps)

    # ── Muon (Jordan et al., 2024) ──
    # Momentum + Orthogonalization 优化器，完全不同于 Adam 系列的思路。
    # 核心机制：先用 Nesterov 动量计算更新方向，再通过 Newton-Schulz 迭代
    # 对更新矩阵做正交化（近似 SVD 后将所有奇异值置为1），确保更新方向在
    # 参数空间中各方向上均匀分布。不需要维护二阶动量，内存占用更低。
    # 在小模型上表现出色，但需要专门的超参数调优（不使用 betas/eps）。
    elif opt_cfg.type == "muon":
        try:
            from muon import Muon
            import torch.distributed as dist
            # Muon uses distributed communication internally; init if needed
            if not dist.is_initialized():
                dist.init_process_group(backend="gloo", init_method="tcp://127.0.0.1:29500",
                                        rank=0, world_size=1)
            # Muon only handles 2D+ params (matrices via Newton-Schulz orthogonalization).
            # 1D params (bias, LayerNorm) use a separate AdamW.
            # We wrap both into a _CombinedOptimizer so the training loop can call
            # optimizer.step() and optimizer.zero_grad() as usual.
            muon_params = [p for n, p in model.named_parameters() if p.requires_grad and p.ndim >= 2]
            adamw_1d = [p for n, p in model.named_parameters() if p.requires_grad and p.ndim < 2]
            muon_opt = Muon(muon_params, lr=training_cfg.learning_rate)
            adamw_opt = AdamW([{"params": adamw_1d, "weight_decay": 0.0}],
                              lr=training_cfg.learning_rate * 0.1,
                              betas=opt_cfg.betas, eps=opt_cfg.eps)
            optimizer = _CombinedOptimizer(muon_opt, adamw_opt)
            print(f"  Muon: {len(muon_params)} matrix params + {len(adamw_1d)} scalar params (AdamW)")
        except ImportError:
            print("WARNING: muon not installed, falling back to AdamW")
            optimizer = AdamW(params, lr=training_cfg.learning_rate,
                            betas=opt_cfg.betas, eps=opt_cfg.eps)

    else:
        raise ValueError(f"Unknown optimizer type: {opt_cfg.type}")

    return optimizer


# ═══════════════════════════════════════════════════════
# HuggingFace Model Export (for evaluation pipeline)
# ═══════════════════════════════════════════════════════

def export_for_eval(model, tokenizer, cfg, export_dir: str):
    """
    Export trained model in HuggingFace format for the evaluation pipeline.
    Creates a self-contained directory loadable via:
        AutoModelForCausalLM.from_pretrained(path, trust_remote_code=True)

    Uses the hf_export module which handles all 5 architectures.
    """
    from .config import config_to_dict
    from .hf_export.export import export_state_dict

    cfg_dict = config_to_dict(cfg)
    model_cfg = cfg_dict.get("model", {})
    arch = model_cfg.get("arch", "gpt_bert")
    tokenizer_dir = cfg_dict.get("data", {}).get("tokenizer_path", "models/tokenizer")

    export_state_dict(
        state_dict=model.state_dict(),
        arch=arch,
        model_cfg=model_cfg,
        output_dir=export_dir,
        tokenizer_dir=tokenizer_dir,
    )