train.py

from __future__ import annotations

import os
from pathlib import Path
import time

os.environ['KERAS_BACKEND']                  = 'jax'
os.environ['XLA_PYTHON_CLIENT_PREALLOCATE']  = 'true'
os.environ['XLA_PYTHON_CLIENT_MEM_FRACTION'] = '0.90'
os.environ['TF_CPP_MIN_LOG_LEVEL']           = '3'
os.environ['PYTHONWARNINGS']                 = 'ignore'
os.environ['TF_DATA_EXPERIMENTAL_SLACK']     = '1'

import numpy as np
import tensorflow as tf
import jax
import keras

from tokenizer import TokenizerWrapper
from veylon_model import create_llm

try:
    from config import (EPOCHS, CONTEXT, vocab_size, D_MODEL, numberoflayers,
                        numberofheads, d_Latent, ffn_mult, swa_window,
                        num_kv_heads, learning_rate, weight_decay, batch_size,
                        data_path)
except Exception:
    EPOCHS         = 20
    CONTEXT        = 2048
    vocab_size     = 32000
    D_MODEL        = 512
    numberoflayers = 8
    numberofheads  = 8
    d_Latent       = 128
    ffn_mult       = 3.5
    swa_window     = 1024
    num_kv_heads   = 2
    learning_rate  = 1e-4
    weight_decay   = 0.01
    batch_size     = 8
    data_path      = './training_data'

keras.mixed_precision.set_global_policy('mixed_bfloat16')


# ─────────────────────────────────────────────────────────────────────────────
# PIPELINE
# ─────────────────────────────────────────────────────────────────────────────

def precompute_windows(tokens: np.ndarray, seq_len: int, stride: int) -> np.ndarray:
    """
    Slice the token array into overlapping (seq_len+1) windows using
    sliding_window_view, then stride-subsample.

    sliding_window_view is preferred over as_strided: it validates bounds
    internally and raises a clear error instead of silently reading garbage
    memory past the array end.

    The [::stride] slice is a zero-copy view.  ascontiguousarray forces one
    real allocation — a dense (N, seq_len+1) int32 buffer — which TensorFlow
    can wrap without copying again.

    Memory: for 5.7M tokens, CONTEXT=2048, stride=1024 → ~45 MB. Trivial
    against the 42 GB host RAM on Colab/Kaggle TPU instances.
    """
    window_len = seq_len + 1
    view = np.lib.stride_tricks.sliding_window_view(tokens, window_len)[::stride]

    if len(view) == 0:
        raise ValueError(
            f"Token array ({len(tokens):,}) is too short for "
            f"seq_len={seq_len}, stride={stride}."
        )

    return np.ascontiguousarray(view, dtype=np.int32)


def _make_tf_data_options() -> tf.data.Options:
    """
    Bundle all tf.data graph-level knobs in one place.

    map_parallelization: lets the runtime fuse and parallelize map ops across
      the thread pool automatically — no manual num_parallel_calls needed for
      simple element-wise transforms.

    parallel_batch: batching itself becomes a parallel operation; each batch
      slot is filled concurrently instead of serially.

    experimental_slack: introduces a one-step slack between the prefetch
      stage and the training loop so the pipeline stays one batch ahead
      without blocking on the next step.  Redundant with the env var above
      but belt-and-suspenders costs nothing.
    """
    opts = tf.data.Options()
    opts.experimental_optimization.map_parallelization = True
    opts.experimental_optimization.parallel_batch      = True
    opts.experimental_slack                            = True
    return opts


def create_tf_dataset(
    windows: np.ndarray,
    batch_size: int,
    shuffle: bool = True,
    shuffle_seed: int = 42,
    cache_path: str = "",          # "" → in-memory cache; "/tmp/..." → disk
) -> tf.data.Dataset:
    """
    High-throughput tf.data pipeline.

    Ordering rationale (each decision is load-bearing):

    1. from_tensor_slices(windows)
       TensorFlow wraps the numpy buffer directly — no tf.constant() copy.
       Cardinality is exactly N (known), enabling optimal prefetch depth and
       shuffle buffer auto-sizing downstream.

    2. cache()  ← BEFORE shuffle and map
       Caches individual (seq_len+1,) sequence tensors.  On epoch 2+ the
       entire dataset is served from RAM, eliminating all numpy I/O.
       Caching before shuffle means the cache stores the canonical data once;
       shuffle order changes every epoch without invalidating the cache.
       Caching after batching would freeze batch composition forever, which
       breaks per-epoch shuffle semantics.

    3. shuffle()  ← AFTER cache, BEFORE repeat
       Operates on individual sequences (correct granularity).
       buffer=min(N, 20_000) bounds peak memory while covering the full
       dataset for corpora this size.

    4. repeat()  ← AFTER shuffle, BEFORE map/batch
       Placed here so the shuffle buffer can draw across epoch boundaries,
       preventing an artificial ordering reset at each epoch edge.

    5. map(split_xy)
       Slices each (seq_len+1,) window into x=(seq_len,) and y=(seq_len,).
       No @tf.function decorator needed — tf.data traces the lambda itself.
       num_parallel_calls=AUTOTUNE parallelises across the CPU thread pool.

    6. batch(drop_remainder=True)
       Static shape [batch_size, seq_len] — mandatory for XLA/TPU.
       No symbolic/dynamic shapes ever reach the device.

    7. prefetch(AUTOTUNE)
       Keeps the device-transfer queue full.  Always last so it buffers
       fully-formed batches, not individual sequences.

    8. with_options(...)
       Applies graph-level optimizations in one shot after the pipeline is
       fully defined.
    """
    n = len(windows)

    ds = tf.data.Dataset.from_tensor_slices(windows)           # (1) wrap

    ds = ds.cache(cache_path)                                  # (2) cache

    if shuffle:
        buf = min(n, 20_000)
        ds  = ds.shuffle(buffer_size=buf, seed=shuffle_seed,
                         reshuffle_each_iteration=True)        # (3) shuffle

    ds = ds.repeat()                                           # (4) repeat

    ds = ds.map(                                               # (5) split
        lambda w: (w[:-1], w[1:]),
        num_parallel_calls=tf.data.AUTOTUNE,
    )

    ds = ds.batch(batch_size, drop_remainder=True)             # (6) batch

    ds = ds.prefetch(tf.data.AUTOTUNE)                        # (7) prefetch

    ds = ds.with_options(_make_tf_data_options())              # (8) opts

    return ds


# ─────────────────────────────────────────────────────────────────────────────
# BENCHMARK
# ─────────────────────────────────────────────────────────────────────────────

def benchmark_dataset(
    ds: tf.data.Dataset,
    batch_size: int,
    seq_len: int,
    n_batches: int = 200,
    label: str = "dataset",
) -> dict:
    """
    Measure pipeline throughput independently of TPU compute.

    We do NOT call .numpy() inside the loop.  That forces a device→host
    transfer and a Python GIL acquisition on every batch, serialising what
    should be a fully-async prefetch chain.  Instead we consume the iterator
    and let TensorFlow materialise tensors lazily.  The end-to-end wall time
    is what the TPU will actually experience.

    Interpret results:
      pipeline_tok/s > 2× tpu_tok/s  →  pipeline is not the bottleneck
      pipeline_tok/s < 1.5× tpu_tok/s →  still starving the TPU
    """
    print(f"\n{'─'*52}")
    print(f"Benchmarking: {label}")
    print(f"{'─'*52}")

    # Warm up: populate cache + prefetch buffers before timing starts.
    for _ in ds.take(5):
        pass
    print("Warm-up (5 batches) complete.")

    start = time.perf_counter()
    for _ in ds.take(n_batches):
        pass
    elapsed = time.perf_counter() - start

    tokens_per_batch = batch_size * seq_len
    total_tokens     = n_batches * tokens_per_batch
    batches_sec      = n_batches / elapsed
    tokens_sec       = total_tokens / elapsed

    print(f"Batches measured : {n_batches}")
    print(f"Elapsed          : {elapsed:.2f} s")
    print(f"Throughput       : {batches_sec:.1f} batches/s")
    print(f"Throughput       : {tokens_sec:,.0f} tokens/s")

    tpu_target = 30_000
    headroom   = tokens_sec / tpu_target
    if headroom >= 2.0:
        verdict = f"✓  {headroom:.1f}× headroom over TPU — pipeline is NOT the bottleneck."
    elif headroom >= 1.2:
        verdict = f"⚠  {headroom:.1f}× headroom — marginal; increase batch_size or reduce stride."
    else:
        verdict = f"✗  {tokens_sec:,.0f} tok/s < TPU target {tpu_target:,} — still starving the TPU."

    print(verdict)
    print(f"{'─'*52}\n")

    return {"batches_sec": batches_sec, "tokens_sec": tokens_sec,
            "headroom_vs_30k": headroom}


# ─────────────────────────────────────────────────────────────────────────────
# THROUGHPUT CALLBACK
# ─────────────────────────────────────────────────────────────────────────────

class ThroughputCallback(keras.callbacks.Callback):
    """
    Logs throughput every LOG_INTERVAL steps and at epoch end.

    Why coarse logging matters on TPU:
    Keras callbacks run on the Python host.  JAX dispatches training steps
    asynchronously — Python returns before XLA finishes the kernel.  A
    time.perf_counter() call on every step forces a host sync, equivalent
    to inserting jax.block_until_ready() every step and destroying async
    pipelining.  At LOG_INTERVAL=50 only ~2% of steps incur this cost.
    """
    LOG_INTERVAL = 50

    def on_epoch_begin(self, epoch, logs=None):
        self._epoch_start    = time.perf_counter()
        self._interval_start = self._epoch_start
        self._step_count     = 0

    def on_train_batch_end(self, batch, logs=None):
        self._step_count += 1
        if batch > 0 and batch % self.LOG_INTERVAL == 0:
            now     = time.perf_counter()
            elapsed = now - self._interval_start
            tps     = (self.LOG_INTERVAL * batch_size * CONTEXT) / elapsed
            loss    = logs.get('loss', float('nan'))
            print(f"  step {batch:5d} | loss {loss:.4f} | {tps:>10,.0f} tok/s")
            self._interval_start = now

    def on_epoch_end(self, epoch, logs=None):
        elapsed      = time.perf_counter() - self._epoch_start
        total_tokens = self._step_count * batch_size * CONTEXT
        val_loss     = logs.get('val_loss', float('nan'))
        print(f"\nEpoch {epoch+1} | {elapsed:.1f}s | "
              f"{total_tokens / elapsed:,.0f} tok/s (avg) | "
              f"val_loss={val_loss:.4f}\n")


# ─────────────────────────────────────────────────────────────────────────────
# HELPERS
# ─────────────────────────────────────────────────────────────────────────────

def load_text_corpus(path: str) -> str:
    p     = Path(path)
    files = [p] if p.is_file() else sorted(p.glob('*.txt'))
    if not files:
        raise FileNotFoundError(f'No .txt files found in {path}')
    return ''.join(fp.read_text(encoding='utf-8') for fp in files)


# ─────────────────────────────────────────────────────────────────────────────
# MAIN
# ─────────────────────────────────────────────────────────────────────────────
class MemoryCallback(keras.callbacks.Callback):
    def on_epoch_end(self, epoch, logs=None):
        stats = jax.devices()[0].memory_stats()

        used = stats["bytes_in_use"] / 1024**3
        peak = stats["peak_bytes_in_use"] / 1024**3
        limit = stats["bytes_limit"] / 1024**3

        print(
            f"\nHBM: {used:.2f}/{limit:.2f} GB "
            f"(Peak: {peak:.2f} GB)"
        )
def main():
    print('\n' + '=' * 60)
    print('BACKEND & DEVICE VERIFICATION')
    print('=' * 60)
    print(f'Keras backend  : {keras.backend.backend()}')
    print(f'JAX devices    : {jax.devices()}')
    print('=' * 60 + '\n')

    # ── Tokenize ──────────────────────────────────────────────────────────
    tokenizer  = TokenizerWrapper('tokenizer.model')
    text       = load_text_corpus(data_path)
    raw_tokens = np.asarray(
        tokenizer.encode(text, add_bos=True, add_eos=False), dtype=np.int32
    )
    print(f'Total tokens: {len(raw_tokens):,}')

    split_idx    = int(len(raw_tokens) * 0.98)
    train_tokens = raw_tokens[:split_idx]
    val_tokens   = raw_tokens[split_idx:]

    stride = max(1, CONTEXT // 2)

    # ── Pre-compute windows ────────────────────────────────────────────────
    print("Pre-computing windows...", end=" ", flush=True)
    t0            = time.perf_counter()
    train_windows = precompute_windows(train_tokens, CONTEXT, stride)
    val_windows   = precompute_windows(val_tokens,   CONTEXT, stride)
    print(f"done in {(time.perf_counter() - t0) * 1000:.1f} ms")
    print(f"  train: {train_windows.shape}  ({train_windows.nbytes / 1e6:.1f} MB)")
    print(f"  val:   {val_windows.shape}  ({val_windows.nbytes / 1e6:.1f} MB)\n")

    # ── Build datasets ─────────────────────────────────────────────────────
    train_ds = create_tf_dataset(train_windows, batch_size, shuffle=True)
    val_ds   = create_tf_dataset(val_windows,   batch_size, shuffle=False)

    # ── Benchmark before training ──────────────────────────────────────────
    # Verifies the pipeline can outpace the TPU before we waste a run.
    benchmark_dataset(train_ds, batch_size, CONTEXT, n_batches=200,
                      label="train_ds")

    # ── Build model ────────────────────────────────────────────────────────
    os.makedirs('checkpoints', exist_ok=True)

    model = create_llm(
        vocab_size   = tokenizer.vocab_size,
        d_model      = D_MODEL,
        n_layers     = numberoflayers,
        n_heads      = numberofheads,
        d_latent     = d_Latent,
        ffn_mult     = ffn_mult,
        max_seq_len  = CONTEXT,
        use_moe      = False,
        num_kv_heads = num_kv_heads,
        swa_window   = swa_window,
    )

    optimizer = keras.optimizers.AdamW(
        learning_rate   = learning_rate,
        weight_decay    = weight_decay,
        global_clipnorm = 1.0,
    )
    model.compile(
        optimizer  = optimizer,
        loss       = keras.losses.SparseCategoricalCrossentropy(from_logits=True),
        jit_compile = True,
    )
    
    # ── Warmup with exact training shape ──────────────────────────────────
    # XLA traces a kernel per unique input shape.  Warming up with any other
    # shape (e.g. the original (1, min(4, CONTEXT))) causes a full recompile
    # on the first real training batch — wasting 30-120 s on the TPU.
    print("Warming up XLA with training shape...", end=" ", flush=True)
    t0      = time.perf_counter()
    dummy_x = np.zeros((batch_size, CONTEXT), dtype=np.int32)
    _       = model(dummy_x, training=False)
    _       = model(dummy_x, training=True)
    print(f"done in {time.perf_counter() - t0:.2f}s\n")

    # ── Forward pass sanity check ──────────────────────────────────────────
    # One explicit transfer (np.array) then work in NumPy — avoids the 4
    # implicit device→host round trips of the original float(ops.min(...))
    # scalar-cast pattern.
    print('─' * 52)
    print('Forward pass validation')
    print('─' * 52)
    for x_batch, _ in train_ds.take(1):
        logits    = model(x_batch[:1], training=False)
        logits_np = np.array(logits)
        print(f'Input  : {x_batch.shape}')
        print(f'Logits : {logits.shape}')
        print(f'min={logits_np.min():.4f}  max={logits_np.max():.4f}')
        if np.isnan(logits_np).any() or np.isinf(logits_np).any():
            raise ValueError('Forward pass produced NaN/Inf — check model init.')
    print('✓ Forward pass clean.\n')
    
    # ── Steps per epoch from actual window count ───────────────────────────
    # Original used len(tokens) // (batch*seq) which undercounts overlapping
    # windows and causes epochs to terminate prematurely.
    print(f"Total params: {model.count_params():,}")
    steps_per_epoch  = max(1, len(train_windows) // batch_size)
    validation_steps = max(1, len(val_windows)   // batch_size)
    print(f"steps_per_epoch  : {steps_per_epoch}")
    print(f"validation_steps : {validation_steps}\n")

    # ── Train ──────────────────────────────────────────────────────────────
    model.fit(
        train_ds,
        validation_data  = val_ds,
        epochs           = EPOCHS,
        steps_per_epoch  = steps_per_epoch,
        validation_steps = validation_steps,
        callbacks        = [
            keras.callbacks.TerminateOnNaN(),
            keras.callbacks.ModelCheckpoint(
                filepath          = 'checkpoints/veylon_{epoch:02d}.weights.h5',
                save_freq         = 'epoch',
                save_weights_only = True,
            ),
            ThroughputCallback(),
            MemoryCallback(),
        ],
    )

    model.save_weights('veylon_final.weights.h5')
    print('Saved veylon_final.weights.h5')


if __name__ == '__main__':
    main()