#!/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 IGNORE_INDEX = -100 DEFAULT_BOS_MARKER = "<|BOS|>" DEFAULT_EOS_MARKER = "<|EOS|>" DEFAULT_BOS_TOKEN = "[BOS]" DEFAULT_EOS_TOKEN = "[EOS]" BOUNDARY_MODE = "marker_aware_generic_special_markers_v11_window_mode_rope" _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, window_mode: str = "stream", ctx_len: int = 512, ): 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") if window_mode not in ("stream", "row"): raise ValueError("--window-mode must be 'stream' or 'row'") if ctx_len <= 0: raise ValueError("--ctx-len 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.window_mode = str(window_mode) self.ctx_len = int(ctx_len) self.need = self.ctx_len + 1 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.train_windows_bin = self.cache_dir / "train_windows.uint64.bin" self.val_windows_bin = self.cache_dir / "val_windows.uint64.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 self.window_mode == "row": if not self.train_windows_bin.exists(): return False if not self.val_windows_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, "window_mode": self.window_mode, "ctx_len": self.ctx_len, "need": self.need, } 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, train_windows: int, val_windows: int, rows_dropped_window: int, rows_dropped_too_short: int, rows_dropped_too_long: 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, "window_mode": self.window_mode, "ctx_len": self.ctx_len, "need": self.need, "train_tokens": int(train_tokens), "val_tokens": int(val_tokens), "train_windows": int(train_windows), "val_windows": int(val_windows), "rows_seen": int(rows_seen), "rows_used": int(rows_used), "rows_dropped_window": int(rows_dropped_window), "rows_dropped_too_short": int(rows_dropped_too_short), "rows_dropped_too_long": int(rows_dropped_too_long), "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), "token_dtype": "uint32", "window_dtype": "uint64_pair_start_length", "row_mode_rule": "row mode keeps rows with 2 <= token_count <= ctx_len+1, pads shorter rows in dataloader, drops rows longer than ctx_len+1", } 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], token_file_obj, window_file_obj, desc: str, ) -> Tuple[int, int, int, int, int, int, int, int]: written_tokens = 0 used_texts = 0 mapped_marker_rows = 0 explicit_boundary_rows = 0 legacy_boundary_rows = 0 windows_written = 0 row_drop = 0 row_drop_short = 0 row_drop_long = 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) token_count = len(ids) if self.window_mode == "row": if token_count < 2: row_drop += 1 row_drop_short += 1 pbar.update(1) continue if token_count > self.need: row_drop += 1 row_drop_long += 1 pbar.update(1) continue start = written_tokens arr = np.asarray(ids, dtype=np.uint32) arr.tofile(token_file_obj) if window_file_obj is None: raise RuntimeError("window_file_obj is required in row mode") win = np.asarray([start, token_count], dtype=np.uint64) win.tofile(window_file_obj) written_tokens += int(arr.size) windows_written += 1 used_texts += 1 else: if token_count >= 2: arr = np.asarray(ids, dtype=np.uint32) arr.tofile(token_file_obj) written_tokens += int(arr.size) used_texts += 1 if token_count >= 2 and not (self.window_mode == "row" and token_count > self.need): 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) or (row_drop > 0 and row_drop % 10_000 == 0): 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:,}", } if self.window_mode == "row": postfix["windows"] = f"{windows_written:,}" postfix["row_drop"] = f"{row_drop:,}" postfix["too_long"] = f"{row_drop_long:,}" pbar.set_postfix(**postfix) pbar.close() return ( written_tokens, mapped_marker_rows, explicit_boundary_rows, legacy_boundary_rows, windows_written, row_drop, row_drop_short, row_drop_long, ) def _flush_text_buffer( self, buffer: List[str], token_file_obj, window_file_obj, rng: random.Random, name: str, ) -> Tuple[int, int, int, int, int, int, int, int]: if not buffer: return 0, 0, 0, 0, 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, windows_written, row_drop, row_drop_short, row_drop_long, ) = self._tokenize_to_file( texts=buffer, token_file_obj=token_file_obj, window_file_obj=window_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:,}") if self.window_mode == "row": print(f"[FLUSH] {name} windows written: {windows_written:,}") print(f"[FLUSH] {name} row_drop: {row_drop:,}") print(f"[FLUSH] {name} row_drop_short: {row_drop_short:,}") print(f"[FLUSH] {name} row_drop_long: {row_drop_long:,}") print() buffer.clear() return ( written_tokens, mapped_marker_rows, explicit_boundary_rows, legacy_boundary_rows, windows_written, row_drop, row_drop_short, row_drop_long, ) 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] Train windows: {meta.get('train_windows', 0):,}") print(f"[INFO] Val windows: {meta.get('val_windows', 0):,}") print(f"[INFO] Rows seen: {meta.get('rows_seen', 0):,}") print(f"[INFO] Rows used: {meta.get('rows_used', 0):,}") print(f"[INFO] Rows dropped/window: {meta.get('rows_dropped_window', 0):,}") print(f"[INFO] Rows dropped too short: {meta.get('rows_dropped_too_short', 0):,}") print(f"[INFO] Rows dropped too long: {meta.get('rows_dropped_too_long', 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] Window mode: {meta.get('window_mode')}") print(f"[INFO] Ctx len in cache: {meta.get('ctx_len')}") 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] Window mode: {self.window_mode}") print(f"[INFO] Ctx len: {self.ctx_len}") print(f"[INFO] Need tokens/window: {self.need}") 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") if self.window_mode == "row": print(f"[INFO] Row window rule: keep rows with 2 <= tokens <= ctx_len+1; pad shorter rows in loader; drop longer rows") else: print(f"[INFO] Stream window rule: old behavior, continuous token stream split into ctx_len+1 blocks") 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") train_windows_tmp = self.train_windows_bin.with_suffix(".tmp") val_windows_tmp = self.val_windows_bin.with_suffix(".tmp") for p in [train_tmp, val_tmp, train_windows_tmp, val_windows_tmp]: if p.exists(): p.unlink() train_tokens = 0 val_tokens = 0 train_windows = 0 val_windows = 0 rows_seen = 0 rows_used = 0 rows_with_mapped_markers = 0 rows_with_explicit_boundaries = 0 rows_with_legacy_boundaries = 0 rows_dropped_window = 0 rows_dropped_too_short = 0 rows_dropped_too_long = 0 train_text_buffer: List[str] = [] val_text_buffer: List[str] = [] if self.window_mode == "row": train_windows_cm = train_windows_tmp.open("wb") val_windows_cm = val_windows_tmp.open("wb") else: train_windows_cm = None val_windows_cm = None try: 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, windows, row_drop, row_drop_short, row_drop_long, ) = self._flush_text_buffer( buffer=train_text_buffer, token_file_obj=f_train, window_file_obj=train_windows_cm, rng=self.shuffle_rng_train, name="train buffer", ) train_tokens += written train_windows += windows rows_with_mapped_markers += marker_rows rows_with_explicit_boundaries += explicit_rows rows_with_legacy_boundaries += legacy_rows rows_dropped_window += row_drop rows_dropped_too_short += row_drop_short rows_dropped_too_long += row_drop_long if len(val_text_buffer) >= max(1_000, self.shuffle_buffer_size // 10): ( written, marker_rows, explicit_rows, legacy_rows, windows, row_drop, row_drop_short, row_drop_long, ) = self._flush_text_buffer( buffer=val_text_buffer, token_file_obj=f_val, window_file_obj=val_windows_cm, rng=self.shuffle_rng_val, name="val buffer", ) val_tokens += written val_windows += windows rows_with_mapped_markers += marker_rows rows_with_explicit_boundaries += explicit_rows rows_with_legacy_boundaries += legacy_rows rows_dropped_window += row_drop rows_dropped_too_short += row_drop_short rows_dropped_too_long += row_drop_long pbar.update(1) if rows_used % 10_000 == 0: 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:,}", } if self.window_mode == "row": postfix["tr_win"] = f"{train_windows:,}" postfix["va_win"] = f"{val_windows:,}" postfix["row_drop"] = f"{rows_dropped_window:,}" pbar.set_postfix(**postfix) ( written, marker_rows, explicit_rows, legacy_rows, windows, row_drop, row_drop_short, row_drop_long, ) = self._flush_text_buffer( buffer=train_text_buffer, token_file_obj=f_train, window_file_obj=train_windows_cm, rng=self.shuffle_rng_train, name="final train buffer", ) train_tokens += written train_windows += windows rows_with_mapped_markers += marker_rows rows_with_explicit_boundaries += explicit_rows rows_with_legacy_boundaries += legacy_rows rows_dropped_window += row_drop rows_dropped_too_short += row_drop_short rows_dropped_too_long += row_drop_long ( written, marker_rows, explicit_rows, legacy_rows, windows, row_drop, row_drop_short, row_drop_long, ) = self._flush_text_buffer( buffer=val_text_buffer, token_file_obj=f_val, window_file_obj=val_windows_cm, rng=self.shuffle_rng_val, name="final val buffer", ) val_tokens += written val_windows += windows rows_with_mapped_markers += marker_rows rows_with_explicit_boundaries += explicit_rows rows_with_legacy_boundaries += legacy_rows rows_dropped_window += row_drop rows_dropped_too_short += row_drop_short rows_dropped_too_long += row_drop_long pbar.close() finally: if train_windows_cm is not None: train_windows_cm.close() if val_windows_cm is not None: val_windows_cm.close() train_tmp.replace(self.train_bin) val_tmp.replace(self.val_bin) if self.window_mode == "row": train_windows_tmp.replace(self.train_windows_bin) val_windows_tmp.replace(self.val_windows_bin) else: if train_windows_tmp.exists(): train_windows_tmp.unlink() if val_windows_tmp.exists(): val_windows_tmp.unlink() 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, train_windows=train_windows, val_windows=val_windows, rows_dropped_window=rows_dropped_window, rows_dropped_too_short=rows_dropped_too_short, rows_dropped_too_long=rows_dropped_too_long, ) print() print("[INFO] Token cache built.") print(f"[INFO] Rows seen: {rows_seen:,}") print(f"[INFO] Rows used: {rows_used:,}") print(f"[INFO] Rows dropped/window: {rows_dropped_window:,}") print(f"[INFO] Rows dropped too short: {rows_dropped_too_short:,}") print(f"[INFO] Rows dropped too long: {rows_dropped_too_long:,}") 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(f"[INFO] Train windows: {train_windows:,}") print(f"[INFO] Val windows: {val_windows:,}") print() class LocalUint32BlockStream(IterableDataset): def __init__( self, bin_path: str, block_size: int, seed: int = 42, shuffle_blocks: bool = False, max_tokens: int = 0, window_mode: str = "stream", windows_path: Optional[str] = None, label_only_loss: bool = False, loss_delimiter_ids: Optional[List[int]] = None, bos_id: Optional[int] = None, eos_id: Optional[int] = None, ): super().__init__() if window_mode not in ("stream", "row"): raise ValueError("window_mode must be 'stream' or 'row'") 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.window_mode = str(window_mode) self.windows_path = str(windows_path) if windows_path is not None else None self.label_only_loss = bool(label_only_loss) self.loss_delimiter_ids = [int(x) for x in (loss_delimiter_ids or [])] self.bos_id = int(bos_id) if bos_id is not None else None self.eos_id = int(eos_id) if eos_id is not None else None 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.window_mode == "row": if self.windows_path is None: raise ValueError("windows_path is required when window_mode='row'") window_file_size = os.path.getsize(self.windows_path) if window_file_size % 16 != 0: raise ValueError(f"Window file size is not divisible by 16: {self.windows_path}") self.num_windows_total = window_file_size // 16 if self.num_windows_total <= 0: raise ValueError(f"No row windows available in {self.windows_path}") self.num_blocks = self._count_valid_row_windows_for_budget() if self.num_blocks <= 0: raise ValueError("No valid row windows available for current max_tokens budget.") else: 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_windows_total = 0 self.num_blocks = self.num_tokens // (self.block_size + 1) if self.num_blocks <= 0: raise ValueError("No full blocks available.") def _count_valid_row_windows_for_budget(self) -> int: if self.max_tokens <= 0: return int(self.num_windows_total) windows = np.memmap( self.windows_path, dtype=np.uint64, mode="r", shape=(self.num_windows_total, 2), ) count = 0 for i in range(self.num_windows_total): start = int(windows[i, 0]) length = int(windows[i, 1]) if start + length <= self.num_tokens: count += 1 return count def set_epoch(self, epoch: int): self._epoch = int(epoch) def _make_block_ids(self, total: int) -> List[int]: block_ids = list(range(total)) if self.shuffle_blocks: rng = random.Random(self.seed + 1_000_003 * self._epoch) rng.shuffle(block_ids) return block_ids @staticmethod def _find_subsequence(seq: List[int], sub: List[int]) -> int: if not sub: return -1 n = len(seq) m = len(sub) if m > n: return -1 for i in range(0, n - m + 1): if seq[i:i + m] == sub: return i return -1 def _apply_label_only_loss_mask( self, raw_tokens: np.ndarray, tgt_arr: np.ndarray, ) -> np.ndarray: if not self.label_only_loss: return tgt_arr masked = np.full(tgt_arr.shape, IGNORE_INDEX, dtype=np.int64) raw = [int(x) for x in raw_tokens.tolist()] eos_id = self.eos_id bos_id = self.bos_id usable = min(len(tgt_arr), max(0, len(raw) - 1)) # Compatible stream : # une fenêtre peut contenir plusieurs lignes : # [BOS] texte === Label [EOS] [BOS] texte === Label [EOS] ... # # On applique donc le masque par segment BOS/EOS. # Pour chaque segment, on garde seulement : # - les tokens cible après le delimiter # - EOS # - BOS, pour apprendre proprement la reprise de ligne en stream segment_start = 0 while segment_start < len(raw): if bos_id is not None: next_bos = -1 for i in range(segment_start, len(raw)): if int(raw[i]) == bos_id: next_bos = i break if next_bos == -1: seg_start = segment_start else: seg_start = next_bos else: seg_start = segment_start if eos_id is not None: seg_end = len(raw) for i in range(seg_start + 1, len(raw)): if int(raw[i]) == eos_id: seg_end = i + 1 break else: seg_end = len(raw) if seg_start >= seg_end: break segment = raw[seg_start:seg_end] local_delim_pos = self._find_subsequence(segment, self.loss_delimiter_ids) if local_delim_pos == -1: keep_start = None else: keep_start = seg_start + local_delim_pos + len(self.loss_delimiter_ids) for j in range(usable): target_raw_index = j + 1 if target_raw_index < seg_start or target_raw_index >= seg_end: continue target_id = int(raw[target_raw_index]) keep = False if keep_start is not None and target_raw_index >= keep_start: keep = True if eos_id is not None and target_id == eos_id: keep = True if bos_id is not None and target_id == bos_id: keep = True if keep: masked[j] = int(tgt_arr[j]) if seg_end <= segment_start: break segment_start = seg_end return masked def _iter_stream(self, worker_id: int, num_workers: int) -> Iterator[Dict[str, torch.Tensor]]: mm = np.memmap( self.bin_path, dtype=np.uint32, mode="r", shape=(self.num_tokens_total,), ) block_ids = self._make_block_ids(self.num_blocks) 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_arr = window[:-1].astype(np.int64, copy=False) tgt_arr = window[1:].astype(np.int64, copy=False) if self.label_only_loss: tgt_arr = self._apply_label_only_loss_mask( raw_tokens=window, tgt_arr=tgt_arr, ) src = torch.from_numpy(src_arr) tgt = torch.from_numpy(tgt_arr) padding_mask = torch.zeros((self.block_size,), dtype=torch.bool) yield { "src": src, "tgt": tgt, "padding_mask": padding_mask, "length": torch.tensor(self.block_size, dtype=torch.long), } def _iter_row(self, worker_id: int, num_workers: int) -> Iterator[Dict[str, torch.Tensor]]: mm = np.memmap( self.bin_path, dtype=np.uint32, mode="r", shape=(self.num_tokens_total,), ) windows = np.memmap( self.windows_path, dtype=np.uint64, mode="r", shape=(self.num_windows_total, 2), ) block_ids = self._make_block_ids(self.num_windows_total) block_ids = block_ids[worker_id::num_workers] max_len = self.block_size + 1 for block_id in block_ids: start = int(windows[block_id, 0]) length = int(windows[block_id, 1]) if length < 2 or length > max_len: continue end = start + length if end > self.num_tokens: continue raw = np.asarray(mm[start:end], dtype=np.uint32) real_len = int(raw.size) - 1 if real_len <= 0: continue src_arr = np.full((self.block_size,), PAD_ID, dtype=np.int64) tgt_arr = np.full((self.block_size,), IGNORE_INDEX, dtype=np.int64) pad_arr = np.ones((self.block_size,), dtype=np.bool_) src_arr[:real_len] = raw[:-1].astype(np.int64, copy=False) tgt_arr[:real_len] = raw[1:].astype(np.int64, copy=False) pad_arr[:real_len] = False if self.label_only_loss: tgt_arr = self._apply_label_only_loss_mask( raw_tokens=raw, tgt_arr=tgt_arr, ) src = torch.from_numpy(src_arr) tgt = torch.from_numpy(tgt_arr) padding_mask = torch.from_numpy(pad_arr) yield { "src": src, "tgt": tgt, "padding_mask": padding_mask, "length": torch.tensor(real_len, dtype=torch.long), } def __iter__(self) -> Iterator[Dict[str, torch.Tensor]]: wi = get_worker_info() if wi is None: worker_id = 0 num_workers = 1 else: worker_id = wi.id num_workers = wi.num_workers if self.window_mode == "row": yield from self._iter_row(worker_id=worker_id, num_workers=num_workers) else: yield from self._iter_stream(worker_id=worker_id, num_workers=num_workers) 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) if "padding_mask" in batch[0]: padding_mask = torch.stack([item["padding_mask"] for item in batch], dim=0) else: 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", rope_base: float = 10000.0, ): 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 head_dim % 2 != 0: raise ValueError( f"RoPE requires even head_dim, got {head_dim}. " "Use n_embd/n_head producing an even head dimension." ) if rope_base <= 0: raise ValueError("--rope-base must be > 0") 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) self.rope_base = float(rope_base) self.positional_encoding = "rope" def rotate_half(x: torch.Tensor) -> torch.Tensor: x_even = x[..., ::2] x_odd = x[..., 1::2] x_rot = torch.stack((-x_odd, x_even), dim=-1) return x_rot.flatten(start_dim=-2) def apply_rope(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor: return (x * cos) + (rotate_half(x) * sin) class RotaryEmbedding(nn.Module): def __init__( self, dim: int, max_position_embeddings: int, base: float = 10000.0, ): super().__init__() if dim % 2 != 0: raise ValueError(f"RoPE dim must be even, got {dim}") self.dim = int(dim) self.max_position_embeddings = int(max_position_embeddings) self.base = float(base) inv_freq = 1.0 / ( self.base ** ( torch.arange( 0, self.dim, 2, dtype=torch.float32, ) / self.dim ) ) self.register_buffer( "inv_freq", inv_freq, persistent=False, ) self._cos_cached = None self._sin_cached = None self._seq_len_cached = 0 self._device_cached = None self._dtype_cached = None def _build_cache( self, seq_len: int, device: torch.device, dtype: torch.dtype, ): t = torch.arange( seq_len, device=device, dtype=torch.float32, ) freqs = torch.einsum( "i,j->ij", t, self.inv_freq.to(device=device, dtype=torch.float32), ) emb = torch.repeat_interleave(freqs, repeats=2, dim=-1) cos = emb.cos().to(dtype=dtype).view(1, 1, seq_len, self.dim) sin = emb.sin().to(dtype=dtype).view(1, 1, seq_len, self.dim) self._cos_cached = cos self._sin_cached = sin self._seq_len_cached = int(seq_len) self._device_cached = device self._dtype_cached = dtype def forward( self, seq_len: int, device: torch.device, dtype: torch.dtype, ) -> Tuple[torch.Tensor, torch.Tensor]: if ( self._cos_cached is None or self._sin_cached is None or self._seq_len_cached < seq_len or self._device_cached != device or self._dtype_cached != dtype ): self._build_cache( seq_len=seq_len, device=device, dtype=dtype, ) return ( self._cos_cached[:, :, :seq_len, :], self._sin_cached[:, :, :seq_len, :], ) 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) self.rope = RotaryEmbedding( dim=self.head_dim, max_position_embeddings=cfg.ctx_len, base=cfg.rope_base, ) 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() cos, sin = self.rope( seq_len=t, device=q.device, dtype=q.dtype, ) q = apply_rope(q, cos, sin) k = apply_rope(k, cos, sin) 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.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}") x = self.tok_emb(idx) 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(ignore_index=IGNORE_INDEX) 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, "positional_encoding": raw_model.cfg.positional_encoding, "rope_base": raw_model.cfg.rope_base, "PAD_ID": PAD_ID, "IGNORE_INDEX": IGNORE_INDEX, "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-SAGE11-WINDOWS-ROPE", 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(tgt.ne(IGNORE_INDEX).sum().item()) 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 pretokenization cache, generic marker->special-token mapping, BOS/EOS boundaries, row/stream window modes, RoPE positional encoding, 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_SAGE11_ROPE") parser.add_argument("--cache-dir", default="lm_token_cache_sage11_marker_special_windows_rope") 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( "--window-mode", default="stream", choices=["stream", "row"], help="stream = old continuous-token behavior. row = one parquet row cannot cross context; rows longer than ctx_len+1 are dropped; shorter rows are padded and ignored in loss.", ) 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( "--rope-base", type=float, default=10000.0, help="RoPE base theta. Default 10000.0.", ) 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( "--label-only-loss", action="store_true", help="Mask loss to -100 everywhere except the label segment after the delimiter and EOS. Designed for rows like: <|BOS|> text === Label <|EOS|>.", ) parser.add_argument( "--target-delimiter", default="===", help="Delimiter before the supervised target label. Default: ===", ) 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, window_mode=args.window_mode, ctx_len=args.ctx_len, ) token_cache.build_if_needed() train_windows_path = str(token_cache.train_windows_bin) if args.window_mode == "row" else None val_windows_path = str(token_cache.val_windows_bin) if args.window_mode == "row" else None loss_delimiter_ids: List[int] = [] if args.label_only_loss: loss_delimiter_ids = token_cache._encode_plain_chunk(args.target_delimiter) if not loss_delimiter_ids: raise RuntimeError(f"Impossible de tokenizer le delimiter: {repr(args.target_delimiter)}") print(f"[INFO] Label-only loss: enabled") print(f"[INFO] Target delimiter: {repr(args.target_delimiter)}") print(f"[INFO] Target delimiter ids: {loss_delimiter_ids}") print(f"[INFO] Loss rule: stream-safe per BOS/EOS segment; -100 before and including delimiter; keep label target + EOS/BOS") else: print(f"[INFO] Label-only loss: disabled") 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, window_mode=args.window_mode, windows_path=train_windows_path, label_only_loss=args.label_only_loss, loss_delimiter_ids=loss_delimiter_ids, bos_id=token_cache.bos_id, eos_id=token_cache.eos_id, ) 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, window_mode=args.window_mode, windows_path=val_windows_path, label_only_loss=args.label_only_loss, loss_delimiter_ids=loss_delimiter_ids, bos_id=token_cache.bos_id, eos_id=token_cache.eos_id, ) 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, rope_base=args.rope_base, ) 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 SAGE11 GENERIC SPECIAL MARKERS + FLASH KERNELS + WINDOW MODE + ROPE") 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] Window mode: {args.window_mode}") print(f"[INFO] Positional encoding: RoPE") print(f"[INFO] RoPE base: {args.rope_base}") 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") if args.window_mode == "row": print(f"[INFO] Row rule: no crossing rows; rows longer than ctx_len+1 dropped; shorter rows padded + loss ignored") else: print(f"[INFO] Stream rule: continuous token stream, old behavior") 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 max: {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/windows: {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 max: {target_tokens:,}") print(f"[INFO] Target compact max: {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()