Ml Engine

ml-engine

[!IMPORTANT] This skill has reference files in the references/ directory. You MUST read them at least once to understand the deep-dive content (TPU Setup, Training Loops, FSDPv2, SPMD, etc.) and call them whenever you need specific information from there.

TPU-first ML research engine. Builds reproducible distributed training pipelines for novel architecture ablations. PyTorch-XLA (torchtpu) is primary; JAX/torchax is first-class secondary. Targets v5e-8 (primary, Kaggle + GCP), v3-8 (legacy quota), v3-64 (multi-pod scaling). Uses modern torch_xla APIs: torch_xla.step(), torch_xla.sync(), torch_xla.launch(), torch_xla.compile(), FSDPv2, SPMD, @assume_pure, scan_layers.

When to Use

• Distributed training on TPU with PyTorch-XLA or JAX
• Novel architecture components: MoE, router, custom FFN, custom LayerNorm
• Reproducible ablation studies with deterministic data pipelines
• SPMD / GSPMD sharding setup and debugging
• Attention kernel selection: Splash, Flash, SDPA on TPU
• Raw optimizer integration: torch.optim, optax, SyncFree AdamW
• datasets streaming with sharded loaders
• wandb logging on multi-host TPU runs
• Offline/online evaluation loops with checkpoint resumption
• Tensor manipulation with einops on TPU
• Multi-pod training: v3-32, v3-64, v4-128
• Pallas custom TPU kernels via call_jax
• Modern torch_xla APIs: torch_xla.step(), torch_xla.sync(), torch_xla.launch(), torch_xla.compile()
• FSDPv2 (GSPMD-based fully sharded data parallelism)
• @assume_pure for tracing optimization
• scan / scan_layers for compile-time reduction on homogeneous layers
• torchax — PyTorch frontend for JAX interoperability
• torch.compile(model, backend='openxla') — Dynamo integration
• Compilation caching via xr.initialize_cache()
• Eager mode via torch_xla.experimental.eager_mode()
• Checkpoint saving, resumption, distributed checkpointing
• Profiling training bottlenecks with XLA profiler

Commands

See README.md for full command documentation, parameters, and examples.

Godmode: /ml

/ml [research idea] — full research scaffold in one command. Runs the full cycle: analyze → plan → implement → verify.

What it produces:

1. Research hypothesis + experiment matrix
2. Mesh setup for target TPU (auto-detected or specified)
3. Training script with correct strategy (SPMD/FSDPv2/DDP)
4. Data pipeline (streaming + sharding)
5. Attention kernel selection (Splash/Flash/SDPA)
6. wandb config + logging hooks
7. Eval loop + perplexity/metric setup
8. Reproducibility seed lock + config dict
9. Ablation matrix if research involves multiple variables
10. Pallas optimization plan if bottleneck detected

When to use /ml vs individual commands:

• /ml = new research idea, greenfield. Get full scaffold.
• Individual commands = targeted work on existing code. Get surgical output.

Core Instructions

1. TPU Detection and Mesh Initialization

Detect TPU version and initialize mesh before any model or data setup. Use modern torch_xla APIs — never legacy xm.xla_device() — because the new tpu_device() context manager handles initialization, mesh setup, and error recovery automatically.

import torch_xla
import torch_xla.core.xla_model as xm
import torch_xla.distributed.spmd as xs
from torch_xla import runtime as xr
import numpy as np

# Modern device access
device = torch_xla.device()
tpu_version = xm.get_tpu_env("TYPE", "v5litepod-8")  # v5litepod-8, v3-8, v2-8, etc.

# 1D mesh for data parallelism on v5e-8 / v3-8 / v2-8
num_devices = xr.global_runtime_device_count()  # SPMD-aware
mesh = xs.Mesh(np.arange(num_devices), (num_devices,), axis_names=("data",))
xs.set_global_mesh(mesh)

See references/tpu-setup.md for single-host mesh, 2D sharding, and device placement. See references/tpu-v5e.md for v5e-specific setup, Kaggle integration, and best practices. See references/torchtpu.md for modern PyTorch XLA runtime (PJRT, torch.compile). See references/multi-pod-training.md for v3-64, torchrun, and distributed checkpointing.

2. Modern Training Loop

Use torch_xla.step() context manager. Replaces manual xm.mark_step() — because manual mark_step() is the #1 source of silent XLA recompilation bugs. Never use xm.mark_step() in new code.

import torch_xla

def train(model, train_loader, optimizer, loss_fn):
    model.to('xla')
    for inputs, labels in train_loader:
        with torch_xla.step():
            inputs, labels = inputs.to('xla'), labels.to('xla')
            optimizer.zero_grad()
            outputs = model(inputs)
            loss = loss_fn(outputs, labels)
            loss.backward()
            optimizer.step()
    torch_xla.sync()

See references/training-loop-patterns.md for multi-process (torch_xla.launch), DDP with XLA backend, compiled step functions, eager mode, and MpDeviceLoader.

3. SPMD Data Parallelism

Single-process multi-device training. Better throughput than DDP for most TPU workloads.

import torch_xla.distributed.spmd as xs
from torch_xla import runtime as xr
import torch_xla.distributed.parallel_loader as pl

xr.use_spmd()

num_devices = xr.global_runtime_device_count()
mesh = xs.Mesh(np.arange(num_devices), (num_devices,), ('data',))
xs.set_global_mesh(mesh)

# Scale batch size — one process handles all devices
batch_size *= num_devices

train_loader = pl.MpDeviceLoader(
    train_loader,
    device,
    input_sharding=xs.ShardingSpec(mesh, ('data', None, None, None))
)

See references/spmd.md for mesh setup, partition specs, and advanced sharding.

4. FSDPv2

Large model training. Requires SPMD mode and mesh with fsdp axis.

import torch_xla.distributed.spmd as xs
from torch_xla.experimental.spmd_fully_sharded_data_parallel import SpmdFullyShardedDataParallel as FSDPv2
from torch_xla import runtime as xr
from torch_xla.distributed.fsdp.wrap import transformer_auto_wrap_policy
import functools

xr.use_spmd()
num_devices = xr.global_runtime_device_count()
mesh = xs.Mesh(np.arange(num_devices), (num_devices, 1), ('fsdp', 'model'))
xs.set_global_mesh(mesh)

auto_wrap_policy = functools.partial(
    transformer_auto_wrap_policy,
    transformer_layer_cls={DecoderLayer},
)
model = FSDPv2(model, auto_wrap_policy=auto_wrap_policy)

Gradient checkpointing (apply before FSDPv2 wrapper):

from torch_xla.distributed.fsdp import checkpoint_module
model = FSDPv2(checkpoint_module(my_module), mesh)

See references/fsdp-v2.md for full FSDPv2 setup, sharding output, and distributed checkpointing.

5. Dual-Framework Runtime (PT-XLA + JAX/torchax)

PT-XLA = primary. JAX = auxiliary ops (custom loss, routing, differentiable sort). torchax = PyTorch syntax on JAX backend.

See references/dual-framework.md for DLPack zero-copy interop patterns and memory management. See references/torchax.md for torchax setup and PyTorch-on-JAX patterns.

6. Model Architecture Patterns

MoE routers, custom FFN (SwiGLU), RMSNorm — all with SPMD sharding annotations baked in.

See references/model-architecture.md for full MoE, custom FFN, LayerNorm, router patterns.

7. Attention Kernel Selection

Auto-detect and select explicitly. Priority: Splash (v3+) → Flash → SDPA. Never silently fall back — because silent fallback to SDPA can mask a 2-3x throughput regression that only surfaces under production load.

See references/attention.md for kernel benchmarking and fallback rules. See references/pallas-custom-kernels.md for Splash Attention via Pallas.

8. Data Pipeline

Stream from datasets, shard across TPU cores, deterministic sampling.

See references/data-pipeline.md for multi-epoch streaming, XLA-Trainer patterns, and JAX pipelines.

9. Optimizer Compatibility

SyncFree AdamW preferred on TPU (~20% faster) — because the sync barrier in standard AdamW wastes TPU cycles waiting for all replicas. Raw instances only -- no wrapper magic.

See references/optimizers.md for LR schedules and custom state handling.

10. Compilation Optimization

Three tools: compilation cache, @assume_pure, scan_layers.

See references/tracing-optimization.md for @assume_pure, scan_layers, benchmarks, and AOTAutograd compatibility. See references/compilation-cache.md for xr.initialize_cache() and shared mounts.

11. torch.compile with OpenXLA

torch.compile(model, backend='openxla') for inference (2-18x speedup) and training steps (experimental).

See references/torchtpu.md for Dynamo integration and performance benchmarks.

12. Eager Mode

Debugging only — because eager mode disables XLA graph fusion, dropping throughput 10-50x. Operations execute immediately -- no lazy graph. Combine with @assume_pure for selective compilation.

13. wandb Logging on TPU

Master process only. Aggregate before logging.

See references/wandb-logging.md for initialization, metric aggregation, checkpoint artifacts, multi-host sync, and JAX logging patterns.

14. Evaluation Loops

Sharded inference, aggregate across devices.

See references/evaluation.md for online/offline eval, perplexity calculation, accuracy, JAX eval, checkpoint resumption, and sharded inference.

15. Reproducible Ablations

Lock all RNGs. Version custom modules with string IDs. Log full config + git hash.

See references/ablations.md for seed locking, config dicts, config hashing, module versioning, deterministic data, and environment capture.

16. Checkpointing

See references/fsdp-v2.md for save/resume patterns and distributed checkpointing with FSDPv2.

17. Tensor Manipulation with einops

Use einops for all reshaping and einsum on TPU. XLA fuses better than manual view + permute.

See references/einops.md for rearrange, einsum, reduce, JAX compatibility, and TPU-specific patterns.

Boundaries

• Does NOT handle hyperparameter search (use Optuna / Ray Tune separately)
• Does NOT optimize inference serving (training focus only)
• Does NOT design model architectures from scratch (provides patterns for common novel components)
• Does NOT replace framework documentation (provides TPU-specific integration patterns)

Self-Verification

Before declaring a training setup complete, verify:

• TPU version detected and mesh shape matches model sharding
• No XLA recompilation warnings in logs
• Loss is decreasing after 100 steps
• Memory usage is within TPU HBM limits
• All xm.mark_step() calls replaced with tpu_step() context manager
• Checkpoint save/load tested at least once

Composability — Working With Other Skills

See PROTOCOL.md (SIP) at skills root for full interop contract.

Domain Declaration

domain: process
composable: true
yields_to: [craft, voice]

ml-engine owns process (training workflow, distributed setup, evaluation loops) and content (ML substance: tensor ops, sharding, attention kernels, optimizer state).

When ml-engine Leads

• Any distributed training setup on TPU
• Model architecture component implementation (MoE, router, FFN, LN)
• Data pipeline with datasets streaming and sharding
• Optimizer integration and custom state handling
• wandb logging aggregation on multi-host TPU
• Evaluation loop construction and metric computation
• Reproducibility guardrails for ablation studies
• FSDPv2 setup and sharding configuration
• Tracing optimization (@assume_pure, scan_layers, compilation caching)
• torch.compile with OpenXLA backend integration
• SPMD mesh and partition spec configuration
• Checkpoint save, resume, and distributed checkpoint patterns
• Training profiling and bottleneck identification

When ml-engine Defers

Layered Composition Rules

1. Process + Content: ml-engine is both process and content. When another content skill provides domain knowledge (e.g., NLP research findings), ml-engine shapes it into training code and config.

2. Process + Craft: ml-engine generates training scripts and metric logs. Craft skill consumes metrics to build visual dashboards. Boundary: ml-engine stops at raw CSV/tensor outputs; craft starts at visualization.

3. Process + Voice: ml-engine writes technical config and code. Voice skill rewrites user-facing explanations. Boundary: code blocks and config dicts are voice-agnostic.

Pipeline Behavior

• Upstream (receives output from another skill): If a research skill provides novel architecture ideas, ml-engine implements them as sharded modules with TPU-compatible attention.
• Downstream (output feeds into another skill): ml-engine outputs training scripts, config files, and metric logs. Downstream skills (craft, voice) consume these for reporting or visualization.

Conflict Signal

If a voice or craft skill attempts to modify code structure or sharding logic:

⚠️ Process conflict: voice/craft skill is modifying training loop structure or sharding annotations. Code structure preserved; prose/visual style applied to surrounding explanations only.