train_llm.py

#!/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], <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 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()