cortex.6.sol / model /README.md
openhands
openhands
feat(cortex): add CORTEX training pipeline and model audit
6fbe100
|
Raw History Blame Contribute Delete
2.66 kB

CORTEX — model

Définition du modèle CORTEX entraînable et gestion des checkpoints.

Fichier Rôle
cortex_model.py Architecture : décodeur Transformer, RMSNorm, RoPE, GQA, SwiGLU
init.py Initialisation des poids, écriture et lecture des checkpoints

Architecture

CortexConfig porte tous les hyper-paramètres. parameter_breakdown() calcule le nombre de paramètres de façon analytique, et count_parameters() le recompte sur le modèle réel : les deux doivent être égaux, ce que vérifie un test.

Le modèle est un décodeur causal standard :

tokens -> embed_tokens -> [N x (RMSNorm -> GQA+RoPE -> résidu -> RMSNorm -> SwiGLU -> résidu)]
       -> RMSNorm -> lm_head -> logits

Points d'implémentation à connaître :

  • Perte par blocs. Le vocabulaire CORTEX compte 129 280 entrées. Un tenseur de logits complet en float32 pour un lot de 2×512 tokens pèserait ~0,5 Go. La perte est donc calculée par tranches (CortexForCausalLM._loss), ce qui évite une pointe mémoire qui faisait redémarrer l'environnement de développement.
  • Poids liés. Avec tie_word_embeddings: true, lm_head.weight est embed_tokens.weight. count_parameters déduplique par identité pour ne pas compter deux fois la même matrice.
  • Masque. En l'absence de attention_mask, le masque causal est délégué à scaled_dot_product_attention. Un masque de padding est combiné au masque causal.

Checkpoints

Un checkpoint CORTEX est un dossier auto-descriptif :

<checkpoint>/
    config.json         hyper-paramètres du modèle
    weights.safetensors
    provenance.json     comment les poids ont été produits
    optimizer.pt        si un optimiseur est fourni
    manifest.json       tailles et SHA-256 de chaque fichier

load_cortex_checkpoint() vérifie le manifeste (taille et SHA-256) avant de charger les poids : un fichier corrompu est refusé, pas silencieusement accepté.

provenance.json contient toujours base_model: null pour les poids initialisés ici, avec une note explicite. Ce champ ne doit être renseigné que si les poids dérivent réellement d'un autre modèle.

Exemple

from model.cortex_model import CortexConfig
from model.init import init_cortex_model, save_cortex_checkpoint, load_cortex_checkpoint

config = CortexConfig(model_name="cortex-dev-1")
model = init_cortex_model(config, seed=1337)
save_cortex_checkpoint(model, config, "checkpoints/demo", step=0)

model, config, provenance = load_cortex_checkpoint("checkpoints/demo")
print(provenance["developer"], provenance["base_model"])