| |
| 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], <bos>, <BOS>, <s>, [CLS] ou équivalent." |
| ) |
|
|
| if self.eos_id is None: |
| raise RuntimeError( |
| "EOS token introuvable. Le tokenizer doit contenir [EOS], <eos>, <EOS>, </s>, [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]", |
| "<bos>", |
| "<BOS>", |
| "<s>", |
| "[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]", |
| "<eos>", |
| "<EOS>", |
| "</s>", |
| "[SEP]", |
| "<sep>", |
| "<SEP>", |
| DEFAULT_EOS_MARKER, |
| ] |
| ) |
|
|
| def _find_sep_id(self) -> Optional[int]: |
| return self._find_token_id( |
| [ |
| "[SEP]", |
| "</s>", |
| "<eos>", |
| "<EOS>", |
| "[EOS]", |
| "<sep>", |
| "<SEP>", |
| 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() |
