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()