#!/usr/bin/env python3 import argparse import json import math import os import random import time from pathlib import Path from typing import Any, Dict, Iterator, List, Optional, Tuple import numpy as np import pyarrow.parquet as pq import torch import torch.nn as nn import torch.nn.functional as F from tokenizers import Tokenizer from torch.utils.data import DataLoader, IterableDataset, get_worker_info from tqdm import tqdm PAD_ID = 0 DEFAULT_BOS_MARKER = "<|BOS|>" DEFAULT_EOS_MARKER = "<|EOS|>" DEFAULT_BOS_TOKEN = "[BOS]" DEFAULT_EOS_TOKEN = "[EOS]" BOUNDARY_MODE = "marker_aware_generic_special_markers_v9" _FLASH2_KERNEL = None _FLASH3_KERNEL = None def get_flash2_kernel(): global _FLASH2_KERNEL if _FLASH2_KERNEL is None: from kernels import get_kernel _FLASH2_KERNEL = get_kernel( "kernels-community/flash-attn2", version=1, ) return _FLASH2_KERNEL def get_flash3_kernel(): global _FLASH3_KERNEL if _FLASH3_KERNEL is None: from kernels import get_kernel _FLASH3_KERNEL = get_kernel( "kernels-community/flash-attn3", version=1, ) return _FLASH3_KERNEL def format_tokens(n: int) -> str: if n >= 1_000_000_000: return f"{n / 1_000_000_000:.2f}B" if n >= 1_000_000: return f"{n / 1_000_000:.2f}M" if n >= 1_000: return f"{n / 1_000:.2f}K" return str(n) def resolve_tokenizer_path(path: str) -> str: p = Path(path) if p.is_dir(): candidate = p / "tokenizer.json" if candidate.exists(): return str(candidate) return str(p) def stable_row_score(row_index: int, seed: int) -> float: x = (row_index + 1) & 0xFFFFFFFFFFFFFFFF x ^= (seed + 0x9E3779B97F4A7C15) & 0xFFFFFFFFFFFFFFFF x = (x * 0xBF58476D1CE4E5B9) & 0xFFFFFFFFFFFFFFFF x ^= x >> 30 x = (x * 0x94D049BB133111EB) & 0xFFFFFFFFFFFFFFFF x ^= x >> 31 return (x & 0xFFFFFFFF) / 0x100000000 def normalize_activity_value(value: Any) -> Optional[str]: if value is None: return None if isinstance(value, str): text = value.strip() return text if text else None if isinstance(value, (list, tuple)): parts = [] for item in value: if item is None: continue s = str(item).strip() if s: parts.append(s) text = " ; ".join(parts).strip() return text if text else None if isinstance(value, dict): text = json.dumps( value, ensure_ascii=False, sort_keys=True, ).strip() return text if text else None text = str(value).strip() return text if text else None def canonical_special_token(value: str) -> str: value = str(value).strip() if not value: raise ValueError("Special token vide.") if value.startswith("[") and value.endswith("]"): inner = value[1:-1].strip() if not inner: raise ValueError(f"Token spécial invalide: {value}") return "[" + inner.upper() + "]" return "[" + value.upper() + "]" def parse_special_marker_spec(spec: str) -> Tuple[str, str]: spec = str(spec).strip() if "=" not in spec: raise ValueError( f"Format --special-marker invalide: {spec}. Format attendu: '<|BOC|>=[BOC]'" ) marker, token = spec.split("=", 1) marker = marker.strip() token = token.strip() if not marker: raise ValueError(f"Marker vide dans: {spec}") if not token: raise ValueError(f"Token vide dans: {spec}") token = canonical_special_token(token) return marker, token def build_marker_token_map(custom_specs: List[str]) -> Dict[str, str]: marker_token_map: Dict[str, str] = { DEFAULT_BOS_MARKER: DEFAULT_BOS_TOKEN, DEFAULT_EOS_MARKER: DEFAULT_EOS_TOKEN, } for spec in custom_specs: marker, token = parse_special_marker_spec(spec) marker_token_map[marker] = token return marker_token_map class RNETokenCache: def __init__( self, src: str, tokenizer_path: str, cache_dir: str, activity_column: str = "activites", row_batch_size: int = 100_000, val_ratio: float = 0.01, seed: int = 42, lowercase: bool = False, append_special_tokens: bool = True, rebuild_cache: bool = False, shuffle_before_tokenize: bool = True, shuffle_buffer_size: int = 500_000, special_marker_specs: Optional[List[str]] = None, ): if not 0.0 < val_ratio < 0.5: raise ValueError("--val-ratio must be > 0 and < 0.5") if shuffle_buffer_size <= 0: raise ValueError("--shuffle-buffer-size must be > 0") self.src = str(src) self.tokenizer_path = resolve_tokenizer_path(tokenizer_path) self.cache_dir = Path(cache_dir) self.activity_column = activity_column self.row_batch_size = int(row_batch_size) self.val_ratio = float(val_ratio) self.seed = int(seed) self.lowercase = bool(lowercase) self.append_special_tokens = bool(append_special_tokens) self.rebuild_cache = bool(rebuild_cache) self.shuffle_before_tokenize = bool(shuffle_before_tokenize) self.shuffle_buffer_size = int(shuffle_buffer_size) self.special_marker_specs = list(special_marker_specs or []) self.cache_dir.mkdir(parents=True, exist_ok=True) self.train_bin = self.cache_dir / "train_tokens.uint32.bin" self.val_bin = self.cache_dir / "val_tokens.uint32.bin" self.meta_path = self.cache_dir / "meta.json" self.tokenizer = Tokenizer.from_file(self.tokenizer_path) self.vocab_size = self.tokenizer.get_vocab_size() self.marker_token_map = build_marker_token_map(self.special_marker_specs) self.marker_id_map = self._build_marker_id_map() self.bos_id = self._find_bos_id() if self.append_special_tokens else None self.eos_id = self._find_eos_id() if self.append_special_tokens else None self.sep_id = self._find_sep_id() if self.append_special_tokens else None if self.append_special_tokens: if self.bos_id is None: raise RuntimeError( "BOS token introuvable. Le tokenizer doit contenir [BOS], , , , [CLS] ou équivalent." ) if self.eos_id is None: raise RuntimeError( "EOS token introuvable. Le tokenizer doit contenir [EOS], , , , [SEP] ou équivalent." ) self.shuffle_rng_train = random.Random(self.seed + 123_456_789) self.shuffle_rng_val = random.Random(self.seed + 987_654_321) def _find_token_id(self, candidates: List[str]) -> Optional[int]: for token in candidates: token_id = self.tokenizer.token_to_id(token) if token_id is not None: return int(token_id) return None def _find_bos_id(self) -> Optional[int]: explicit_token = self.marker_token_map.get(DEFAULT_BOS_MARKER, DEFAULT_BOS_TOKEN) return self._find_token_id( [ explicit_token, "[BOS]", "", "", "", "[CLS]", DEFAULT_BOS_MARKER, ] ) def _find_eos_id(self) -> Optional[int]: explicit_token = self.marker_token_map.get(DEFAULT_EOS_MARKER, DEFAULT_EOS_TOKEN) return self._find_token_id( [ explicit_token, "[EOS]", "", "", "", "[SEP]", "", "", DEFAULT_EOS_MARKER, ] ) def _find_sep_id(self) -> Optional[int]: return self._find_token_id( [ "[SEP]", "", "", "", "[EOS]", "", "", DEFAULT_EOS_MARKER, ] ) def _build_marker_id_map(self) -> Dict[str, int]: marker_id_map: Dict[str, int] = {} for marker, token in self.marker_token_map.items(): token_id = self.tokenizer.token_to_id(token) if token_id is None: raise RuntimeError( f"Token spécial introuvable dans le tokenizer: marker {repr(marker)} -> token {repr(token)}. " f"Ajoute-le au tokenizer avec --add-special-token." ) marker_id_map[marker] = int(token_id) return marker_id_map def _cache_is_valid(self) -> bool: if self.rebuild_cache: return False if not self.train_bin.exists(): return False if not self.val_bin.exists(): return False if not self.meta_path.exists(): return False try: meta = json.loads(self.meta_path.read_text(encoding="utf-8")) except Exception: return False expected = { "src": os.path.abspath(self.src), "tokenizer_path": os.path.abspath(self.tokenizer_path), "activity_column": self.activity_column, "val_ratio": self.val_ratio, "seed": self.seed, "lowercase": self.lowercase, "append_special_tokens": self.append_special_tokens, "bos_id": self.bos_id, "eos_id": self.eos_id, "sep_id": self.sep_id, "vocab_size": self.vocab_size, "shuffle_before_tokenize": self.shuffle_before_tokenize, "shuffle_buffer_size": self.shuffle_buffer_size, "boundary_mode": BOUNDARY_MODE, "default_bos_marker": DEFAULT_BOS_MARKER, "default_eos_marker": DEFAULT_EOS_MARKER, "marker_token_map": self.marker_token_map, "marker_id_map": self.marker_id_map, } for key, value in expected.items(): if meta.get(key) != value: return False return True def _write_meta( self, train_tokens: int, val_tokens: int, rows_seen: int, rows_used: int, rows_with_mapped_markers: int, rows_with_explicit_boundaries: int, rows_with_legacy_boundaries: int, ): payload = { "src": os.path.abspath(self.src), "tokenizer_path": os.path.abspath(self.tokenizer_path), "activity_column": self.activity_column, "val_ratio": self.val_ratio, "seed": self.seed, "lowercase": self.lowercase, "append_special_tokens": self.append_special_tokens, "bos_id": self.bos_id, "eos_id": self.eos_id, "sep_id": self.sep_id, "vocab_size": self.vocab_size, "shuffle_before_tokenize": self.shuffle_before_tokenize, "shuffle_buffer_size": self.shuffle_buffer_size, "boundary_mode": BOUNDARY_MODE, "default_bos_marker": DEFAULT_BOS_MARKER, "default_eos_marker": DEFAULT_EOS_MARKER, "marker_token_map": self.marker_token_map, "marker_id_map": self.marker_id_map, "train_tokens": int(train_tokens), "val_tokens": int(val_tokens), "rows_seen": int(rows_seen), "rows_used": int(rows_used), "rows_with_mapped_markers": int(rows_with_mapped_markers), "rows_with_explicit_boundaries": int(rows_with_explicit_boundaries), "rows_with_legacy_boundaries": int(rows_with_legacy_boundaries), "dtype": "uint32", } self.meta_path.write_text( json.dumps(payload, ensure_ascii=False, indent=2), encoding="utf-8", ) def _shuffle_buffer_with_progress( self, buffer: List[str], rng: random.Random, desc: str, ): n = len(buffer) if n <= 1: return pbar = tqdm( total=n - 1, desc=desc, dynamic_ncols=True, unit="swap", ) for i in range(n - 1, 0, -1): j = rng.randint(0, i) buffer[i], buffer[j] = buffer[j], buffer[i] pbar.update(1) pbar.close() def _has_explicit_bos_and_eos_markers(self, text: str) -> bool: return DEFAULT_BOS_MARKER in text and DEFAULT_EOS_MARKER in text def _has_any_mapped_marker(self, text: str) -> bool: for marker in self.marker_id_map.keys(): if marker in text: return True return False def _encode_plain_chunk(self, text: str) -> List[int]: if not text: return [] if self.lowercase: text = text.lower() ids = self.tokenizer.encode( text, add_special_tokens=False, ).ids return [int(x) for x in ids] def _find_next_marker(self, text: str, start: int) -> Tuple[int, Optional[str], Optional[int]]: best_pos = -1 best_marker = None best_id = None for marker, marker_id in self.marker_id_map.items(): pos = text.find(marker, start) if pos == -1: continue if best_pos == -1 or pos < best_pos: best_pos = pos best_marker = marker best_id = marker_id return best_pos, best_marker, best_id def _encode_text_replacing_markers(self, text: str) -> Tuple[List[int], bool]: ids: List[int] = [] i = 0 n = len(text) used_marker = False while i < n: marker_pos, marker, marker_id = self._find_next_marker(text, i) if marker_pos == -1 or marker is None or marker_id is None: chunk = text[i:] ids.extend(self._encode_plain_chunk(chunk)) break chunk = text[i:marker_pos] ids.extend(self._encode_plain_chunk(chunk)) ids.append(int(marker_id)) used_marker = True i = marker_pos + len(marker) return ids, used_marker def _encode_text_with_boundaries(self, text: str) -> Tuple[List[int], bool, bool]: has_explicit_boundaries = self._has_explicit_bos_and_eos_markers(text) has_any_marker = self._has_any_mapped_marker(text) if not self.append_special_tokens: ids, used_marker = self._encode_text_replacing_markers(text) return ids, used_marker, has_explicit_boundaries if has_any_marker: ids, used_marker = self._encode_text_replacing_markers(text) if has_explicit_boundaries: return ids, used_marker, True ids = [int(self.bos_id)] + ids + [int(self.eos_id)] return ids, used_marker, False if self.lowercase: text = text.lower() ids = self.tokenizer.encode( text, add_special_tokens=False, ).ids ids = [int(x) for x in ids] ids = [int(self.bos_id)] + ids + [int(self.eos_id)] return ids, False, False def _tokenize_to_file( self, texts: List[str], file_obj, desc: str, ) -> Tuple[int, int, int, int]: written_tokens = 0 used_texts = 0 mapped_marker_rows = 0 explicit_boundary_rows = 0 legacy_boundary_rows = 0 pbar = tqdm( total=len(texts), desc=desc, dynamic_ncols=True, unit="texts", ) for text in texts: ids, used_mapped_marker, used_explicit_boundaries = self._encode_text_with_boundaries(text) if len(ids) >= 2: arr = np.asarray(ids, dtype=np.uint32) arr.tofile(file_obj) written_tokens += int(arr.size) used_texts += 1 if used_mapped_marker: mapped_marker_rows += 1 if used_explicit_boundaries: explicit_boundary_rows += 1 else: legacy_boundary_rows += 1 pbar.update(1) if used_texts > 0 and used_texts % 10_000 == 0: pbar.set_postfix( used=f"{used_texts:,}", tokens=format_tokens(written_tokens), markers=f"{mapped_marker_rows:,}", explicit=f"{explicit_boundary_rows:,}", legacy=f"{legacy_boundary_rows:,}", ) pbar.close() return written_tokens, mapped_marker_rows, explicit_boundary_rows, legacy_boundary_rows def _flush_text_buffer( self, buffer: List[str], file_obj, rng: random.Random, name: str, ) -> Tuple[int, int, int, int]: if not buffer: return 0, 0, 0, 0 print() print(f"[FLUSH] {name}") print(f"[FLUSH] texts in buffer: {len(buffer):,}") if self.shuffle_before_tokenize: self._shuffle_buffer_with_progress( buffer=buffer, rng=rng, desc=f"Shuffling {name}", ) written_tokens, mapped_marker_rows, explicit_boundary_rows, legacy_boundary_rows = self._tokenize_to_file( texts=buffer, file_obj=file_obj, desc=f"Tokenizing {name}", ) print(f"[FLUSH] {name} tokens written: {written_tokens:,}") print(f"[FLUSH] {name} mapped marker rows: {mapped_marker_rows:,}") print(f"[FLUSH] {name} explicit boundary rows: {explicit_boundary_rows:,}") print(f"[FLUSH] {name} legacy boundary rows: {legacy_boundary_rows:,}") print() buffer.clear() return written_tokens, mapped_marker_rows, explicit_boundary_rows, legacy_boundary_rows def build_if_needed(self): if self._cache_is_valid(): print("[INFO] Token cache found.") meta = json.loads(self.meta_path.read_text(encoding="utf-8")) print(f"[INFO] Train tokens: {meta['train_tokens']:,}") print(f"[INFO] Val tokens: {meta['val_tokens']:,}") print(f"[INFO] Rows seen: {meta.get('rows_seen', 0):,}") print(f"[INFO] Rows used: {meta.get('rows_used', 0):,}") print(f"[INFO] Mapped marker rows: {meta.get('rows_with_mapped_markers', 0):,}") print(f"[INFO] Explicit boundary rows: {meta.get('rows_with_explicit_boundaries', 0):,}") print(f"[INFO] Legacy boundary rows: {meta.get('rows_with_legacy_boundaries', 0):,}") print(f"[INFO] Vocab size: {meta['vocab_size']:,}") print(f"[INFO] BOS id: {meta.get('bos_id')}") print(f"[INFO] EOS id: {meta.get('eos_id')}") print(f"[INFO] SEP id: {meta.get('sep_id')}") print(f"[INFO] Boundary mode: {meta.get('boundary_mode')}") print(f"[INFO] Marker token map: {meta.get('marker_token_map')}") print(f"[INFO] Marker id map: {meta.get('marker_id_map')}") print(f"[INFO] Shuffle before tok: {meta.get('shuffle_before_tokenize')}") print(f"[INFO] Shuffle buffer: {meta.get('shuffle_buffer_size'):,}") return print("[INFO] Building token cache from parquet.") print(f"[INFO] Source: {self.src}") print(f"[INFO] Column: {self.activity_column}") print(f"[INFO] Tokenizer: {self.tokenizer_path}") print(f"[INFO] Cache dir: {self.cache_dir}") print(f"[INFO] Vocab size: {self.vocab_size:,}") print(f"[INFO] Append special: {self.append_special_tokens}") print(f"[INFO] BOS id: {self.bos_id}") print(f"[INFO] EOS id: {self.eos_id}") print(f"[INFO] SEP id: {self.sep_id}") print(f"[INFO] Boundary mode: {BOUNDARY_MODE}") print(f"[INFO] Marker token map: {self.marker_token_map}") print(f"[INFO] Marker id map: {self.marker_id_map}") print(f"[INFO] Explicit boundary rule: if <|BOS|> and <|EOS|> are present, no auto BOS/EOS") print(f"[INFO] Legacy boundary rule: otherwise BOS + text + EOS") print(f"[INFO] Shuffle before tok: {self.shuffle_before_tokenize}") print(f"[INFO] Shuffle buffer size: {self.shuffle_buffer_size:,}") print() pf = pq.ParquetFile(self.src) if self.activity_column not in pf.schema.names: raise ValueError( f"Column '{self.activity_column}' not found. Available columns: {pf.schema.names}" ) total_rows = pf.metadata.num_rows train_tmp = self.train_bin.with_suffix(".tmp") val_tmp = self.val_bin.with_suffix(".tmp") if train_tmp.exists(): train_tmp.unlink() if val_tmp.exists(): val_tmp.unlink() train_tokens = 0 val_tokens = 0 rows_seen = 0 rows_used = 0 rows_with_mapped_markers = 0 rows_with_explicit_boundaries = 0 rows_with_legacy_boundaries = 0 train_text_buffer: List[str] = [] val_text_buffer: List[str] = [] with train_tmp.open("wb") as f_train, val_tmp.open("wb") as f_val: pbar = tqdm( total=total_rows, desc="Reading + shuffling + tokenizing rows", dynamic_ncols=True, unit="rows", ) for batch in pf.iter_batches( batch_size=self.row_batch_size, columns=[self.activity_column], ): d = batch.to_pydict() values = d[self.activity_column] for value in values: row_index = rows_seen rows_seen += 1 text = normalize_activity_value(value) if text is None: pbar.update(1) continue if not text: pbar.update(1) continue if stable_row_score(row_index, self.seed) < self.val_ratio: val_text_buffer.append(text) else: train_text_buffer.append(text) rows_used += 1 if len(train_text_buffer) >= self.shuffle_buffer_size: written, marker_rows, explicit_rows, legacy_rows = self._flush_text_buffer( buffer=train_text_buffer, file_obj=f_train, rng=self.shuffle_rng_train, name="train buffer", ) train_tokens += written rows_with_mapped_markers += marker_rows rows_with_explicit_boundaries += explicit_rows rows_with_legacy_boundaries += legacy_rows if len(val_text_buffer) >= max(1_000, self.shuffle_buffer_size // 10): written, marker_rows, explicit_rows, legacy_rows = self._flush_text_buffer( buffer=val_text_buffer, file_obj=f_val, rng=self.shuffle_rng_val, name="val buffer", ) val_tokens += written rows_with_mapped_markers += marker_rows rows_with_explicit_boundaries += explicit_rows rows_with_legacy_boundaries += legacy_rows pbar.update(1) if rows_used % 10_000 == 0: pbar.set_postfix( used=f"{rows_used:,}", train_tok=format_tokens(train_tokens), val_tok=format_tokens(val_tokens), tr_buf=f"{len(train_text_buffer):,}", va_buf=f"{len(val_text_buffer):,}", markers=f"{rows_with_mapped_markers:,}", explicit=f"{rows_with_explicit_boundaries:,}", legacy=f"{rows_with_legacy_boundaries:,}", ) written, marker_rows, explicit_rows, legacy_rows = self._flush_text_buffer( buffer=train_text_buffer, file_obj=f_train, rng=self.shuffle_rng_train, name="final train buffer", ) train_tokens += written rows_with_mapped_markers += marker_rows rows_with_explicit_boundaries += explicit_rows rows_with_legacy_boundaries += legacy_rows written, marker_rows, explicit_rows, legacy_rows = self._flush_text_buffer( buffer=val_text_buffer, file_obj=f_val, rng=self.shuffle_rng_val, name="final val buffer", ) val_tokens += written rows_with_mapped_markers += marker_rows rows_with_explicit_boundaries += explicit_rows rows_with_legacy_boundaries += legacy_rows pbar.close() train_tmp.replace(self.train_bin) val_tmp.replace(self.val_bin) self._write_meta( train_tokens=train_tokens, val_tokens=val_tokens, rows_seen=rows_seen, rows_used=rows_used, rows_with_mapped_markers=rows_with_mapped_markers, rows_with_explicit_boundaries=rows_with_explicit_boundaries, rows_with_legacy_boundaries=rows_with_legacy_boundaries, ) print() print("[INFO] Token cache built.") print(f"[INFO] Rows seen: {rows_seen:,}") print(f"[INFO] Rows used: {rows_used:,}") print(f"[INFO] Mapped marker rows: {rows_with_mapped_markers:,}") print(f"[INFO] Explicit boundary rows: {rows_with_explicit_boundaries:,}") print(f"[INFO] Legacy boundary rows: {rows_with_legacy_boundaries:,}") print(f"[INFO] Train tokens: {train_tokens:,}") print(f"[INFO] Val tokens: {val_tokens:,}") print() class LocalUint32BlockStream(IterableDataset): def __init__( self, bin_path: str, block_size: int, seed: int = 42, shuffle_blocks: bool = False, max_tokens: int = 0, ): super().__init__() self.bin_path = str(bin_path) self.block_size = int(block_size) self.seed = int(seed) self.shuffle_blocks = bool(shuffle_blocks) self.max_tokens = int(max_tokens) self._epoch = 0 file_size = os.path.getsize(self.bin_path) if file_size % 4 != 0: raise ValueError(f"Token file size is not divisible by 4: {self.bin_path}") self.num_tokens_total = file_size // 4 if self.max_tokens > 0: self.num_tokens = min(self.num_tokens_total, self.max_tokens) else: self.num_tokens = self.num_tokens_total if self.num_tokens <= self.block_size + 1: raise ValueError( f"Not enough tokens in {self.bin_path}: " f"{self.num_tokens} <= block_size+1={self.block_size + 1}" ) self.num_blocks = self.num_tokens // (self.block_size + 1) if self.num_blocks <= 0: raise ValueError("No full blocks available.") def set_epoch(self, epoch: int): self._epoch = int(epoch) def __iter__(self) -> Iterator[Dict[str, torch.Tensor]]: mm = np.memmap( self.bin_path, dtype=np.uint32, mode="r", shape=(self.num_tokens_total,), ) wi = get_worker_info() if wi is None: worker_id = 0 num_workers = 1 else: worker_id = wi.id num_workers = wi.num_workers block_ids = list(range(self.num_blocks)) if self.shuffle_blocks: rng = random.Random(self.seed + 1_000_003 * self._epoch) rng.shuffle(block_ids) block_ids = block_ids[worker_id::num_workers] need = self.block_size + 1 for block_id in block_ids: start = block_id * need end = start + need if end > self.num_tokens: continue window = np.asarray(mm[start:end], dtype=np.uint32) src = torch.from_numpy(window[:-1].astype(np.int64, copy=False)) tgt = torch.from_numpy(window[1:].astype(np.int64, copy=False)) yield { "src": src, "tgt": tgt, "length": torch.tensor(self.block_size, dtype=torch.long), } def collate_lm_fixed(batch): src = torch.stack([item["src"] for item in batch], dim=0) tgt = torch.stack([item["tgt"] for item in batch], dim=0) padding_mask = torch.zeros( src.shape, dtype=torch.bool, ) return src, tgt, padding_mask class GPTConfig: def __init__( self, vocab_size: int, ctx_len: int = 512, n_layer: int = 4, n_head: int = 4, n_embd: int = 384, dropout: float = 0.0, attention_backend: str = "sage", ): if attention_backend not in ("sage", "torch", "flash2", "flash3"): raise ValueError("--attention-backend must be 'sage', 'torch', 'flash2' or 'flash3'") if n_embd % n_head != 0: raise ValueError("n_embd must be divisible by n_head") head_dim = n_embd // n_head if attention_backend == "sage" and head_dim not in (64, 96, 128): raise ValueError( f"SageAttention requires head_dim in [64, 96, 128], got {head_dim}. " "Examples: 384/4=96, 384/6=64, 256/4=64, 128/2=64." ) if attention_backend == "sage" and dropout != 0.0: raise ValueError("SageAttention strict mode requires --dropout 0.0") if attention_backend == "flash3" and dropout != 0.0: raise ValueError("FlashAttention3 backend requires --dropout 0.0") if attention_backend in ("flash2", "flash3") and head_dim % 8 != 0: raise ValueError( f"FlashAttention requires head_dim multiple of 8, got {head_dim}." ) self.vocab_size = int(vocab_size) self.ctx_len = int(ctx_len) self.n_layer = int(n_layer) self.n_head = int(n_head) self.n_embd = int(n_embd) self.dropout = float(dropout) self.attention_backend = str(attention_backend) class CausalSelfAttention(nn.Module): def __init__(self, cfg: GPTConfig): super().__init__() self.n_head = cfg.n_head self.head_dim = cfg.n_embd // cfg.n_head self.attention_backend = cfg.attention_backend self.dropout_p = float(cfg.dropout) self.qkv = nn.Linear(cfg.n_embd, 3 * cfg.n_embd, bias=False) self.proj = nn.Linear(cfg.n_embd, cfg.n_embd, bias=False) self.dropout = nn.Dropout(cfg.dropout) mask = torch.tril(torch.ones(cfg.ctx_len, cfg.ctx_len)) self.register_buffer( "mask", mask.view(1, 1, cfg.ctx_len, cfg.ctx_len), persistent=False, ) self.sageattn = None self.flash_kernel = None if self.attention_backend == "sage": try: from sageattention import sageattn except Exception as exc: raise RuntimeError( "SageAttention demandé, mais impossible d'importer : " "from sageattention import sageattn" ) from exc self.sageattn = sageattn if self.attention_backend == "flash2": try: self.flash_kernel = get_flash2_kernel() except Exception as exc: raise RuntimeError( "FlashAttention2 demandé, mais impossible de charger : " 'get_kernel("kernels-community/flash-attn2", version=1)' ) from exc if self.attention_backend == "flash3": try: self.flash_kernel = get_flash3_kernel() except Exception as exc: raise RuntimeError( "FlashAttention3 demandé, mais impossible de charger : " 'get_kernel("kernels-community/flash-attn3", version=1)' ) from exc def _torch_attention( self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, t: int, ) -> torch.Tensor: scores = (q @ k.transpose(-2, -1)) / math.sqrt(self.head_dim) scores = scores.masked_fill( self.mask[:, :, :t, :t] == 0, float("-inf"), ) att = F.softmax(scores.float(), dim=-1).to(q.dtype) att = self.dropout(att) y = att @ v return y def _sage_attention( self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, ) -> torch.Tensor: if self.sageattn is None: raise RuntimeError("SageAttention demandé mais sageattn est None") if not q.is_cuda: raise RuntimeError("SageAttention exige CUDA") q = q.contiguous() k = k.contiguous() v = v.contiguous() y = self.sageattn( q, k, v, tensor_layout="HND", is_causal=True, ) return y def _flash2_attention( self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, ) -> torch.Tensor: if self.flash_kernel is None: raise RuntimeError("FlashAttention2 demandé mais flash_kernel est None") if not q.is_cuda: raise RuntimeError("FlashAttention2 exige CUDA") q = q.transpose(1, 2).contiguous() k = k.transpose(1, 2).contiguous() v = v.transpose(1, 2).contiguous() dropout_p = self.dropout_p if self.training else 0.0 y = self.flash_kernel.flash_attn_func( q, k, v, dropout_p=dropout_p, causal=True, ) y = y.transpose(1, 2).contiguous() return y def _flash3_attention( self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, ) -> torch.Tensor: if self.flash_kernel is None: raise RuntimeError("FlashAttention3 demandé mais flash_kernel est None") if not q.is_cuda: raise RuntimeError("FlashAttention3 exige CUDA") q = q.transpose(1, 2).contiguous() k = k.transpose(1, 2).contiguous() v = v.transpose(1, 2).contiguous() y = self.flash_kernel.flash_attn_func( q, k, v, causal=True, ) y = y.transpose(1, 2).contiguous() return y def forward(self, x: torch.Tensor) -> torch.Tensor: b, t, c = x.shape qkv = self.qkv(x) q, k, v = qkv.chunk(3, dim=-1) q = q.view(b, t, self.n_head, self.head_dim).transpose(1, 2).contiguous() k = k.view(b, t, self.n_head, self.head_dim).transpose(1, 2).contiguous() v = v.view(b, t, self.n_head, self.head_dim).transpose(1, 2).contiguous() if self.attention_backend == "sage": y = self._sage_attention(q, k, v) elif self.attention_backend == "flash2": y = self._flash2_attention(q, k, v) elif self.attention_backend == "flash3": y = self._flash3_attention(q, k, v) else: y = self._torch_attention(q, k, v, t) y = y.transpose(1, 2).contiguous().view(b, t, c) y = self.proj(y) return y class MLP(nn.Module): def __init__(self, cfg: GPTConfig): super().__init__() self.fc = nn.Linear(cfg.n_embd, 4 * cfg.n_embd, bias=False) self.proj = nn.Linear(4 * cfg.n_embd, cfg.n_embd, bias=False) self.dropout = nn.Dropout(cfg.dropout) def forward(self, x: torch.Tensor) -> torch.Tensor: x = self.fc(x) x = F.gelu(x) x = self.proj(x) x = self.dropout(x) return x class Block(nn.Module): def __init__(self, cfg: GPTConfig): super().__init__() self.ln1 = nn.LayerNorm(cfg.n_embd) self.attn = CausalSelfAttention(cfg) self.ln2 = nn.LayerNorm(cfg.n_embd) self.mlp = MLP(cfg) def forward(self, x: torch.Tensor) -> torch.Tensor: x = x + self.attn(self.ln1(x)) x = x + self.mlp(self.ln2(x)) return x class TinyGPT(nn.Module): def __init__(self, cfg: GPTConfig): super().__init__() self.cfg = cfg self.tok_emb = nn.Embedding(cfg.vocab_size, cfg.n_embd) self.pos_emb = nn.Embedding(cfg.ctx_len, cfg.n_embd) self.drop = nn.Dropout(cfg.dropout) self.blocks = nn.ModuleList( [Block(cfg) for _ in range(cfg.n_layer)] ) self.ln_f = nn.LayerNorm(cfg.n_embd) self.head = nn.Linear(cfg.n_embd, cfg.vocab_size, bias=False) self.head.weight = self.tok_emb.weight self.apply(self._init_weights) def _init_weights(self, module): if isinstance(module, nn.Linear): nn.init.normal_( module.weight, mean=0.0, std=0.02, ) if isinstance(module, nn.Embedding): nn.init.normal_( module.weight, mean=0.0, std=0.02, ) def forward( self, idx: torch.Tensor, return_hidden: bool = False, ): b, t = idx.shape if t > self.cfg.ctx_len: raise ValueError(f"Input length {t} > ctx_len {self.cfg.ctx_len}") pos = torch.arange( 0, t, dtype=torch.long, device=idx.device, ).unsqueeze(0) x = self.tok_emb(idx) + self.pos_emb(pos) x = self.drop(x) for block in self.blocks: x = block(x) hidden = self.ln_f(x) logits = self.head(hidden) if return_hidden: return logits, hidden return logits def embed_mean_pool(self, idx: torch.Tensor) -> torch.Tensor: _, hidden = self.forward(idx, return_hidden=True) mask = idx.ne(PAD_ID).unsqueeze(-1).to(hidden.dtype) summed = (hidden * mask).sum(dim=1) denom = mask.sum(dim=1).clamp(min=1.0) emb = summed / denom emb = F.normalize(emb, p=2, dim=-1) return emb def param_count(model: nn.Module) -> int: return int(sum(p.numel() for p in model.parameters())) class RNETrainer: def __init__( self, model: TinyGPT, train_loader: DataLoader, val_loader: DataLoader, out_dir: str, max_steps: int, lr: float, weight_decay: float, save_every: int, log_every: int, val_every: int, val_batches: int, dtype: str, grad_clip: float, device: torch.device, compile_model: bool = False, ): self.model = model self.train_loader = train_loader self.val_loader = val_loader self.out_dir = Path(out_dir) self.max_steps = int(max_steps) self.lr = float(lr) self.weight_decay = float(weight_decay) self.save_every = int(save_every) self.log_every = int(log_every) self.val_every = int(val_every) self.val_batches = int(val_batches) self.dtype = dtype self.grad_clip = float(grad_clip) self.device = device if dtype == "float16": self.amp_dtype = torch.float16 elif dtype == "bfloat16": self.amp_dtype = torch.bfloat16 else: self.amp_dtype = torch.float32 self.use_amp = self.device.type == "cuda" and dtype in ("float16", "bfloat16") self.optimizer = torch.optim.AdamW( self.model.parameters(), lr=self.lr, betas=(0.9, 0.95), weight_decay=self.weight_decay, ) self.scaler = torch.amp.GradScaler( "cuda", enabled=self.use_amp, ) self.criterion = nn.CrossEntropyLoss() if compile_model: self.model = torch.compile(self.model) self.tokens_seen_total = 0 self.tokens_seen_since = 0 self.steps_since = 0 self.amp_overflow_count = 0 self.rate_t0 = time.perf_counter() def _set_lr(self, lr: float): for group in self.optimizer.param_groups: group["lr"] = lr def _get_lr(self, step: int) -> float: return self.lr def _reset_rate_window(self): self.rate_t0 = time.perf_counter() self.tokens_seen_since = 0 self.steps_since = 0 self.amp_overflow_count = 0 def _rate_info(self) -> Tuple[float, float]: now = time.perf_counter() dt = max(now - self.rate_t0, 1e-9) tok_s = self.tokens_seen_since / dt step_s = self.steps_since / dt return tok_s, step_s def _save(self, step: int): self.out_dir.mkdir(parents=True, exist_ok=True) raw_model = self.model._orig_mod if hasattr(self.model, "_orig_mod") else self.model payload = { "step": int(step), "model": raw_model.state_dict(), "optimizer": self.optimizer.state_dict(), "config": { "vocab_size": raw_model.cfg.vocab_size, "ctx_len": raw_model.cfg.ctx_len, "n_layer": raw_model.cfg.n_layer, "n_head": raw_model.cfg.n_head, "n_embd": raw_model.cfg.n_embd, "dropout": raw_model.cfg.dropout, "attention_backend": raw_model.cfg.attention_backend, "PAD_ID": PAD_ID, "boundary_mode": BOUNDARY_MODE, "default_bos_marker": DEFAULT_BOS_MARKER, "default_eos_marker": DEFAULT_EOS_MARKER, }, "tokens_seen_total": int(self.tokens_seen_total), } ckpt = self.out_dir / f"checkpoint_step_{step}.pt" latest = self.out_dir / "latest.pt" torch.save(payload, ckpt) torch.save(payload, latest) print(f"\n[SAVE] {ckpt}") def evaluate(self) -> float: self.model.eval() total_loss = 0.0 seen = 0 with torch.no_grad(): for batch in self.val_loader: src, tgt, padding_mask = batch src = src.to(self.device, non_blocking=True) tgt = tgt.to(self.device, non_blocking=True) with torch.autocast( device_type="cuda", dtype=self.amp_dtype, enabled=self.use_amp, ): logits = self.model(src) loss = self.criterion( logits.reshape(-1, logits.size(-1)).float(), tgt.reshape(-1), ) total_loss += float(loss.item()) seen += 1 if seen >= self.val_batches: break self.model.train() return total_loss / max(1, seen) def train(self): self.model.train() step = 0 running_loss = 0.0 running_count = 0 last_val_loss = None train_iter = iter(self.train_loader) self._reset_rate_window() pbar = tqdm( total=self.max_steps, desc="Training/LM-SAGE9-FLASH-KERNELS", dynamic_ncols=True, ) while step < self.max_steps: try: src, tgt, padding_mask = next(train_iter) except StopIteration: train_iter = iter(self.train_loader) src, tgt, padding_mask = next(train_iter) src = src.to(self.device, non_blocking=True) tgt = tgt.to(self.device, non_blocking=True) batch_tokens = int(src.numel()) lr = self._get_lr(step + 1) self._set_lr(lr) self.optimizer.zero_grad(set_to_none=True) with torch.autocast( device_type="cuda", dtype=self.amp_dtype, enabled=self.use_amp, ): logits = self.model(src) loss = self.criterion( logits.reshape(-1, logits.size(-1)).float(), tgt.reshape(-1), ) if not torch.isfinite(loss): raise RuntimeError(f"Non-finite loss detected: {loss.item()}") self.scaler.scale(loss).backward() self.scaler.unscale_(self.optimizer) if self.grad_clip > 0: nn.utils.clip_grad_norm_( self.model.parameters(), max_norm=self.grad_clip, ) scale_before = float(self.scaler.get_scale()) self.scaler.step(self.optimizer) self.scaler.update() scale_after = float(self.scaler.get_scale()) if self.use_amp and scale_after < scale_before: self.amp_overflow_count += 1 self.optimizer.zero_grad(set_to_none=True) if self.amp_overflow_count <= 3: print( f"[amp] overflow detected: scale {scale_before:.1f} -> {scale_after:.1f}; skipping update" ) continue step += 1 pbar.update(1) self.tokens_seen_total += batch_tokens self.tokens_seen_since += batch_tokens self.steps_since += 1 running_loss += float(loss.item()) running_count += 1 if step % self.val_every == 0: last_val_loss = self.evaluate() if step % self.log_every == 0: avg_loss = running_loss / max(1, running_count) ppl = math.exp(min(avg_loss, 20.0)) tok_s, step_s = self._rate_info() postfix = { "loss": f"{avg_loss:.4f}", "ppl": f"{ppl:.2f}", "lr": f"{lr:.2e}", "seen": format_tokens(self.tokens_seen_total), "tok_s": f"{tok_s:,.0f}", "step_s": f"{step_s:.2f}", } if last_val_loss is not None: postfix["val_loss"] = f"{last_val_loss:.4f}" postfix["val_ppl"] = f"{math.exp(min(last_val_loss, 20.0)):.2f}" if self.amp_overflow_count > 0: postfix["amp_of"] = str(self.amp_overflow_count) pbar.set_postfix(**postfix) running_loss = 0.0 running_count = 0 self._reset_rate_window() if step % self.save_every == 0: self._save(step) pbar.close() self._save(step) print() print("[DONE] Training finished.") print(f"[DONE] Steps: {step:,}") print(f"[DONE] Tokens seen: {self.tokens_seen_total:,}") print(f"[DONE] Tokens compact: {format_tokens(self.tokens_seen_total)}") if last_val_loss is not None: print(f"[DONE] Last val loss: {last_val_loss:.6f}") print(f"[DONE] Last val ppl: {math.exp(min(last_val_loss, 20.0)):.6f}") def parse_args(): parser = argparse.ArgumentParser( description="LM trainer with shuffled pretokenization cache, generic marker->special-token mapping, BOS/EOS boundaries, SageAttention, torch attention, FlashAttention2 and FlashAttention3 via HF kernels." ) parser.add_argument("--src", required=True) parser.add_argument("--tokenizer", required=True) parser.add_argument("--out-dir", default="LM_SAGE9") parser.add_argument("--cache-dir", default="lm_token_cache_sage9_marker_special") parser.add_argument("--activity-column", default="activites") parser.add_argument("--row-batch-size", type=int, default=100_000) parser.add_argument("--rebuild-cache", action="store_true") parser.add_argument("--shuffle-before-tokenize", action="store_true") parser.add_argument("--no-shuffle-before-tokenize", action="store_true") parser.add_argument("--shuffle-buffer-size", type=int, default=500_000) parser.add_argument("--ctx-len", type=int, default=512) parser.add_argument("--batch-size", type=int, default=4) parser.add_argument("--num-workers", type=int, default=0) parser.add_argument("--shuffle-blocks", action="store_true") parser.add_argument("--max-train-tokens", type=int, default=0) parser.add_argument("--max-val-tokens", type=int, default=0) parser.add_argument("--val-ratio", type=float, default=0.01) parser.add_argument("--val-every", type=int, default=2000) parser.add_argument("--val-batches", type=int, default=10) parser.add_argument("--n-layer", type=int, default=4) parser.add_argument("--n-head", type=int, default=4) parser.add_argument("--n-embd", type=int, default=384) parser.add_argument("--dropout", type=float, default=0.0) parser.add_argument( "--attention-backend", default="sage", choices=["sage", "torch", "flash2", "flash3"], ) parser.add_argument("--lr", type=float, default=3e-4) parser.add_argument("--weight-decay", type=float, default=0.1) parser.add_argument("--max-steps", type=int, default=50_000) parser.add_argument("--save-every", type=int, default=10_000) parser.add_argument("--log-every", type=int, default=20) parser.add_argument("--grad-clip", type=float, default=1.0) parser.add_argument("--dtype", default="bfloat16", choices=["float32", "float16", "bfloat16"]) parser.add_argument("--device", default="cuda") parser.add_argument("--seed", type=int, default=42) parser.add_argument("--lowercase", action="store_true") parser.add_argument( "--special-marker", action="append", default=[], help='Map a dataset marker to a tokenizer special token. Example: --special-marker "<|BOC|>=[BOC]". Can be repeated.', ) parser.add_argument( "--no-special-boundaries", action="store_true", help="Disable BOS/EOS insertion and marker replacement during pretokenization.", ) parser.add_argument( "--no-append-sep", action="store_true", help="Legacy alias: disables BOS/EOS insertion too.", ) parser.add_argument("--compile", action="store_true") return parser.parse_args() def main(): args = parse_args() random.seed(args.seed) np.random.seed(args.seed) torch.manual_seed(args.seed) if args.device == "cuda" and not torch.cuda.is_available(): print("[WARN] CUDA unavailable, using CPU.") args.device = "cpu" if args.attention_backend in ("sage", "flash2", "flash3") and args.device != "cuda": raise RuntimeError(f"--attention-backend {args.attention_backend} requires --device cuda") if args.no_shuffle_before_tokenize: shuffle_before_tokenize = False else: shuffle_before_tokenize = True if args.shuffle_before_tokenize: shuffle_before_tokenize = True append_special_tokens = True if args.no_special_boundaries: append_special_tokens = False if args.no_append_sep: append_special_tokens = False token_cache = RNETokenCache( src=args.src, tokenizer_path=args.tokenizer, cache_dir=args.cache_dir, activity_column=args.activity_column, row_batch_size=args.row_batch_size, val_ratio=args.val_ratio, seed=args.seed, lowercase=args.lowercase, append_special_tokens=append_special_tokens, rebuild_cache=args.rebuild_cache, shuffle_before_tokenize=shuffle_before_tokenize, shuffle_buffer_size=args.shuffle_buffer_size, special_marker_specs=args.special_marker, ) token_cache.build_if_needed() train_ds = LocalUint32BlockStream( bin_path=str(token_cache.train_bin), block_size=args.ctx_len, seed=args.seed, shuffle_blocks=args.shuffle_blocks, max_tokens=args.max_train_tokens, ) val_ds = LocalUint32BlockStream( bin_path=str(token_cache.val_bin), block_size=args.ctx_len, seed=args.seed + 10_000_000, shuffle_blocks=False, max_tokens=args.max_val_tokens, ) train_loader = DataLoader( train_ds, batch_size=args.batch_size, num_workers=args.num_workers, collate_fn=collate_lm_fixed, drop_last=True, pin_memory=(args.device == "cuda"), persistent_workers=(args.num_workers > 0), ) val_loader = DataLoader( val_ds, batch_size=args.batch_size, num_workers=max(0, args.num_workers // 2), collate_fn=collate_lm_fixed, drop_last=True, pin_memory=(args.device == "cuda"), persistent_workers=(args.num_workers > 1), ) cfg = GPTConfig( vocab_size=token_cache.vocab_size, ctx_len=args.ctx_len, n_layer=args.n_layer, n_head=args.n_head, n_embd=args.n_embd, dropout=args.dropout, attention_backend=args.attention_backend, ) device = torch.device(args.device) model = TinyGPT(cfg).to(device) params = param_count(model) target_tokens = args.max_steps * args.batch_size * args.ctx_len train_epoch_steps = max(1, train_ds.num_blocks // max(1, args.batch_size)) approx_epochs = args.max_steps / train_epoch_steps print("[INFO] LM SAGE9 GENERIC SPECIAL MARKERS + FLASH KERNELS") print(f"[INFO] Source: {args.src}") print(f"[INFO] Activity column: {args.activity_column}") print(f"[INFO] Tokenizer: {token_cache.tokenizer_path}") print(f"[INFO] Cache dir: {args.cache_dir}") print(f"[INFO] Vocab size: {token_cache.vocab_size:,}") print(f"[INFO] Append special tokens: {append_special_tokens}") print(f"[INFO] BOS id: {token_cache.bos_id}") print(f"[INFO] EOS id: {token_cache.eos_id}") print(f"[INFO] SEP id: {token_cache.sep_id}") print(f"[INFO] Boundary mode: {BOUNDARY_MODE}") print(f"[INFO] Marker token map: {token_cache.marker_token_map}") print(f"[INFO] Marker id map: {token_cache.marker_id_map}") print(f"[INFO] Boundary rule: explicit <|BOS|> + <|EOS|> => no auto BOS/EOS") print(f"[INFO] Legacy rule: otherwise BOS + text + EOS") print(f"[INFO] Shuffle before tok: {shuffle_before_tokenize}") print(f"[INFO] Shuffle buffer size: {args.shuffle_buffer_size:,}") print(f"[INFO] Ctx len: {args.ctx_len}") print(f"[INFO] Batch size: {args.batch_size}") print(f"[INFO] Num workers: {args.num_workers}") print(f"[INFO] Shuffle blocks: {args.shuffle_blocks}") print(f"[INFO] Tokens / step: {args.batch_size * args.ctx_len:,}") print(f"[INFO] Train tokens file: {train_ds.num_tokens:,}") print(f"[INFO] Val tokens file: {val_ds.num_tokens:,}") print(f"[INFO] Train blocks: {train_ds.num_blocks:,}") print(f"[INFO] Steps / epoch: {train_epoch_steps:,}") print(f"[INFO] Approx epochs: {approx_epochs:.2f}") print(f"[INFO] Target tokens seen: {target_tokens:,}") print(f"[INFO] Target compact: {format_tokens(target_tokens)}") print(f"[INFO] Val ratio: {args.val_ratio}") print(f"[INFO] Val every: {args.val_every}") print(f"[INFO] Val batches: {args.val_batches}") print(f"[INFO] Params: {params:,}") print(f"[INFO] Device: {device}") print(f"[INFO] Dtype: {args.dtype}") print(f"[INFO] Attention backend: {args.attention_backend}") print(f"[INFO] Head dim: {args.n_embd // args.n_head}") print(f"[INFO] LR fixed: {args.lr}") print(f"[INFO] Output dir: {args.out_dir}") print() trainer = RNETrainer( model=model, train_loader=train_loader, val_loader=val_loader, out_dir=args.out_dir, max_steps=args.max_steps, lr=args.lr, weight_decay=args.weight_decay, save_every=args.save_every, log_every=args.log_every, val_every=args.val_every, val_batches=args.val_batches, dtype=args.dtype, grad_clip=args.grad_clip, device=device, compile_model=args.compile, ) trainer.train() if __name__ == "__main__": main()