Update main_medium.py

main_medium.py

@@ -240,69 +240,148 @@ HF_BACKUP_PRIVATE = True
 # "meilleur") existait juste avant le crash.
 HF_BACKUP_LATEST_INTERVAL_SECONDS = 600  # 10 min
 
-_hf_backup_api = None
-
-
-def _get_hf_backup_api():
-    """Instancie paresseusement le client HfApi (une seule fois), et
-    crée le repo de sauvegarde s'il n'existe pas déjà."""
-    global _hf_backup_api
-    if _hf_backup_api is not None:
-        return _hf_backup_api
+# --- Sauvegarde du CORPUS tokenisé (gros .bin, plusieurs heures de construction) ---
+# Poussé UNE fois vers un repo HF de type "dataset" (distinct du repo de modèles),
+# puis re-téléchargé automatiquement si le .bin local est absent ou vide
+# (session Colab/Kaggle neuve, Drive corrompu...). Public par défaut: ce sont des
+# données web publiques (Wikipedia / FineWeb-2) et ça évite les soucis de quota
+# du stockage privé déjà rencontrés (403). Mets True pour un repo privé.
+HF_CORPUS_BACKUP_ENABLED = True
+HF_CORPUS_REPO_ID = "JuloEco/opsiom-fr-corpus"  # ⚠️ adapte à ton propre namespace HF
+HF_CORPUS_PRIVATE = False
+
+_hf_apis = {}
+
+
+def _get_hf_api(repo_id: str, repo_type: str, private: bool):
+    """Instancie paresseusement un client HfApi et crée le repo cible s'il
+    n'existe pas encore (un client mis en cache par couple repo_id/repo_type)."""
+    key = (repo_id, repo_type)
+    if key in _hf_apis:
+        return _hf_apis[key]
     from huggingface_hub import HfApi
     api = HfApi()
     try:
-        api.create_repo(repo_id=HF_BACKUP_REPO_ID, private=HF_BACKUP_PRIVATE, repo_type="model", exist_ok=True)
+        api.create_repo(repo_id=repo_id, private=private, repo_type=repo_type, exist_ok=True)
     except Exception as e:
-        print(f"⚠️ Impossible de créer/vérifier le repo de sauvegarde HF '{HF_BACKUP_REPO_ID}' ({e}). "
-              f"La sauvegarde externe est désactivée pour cette session.")
+        print(f"⚠️ Impossible de créer/vérifier le repo HF '{repo_id}' ({repo_type}) ({e}). "
+              f"La sauvegarde externe correspondante est désactivée pour cette session.")
         return None
-    _hf_backup_api = api
+    _hf_apis[key] = api
     return api
 
 
-def push_backup_to_hf(local_path: str, path_in_repo: str) -> None:
-    """Pousse un fichier vers le repo de sauvegarde HF Hub. Ne lève JAMAIS
-    d'exception: un problème réseau ponctuel ne doit pas interrompre
-    l'entraînement — on log juste un avertissement et on continue."""
+def push_backup_to_hf(local_path: str, path_in_repo: str, repo_id: str | None = None,
+                      repo_type: str = "model", private: bool | None = None) -> None:
+    """Pousse un fichier vers un repo HF Hub (par défaut: le repo de checkpoints).
+    Ne lève JAMAIS d'exception: un problème réseau ponctuel ne doit pas
+    interrompre l'entraînement — on log juste un avertissement et on continue."""
     if not HF_BACKUP_ENABLED or not os.path.exists(local_path):
         return
-    api = _get_hf_backup_api()
+    repo_id = repo_id or HF_BACKUP_REPO_ID
+    private = HF_BACKUP_PRIVATE if private is None else private
+    api = _get_hf_api(repo_id, repo_type, private)
     if api is None:
         return
     try:
         api.upload_file(
             path_or_fileobj=local_path,
             path_in_repo=path_in_repo,
-            repo_id=HF_BACKUP_REPO_ID,
-            repo_type="model",
+            repo_id=repo_id,
+            repo_type=repo_type,
         )
         size_mb = os.path.getsize(local_path) / 1e6
-        print(f"   ↳ 🛟 Sauvegarde externe: '{path_in_repo}' poussé vers {HF_BACKUP_REPO_ID} ({size_mb:.1f} Mo)")
+        print(f"   ↳ 🛟 Sauvegarde externe: '{path_in_repo}' poussé vers {repo_id} ({size_mb:.1f} Mo)")
     except Exception as e:
         print(f"   ↳ ⚠️ Échec de la sauvegarde externe de '{path_in_repo}' ({e}) — entraînement non interrompu.")
 
 
-def try_restore_backup_from_hf(local_path: str, path_in_repo: str) -> bool:
-    """Tente de restaurer un fichier depuis le repo de sauvegarde HF Hub si le
-    fichier local est absent — utile après une session Kaggle perdue sans
-    'Save Version', pour reprendre automatiquement sans réimport manuel.
-    Retourne True si un fichier a été restauré."""
+def try_restore_backup_from_hf(local_path: str, path_in_repo: str, repo_id: str | None = None,
+                               repo_type: str = "model", direct: bool = False) -> bool:
+    """Tente de restaurer un fichier depuis un repo HF Hub si le fichier local
+    est absent — utile après une session perdue, pour reprendre automatiquement
+    sans réimport manuel. `direct=True` télécharge directement dans le dossier
+    de destination (pas de copie via le cache HF): indispensable pour un fichier
+    de plusieurs Go. Retourne True si un fichier a été restauré."""
     if not HF_BACKUP_ENABLED or os.path.exists(local_path):
         return False
+    repo_id = repo_id or HF_BACKUP_REPO_ID
     try:
         from huggingface_hub import hf_hub_download
-        downloaded = hf_hub_download(repo_id=HF_BACKUP_REPO_ID, filename=path_in_repo, repo_type="model")
-        import shutil
         os.makedirs(os.path.dirname(local_path) or ".", exist_ok=True)
-        shutil.copy2(downloaded, local_path)
-        print(f"♻️  Restauré depuis la sauvegarde externe HF Hub: '{path_in_repo}' -> '{local_path}'")
+        if direct:
+            target_dir = os.path.dirname(local_path) or "."
+            downloaded = hf_hub_download(repo_id=repo_id, filename=path_in_repo,
+                                         repo_type=repo_type, local_dir=target_dir)
+            if os.path.abspath(downloaded) != os.path.abspath(local_path):
+                os.replace(downloaded, local_path)
+        else:
+            downloaded = hf_hub_download(repo_id=repo_id, filename=path_in_repo, repo_type=repo_type)
+            import shutil
+            shutil.copy2(downloaded, local_path)
+        print(f"♻️  Restauré depuis HF Hub ({repo_id}): '{path_in_repo}' -> '{local_path}'")
         return True
     except Exception as e:
-        print(f"ℹ️ Pas de sauvegarde externe utilisable pour '{path_in_repo}' ({e}).")
+        print(f"ℹ️ Pas de sauvegarde externe utilisable pour '{path_in_repo}' sur {repo_id} ({e}).")
         return False
 
 
+def restore_corpus_from_hf(bin_path: str, meta_path: str) -> bool:
+    """Si le corpus .bin est absent ou vide en local, le re-télécharge depuis le
+    repo HF 'dataset' (avec son meta). Ne touche à rien si un .bin non vide existe."""
+    if not (HF_BACKUP_ENABLED and HF_CORPUS_BACKUP_ENABLED):
+        return False
+    if os.path.exists(bin_path) and os.path.getsize(bin_path) > 0:
+        return False
+    if os.path.exists(bin_path):
+        os.remove(bin_path)  # fichier vide: sinon try_restore le considérerait "présent"
+    print(f"⬇️ Corpus local absent/vide — tentative de téléchargement depuis {HF_CORPUS_REPO_ID} "
+          f"(plusieurs Go, sans barre de progression, patiente)...")
+    ok = try_restore_backup_from_hf(bin_path, os.path.basename(bin_path),
+                                    repo_id=HF_CORPUS_REPO_ID, repo_type="dataset", direct=True)
+    if not ok:
+        return False
+    old_meta = None
+    if os.path.exists(meta_path):
+        with open(meta_path, "rb") as f:
+            old_meta = f.read()
+        os.remove(meta_path)
+    if not try_restore_backup_from_hf(meta_path, os.path.basename(meta_path),
+                                      repo_id=HF_CORPUS_REPO_ID, repo_type="dataset", direct=True):
+        if old_meta is not None:
+            with open(meta_path, "wb") as f:
+                f.write(old_meta)
+    return True
+
+
+def push_corpus_to_hf_if_needed(bin_path: str, meta_path: str) -> None:
+    """Pousse le corpus (.bin puis .meta) vers le repo HF 'dataset'. Le .bin
+    (~3,5 Go) n'est renvoyé que si la taille distante diffère de la taille
+    locale: pas de ré-upload inutile à chaque lancement. Ne lève jamais."""
+    if not (HF_BACKUP_ENABLED and HF_CORPUS_BACKUP_ENABLED):
+        return
+    if not (os.path.exists(bin_path) and os.path.getsize(bin_path) > 0 and os.path.exists(meta_path)):
+        return
+    api = _get_hf_api(HF_CORPUS_REPO_ID, "dataset", HF_CORPUS_PRIVATE)
+    if api is None:
+        return
+    bin_name = os.path.basename(bin_path)
+    try:
+        infos = api.get_paths_info(HF_CORPUS_REPO_ID, [bin_name], repo_type="dataset")
+        if infos and infos[0].size == os.path.getsize(bin_path):
+            print(f"✅ Corpus déjà sauvegardé sur {HF_CORPUS_REPO_ID} (même taille) — pas de ré-envoi.")
+        else:
+            raise FileNotFoundError("absent ou taille différente")
+    except Exception:
+        print(f"⬆️ Envoi du corpus ({os.path.getsize(bin_path) / 1e9:.2f} Go) vers {HF_CORPUS_REPO_ID} — "
+              f"quelques minutes la première fois...")
+        push_backup_to_hf(bin_path, bin_name, repo_id=HF_CORPUS_REPO_ID,
+                          repo_type="dataset", private=HF_CORPUS_PRIVATE)
+    # Le meta (quelques centaines d'octets) est toujours repoussé, après le .bin.
+    push_backup_to_hf(meta_path, os.path.basename(meta_path), repo_id=HF_CORPUS_REPO_ID,
+                      repo_type="dataset", private=HF_CORPUS_PRIVATE)
+
+
 # ============================================================================
 # Configuration
 # ============================================================================
@@ -362,6 +441,7 @@ EVAL_INTERVAL = 200        # Évaluation + sauvegarde du meilleur modèle tous l
 GEN_INTERVAL = 400         # Génération d'un échantillon de contrôle tous les N steps
 CHECKPOINT_PATH = os.path.join(_SAVE_DIR, "best_model_medium.pt")
 LATEST_CHECKPOINT_PATH = os.path.join(_SAVE_DIR, "latest_model_medium.pt")
+FINAL_CHECKPOINT_PATH = os.path.join(_SAVE_DIR, "final_model_medium.pt")
 SEED = 1337
 
 # Cible de tokens pour le pré-entraînement à grande échelle (modèle "Large").
@@ -1181,6 +1261,10 @@ def build_pretraining_corpus_bin(
             os.remove(bin_path)
         os.replace(alt_path, bin_path)
 
+    # ⬇️ Dernier recours avant une reconstruction de plusieurs heures: si le .bin
+    # est toujours absent/vide, on tente de le re-télécharger depuis HF Hub.
+    restore_corpus_from_hf(bin_path, meta_path)
+
     if os.path.exists(bin_path) and os.path.exists(meta_path):
         with open(meta_path, "r", encoding="utf-8") as f:
             meta = json.load(f)
@@ -1438,9 +1522,30 @@ def prepare_data_once() -> tuple[str, int]:
     bin_path, total_tokens = build_pretraining_corpus_bin(
         tokenizer, target_tokens=TARGET_PRETRAIN_TOKENS, extra_text_once=seed_text,
     )
+    # 🛟 Sauvegarde du corpus sur HF Hub (no-op si déjà présent avec la même taille).
+    push_corpus_to_hf_if_needed(bin_path, PRETRAIN_META_PATH)
     return bin_path, total_tokens
 
 
+def _checkpoint_progress(path: str) -> int:
+    """Nombre de tokens vus enregistré dans un checkpoint (0 si absent/illisible/
+    ancien format sans 'tokens_seen')."""
+    try:
+        c = torch.load(path, map_location="cpu", weights_only=False)
+        return int(c.get("tokens_seen") or 0)
+    except Exception:
+        return -1
+
+
+def _pick_resume_checkpoint(paths: list[str]) -> str | None:
+    """Parmi les checkpoints locaux existants, renvoie celui qui a vu le plus de
+    tokens (en cas d'égalité, le premier de la liste, donc 'best')."""
+    existing = [p for p in paths if os.path.exists(p)]
+    if not existing:
+        return None
+    return max(existing, key=lambda p: (_checkpoint_progress(p), -existing.index(p)))
+
+
 def main_worker(rank: int, world_size: int, bin_path: str, total_tokens: int) -> None:
 
     if world_size > 1:
@@ -1501,15 +1606,21 @@ def main_worker(rank: int, world_size: int, bin_path: str, total_tokens: int) ->
     # depuis la sauvegarde externe HF Hub (utile après une session Kaggle
     # perdue sans 'Save Version'), PUIS le comportement RESUME_FROM_CHECKPOINT
     # habituel sur le fichier local (potentiellement fraîchement restauré).
-    if accelerator.is_main_process and RESUME_FROM_CHECKPOINT and not os.path.exists(CHECKPOINT_PATH):
-        try_restore_backup_from_hf(CHECKPOINT_PATH, "best_model_medium.pt")
+    if accelerator.is_main_process and RESUME_FROM_CHECKPOINT:
+        if not os.path.exists(CHECKPOINT_PATH):
+            try_restore_backup_from_hf(CHECKPOINT_PATH, "best_model_medium.pt")
+        if not os.path.exists(LATEST_CHECKPOINT_PATH):
+            try_restore_backup_from_hf(LATEST_CHECKPOINT_PATH, "latest_model_medium.pt")
     if accelerator.num_processes > 1:
         accelerator.wait_for_everyone()  # tous les rangs attendent la restauration éventuelle du rang 0
 
     resumed_step = 0
     resumed_step_is_approximate = False
-    if RESUME_FROM_CHECKPOINT and os.path.exists(CHECKPOINT_PATH):
-        _ckpt = torch.load(CHECKPOINT_PATH, map_location="cpu", weights_only=False)
+    # On reprend du checkpoint le plus avancé entre "best" et "latest" (le latest
+    # peut être plus récent que le best s'il n'y a pas eu d'amélioration de val_loss).
+    _resume_path = _pick_resume_checkpoint([CHECKPOINT_PATH, LATEST_CHECKPOINT_PATH]) if RESUME_FROM_CHECKPOINT else None
+    if _resume_path is not None:
+        _ckpt = torch.load(_resume_path, map_location="cpu", weights_only=False)
         try:
             model.load_state_dict(_ckpt["model_state_dict"])
             _tokens_seen = _ckpt.get("tokens_seen")
@@ -1533,9 +1644,11 @@ def main_worker(rank: int, world_size: int, bin_path: str, total_tokens: int) ->
                 resumed_step = _old_step // GRAD_ACCUM_STEPS
                 resumed_step_is_approximate = _old_step > 0
             if accelerator.is_main_process:
-                print(f"♻️  Reprise depuis '{CHECKPOINT_PATH}' (au lieu de repartir de poids aléatoires)...")
+                _pv = _ckpt.get("val_loss")
+                _pv_txt = f"{_pv:.4f}" if _pv is not None else "n/a (checkpoint 'latest')"
+                print(f"♻️  Reprise depuis '{_resume_path}' (au lieu de repartir de poids aléatoires)...")
                 print(f"   ↳ poids chargés (step réel estimé: {resumed_step}, "
-                      f"val_loss précédent: {_ckpt.get('val_loss', float('nan')):.4f})")
+                      f"val_loss précédent: {_pv_txt})")
                 if resumed_step_is_approximate:
                     print(f"   ↳ ⚠️ Checkpoint antérieur à l'ajout de l'accumulation de gradient — "
                           f"step converti par approximation (÷{GRAD_ACCUM_STEPS}), pas une valeur exacte.")
@@ -1548,7 +1661,7 @@ def main_worker(rank: int, world_size: int, bin_path: str, total_tokens: int) ->
             # repart de poids aléatoires — avec un avertissement explicite,
             # pour que ça ne passe jamais inaperçu.
             if accelerator.is_main_process:
-                print(f"⚠️  '{CHECKPOINT_PATH}' existe mais ne correspond pas à l'architecture actuelle "
+                print(f"⚠️  '{_resume_path}' existe mais ne correspond pas à l'architecture actuelle "
                       f"(dim={args.dim}, n_layers={args.n_layers}) — probablement un checkpoint d'une "
                       f"AUTRE taille de modèle. Poids ALÉATOIRES utilisés pour ce run (le fichier "
                       f"n'est pas touché). Renomme/déplace-le si ce n'est pas voulu.\n   Détail: {e}")
@@ -1602,7 +1715,12 @@ def main_worker(rank: int, world_size: int, bin_path: str, total_tokens: int) ->
 
     data_iter = infinite_loader(loader)
 
-    best_val_loss = float(_ckpt.get("val_loss", float("inf"))) if resumed_step else float("inf")
+    # Meilleur val_loss connu: si on reprend depuis "latest" (val_loss=None), on le
+    # relit dans le fichier "best" pour ne PAS écraser un meilleur modèle par un pire.
+    _prev_best = _ckpt.get("val_loss") if resumed_step else None
+    if resumed_step and _prev_best is None and os.path.exists(CHECKPOINT_PATH):
+        _prev_best = torch.load(CHECKPOINT_PATH, map_location="cpu", weights_only=False).get("val_loss")
+    best_val_loss = float(_prev_best) if _prev_best is not None else float("inf")
     _last_latest_backup = time.time()
     # ⚠️ Throttle RETIRÉ : il causait un problème pire que celui qu'il évitait.
     # En repo PRIVÉ, on avait dépassé le quota de stockage (403) à force de
@@ -1683,6 +1801,14 @@ def main_worker(rank: int, world_size: int, bin_path: str, total_tokens: int) ->
     if accelerator.is_main_process:
         print("\n✅ Entraînement terminé.")
 
+        # 🛟 État final (dernier step, pas forcément le "best") sauvegardé et poussé.
+        torch.save(
+            {"model_state_dict": raw_model.state_dict(), "args": args, "step": effective_max_steps - 1,
+             "val_loss": None, "tokens_seen": effective_max_steps * tokens_per_micro_step * GRAD_ACCUM_STEPS},
+            FINAL_CHECKPOINT_PATH,
+        )
+        push_backup_to_hf(FINAL_CHECKPOINT_PATH, "final_model_medium.pt")
+
         if os.path.exists(CHECKPOINT_PATH):
             checkpoint = torch.load(CHECKPOINT_PATH, map_location=device, weights_only=False)
             raw_model.load_state_dict(checkpoint["model_state_dict"])