src/train/train_moe.py

"""Phase 10 — entraînement du modèle texte avec FFN Mixture-of-Experts.

Remplace les FFN denses du transformer texte (Phase 2) par des **MoE 20 experts,
top-2**. La perte totale = perte de langage + ``aux_loss_coef`` * loss d'équilibrage.

Usage :

    python -m src.train.train_moe --out-dir checkpoints/moe_toy \\
        --data-dir data/processed/tiny_fr --max-steps 400
"""

from __future__ import annotations

import argparse
import json
import math
import time
from dataclasses import asdict, dataclass
from pathlib import Path

import numpy as np
import torch

from src.data.prepare_text import load_bin, prepare
from src.model.moe import CortexMoELM, MoEConfig, load_balance_metrics
from src.model.text_model import TextConfig
from src.tokenizer import load_tokenizer
from src.train.train_text import get_lr, pick_device

ROOT = Path(__file__).resolve().parents[2]

TINY_FR = [
    "Le chat dort sur le tapis du salon.",
    "La petite fille mange une pomme rouge.",
    "Il fait beau aujourd'hui dans la ville.",
    "Le chien court rapidement dans le jardin.",
    "Nous lisons un livre intéressant le soir.",
    "La voiture roule doucement sur la route.",
    "Les enfants jouent ensemble dans la cour.",
    "Le soleil brille au-dessus des montagnes.",
] * 40


@dataclass
class MoETrainConfig:
    data_dir: str = "data/processed/tiny_fr"
    out_dir: str = "checkpoints/moe_toy"
    max_steps: int = 400
    batch_size: int = 16
    ctx_len: int = 128
    lr: float = 3e-3
    warmup_steps: int = 40
    min_lr_ratio: float = 0.1
    weight_decay: float = 0.1
    grad_clip: float = 1.0
    n_layer: int = 4
    n_head: int = 4
    n_embd: int = 128
    n_experts: int = 20
    top_k: int = 2
    aux_loss_coef: float = 0.01
    jitter: float = 0.01
    seed: int = 1337
    device: str = "auto"


def get_batch(data: np.ndarray, ctx_len: int, batch_size: int, rng) -> tuple[torch.Tensor, torch.Tensor]:
    ix = rng.integers(0, len(data) - ctx_len - 1, batch_size)
    x = np.stack([data[i:i + ctx_len] for i in ix]).astype(np.int64)
    y = np.stack([data[i + 1:i + 1 + ctx_len] for i in ix]).astype(np.int64)
    return torch.from_numpy(x), torch.from_numpy(y)


def load_or_prepare(config_data_dir: str, tokenizer):
    d = Path(config_data_dir)
    if not (d / "train.bin").is_file():
        prepare(TINY_FR, d, tokenizer, val_frac=0.05, extra_meta={"preset": "tiny_fr"})
    return load_bin(d / "train.bin"), load_bin(d / "val.bin")


def build_model(cfg: MoETrainConfig, tokenizer) -> CortexMoELM:
    text = TextConfig(vocab_size=tokenizer.vocab_size, n_layer=cfg.n_layer, n_head=cfg.n_head,
                      n_embd=cfg.n_embd, ctx_len=cfg.ctx_len, dropout=0.0,
                      tie_embeddings=True)
    moe = MoEConfig(n_experts=cfg.n_experts, top_k=cfg.top_k, aux_loss_coef=cfg.aux_loss_coef,
                    jitter=cfg.jitter)
    return CortexMoELM(text, moe)


def train(cfg: MoETrainConfig) -> dict:
    torch.manual_seed(cfg.seed)
    device = pick_device(cfg.device)
    tokenizer = load_tokenizer()
    out_dir = Path(cfg.out_dir)
    out_dir.mkdir(parents=True, exist_ok=True)

    train_data, val_data = load_or_prepare(cfg.data_dir, tokenizer)
    model = build_model(cfg, tokenizer).to(device)

    decay, no_decay = [], []
    for _n, p in model.named_parameters():
        if p.requires_grad:
            (decay if p.dim() >= 2 else no_decay).append(p)
    opt = torch.optim.AdamW([{"params": decay, "weight_decay": cfg.weight_decay},
                             {"params": no_decay, "weight_decay": 0.0}], lr=cfg.lr,
                            betas=(0.9, 0.95))

    rng = np.random.default_rng(cfg.seed)
    t0 = time.time()
    log_loss = torch.tensor(float("nan"))
    for step in range(cfg.max_steps):
        lr = get_lr(step, _LRShim(cfg))
        for g in opt.param_groups:
            g["lr"] = lr
        x, y = get_batch(train_data, cfg.ctx_len, cfg.batch_size, rng)
        x, y = x.to(device), y.to(device)
        emb = model.token_embedding(x)
        _logits, loss, aux, load = model(emb, y)
        total = loss + cfg.aux_loss_coef * aux
        opt.zero_grad()
        total.backward()
        if cfg.grad_clip > 0:
            torch.nn.utils.clip_grad_norm_(model.parameters(), cfg.grad_clip)
        opt.step()
        log_loss = loss.detach()
        if (step + 1) % 50 == 0 or step == 0:
            print(f"step {step:4d} | lm {loss.item():.3f} | aux {aux.item():.3f} | "
                  f"load-H {load_balance_metrics(load)['entropy_norm']:.2f} | {time.time() - t0:.0f}s", flush=True)

    val_loss = estimate_val(model, val_data, cfg, rng, device)
    torch.save({"model": model.state_dict(), "train_config": asdict(cfg)},
               out_dir / "best.pt")
    summary = {
        "final_train_loss": round(float(log_loss), 4),
        "val_loss": round(val_loss, 4),
        "val_perplexity": round(math.exp(val_loss), 3),
        "expert_load_entropy": round(load_balance_metrics(load)["entropy_norm"], 4),
        "total_parameters": model.total_parameters,
        "active_parameters": model.active_parameters,
        "steps": cfg.max_steps,
        "elapsed_s": round(time.time() - t0, 1),
    }
    (out_dir / "summary.json").write_text(json.dumps(summary, ensure_ascii=False, indent=2),
                                         encoding="utf-8")
    print("Résumé :", json.dumps(summary, ensure_ascii=False), flush=True)
    return summary


class _LRShim:
    """Adapte MoETrainConfig pour réutiliser `get_lr` de Phase 3."""

    def __init__(self, cfg: MoETrainConfig):
        self.lr = cfg.lr
        self.warmup_steps = cfg.warmup_steps
        self.max_steps = cfg.max_steps
        self.min_lr_ratio = cfg.min_lr_ratio


@torch.no_grad()
def estimate_val(model, val_data, cfg, rng, device, n_batches: int = 10) -> float:
    model.eval()
    losses = []
    for _ in range(n_batches):
        x, y = get_batch(val_data, cfg.ctx_len, cfg.batch_size, rng)
        emb = model.token_embedding(x.to(device))
        _logits, loss, _aux, _load = model(emb, y.to(device))
        losses.append(float(loss))
    model.train()
    return float(np.mean(losses))


def main() -> None:
    parser = argparse.ArgumentParser(description="Entraîner le modèle MoE (Phase 10).")
    for field, value in asdict(MoETrainConfig()).items():
        parser.add_argument(f"--{field.replace('_', '-')}", type=type(value), default=value)
    ns = parser.parse_args()
    train(MoETrainConfig(**{k.replace("-", "_"): v for k, v in vars(ns).items()}))


if __name__ == "__main__":
    main()