"""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()