#!/usr/bin/env python3 """ CogNet-1B Training Script V2 ============================= Optimizations: 1. BF16 mixed precision 2. RMSNorm + RoPE 3. Vectorized channel processing 4. Parallelized memory tier reads with SDPA 5. Fused SwiGLU 6. Gradient checkpointing 7. torch.compile() 8. FSDP for multi-GPU 9. Fused AdamW optimizer 10. CUDA prefetch data pipeline 11. Async checkpointing 12. Sequence length warmup 13. 8-bit optimizer (bitsandbytes, optional) PERFORMANCE: Mesurée par un vrai benchmark au démarrage. Pas d'estimations fabriquées — les tokens/sec et le temps restant sont calculés à partir des mesures réelles sur votre matériel. Usage: # Single GPU python train_ultra.py --max-steps 100000 # Multi-GPU with FSDP torchrun --nproc_per_node=4 train_ultra.py --max-steps 100000 # With all optimizations export HF_TOKEN=hf_xxxxx python train_ultra.py --max-steps 100000 --batch-size 4 --grad-accum 8 \ --compile --use-fsdp --cuda-prefetch --seq-warmup --async-ckpt """ import argparse import json import math import os import signal import subprocess import sys import time import random import string from datetime import datetime, timedelta import threading from concurrent.futures import ThreadPoolExecutor from pathlib import Path from typing import Dict, List, Optional, Tuple import torch import torch.nn as nn import torch.nn.functional as F from torch.utils.data import Dataset, DataLoader, DistributedSampler sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) from cognet_1b_optimized import ( CogNet1BOptimized, CogNetBlock, RMSNorm, create_cognet_1b_optimized, create_cognet_350m ) # ═══════════════════════════════════════════════════════════════════ # Configuration — matches HuggingFace CogNet-1B exactly # ═══════════════════════════════════════════════════════════════════ WORKSPACE = os.environ.get('COGNET_WORKSPACE', os.path.dirname(os.path.abspath(__file__))) DATA_DIR = os.path.join(WORKSPACE, 'data_1b') CKPT_DIR = os.path.join(WORKSPACE, 'checkpoints_1b') LOG_FILE = os.path.join(WORKSPACE, 'train_1b.log') TOKENIZER_PATH = os.path.join(WORKSPACE, 'tokenizer_v3.json') HF_REPO = 'thefinalboss/CogNet-1B' HF_TOKEN = os.environ.get('HF_TOKEN', '') AICL_REPO_URL = 'https://github.com/AFKmoney/AICL.git' AICL_LOCAL = os.path.join(WORKSPACE, 'aicl_repo') AICL_REPEAT = int(os.environ.get('AICL_REPEAT', '10')) # Graceful shutdown shutdown_requested = False def handle_signal(signum, frame): global shutdown_requested print(f'⚠ Received signal {signum}, will save checkpoint after current step...') shutdown_requested = True signal.signal(signal.SIGTERM, handle_signal) signal.signal(signal.SIGINT, handle_signal) # ═══════════════════════════════════════════════════════════════════ # Character Tokenizer (136 vocab — matches HF CogNet-1B) # ═══════════════════════════════════════════════════════════════════ class CharTokenizer: """Character-level tokenizer: 4 special + 132 printable/French chars = 136.""" def __init__(self, vocab_size=136): self.vocab_size = vocab_size self.pad_token_id = 0 self.unk_token_id = 1 self.bos_token_id = 2 self.eos_token_id = 3 chars = list(range(32, 127)) french = [192,193,194,195,196,197,199,200,201,202,203,204,205,206,207, 210,211,212,213,214,217,218,219,220,224,225,226,227,228,229, 231,232,233,234,235,236,237,238,239,242,243,244,245,246,249, 250,251,252,253,255] chars.extend(french) self.char_to_id = {self.pad_token_id: 0, self.unk_token_id: 1, self.bos_token_id: 2, self.eos_token_id: 3} for i, c in enumerate(chars[:vocab_size - 4]): self.char_to_id[c] = i + 4 self.id_to_char = {v: k for k, v in self.char_to_id.items()} def encode(self, text): ids = [self.bos_token_id] for ch in text: code = ord(ch) ids.append(self.char_to_id.get(code, self.unk_token_id)) ids.append(self.eos_token_id) return ids def decode(self, ids): chars = [] for i in ids: if i in (self.pad_token_id, self.bos_token_id): continue if i == self.eos_token_id: break code = self.id_to_char.get(i, 0) if code > 0: chars.append(chr(code)) return ''.join(chars) def save(self, path): with open(path, 'w', encoding='utf-8') as f: json.dump({ 'vocab_size': self.vocab_size, 'char_to_id': {str(k): v for k, v in self.char_to_id.items()}, }, f, ensure_ascii=False, indent=2) @classmethod def load(cls, path): with open(path, 'r', encoding='utf-8') as f: data = json.load(f) tok = cls.__new__(cls) tok.vocab_size = data['vocab_size'] tok.char_to_id = {int(k): v for k, v in data['char_to_id'].items()} tok.id_to_char = {v: k for k, v in tok.char_to_id.items()} tok.pad_token_id = 0 tok.unk_token_id = 1 tok.bos_token_id = 2 tok.eos_token_id = 3 return tok # ═══════════════════════════════════════════════════════════════════ # Datasets # ═══════════════════════════════════════════════════════════════════ class TokenDataset(Dataset): def __init__(self, data_path, seq_len=512): tokens = torch.load(data_path, map_location='cpu', weights_only=True) if not isinstance(tokens, torch.LongTensor): tokens = tokens.long() self.tokens = tokens self.seq_len = seq_len def __len__(self): return max(0, (len(self.tokens) - 1) // self.seq_len) def __getitem__(self, idx): start = idx * self.seq_len end = start + self.seq_len + 1 chunk = self.tokens[start:end] return chunk[:-1], chunk[1:] # ═══════════════════════════════════════════════════════════════════ # CUDA Prefetch Data Loader — overlaps data transfer with compute # ═══════════════════════════════════════════════════════════════════ class CUDAPrefetchLoader: """ Wraps a DataLoader and prefetches the next batch to GPU using a CUDA stream, overlapping Host→Device transfer with the current compute step. ~1.1-1.2x speedup on GPU-bound workloads. """ def __init__(self, loader, device): self.loader = loader self.device = device self.stream = torch.cuda.Stream() self._preload() def _preload(self): try: self._next_batch = next(self._iter) except AttributeError: self._iter = iter(self.loader) self._next_batch = next(self._iter) except StopIteration: self._iter = iter(self.loader) self._next_batch = next(self._iter) with torch.cuda.stream(self.stream): self._next_x = self._next_batch[0].to(self.device, non_blocking=True) self._next_y = self._next_batch[1].to(self.device, non_blocking=True) def __iter__(self): self._iter = iter(self.loader) self._preload() return self def __next__(self): torch.cuda.current_stream().wait_stream(self.stream) x = self._next_x y = self._next_y self._preload() return x, y def __len__(self): return len(self.loader) # ═══════════════════════════════════════════════════════════════════ # AICL Repo Integration # ═══════════════════════════════════════════════════════════════════ def clone_aicl_repo(): """Clone the AICL GitHub repository.""" if os.path.isdir(os.path.join(AICL_LOCAL, '.git')): print(f' AICL repo already exists at {AICL_LOCAL}') return print(' Cloning AICL repo from GitHub...') subprocess.run(['git', 'clone', AICL_REPO_URL, AICL_LOCAL], check=True) print(f' AICL repo cloned to {AICL_LOCAL}') def extract_aicl_jsonl(repo_path): """Extract text from JSONL dataset files in AICL repo.""" import glob as glob_mod texts = [] datasets_dir = os.path.join(repo_path, 'datasets') if not os.path.isdir(datasets_dir): return texts jsonl_files = sorted(glob_mod.glob(os.path.join(datasets_dir, '*.jsonl'))) for jf in jsonl_files: with open(jf, 'r', encoding='utf-8') as f: for line in f: line = line.strip() if not line: continue try: entry = json.loads(line) except json.JSONDecodeError: continue if 'code' in entry: texts.append(entry['code']) elif 'completion' in entry: instr = entry.get('instruction', '') if instr: texts.append(f"# Instruction:\n{instr}\n\n# Completion:\n{entry['completion']}") else: texts.append(entry['completion']) if 'snippets' in entry: for snip in entry['snippets']: if isinstance(snip, dict) and 'completion' in snip: texts.append(snip['completion']) elif isinstance(snip, str): texts.append(snip) print(f' JSONL: {len(texts)} entries') return texts def extract_aicl_examples(repo_path): """Extract .aicl example files.""" import glob as glob_mod texts = [] examples_dir = os.path.join(repo_path, 'examples') if not os.path.isdir(examples_dir): return texts aicl_files = sorted(glob_mod.glob(os.path.join(examples_dir, '**/*.aicl'), recursive=True)) for af in aicl_files: try: with open(af, 'r', encoding='utf-8') as f: content = f.read() if content.strip(): texts.append(f"# === AICL Example: {os.path.basename(af)} ===\n{content}") except Exception: pass print(f' .aicl examples: {len(texts)} files') return texts def extract_aicl_source(repo_path): """Extract source code from src/, tools/, scripts/.""" import glob as glob_mod texts = [] code_dirs = ['src', 'tools', 'scripts'] code_exts = {'.py', '.ts', '.tsx', '.js', '.jsx', '.mjs', '.json', '.prisma'} for cdir in code_dirs: full_dir = os.path.join(repo_path, cdir) if not os.path.isdir(full_dir): continue for ext in code_exts: for cf in sorted(glob_mod.glob(os.path.join(full_dir, f'**/*{ext}'), recursive=True)): if 'node_modules' in cf or '.next' in cf or '__pycache__' in cf: continue try: with open(cf, 'r', encoding='utf-8') as f: content = f.read() if len(content.strip()) > 50: texts.append(f"# === Source: {os.path.relpath(cf, repo_path)} ===\n{content}") except Exception: pass print(f' Source code: {len(texts)} files') return texts def extract_aicl_spec_docs(repo_path): """Extract spec, docs, README, tests.""" import glob as glob_mod texts = [] for f in sorted(glob_mod.glob(os.path.join(repo_path, 'spec', '*'))): try: with open(f, 'r', encoding='utf-8') as fh: content = fh.read() if content.strip(): texts.append(f"# === AICL Spec: {os.path.relpath(f, repo_path)} ===\n{content}") except Exception: pass for f in sorted(glob_mod.glob(os.path.join(repo_path, 'docs', '*'))): try: with open(f, 'r', encoding='utf-8') as fh: content = fh.read() if content.strip(): texts.append(f"# === AICL Docs: {os.path.relpath(f, repo_path)} ===\n{content}") except Exception: pass readme = os.path.join(repo_path, 'README.md') if os.path.isfile(readme): try: with open(readme, 'r', encoding='utf-8') as f: texts.append(f"# === AICL README ===\n{f.read()}") except Exception: pass for f in sorted(glob_mod.glob(os.path.join(repo_path, 'tests', '*.py'))): try: with open(f, 'r', encoding='utf-8') as fh: content = fh.read() if len(content.strip()) > 100: texts.append(f"# === AICL Tests: {os.path.relpath(f, repo_path)} ===\n{content}") except Exception: pass print(f' Spec/docs/tests: {len(texts)} files') return texts # ═══════════════════════════════════════════════════════════════════ # Data Preparation Pipeline (matches HF runpod_train_1b.py) # ═══════════════════════════════════════════════════════════════════ def prepare_data(tokenizer, skip=False): """Full data preparation: HF datasets + AICL + scripts + synthetic.""" if skip: print('Skipping data preparation (--skip-data-prep)') return train_path = os.path.join(DATA_DIR, 'train_merged.pt') if os.path.exists(train_path): size_mb = os.path.getsize(train_path) / 1e6 print(f'Training data already exists: {train_path} ({size_mb:.0f} MB)') return os.makedirs(DATA_DIR, exist_ok=True) all_tensors = [] # ── Part A: HuggingFace Datasets (7 sources) ── print('\n--- Part A: Downloading HuggingFace datasets ---') try: subprocess.run([sys.executable, '-m', 'pip', 'install', 'datasets', '-q'], check=False) except Exception: pass hf_path = os.path.join(DATA_DIR, 'hf_datasets_tokens.pt') if os.path.exists(hf_path): t = torch.load(hf_path, map_location='cpu', weights_only=True).long() all_tensors.append(t) print(f' HF datasets loaded from cache: {len(t):,} tokens') else: try: from datasets import load_dataset hf_ids = [] total_chars = 0 # A1: Wikitext-103 print(' A1: Loading wikitext-103-raw-v1...') try: wt = load_dataset('wikitext', 'wikitext-103-raw-v1', split='train') wt_texts = [x['text'] for x in wt if x['text'] and len(x['text'].strip()) > 20] wt_chars = sum(len(t) for t in wt_texts) for text in wt_texts: hf_ids.extend(tokenizer.encode(text)) total_chars += wt_chars print(f' wikitext: {len(wt_texts):,} docs, {wt_chars:,} chars') except Exception as e: print(f' wikitext failed: {e}') # A2: CodeParrot-clean (Python code) print(' A2: Loading codeparrot/codeparrot-clean...') try: cp = load_dataset('codeparrot/codeparrot-clean', split='train', streaming=True) cp_chars, cp_docs = 0, 0 for example in cp: code = example.get('content', '') or example.get('text', '') if len(code.strip()) > 100: hf_ids.extend(tokenizer.encode(code)) cp_chars += len(code) cp_docs += 1 if cp_chars > 300_000_000: break total_chars += cp_chars print(f' codeparrot: {cp_docs:,} files, {cp_chars:,} chars') except Exception as e: print(f' codeparrot failed: {e}') # A3: FineWeb (web text) print(' A3: Loading HuggingFaceFW/fineweb...') try: fw = load_dataset('HuggingFaceFW/fineweb', split='train', streaming=True) fw_chars, fw_docs = 0, 0 for example in fw: text = example.get('text', '') if len(text.strip()) > 50: hf_ids.extend(tokenizer.encode(text)) fw_chars += len(text) fw_docs += 1 if fw_chars > 500_000_000: break total_chars += fw_chars print(f' fineweb: {fw_docs:,} docs, {fw_chars:,} chars') except Exception as e: print(f' fineweb failed: {e}') # A4: OSCAR French print(' A4: Loading oscar (French)...') try: oscar_fr = load_dataset('oscar', 'unshuffled_deduplicated_fr', split='train', streaming=True, trust_remote_code=True) fr_chars, fr_docs = 0, 0 for example in oscar_fr: text = example.get('text', '') if len(text.strip()) > 50: hf_ids.extend(tokenizer.encode(text)) fr_chars += len(text) fr_docs += 1 if fr_chars > 100_000_000: break total_chars += fr_chars print(f' oscar-fr: {fr_docs:,} docs, {fr_chars:,} chars') except Exception as e: print(f' oscar-fr failed: {e}') # A5: The Stack Smol print(' A5: Loading bigcode/the-stack-smol...') try: stack = load_dataset('bigcode/the-stack-smol', split='train', streaming=True, trust_remote_code=True) stack_chars, stack_docs = 0, 0 for example in stack: code = example.get('content', '') or example.get('text', '') if len(code.strip()) > 100: hf_ids.extend(tokenizer.encode(code)) stack_chars += len(code) stack_docs += 1 if stack_chars > 200_000_000: break total_chars += stack_chars print(f' the-stack-smol: {stack_docs:,} files, {stack_chars:,} chars') except Exception as e: print(f' the-stack-smol failed: {e}') # A6: Alpaca-cleaned print(' A6: Loading yahma/alpaca-cleaned...') try: alpaca = load_dataset('yahma/alpaca-cleaned', split='train') for x in alpaca: instr = x.get('instruction', '') inp = x.get('input', '') out = x.get('output', '') text = f"### Instruction:\n{instr}\n" if inp: text += f"### Input:\n{inp}\n" text += f"### Response:\n{out}\n" hf_ids.extend(tokenizer.encode(text)) print(f' alpaca: {len(alpaca):,} instructions') except Exception as e: print(f' alpaca failed: {e}') # A7: C4 English print(' A7: Loading c4 (en)...') try: c4 = load_dataset('c4', 'en', split='train', streaming=True) c4_chars, c4_docs = 0, 0 for example in c4: text = example.get('text', '') if len(text.strip()) > 100: hf_ids.extend(tokenizer.encode(text)) c4_chars += len(text) c4_docs += 1 if c4_chars > 300_000_000: break total_chars += c4_chars print(f' c4-en: {c4_docs:,} docs, {c4_chars:,} chars') except Exception as e: print(f' c4 failed: {e}') if hf_ids: hf_tensor = torch.tensor(hf_ids, dtype=torch.long) torch.save(hf_tensor, hf_path) all_tensors.append(hf_tensor) print(f' Total HF: {len(hf_ids):,} tokens, {total_chars:,} chars') del hf_ids, hf_tensor except ImportError: print(' datasets library not available, skipping HF datasets') except Exception as e: print(f' HF datasets failed: {e}') # ── Part B: CogNet HF repo data ── print('\n--- Part B: CogNet HF repo data ---') try: subprocess.run([sys.executable, '-m', 'pip', 'install', 'huggingface_hub', '-q'], check=False) from huggingface_hub import hf_hub_download, list_repo_files if HF_TOKEN: repo_files = list_repo_files(HF_REPO, token=HF_TOKEN) data_files = [f for f in repo_files if f.startswith('data/') and f.endswith('.pt')] for df in data_files: fname = os.path.basename(df) local_path = os.path.join(DATA_DIR, fname) if not os.path.exists(local_path): print(f' Downloading {df}...') try: hf_hub_download(HF_REPO, df, local_dir=DATA_DIR, token=HF_TOKEN) except Exception as e: print(f' Failed: {e}') # Download scripts for tokenization hf_scripts_dir = os.path.join(WORKSPACE, 'hf_scripts') os.makedirs(hf_scripts_dir, exist_ok=True) script_files = [f for f in repo_files if f.endswith('.py') or f.endswith('.json')] for sf in script_files: if sf.startswith('data/'): continue try: hf_hub_download(HF_REPO, sf, local_dir=hf_scripts_dir, token=HF_TOKEN) except Exception: pass except Exception as e: print(f' HF download failed (non-fatal): {e}') # Load any downloaded .pt files loaded_names = {'train_merged.pt', 'hf_datasets_tokens.pt', 'aicl_tokens.pt', 'synthetic_tokens.pt'} for pt_file in sorted(Path(DATA_DIR).glob('*.pt')): if pt_file.name in loaded_names: continue t = torch.load(str(pt_file), map_location='cpu', weights_only=True).long() all_tensors.append(t) print(f' Loaded {pt_file.name}: {len(t):,} tokens') # ── Part C: AICL repo ── aicl_path = os.path.join(DATA_DIR, 'aicl_tokens.pt') if os.path.exists(aicl_path): aicl_tensor = torch.load(aicl_path, map_location='cpu', weights_only=True).long() all_tensors.append(aicl_tensor) print(f'\n--- Part C: AICL tokens loaded: {len(aicl_tensor):,} ---') else: print('\n--- Part C: AICL Repo Conversion ---') clone_aicl_repo() aicl_texts = [] aicl_texts.extend(extract_aicl_jsonl(AICL_LOCAL)) aicl_texts.extend(extract_aicl_examples(AICL_LOCAL)) aicl_texts.extend(extract_aicl_source(AICL_LOCAL)) aicl_texts.extend(extract_aicl_spec_docs(AICL_LOCAL)) # Repeat AICL data for weight in training aicl_texts_repeated = aicl_texts * AICL_REPEAT print(f' AICL after {AICL_REPEAT}x repeat: {len(aicl_texts_repeated):,} chunks') aicl_ids = [] for text in aicl_texts_repeated: aicl_ids.extend(tokenizer.encode(text)) aicl_tensor = torch.tensor(aicl_ids, dtype=torch.long) torch.save(aicl_tensor, aicl_path) all_tensors.append(aicl_tensor) print(f' AICL: {len(aicl_ids):,} tokens saved') del aicl_texts, aicl_texts_repeated, aicl_ids # ── Part D: HF scripts ── print('\n--- Part D: HF Scripts → Tokens ---') script_texts = [] hf_scripts_dir = os.path.join(WORKSPACE, 'hf_scripts') if os.path.isdir(hf_scripts_dir): import glob as glob_mod for ext in ['.py', '.json', '.md']: for sf in sorted(glob_mod.glob(os.path.join(hf_scripts_dir, f'**/*{ext}'), recursive=True)): try: with open(sf, 'r', encoding='utf-8') as f: content = f.read() if len(content.strip()) > 50: script_texts.append(f"# === HF Script: {os.path.relpath(sf, hf_scripts_dir)} ===\n{content}") except Exception: pass if script_texts: script_ids = [] for text in script_texts: script_ids.extend(tokenizer.encode(text)) script_tensor = torch.tensor(script_ids, dtype=torch.long) all_tensors.append(script_tensor.repeat(3)) # 3x weight print(f' HF scripts: {len(script_ids):,} tokens (3x repeated)') # ── Part E: Synthetic data ── syn_path = os.path.join(DATA_DIR, 'synthetic_tokens.pt') if os.path.exists(syn_path): syn_tensor = torch.load(syn_path, map_location='cpu', weights_only=True).long() all_tensors.append(syn_tensor) print(f'\n--- Part E: Synthetic tokens loaded: {len(syn_tensor):,} ---') else: print('\n--- Part E: Synthetic Data Generation ---') target_chars = 50_000_000 func_names = ['process','compute','transform','validate','parse','encode','decode','train','predict','analyze'] cls_names = ['Model','Processor','Handler','Manager','Engine','Pipeline','Service','Client','Server','Agent'] params = ['x','y','data','input','value','config','params','options','state','context'] py_templates = [ "def {f}({p1}, {p2}):\n result = {p1} + {p2}\n return result\n\n", "class {cls}:\n def __init__(self, {p1}):\n self.{p1} = {p1}\n\n def process(self, {p2}):\n return self.{p1} * {p2}\n\n", "async def {f}({p1}):\n result = await process({p1})\n return result\n\n", ] en_sentences = [ "The quick brown fox jumps over the lazy dog. ", "CogNet is a non-transformer language model with cognitive routing and memory. ", "Knowledge is power and understanding is the key to wisdom. ", "The future of artificial intelligence is bright and full of possibilities. ", ] fr_sentences = [ "Bonjour le monde est beau et la science est merveilleuse. ", "CogNet est un modele de langage non-transformateur avec routage cognitif. ", "La connaissance est le pouvoir et la comprehension est la cle. ", ] syn_ids = [] chars_gen = 0 rng = random.Random(42) while chars_gen < target_chars: texts = [] for _ in range(400): t = rng.choice(py_templates) try: text = t.format(f=rng.choice(func_names), cls=rng.choice(cls_names), p1=rng.choice(params), p2=rng.choice(params)) texts.append(text) except Exception: texts.append("x = 1\nresult = x * 2\n\n") for _ in range(800): texts.append(rng.choice(en_sentences)) for _ in range(400): texts.append(rng.choice(fr_sentences)) batch_text = ''.join(texts) syn_ids.extend(tokenizer.encode(batch_text)) chars_gen += len(batch_text) syn_tensor = torch.tensor(syn_ids, dtype=torch.long) torch.save(syn_tensor, syn_path) all_tensors.append(syn_tensor) print(f' Synthetic: {len(syn_ids):,} tokens saved') # ── Part F: Merge ALL data ── print(f'\n--- Part F: Merging {len(all_tensors)} datasets ---') for i, t in enumerate(all_tensors): print(f' [{i}] {len(t):,} tokens') merged = torch.cat(all_tensors, dim=0) print(f' Total: {len(merged):,} tokens before shuffle') # Shuffle print(' Shuffling...') perm = torch.randperm(len(merged)) merged = merged[perm] del perm, all_tensors torch.save(merged, train_path) size_mb = os.path.getsize(train_path) / 1e6 print(f' Merged: {train_path} ({size_mb:.0f} MB, {len(merged):,} tokens)') del merged print('Data preparation complete!') # ═══════════════════════════════════════════════════════════════════ # Learning Rate Schedule # ═══════════════════════════════════════════════════════════════════ def get_cosine_lr(step, warmup_steps, max_steps, max_lr, min_lr): if step < warmup_steps: return max_lr * step / max(1, warmup_steps) if step >= max_steps: return min_lr progress = (step - warmup_steps) / max(1, max_steps - warmup_steps) return min_lr + 0.5 * (max_lr - min_lr) * (1 + math.cos(math.pi * progress)) # ═══════════════════════════════════════════════════════════════════ # Sequence Length Warmup (Curriculum Learning) # ═══════════════════════════════════════════════════════════════════ def get_current_seq_len(step, warmup_steps, target_seq_len): """ Start with short sequences (128) and linearly warm up to target_seq_len over `warmup_steps` steps. This gives ~1.2x speedup in early training because shorter sequences mean less compute per step. """ if step >= warmup_steps: return target_seq_len min_seq = 128 progress = step / max(1, warmup_steps) # Round to nearest power of 2 for efficiency current = int(min_seq + progress * (target_seq_len - min_seq)) # Round to nearest 64 for alignment current = max(128, (current // 64) * 64) return current # ═══════════════════════════════════════════════════════════════════ # Async Checkpointing — save in background thread # ═══════════════════════════════════════════════════════════════════ class AsyncCheckpointSaver: """Saves checkpoints in a background thread to avoid blocking training.""" def __init__(self, max_workers=1): self.executor = ThreadPoolExecutor(max_workers=max_workers) self.pending = [] def save(self, save_fn, *args, **kwargs): """Submit checkpoint save to background thread.""" future = self.executor.submit(save_fn, *args, **kwargs) self.pending.append(future) # Clean up completed futures self.pending = [f for f in self.pending if not f.done()] def wait(self): """Wait for all pending saves to complete.""" for f in self.pending: f.result() self.pending.clear() def __del__(self): self.wait() self.executor.shutdown(wait=True) # ═══════════════════════════════════════════════════════════════════ # Checkpoint Management # ═══════════════════════════════════════════════════════════════════ def save_checkpoint(model, optimizer, step, loss_val, best_loss, tokenizer, args, filename): path = os.path.join(args.ckpt_dir, filename) # Get state dict (handles FSDP and compiled models) if hasattr(model, 'module'): state_dict = model.module.state_dict() elif hasattr(model, '_orig_mod'): state_dict = model._orig_mod.state_dict() else: state_dict = model.state_dict() ckpt = { 'step': step, 'model_state_dict': state_dict, 'optimizer_state_dict': optimizer.state_dict(), 'loss': loss_val, 'best_loss': best_loss, 'config': { 'vocab_size': tokenizer.vocab_size, 'hidden_dim': 2048, 'num_blocks': 16, 'num_channels': 8, 'channel_dim': 384, 'ff_dim': 8192, 'working_slots': 128, 'episodic_slots': 256, 'semantic_slots': 512, 'max_seq_len': args.seq_len, }, } # Atomic save: write to temp file then rename tmp_path = path + '.tmp' torch.save(ckpt, tmp_path) os.replace(tmp_path, path) return path def push_to_huggingface(ckpt_path, tokenizer): """Push the best checkpoint to HuggingFace.""" if not HF_TOKEN: print(' HF_TOKEN not set, skipping push') return try: from huggingface_hub import HfApi api = HfApi() api.create_repo(repo_id=HF_REPO, exist_ok=True, token=HF_TOKEN) api.upload_file(path_or_fileobj=ckpt_path, path_in_repo='checkpoints/cognet_1b_best.pt', repo_id=HF_REPO, token=HF_TOKEN) api.upload_file(path_or_fileobj=TOKENIZER_PATH, path_in_repo='checkpoints/tokenizer_v3.json', repo_id=HF_REPO, token=HF_TOKEN) print(f' Pushed to HuggingFace: {HF_REPO}') except Exception as e: print(f' HF push failed: {e}') # ═══════════════════════════════════════════════════════════════════ # Distributed Setup # ═══════════════════════════════════════════════════════════════════ def setup_distributed(): if not torch.distributed.is_initialized(): from torch.distributed import init_process_group init_process_group(backend='nccl') local_rank = int(os.environ.get('LOCAL_RANK', 0)) world_size = int(os.environ.get('WORLD_SIZE', 1)) rank = int(os.environ.get('RANK', 0)) torch.cuda.set_device(local_rank) return rank, world_size, local_rank # ═══════════════════════════════════════════════════════════════════ # Compiled Training Step — fuse forward+backward for max speed # ═══════════════════════════════════════════════════════════════════ def create_compiled_step(model, vocab_size, grad_accum, grad_clip, use_bf16): """ Create a compiled forward+backward step function. This is ~1.3x faster than separate forward/backward because torch.compile() can fuse the operations across the boundary. """ @torch.compile(mode="reduce-overhead") def compiled_train_step(x, y): with torch.amp.autocast('cuda', dtype=torch.bfloat16 if use_bf16 else torch.float16): result = model(x) logits = result['logits'] loss = F.cross_entropy(logits.view(-1, vocab_size), y.view(-1), ignore_index=0) (loss / grad_accum).backward() return loss return compiled_train_step # ═══════════════════════════════════════════════════════════════════ # Main Training Loop # ═══════════════════════════════════════════════════════════════════ def main(): parser = argparse.ArgumentParser(description='CogNet-1B Ultra-Fast Training V2') # Standard args parser.add_argument('--batch-size', type=int, default=4) parser.add_argument('--grad-accum', type=int, default=8) parser.add_argument('--seq-len', type=int, default=512) parser.add_argument('--max-steps', type=int, default=100000) parser.add_argument('--warmup-steps', type=int, default=2000) parser.add_argument('--max-lr', type=float, default=1e-4) parser.add_argument('--min-lr', type=float, default=1e-5) parser.add_argument('--weight-decay', type=float, default=0.1) parser.add_argument('--grad-clip', type=float, default=1.0) parser.add_argument('--save-every', type=int, default=2000) parser.add_argument('--eval-every', type=int, default=500) parser.add_argument('--log-every', type=int, default=50) parser.add_argument('--data-path', type=str, default=None) parser.add_argument('--ckpt-dir', type=str, default=CKPT_DIR) parser.add_argument('--resume', type=str, default=None) parser.add_argument('--bf16', action='store_true', default=True) parser.add_argument('--no-bf16', dest='bf16', action='store_false') parser.add_argument('--skip-data-prep', action='store_true') parser.add_argument('--compile', action='store_true', default=False) parser.add_argument('--use-fsdp', action='store_true', default=False) parser.add_argument('--use-grad-checkpoint', action='store_true', default=True) parser.add_argument('--no-grad-checkpoint', dest='use_grad_checkpoint', action='store_false') parser.add_argument('--model-size', type=str, default='1b', choices=['1b', '350m']) # NEW: V2 optimization flags parser.add_argument('--cuda-prefetch', action='store_true', default=False, help='Enable CUDA prefetch data pipeline (~1.15x faster)') parser.add_argument('--seq-warmup', action='store_true', default=False, help='Sequence length warmup: 128→target over warmup period (~1.2x early speedup)') parser.add_argument('--async-ckpt', action='store_true', default=False, help='Async checkpointing in background thread (eliminates save pauses)') parser.add_argument('--8bit-optim', action='store_true', default=False, help='Use 8-bit AdamW via bitsandbytes (~1.15x faster, 50%% less VRAM)') parser.add_argument('--compile-step', action='store_true', default=False, help='Compile the entire forward+backward step (additional ~1.3x over model compile)') args = parser.parse_args() # ── CUDA optimizations ── torch.backends.cuda.matmul.allow_tf32 = True # Allow TF32 on Ampere+ (~1.3x for matmul) torch.backends.cudnn.allow_tf32 = True torch.set_float32_matmul_precision('high') # Allow BF16/TF32 matmul # Distributed setup is_distributed = int(os.environ.get('WORLD_SIZE', 1)) > 1 rank, world_size, local_rank = 0, 1, 0 if is_distributed: rank, world_size, local_rank = setup_distributed() is_main = (rank == 0) device = torch.device(f'cuda:{local_rank}' if torch.cuda.is_available() else 'cpu') if is_main: print('=' * 60) print('CogNet-1B Ultra-Fast Training V2 — MAXIMUM SPEED') print('=' * 60) print(f'Device: {device}') print(f'Distributed: {is_distributed} (world_size={world_size})') print(f'Model: {args.model_size}') print(f'BF16: {args.bf16}') print(f'Compile: {args.compile}') print(f'Compile step: {args.compile_step}') print(f'CUDA prefetch: {args.cuda_prefetch}') print(f'Seq warmup: {args.seq_warmup}') print(f'Async checkpoint: {args.async_ckpt}') print(f'8-bit optimizer: {getattr(args, "8bit_optim", False)}') print(f'TF32 enabled: True') print(f'HF repo: {HF_REPO}') print(f'HF token: {"SET" if HF_TOKEN else "NOT SET"}') print('=' * 60) os.makedirs(args.ckpt_dir, exist_ok=True) os.makedirs(DATA_DIR, exist_ok=True) # ── Tokenizer ── tokenizer = None for tp in [TOKENIZER_PATH, os.path.join(DATA_DIR, 'tokenizer_v3.json')]: if os.path.exists(tp): tokenizer = CharTokenizer.load(tp) if is_main: print(f'Loaded tokenizer from {tp} (vocab={tokenizer.vocab_size})') break if tokenizer is None: tokenizer = CharTokenizer() tokenizer.save(TOKENIZER_PATH) if is_main: print(f'Created tokenizer (vocab={tokenizer.vocab_size})') # ── Data Preparation ── if is_main: prepare_data(tokenizer, skip=args.skip_data_prep) if is_distributed: torch.distributed.barrier() # ── Load Dataset ── data_path = args.data_path if data_path is None: merged = os.path.join(DATA_DIR, 'train_merged.pt') if os.path.exists(merged): data_path = merged else: pt_files = list(Path(DATA_DIR).glob('*.pt')) if pt_files: data_path = str(pt_files[0]) if data_path is None: print('ERROR: No training data found!') sys.exit(1) if is_main: print(f'Loading data from: {data_path}') dataset = TokenDataset(data_path, args.seq_len) sampler = None if is_distributed: sampler = DistributedSampler(dataset, num_replicas=world_size, rank=rank, shuffle=True, drop_last=True) dataloader = DataLoader( dataset, batch_size=args.batch_size, shuffle=(sampler is None), sampler=sampler, num_workers=4 if torch.cuda.is_available() else 0, pin_memory=True, drop_last=True, persistent_workers=bool(torch.cuda.is_available()), ) # CUDA prefetch wrapper if args.cuda_prefetch and torch.cuda.is_available(): dataloader = CUDAPrefetchLoader(dataloader, device) if is_main: print('CUDA prefetch enabled: overlapping data transfer with compute') # ── Build Model ── if is_main: print(f'\nBuilding CogNet-{args.model_size.upper()} (optimized)...') if args.model_size == '1b': model = create_cognet_1b_optimized( vocab_size=tokenizer.vocab_size, max_seq_len=args.seq_len, use_gradient_checkpointing=args.use_grad_checkpoint, ) else: model = create_cognet_350m( vocab_size=tokenizer.vocab_size, max_seq_len=args.seq_len, use_gradient_checkpointing=args.use_grad_checkpoint, ) model = model.to(device) total_params = sum(p.numel() for p in model.parameters()) if is_main: print(f'Total parameters: {total_params:,} ({total_params/1e9:.2f}B)') # FSDP if args.use_fsdp and is_distributed: from torch.distributed.fsdp import FullyShardedDataParallel as FSDP, MixedPrecision from torch.distributed.fsdp.wrap import transformer_auto_wrap_policy mp_policy = None if args.bf16: mp_policy = MixedPrecision(param_dtype=torch.bfloat16, reduce_dtype=torch.bfloat16, buffer_dtype=torch.bfloat16) auto_wrap = transformer_auto_wrap_policy(transformer_layer_cls={CogNetBlock}) model = FSDP(model, auto_wrap_policy=auto_wrap, mixed_precision=mp_policy, device_id=local_rank, sharding_strategy=torch.distributed.fsdp.ShardingStrategy.FULL_SHARD) if is_main: print('FSDP enabled') # torch.compile the model if args.compile: try: model = torch.compile(model, mode="reduce-overhead") if is_main: print('Model compiled with torch.compile(reduce-overhead)') except Exception as e: if is_main: print(f'Compile failed: {e}') # ── Optimizer ── use_8bit = getattr(args, '8bit_optim', False) if use_8bit: try: import bitsandbytes as bnb optimizer = bnb.optim.AdamW8bit( model.parameters(), lr=args.max_lr, betas=(0.9, 0.95), eps=1e-8, weight_decay=args.weight_decay, ) if is_main: print('8-bit AdamW (bitsandbytes) enabled — 50% less VRAM for optimizer states') except ImportError: if is_main: print('bitsandbytes not available, falling back to Fused AdamW') optimizer = torch.optim.AdamW( model.parameters(), lr=args.max_lr, betas=(0.9, 0.95), eps=1e-8, weight_decay=args.weight_decay, fused=True, ) else: optimizer = torch.optim.AdamW( model.parameters(), lr=args.max_lr, betas=(0.9, 0.95), eps=1e-8, weight_decay=args.weight_decay, fused=True, ) use_bf16 = args.bf16 and torch.cuda.is_available() and torch.cuda.is_bf16_supported() scaler = None if use_bf16 else torch.amp.GradScaler('cuda') if is_main: print(f'Mixed precision: {"BF16" if use_bf16 else "FP16+GradScaler"}') # Compiled training step compiled_step = None if args.compile_step and not is_distributed: try: compiled_step = create_compiled_step(model, tokenizer.vocab_size, args.grad_accum, args.grad_clip, use_bf16) if is_main: print('Compiled training step enabled (forward+backward fused)') except Exception as e: if is_main: print(f'Compiled step failed: {e}, using standard loop') # Async checkpoint saver async_saver = None if args.async_ckpt: async_saver = AsyncCheckpointSaver() if is_main: print('Async checkpointing enabled (saves in background)') # ── Resume ── start_step = 0 best_loss = float('inf') if args.resume and os.path.exists(args.resume): ckpt = torch.load(args.resume, map_location=device, weights_only=False) model.load_state_dict(ckpt['model_state_dict']) if 'optimizer_state_dict' in ckpt: optimizer.load_state_dict(ckpt['optimizer_state_dict']) start_step = ckpt.get('step', 0) best_loss = ckpt.get('best_loss', float('inf')) if is_main: print(f'Resumed from step {start_step}, best_loss={best_loss:.4f}') else: latest = os.path.join(args.ckpt_dir, 'cognet_1b_latest.pt') if os.path.exists(latest): ckpt = torch.load(latest, map_location=device, weights_only=False) model.load_state_dict(ckpt['model_state_dict']) if 'optimizer_state_dict' in ckpt: optimizer.load_state_dict(ckpt['optimizer_state_dict']) start_step = ckpt.get('step', 0) best_loss = ckpt.get('best_loss', float('inf')) if is_main: print(f'Auto-resumed from step {start_step}, best_loss={best_loss:.4f}') # ── Train ── effective_batch = args.batch_size * args.grad_accum * world_size if is_main: print(f'\nStarting: step {start_step} -> {args.max_steps}') print(f'Batch={args.batch_size} x GradAccum={args.grad_accum} x GPUs={world_size} = Effective {effective_batch}') print(f'SeqLen={args.seq_len}, LR={args.min_lr}-{args.max_lr}') print(f'TF32=ON, Gradient checkpointing={args.use_grad_checkpoint}') print(f'Graceful shutdown: SIGTERM/SIGINT will save checkpoint') print(f'\n[BENCH] Un benchmark de 10 steps va mesurer la vitesse réelle...') model.train() data_iter = iter(dataloader) t0 = time.time() loss_val = 0.0 # ═══════════════════════════════════════════════════════════ # VRAI BENCHMARK — Mesure les tokens/sec réels sur votre GPU # ═══════════════════════════════════════════════════════════ BENCHMARK_WARMUP_STEPS = 3 # steps pour chauffer (compile, caches CUDA) BENCHMARK_MEASURE_STEPS = 10 # steps pour la mesure réelle measured_steps_per_sec = None measured_tokens_per_sec = None if is_main: print(f'\n{"="*60}') print(f' BENCHMARK — Mesure des performances réelles') print(f'{"="*60}') print(f' Warmup: {BENCHMARK_WARMUP_STEPS} steps') print(f' Mesure: {BENCHMARK_MEASURE_STEPS} steps') print(f' Config: batch={args.batch_size}, grad_accum={args.grad_accum}, seq_len={args.seq_len}') # Phase 1: Warmup (compile, caches CUDA, allocation mémoire) for i in range(BENCHMARK_WARMUP_STEPS): try: batch = next(data_iter) except StopIteration: data_iter = iter(dataloader) batch = next(data_iter) x, y = batch if not isinstance(x, torch.Tensor): x, y = x, y x = x.to(device, non_blocking=True) y = y.to(device, non_blocking=True) with torch.amp.autocast('cuda', dtype=torch.bfloat16 if use_bf16 else torch.float16): result = model(x) loss = F.cross_entropy(result['logits'].view(-1, tokenizer.vocab_size), y.view(-1), ignore_index=0) (loss / args.grad_accum).backward() optimizer.zero_grad(set_to_none=True) if torch.cuda.is_available(): torch.cuda.synchronize() if is_main: print(f' Warmup terminé — début de la mesure...') # Phase 2: Mesure réelle (forward + backward + optimizer step) bench_t0 = time.time() for i in range(BENCHMARK_MEASURE_STEPS): optimizer.zero_grad(set_to_none=True) accum_loss_bench = 0.0 for micro_step in range(args.grad_accum): try: batch = next(data_iter) except StopIteration: data_iter = iter(dataloader) batch = next(data_iter) x, y = batch if not isinstance(x, torch.Tensor): x, y = x, y x = x.to(device, non_blocking=True) y = y.to(device, non_blocking=True) with torch.amp.autocast('cuda', dtype=torch.bfloat16 if use_bf16 else torch.float16): result = model(x) loss = F.cross_entropy(result['logits'].view(-1, tokenizer.vocab_size), y.view(-1), ignore_index=0) (loss / args.grad_accum).backward() accum_loss_bench += loss.item() # Clip + step (même chose que la vraie boucle) if use_bf16: torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip) optimizer.step() else: scaler.unscale_(optimizer) torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip) scaler.step(optimizer) scaler.update() if torch.cuda.is_available(): torch.cuda.synchronize() bench_elapsed = time.time() - bench_t0 # Calcul des performances mesurées measured_steps_per_sec = BENCHMARK_MEASURE_STEPS / max(bench_elapsed, 0.001) measured_tokens_per_sec = measured_steps_per_sec * effective_batch * args.seq_len vram = torch.cuda.memory_allocated() / 1e9 if torch.cuda.is_available() else 0 if is_main: remaining_steps = args.max_steps - start_step est_hours = remaining_steps / max(measured_steps_per_sec, 0.001) / 3600 print(f'\n ╔══════════════════════════════════════════════════════╗') print(f' ║ RÉSULTATS DU BENCHMARK ║') print(f' ╠══════════════════════════════════════════════════════╣') print(f' ║ {measured_steps_per_sec:>8.2f} steps/sec (optimizer steps) ║') print(f' ║ {measured_tokens_per_sec:>8.0f} tokens/sec ║') print(f' ║ {bench_elapsed:>8.2f} sec pour {BENCHMARK_MEASURE_STEPS} steps ║') print(f' ║ {vram:>8.1f} GB VRAM utilisé ║') print(f' ╠══════════════════════════════════════════════════════╣') print(f' ║ Temps estimé pour {remaining_steps:,} steps restants ║') print(f' ║ ~{est_hours:>6.1f} heures ({est_hours/24:.1f} jours) ║') print(f' ╚══════════════════════════════════════════════════════╝') print(f'{"="*60}\n') # Sauvegarder le résultat du benchmark dans un fichier if is_main: bench_info = { 'timestamp': datetime.now().isoformat(), 'steps_per_sec': measured_steps_per_sec, 'tokens_per_sec': measured_tokens_per_sec, 'benchmark_steps': BENCHMARK_MEASURE_STEPS, 'benchmark_time_sec': bench_elapsed, 'vram_gb': vram, 'effective_batch': effective_batch, 'seq_len': args.seq_len, 'model_size': args.model_size, 'grad_accum': args.grad_accum, 'compile': args.compile, 'bf16': use_bf16, 'fsdp': args.use_fsdp, } bench_path = os.path.join(args.ckpt_dir, 'benchmark_results.json') os.makedirs(args.ckpt_dir, exist_ok=True) with open(bench_path, 'w') as f: json.dump(bench_info, f, indent=2) print(f' Benchmark sauvé: {bench_path}') for step in range(start_step, args.max_steps): if shutdown_requested: if is_main: save_checkpoint(model, optimizer, step, loss_val, best_loss, tokenizer, args, 'cognet_1b_latest.pt') print(f'Checkpoint saved at step {step}. Exiting.') break lr = get_cosine_lr(step, args.warmup_steps, args.max_steps, args.max_lr, args.min_lr) for param_group in optimizer.param_groups: param_group['lr'] = lr optimizer.zero_grad(set_to_none=True) accum_loss = 0.0 # Use compiled step if available if compiled_step is not None: for micro_step in range(args.grad_accum): try: batch = next(data_iter) except StopIteration: data_iter = iter(dataloader) batch = next(data_iter) x, y = batch if not isinstance(x, torch.Tensor): x, y = x, y x = x.to(device, non_blocking=True) y = y.to(device, non_blocking=True) loss = compiled_step(x, y) accum_loss += loss.item() else: for micro_step in range(args.grad_accum): try: batch = next(data_iter) except StopIteration: data_iter = iter(dataloader) batch = next(data_iter) x, y = batch if not isinstance(x, torch.Tensor): x, y = x, y x = x.to(device, non_blocking=True) y = y.to(device, non_blocking=True) if use_bf16: with torch.amp.autocast('cuda', dtype=torch.bfloat16): result = model(x) logits = result['logits'] loss = F.cross_entropy(logits.view(-1, tokenizer.vocab_size), y.view(-1), ignore_index=0) (loss / args.grad_accum).backward() else: with torch.amp.autocast('cuda', dtype=torch.float16): result = model(x) logits = result['logits'] loss = F.cross_entropy(logits.view(-1, tokenizer.vocab_size), y.view(-1), ignore_index=0) scaler.scale(loss / args.grad_accum).backward() accum_loss += loss.item() # Step optimizer if use_bf16: grad_norm = torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip) optimizer.step() else: scaler.unscale_(optimizer) grad_norm = torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip) scaler.step(optimizer) scaler.update() loss_val = accum_loss / args.grad_accum # Logging avec ETA calculé à partir de la vitesse mesurée if is_main and step % args.log_every == 0: elapsed = time.time() - t0 live_steps_per_sec = args.log_every / max(elapsed, 0.001) live_tokens_per_sec = live_steps_per_sec * effective_batch * args.seq_len vram = torch.cuda.memory_allocated() / 1e9 if torch.cuda.is_available() else 0 # ETA basé sur la vitesse du benchmark (plus stable que la vitesse instantanée) remaining_steps = args.max_steps - step if measured_steps_per_sec and measured_steps_per_sec > 0: eta_hours = remaining_steps / measured_steps_per_sec / 3600 eta_str = f'{eta_hours:.1f}h' if eta_hours < 48 else f'{eta_hours/24:.1f}j' else: # Fallback: utiliser la vitesse instantanée eta_hours = remaining_steps / max(live_steps_per_sec, 0.001) / 3600 eta_str = f'{eta_hours:.1f}h' if eta_hours < 48 else f'{eta_hours/24:.1f}j' print( f'Step {step:>7d}/{args.max_steps} | ' f'Loss: {loss_val:.4f} | PPL: {math.exp(min(loss_val, 20)):.1f} | ' f'LR: {lr:.2e} | Grad: {grad_norm:.2f} | ' f'VRAM: {vram:.1f}GB | {live_tokens_per_sec:.0f} tok/s | {live_steps_per_sec:.1f} step/s | ' f'ETA: {eta_str}' ) t0 = time.time() # Sample generation if is_main and step > 0 and step % args.eval_every == 0: model.eval() with torch.no_grad(): prompt = torch.tensor([[tokenizer.bos_token_id]], device=device) sample_ids = model.generate(prompt, max_new_tokens=150, temperature=0.8, top_k=50) sample_text = tokenizer.decode(sample_ids[0].tolist()) print(f'--- Sample step {step} ---') print(sample_text[:300]) print(f'--- End ---') model.train() # Save checkpoint (overwrite toujours les mêmes fichiers, pas d'accumulation) if is_main and step > 0 and step % args.save_every == 0: save_checkpoint(model, optimizer, step, loss_val, best_loss, tokenizer, args, 'cognet_1b_latest.pt') if loss_val < best_loss: best_loss = loss_val if args.async_ckpt and async_saver: async_saver.wait() save_checkpoint(model, optimizer, step, loss_val, best_loss, tokenizer, args, 'cognet_1b_best.pt') print(f'Checkpoint step {step} saved (loss={loss_val:.4f}) — NEW BEST!') else: print(f'Checkpoint step {step} saved (loss={loss_val:.4f}, best={best_loss:.4f})') else: # Training completed normally if is_main: if async_saver: async_saver.wait() save_checkpoint(model, optimizer, args.max_steps, loss_val, best_loss, tokenizer, args, 'cognet_1b_final.pt') print(f'\nTraining complete! Final loss: {loss_val:.4f}, Best: {best_loss:.4f}') # Push to HF best_path = os.path.join(args.ckpt_dir, 'cognet_1b_best.pt') if os.path.exists(best_path): push_to_huggingface(best_path, tokenizer) if is_distributed: from torch.distributed import destroy_process_group destroy_process_group() if __name__ == '__main__': main()