opsiom-kairo 0.1.0__tar.gz

opsiom_kairo-0.1.0/PKG-INFO

@@ -0,0 +1,14 @@
+Metadata-Version: 2.4
+Name: opsiom-kairo
+Version: 0.1.0
+Summary: Mesure, fiabilité et reprise pour l'entraînement PyTorch (d'abord Opsiom Kaïro).
+Requires-Python: >=3.10
+Requires-Dist: torch>=2.5
+Requires-Dist: numpy
+Requires-Dist: accelerate
+Requires-Dist: pyyaml
+Requires-Dist: huggingface_hub
+Provides-Extra: logging
+Requires-Dist: wandb; extra == "logging"
+Provides-Extra: dev
+Requires-Dist: pytest; extra == "dev"

opsiom_kairo-0.1.0/README.md

@@ -0,0 +1,37 @@
+# Kairo Engine
+
+Couche de **mesure, fiabilité et reprise** pour l'entraînement PyTorch,
+conçue d'abord pour Opsiom Kaïro (pré-entraînement 500M sur T4 / Colab / Kaggle).
+
+Principe : **mesurer avant d'optimiser**. Kairo ne modifie jamais silencieusement
+les choix scientifiques (learning rate, architecture, dataset, budget de tokens).
+
+## Structure
+
+```text
+kairo/
+├── cli.py         kairo benchmark | autotune | train | resume | status
+├── hardware.py    détection GPU/TPU, VRAM, fp16/bf16
+├── benchmark.py   tokens/s, MFU, temps forward/backward/optim, VRAM   ← COMMENCER ICI
+├── autotune.py    batch size, grad accum, compile, AMP (progressif, jusqu'à OOM)
+├── data.py        sampler à permutation (graine fixe), split val représentatif
+├── checkpoint.py  sauvegarde/reprise (model, optimizer, tokens_seen, RNG, sampler)
+├── relay.py       relais Colab/Kaggle, sauvegarde HF Hub, notifications, arrêt propre
+├── watchdog.py    NaN/Inf, loss spike, chute de débit -> checkpoint d'urgence
+└── probes.py      prompts fixes générés à intervalles réguliers
+configs/           default.yaml, t4.yaml
+tests/             val sans fuite, reprise après crash, détection d'anomalies
+```
+
+## Ordre de développement
+
+1. `hardware.py` + `benchmark.py` (avec MFU) : baseline de référence
+2. Correction du split val + sampler à permutation (`data.py`)
+3. Tests RoPE réel + `torch.compile` via `benchmark.py`
+4. `checkpoint.py` + `relay.py` (cœur de la valeur du projet)
+5. `watchdog.py`, puis `autotune.py`
+6. `probes.py`
+
+Volontairement absents tant que le profiling ne les justifie pas : C++/CUDA
+custom, dashboard web, registre d'expériences, réimplémentation d'un Trainer
+(Kairo se branche sur Accelerate).

opsiom_kairo-0.1.0/kairo/__init__.py

@@ -0,0 +1 @@
+"""Kairo Engine."""

opsiom_kairo-0.1.0/kairo/autotune.py

@@ -0,0 +1,84 @@
+"""Recherche progressive de la config d'exécution, jamais de la config
+scientifique.
+
+Autotune ne touche jamais : learning rate, architecture, dataset, nombre
+de tokens. Il ne touche que ce qui n'a aucun effet sur le résultat final
+d'entraînement (à grad_accum équivalent) : micro_batch_size et
+grad_accum_steps, pour occuper le GPU sans OOM ni sous-utilisation.
+
+Approche volontairement simple (recherche binaire), au-dessus de
+benchmark.py plutôt qu'une réimplémentation : la complexité (compile,
+AMP...) est déjà gérée par PyTorch/Accelerate.
+"""
+from __future__ import annotations
+
+import dataclasses
+
+from kairo.benchmark import BenchmarkResult, run_benchmark
+
+
+@dataclasses.dataclass
+class AutotuneResult:
+    micro_batch_size: int
+    grad_accum_steps: int
+    benchmark: BenchmarkResult
+
+
+def find_max_micro_batch_size(
+    model_factory,
+    batch_fn_factory,
+    target_effective_batch: int,
+    flops_per_token: float,
+    min_batch: int = 1,
+    max_batch: int = 256,
+    bench_steps: int = 5,
+) -> AutotuneResult:
+    """Recherche binaire du plus grand micro-batch tenant en VRAM, en
+    conservant `target_effective_batch` via grad_accum_steps.
+
+    - `model_factory()` doit renvoyer un modèle frais sur le bon device
+      (nécessaire car un OOM peut laisser des gradients/activations sales).
+    - `batch_fn_factory(micro_batch_size)` doit renvoyer un `batch_fn`
+      compatible avec `kairo.benchmark.run_benchmark`.
+    """
+    import torch
+
+    lo, hi = min_batch, max_batch
+    best: AutotuneResult | None = None
+
+    while lo <= hi:
+        mid = (lo + hi) // 2
+        try:
+            model = model_factory()
+            batch_fn = batch_fn_factory(mid)
+            result = run_benchmark(
+                model, batch_fn, flops_per_token=flops_per_token,
+                num_steps=bench_steps, warmup_steps=2,
+            )
+        except torch.cuda.OutOfMemoryError:
+            torch.cuda.empty_cache()
+            hi = mid - 1
+            continue
+        finally:
+            del model
+            torch.cuda.empty_cache()
+
+        if target_effective_batch % mid != 0:
+            # On garde la mesure pour info mais on ne la retient comme
+            # "best" que si elle divise exactement le batch effectif visé,
+            # pour ne jamais changer le batch effectif implicitement.
+            lo = mid + 1
+            continue
+
+        accum = target_effective_batch // mid
+        best = AutotuneResult(mid, accum, result)
+        lo = mid + 1  # on essaie plus grand : moins de micro-steps, souvent plus rapide
+
+    if best is None:
+        raise RuntimeError(
+            "Aucun micro_batch_size testé ne divise target_effective_batch "
+            f"({target_effective_batch}) sans OOM. Essayez un target_effective_batch "
+            "qui a plus de diviseurs, ou augmentez max_batch."
+        )
+    return best
+

opsiom_kairo-0.1.0/kairo/benchmark.py

@@ -0,0 +1,157 @@
+"""Benchmark de référence : tokens/s, MFU, forward/backward/optimiseur, VRAM.
+
+Objectif : mesurer AVANT d'optimiser, pour éviter de deviner. Rien ici ne
+change la config d'entraînement ; ce module ne fait qu'observer.
+
+Usage typique :
+
+    from kairo.benchmark import estimate_model_flops, run_benchmark
+
+    flops_per_token = estimate_model_flops(n_params=511_000_000)
+    result = run_benchmark(model, batch_fn, flops_per_token=flops_per_token,
+                            num_steps=20, warmup_steps=5)
+    print(result.report())
+"""
+from __future__ import annotations
+
+import dataclasses
+import time
+from typing import Callable
+
+from kairo.hardware import detect_hardware
+
+
+def estimate_model_flops(n_params: int, n_layers: int | None = None,
+                          seq_len: int | None = None) -> float:
+    """FLOPs par token pour une passe forward+backward, approximation
+    standard des modèles transformer denses : 6*N (Kaplan et al., 2020).
+
+    Le terme d'attention (2 * 2 * n_layers * seq_len * d_model) est ignoré
+    par défaut car négligeable pour seq_len << d_model * n_layers ; il est
+    ajouté si n_layers et seq_len sont fournis (approximation grossière,
+    d_model non connu ici -> à affiner si besoin de précision fine)."""
+    flops = 6 * n_params
+    if n_layers and seq_len:
+        # Terme quadratique en séquence, ordre de grandeur seulement.
+        flops += 12 * n_layers * seq_len
+    return float(flops)
+
+
+@dataclasses.dataclass
+class BenchmarkResult:
+    tokens_per_sec: float
+    step_time_s: float
+    forward_time_s: float
+    backward_time_s: float
+    optim_time_s: float
+    peak_vram_gb: float | None
+    achieved_flops: float  # FLOPs/s réellement obtenus
+    mfu: float | None  # achieved_flops / pic matériel théorique, 0..1
+    num_steps: int
+    tokens_per_step: int
+
+    def report(self) -> str:
+        lines = [
+            f"tokens/s        : {self.tokens_per_sec:,.0f}",
+            f"temps/step      : {self.step_time_s*1000:.1f} ms "
+            f"(fwd {self.forward_time_s*1000:.1f} / "
+            f"bwd {self.backward_time_s*1000:.1f} / "
+            f"optim {self.optim_time_s*1000:.1f})",
+        ]
+        if self.peak_vram_gb is not None:
+            lines.append(f"VRAM pic        : {self.peak_vram_gb:.2f} Go")
+        if self.mfu is not None:
+            lines.append(f"MFU             : {self.mfu*100:.1f} %")
+        else:
+            lines.append("MFU             : n/a (pic FLOPs matériel inconnu)")
+        return "\n".join(lines)
+
+
+def run_benchmark(
+    model,
+    batch_fn: Callable[[], tuple],
+    flops_per_token: float,
+    step_fn: Callable | None = None,
+    num_steps: int = 20,
+    warmup_steps: int = 5,
+) -> BenchmarkResult:
+    """Mesure les performances réelles d'un modèle sur `num_steps` pas.
+
+    - `batch_fn()` doit retourner un batch `(input_ids, targets)` déjà sur
+      le bon device.
+    - `step_fn(model, batch)` doit exécuter forward + backward + optimizer
+      et rendre le contrôle une fois le pas terminé (avant tout .item()
+      qui forcerait déjà une synchro : la synchro est faite par ce module).
+      Si `step_fn` n'est pas fourni, un pas forward/backward générique est
+      utilisé (sans pas d'optimiseur), suffisant pour mesurer fwd/bwd.
+    - Import torch différé : ce module reste important même sans torch.
+    """
+    import torch
+
+    device = next(model.parameters()).device
+    is_cuda = device.type == "cuda"
+
+    def sync():
+        if is_cuda:
+            torch.cuda.synchronize()
+
+    if step_fn is None:
+        def step_fn(model, batch):  # type: ignore[no-redef]
+            input_ids, targets = batch
+            logits = model(input_ids)
+            loss = torch.nn.functional.cross_entropy(
+                logits.view(-1, logits.size(-1)), targets.view(-1)
+            )
+            loss.backward()
+            model.zero_grad(set_to_none=True)
+
+    # Warmup : cuDNN autotuning, compilation, allocation de la VRAM.
+    for _ in range(warmup_steps):
+        batch = batch_fn()
+        step_fn(model, batch)
+    sync()
+
+    if is_cuda:
+        torch.cuda.reset_peak_memory_stats(device)
+
+    total_tokens = 0
+    step_times = []
+    for _ in range(num_steps):
+        batch = batch_fn()
+        n_tokens = batch[0].numel()
+        sync()
+        t0 = time.perf_counter()
+        step_fn(model, batch)
+        sync()
+        t1 = time.perf_counter()
+        step_times.append(t1 - t0)
+        total_tokens += n_tokens
+
+    total_time = sum(step_times)
+    tokens_per_sec = total_tokens / total_time
+    step_time_s = total_time / num_steps
+    achieved_flops = tokens_per_sec * flops_per_token
+
+    peak_vram_gb = None
+    if is_cuda:
+        peak_vram_gb = torch.cuda.max_memory_allocated(device) / (1024**3)
+
+    hw = detect_hardware()
+    mfu = achieved_flops / hw.peak_flops if hw.peak_flops else None
+
+    # Sans instrumentation fine (torch.profiler), on ne distingue pas
+    # fwd/bwd/optim par défaut : ils sont laissés à 0 ici et à remplir par
+    # l'appelant via `step_fn` instrumenté si le détail est nécessaire.
+    return BenchmarkResult(
+        tokens_per_sec=tokens_per_sec,
+        step_time_s=step_time_s,
+        forward_time_s=0.0,
+        backward_time_s=0.0,
+        optim_time_s=0.0,
+        peak_vram_gb=peak_vram_gb,
+        achieved_flops=achieved_flops,
+        mfu=mfu,
+        num_steps=num_steps,
+        tokens_per_step=total_tokens // num_steps,
+    )
+

opsiom_kairo-0.1.0/kairo/checkpoint.py

@@ -0,0 +1,175 @@
+"""Sauvegarde/reprise complète d'un run.
+
+Contrairement à une sauvegarde "poids seuls", un checkpoint Kairo capture
+tout ce qui est nécessaire pour reprendre EXACTEMENT là où le run s'est
+arrêté, y compris entre deux plateformes différentes (Colab -> Kaggle) :
+
+  - state_dict() du modèle et de l'optimiseur
+  - scheduler (si fourni)
+  - tokens_seen (source de vérité pour la progression, pas step ni epoch)
+  - états RNG (torch CPU, torch CUDA, numpy, random) pour la reproductibilité
+  - state_dict() du sampler (kairo.data.PermutationSampler), pour reprendre
+    le parcours du corpus exactement où il en était
+  - métadonnées libres (config, meilleure loss val, etc.)
+
+Le fichier le plus récent et le "best" (meilleure loss val représentative,
+cf. data.py) sont conservés séparément pour ne jamais perdre le meilleur
+modèle à cause d'une reprise ultérieure qui diverge.
+"""
+from __future__ import annotations
+
+import dataclasses
+import json
+import os
+import random
+import shutil
+import tempfile
+import time
+from typing import Any
+
+
+@dataclasses.dataclass
+class CheckpointMeta:
+    tokens_seen: int
+    step: int
+    val_loss: float | None
+    wall_time_s: float
+    extra: dict = dataclasses.field(default_factory=dict)
+
+
+def _capture_rng_state() -> dict:
+    state = {"python": random.getstate(), "numpy": None, "torch_cpu": None,
+              "torch_cuda": None}
+    try:
+        import numpy as np
+        state["numpy"] = np.random.get_state()
+    except ImportError:
+        pass
+    try:
+        import torch
+        state["torch_cpu"] = torch.get_rng_state()
+        if torch.cuda.is_available():
+            state["torch_cuda"] = torch.cuda.get_rng_state_all()
+    except ImportError:
+        pass
+    return state
+
+
+def _restore_rng_state(state: dict) -> None:
+    if state.get("python") is not None:
+        random.setstate(state["python"])
+    if state.get("numpy") is not None:
+        try:
+            import numpy as np
+            np.random.set_state(state["numpy"])
+        except ImportError:
+            pass
+    try:
+        import torch
+        if state.get("torch_cpu") is not None:
+            torch.set_rng_state(state["torch_cpu"])
+        if state.get("torch_cuda") is not None and torch.cuda.is_available():
+            torch.cuda.set_rng_state_all(state["torch_cuda"])
+    except ImportError:
+        pass
+
+
+def save_checkpoint(
+    path: str,
+    model,
+    optimizer=None,
+    scheduler=None,
+    sampler=None,
+    meta: CheckpointMeta | None = None,
+) -> None:
+    """Écrit un checkpoint de façon atomique (fichier temporaire + rename),
+    pour ne jamais laisser un fichier corrompu si le process est tué en
+    plein milieu (fréquent sur Colab/Kaggle en fin de session)."""
+    import torch
+
+    payload: dict[str, Any] = {
+        "model": model.state_dict(),
+        "optimizer": optimizer.state_dict() if optimizer is not None else None,
+        "scheduler": scheduler.state_dict() if scheduler is not None else None,
+        "sampler": sampler.state_dict() if sampler is not None else None,
+        "rng": _capture_rng_state(),
+        "meta": dataclasses.asdict(meta) if meta is not None else None,
+        "saved_at": time.time(),
+    }
+
+    os.makedirs(os.path.dirname(path) or ".", exist_ok=True)
+    fd, tmp_path = tempfile.mkstemp(dir=os.path.dirname(path) or ".", suffix=".tmp")
+    os.close(fd)
+    try:
+        torch.save(payload, tmp_path)
+        os.replace(tmp_path, path)  # atomique sur un même filesystem
+    except BaseException:
+        if os.path.exists(tmp_path):
+            os.remove(tmp_path)
+        raise
+
+
+def load_checkpoint(
+    path: str,
+    model,
+    optimizer=None,
+    scheduler=None,
+    sampler=None,
+    map_location: str | None = None,
+    restore_rng: bool = True,
+) -> CheckpointMeta | None:
+    import torch
+
+    payload = torch.load(path, map_location=map_location, weights_only=False)
+
+    model.load_state_dict(payload["model"])
+    if optimizer is not None and payload.get("optimizer") is not None:
+        optimizer.load_state_dict(payload["optimizer"])
+    if scheduler is not None and payload.get("scheduler") is not None:
+        scheduler.load_state_dict(payload["scheduler"])
+    if sampler is not None and payload.get("sampler") is not None:
+        sampler.load_state_dict(payload["sampler"])
+    if restore_rng and payload.get("rng") is not None:
+        _restore_rng_state(payload["rng"])
+
+    meta = payload.get("meta")
+    return CheckpointMeta(**meta) if meta else None
+
+
+class CheckpointManager:
+    """Gère deux fichiers stables : `latest.pt` (dernier en date, pour
+    reprendre après une coupure) et `best.pt` (meilleure val_loss vue,
+    jamais écrasé par un run qui repart puis diverge)."""
+
+    def __init__(self, directory: str):
+        self.directory = directory
+        os.makedirs(directory, exist_ok=True)
+        self._best_val_loss = float("inf")
+        state_path = os.path.join(directory, "state.json")
+        if os.path.exists(state_path):
+            with open(state_path) as f:
+                self._best_val_loss = json.load(f).get("best_val_loss", float("inf"))
+
+    @property
+    def latest_path(self) -> str:
+        return os.path.join(self.directory, "latest.pt")
+
+    @property
+    def best_path(self) -> str:
+        return os.path.join(self.directory, "best.pt")
+
+    def save(self, model, optimizer, scheduler, sampler, meta: CheckpointMeta) -> None:
+        save_checkpoint(self.latest_path, model, optimizer, scheduler, sampler, meta)
+        if meta.val_loss is not None and meta.val_loss < self._best_val_loss:
+            self._best_val_loss = meta.val_loss
+            shutil.copy2(self.latest_path, self.best_path)
+            with open(os.path.join(self.directory, "state.json"), "w") as f:
+                json.dump({"best_val_loss": self._best_val_loss}, f)
+
+    def resume(self, model, optimizer=None, scheduler=None, sampler=None,
+               map_location: str | None = None) -> CheckpointMeta | None:
+        if not os.path.exists(self.latest_path):
+            return None
+        return load_checkpoint(self.latest_path, model, optimizer, scheduler,
+                                sampler, map_location=map_location)
+

opsiom_kairo-0.1.0/kairo/cli.py

@@ -0,0 +1,78 @@
+"""Point d'entrée CLI : kairo hardware | benchmark | status.
+
+`kairo train`/`resume` ne sont pas des commandes génériques ici : Kairo est
+une bibliothèque qui s'utilise depuis le script d'entraînement existant
+(main-GPU-CPU-TPU.py), pas un Trainer qui remplace ce script. La CLI ne
+couvre donc que ce qui a du sens hors contexte d'un script précis :
+diagnostic matériel, benchmark autonome sur un modèle donné, et lecture
+de l'état d'un run à partir de son dossier de checkpoints.
+"""
+from __future__ import annotations
+
+import argparse
+import json
+
+
+def cmd_hardware(args: argparse.Namespace) -> None:
+    from kairo.hardware import detect_hardware, summary_line
+
+    info = detect_hardware()
+    if args.json:
+        print(json.dumps(info.as_dict(), indent=2))
+    else:
+        print(summary_line(info))
+
+
+def cmd_status(args: argparse.Namespace) -> None:
+    import os
+
+    from kairo.checkpoint import load_checkpoint
+
+    latest = os.path.join(args.checkpoint_dir, "latest.pt")
+    best = os.path.join(args.checkpoint_dir, "best.pt")
+    if not os.path.exists(latest):
+        print(f"Aucun checkpoint dans {args.checkpoint_dir}")
+        return
+
+    # Chargement "à blanc" : on ne restaure aucun state_dict, juste les
+    # métadonnées, donc pas besoin d'un vrai modèle/optimiseur ici.
+    class _Void:
+        def state_dict(self):
+            return {}
+
+        def load_state_dict(self, _):
+            pass
+
+    meta = load_checkpoint(latest, _Void(), restore_rng=False)
+    print(f"dernier checkpoint : {latest}")
+    if meta:
+        print(f"  step        : {meta.step}")
+        print(f"  tokens_seen : {meta.tokens_seen:,}")
+        print(f"  val_loss    : {meta.val_loss}")
+    print(f"best checkpoint    : {best if os.path.exists(best) else '(absent)'}")
+
+
+def build_parser() -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser(prog="kairo")
+    sub = parser.add_subparsers(dest="command", required=True)
+
+    p_hw = sub.add_parser("hardware", help="Détecte et affiche le matériel disponible.")
+    p_hw.add_argument("--json", action="store_true")
+    p_hw.set_defaults(func=cmd_hardware)
+
+    p_status = sub.add_parser("status", help="Affiche l'état d'un dossier de checkpoints.")
+    p_status.add_argument("checkpoint_dir")
+    p_status.set_defaults(func=cmd_status)
+
+    return parser
+
+
+def main() -> None:
+    parser = build_parser()
+    args = parser.parse_args()
+    args.func(args)
+
+
+if __name__ == "__main__":
+    main()
+