train2.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 jax.sharding as sharding
from jax.sharding import PartitionSpec as P
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
            # Multiply by effective_batch if data parallelism is active
            tps     = (self.LOG_INTERVAL * effective_batch * 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 * effective_batch * 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)


# ─────────────────────────────────────────────────────────────────────────────
# SHARDING (Data Parallelism for v5e-8)
# ─────────────────────────────────────────────────────────────────────────────

def setup_data_parallelism():
    """
    Configure JAX for pure data parallelism across all TPU devices.

    For v5e-8 (8 chips):
      - Each chip runs the full model
      - Different batch shards go to each chip
      - Gradients are AllReduced across chips
      - Effective batch = batch_per_chip × num_chips

    This is the simplest and most efficient sharding strategy for models
    that fit on a single chip.  For 8M parameters, the model easily fits
    on one v5e chip (16 GB HBM each).

    Returns mesh_shape (tuple) and sharding spec dict for inputs/weights.
    """
    devices = jax.devices()
    n_devices = len(devices)

    if n_devices == 1:
        print(f"Single device detected ({devices[0]}). Data parallelism disabled.")
        return None, None

    print(f"\n{'─'*52}")
    print(f"Data Parallelism Setup (v5e-{n_devices})")
    print(f"{'─'*52}")
    print(f"Devices: {n_devices}")
    print(f"Device type: {devices[0].platform}")

    # Create 1D mesh for data parallelism: each axis is one device
    mesh = sharding.Mesh(
        devices=np.array(devices).reshape((n_devices,)),
        axis_names=("batch",)
    )

    # Sharding specs for inputs and weights
    # Input shape: (batch, seq_len)
    #   Shard batch dimension across devices, keep seq_len replicated
    input_spec = P("batch", None)

    # Weights: replicate across all devices (each chip has full model)
    weight_spec = P(None)

    # Activation gradients during backward pass
    #   Same as input: shard batch, replicate seq
    activation_spec = P("batch", None)

    print(f"Mesh shape: {mesh.shape}")
    print(f"Input sharding: {input_spec}  (batch sharded, seq replicated)")
    print(f"Weight sharding: {weight_spec}  (fully replicated)")
    print(f"Expected effective batch: {n_devices} × batch_size")
    print(f"{'─'*52}\n")

    return mesh, {
        "input": input_spec,
        "weight": weight_spec,
        "activation": activation_spec,
    }


def apply_sharding_to_model(model, mesh, sharding_specs):
    """
    Apply JAX sharding annotations to a Keras model compiled with JAX backend.

    WARNING: This is JAX-level sharding and only works if:
    1. Model is built with JAX-native ops (Keras layers with JAX backend)
    2. jit_compile=True is set in model.compile()
    3. No TensorFlow-only ops are used in the model

    For Keras models, we set the sharding via jax.Array.with_sharding_constraint()
    in a custom train step. However, Keras 3 makes this tricky because it wraps
    training in its own jit.

    Simpler approach: set jax.config to globally use this mesh, and let XLA
    infer sharding from the mesh context.
    """
    if mesh is None:
        return

    # Tell JAX to use this mesh for all operations
    with mesh:
        print("Sharding configuration loaded.")
        print(f"  All matmul ops will shard batch dim across {mesh.shape[0]} devices")
        print(f"  Gradient AllReduce will use {mesh.shape[0]}-way ICI collective\n")


# ─────────────────────────────────────────────────────────────────────────────
# 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')

    # ── Setup data parallelism ─────────────────────────────────────────────
    mesh, sharding_specs = setup_data_parallelism()
    if mesh is not None:
        apply_sharding_to_model(None, mesh, sharding_specs)

    # ── 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 ─────────────────────────────────────────────────────
    # If data parallelism is active, the effective batch becomes:
    #   batch_per_device × num_devices
    # But tf.data still sees batch_per_device — JAX handles the replication.
    num_devices = len(jax.devices())
    actual_batch_size = batch_size  # Per-device batch
    effective_batch = actual_batch_size * num_devices if mesh is not None else actual_batch_size

    if mesh is not None:
        print(f"Effective batch size: {actual_batch_size} × {num_devices} devices = {effective_batch}")

    train_ds = create_tf_dataset(train_windows, actual_batch_size, shuffle=True)
    val_ds   = create_tf_dataset(val_windows,   actual_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()