# Ml Engine

> >-

- **Type:** Skill
- **Install:** `agentstack add skill-isnoobgrammer-skills-for-agents-ml-engine`
- **Verified:** Yes — security-reviewed for prompt injection and unsafe behavior
- **Seller:** [IsNoobgrammer](https://agentstack.voostack.com/s/isnoobgrammer)
- **Installs:** 0
- **Category:** [Agent Skills](https://agentstack.voostack.com/c/agent-skills)
- **Latest version:** 0.1.0
- **License:** MIT
- **Upstream author:** [IsNoobgrammer](https://github.com/IsNoobgrammer)
- **Source:** https://github.com/IsNoobgrammer/skills-for-agents/tree/main/skills/ml-engine
- **Website:** https://isnoobgrammer.github.io/skills-for-agents/

## Install

```sh
agentstack add skill-isnoobgrammer-skills-for-agents-ml-engine
```

Requires the [AgentStack CLI](https://agentstack.voostack.com/docs/cli). Works with Claude Code, Cursor, and any MCP-compatible agent.

## About

# 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.

| Command | Purpose | Example |
|---------|---------|---------|
| `/ml [idea]` | **Godmode** — full research scaffold end-to-end | `/ml "sparse MoE with dynamic clamping"` |
| `/ml-train [model] [tpu] [strategy]` | Generate training script | `/ml-train moe v3-8 fsdp` |
| `/ml-mesh [tpu] [parallelism]` | Generate mesh setup | `/ml-mesh v3-64 2d-data-model` |
| `/ml-debug [issue]` | Debug XLA issues | `/ml-debug recompilation` |
| `/ml-benchmark [kernel] [seq-len]` | Benchmark attention kernel | `/ml-benchmark auto 2048` |
| `/ml-migrate [code]` | Migrate old API → modern | `/ml-migrate train.py` |
| `/ml-port [code] [mode]` | Port PyTorch → torch_xla native | `/ml-port model.py native` |
| `/ml-optimize [object] [method]` | Optimize XLA bottleneck | `/ml-optimize router pallas` |
| `/ml-plan [idea] [scope]` | Plan & template research | `/ml-plan "sparse MoE" full` |
| `/ml-ablate [var] [values]` | Run ablation matrix | `/ml-ablate top_k "1,2,4"` |
| `/ml-checkpoint [action] [path]` | Save/resume/inspect checkpoint | `/ml-checkpoint resume ./ckpt` |
| `/ml-profile [target] [steps]` | Profile training bottleneck | `/ml-profile step_fn 100` |

---

## 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.

```python
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.

```python
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.

```python
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.

```python
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):

```python
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

```yaml
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

| Other Skill's Domain | What ml-engine Does |
|---------------------|------------------------|
| **Voice** (e.g. blogger) | Preserves factual content about training setup; voice skill rewrites prose tone. ml-engine owns all code, config, and metric content. |
| **Craft** (e.g. painter) | Preserves training code and data pipeline; craft skill adds UI/visualization for dashboards, plots, or reports. ml-engine provides raw metrics. |
| **Density** (e.g. caveman, compress) | SKILL.md and reference files are source files — not compressible. Conversational explanations about ml-engine can be compressed. |

### 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.`

## Source & license

This open-source skill is cataloged on AgentStack and links to its original source — we do not rehost the code.

- **Author:** [IsNoobgrammer](https://github.com/IsNoobgrammer)
- **Source:** [IsNoobgrammer/skills-for-agents](https://github.com/IsNoobgrammer/skills-for-agents)
- **License:** MIT
- **Homepage:** https://isnoobgrammer.github.io/skills-for-agents/

Install and usage instructions live in the source repository linked above.

## Pricing

- **Free** — Free

## Security capabilities

Automated source analysis of v0.1.0 — what this tool can access:

- **Network access:** no
- **Filesystem access:** no
- **Shell / process execution:** no
- **Environment & secrets:** no
- **Dynamic code execution:** yes

*"Yes" means the capability is present in the source — more access means more to trust, not that it is unsafe.*


## Versions

- **0.1.0** — security scan: passed — Imported from the upstream source.

## Links

- Listing page: https://agentstack.voostack.com/l/skill-isnoobgrammer-skills-for-agents-ml-engine
- Seller: https://agentstack.voostack.com/s/isnoobgrammer
- Browse the marketplace: https://agentstack.voostack.com/browse

---
Listed on AgentStack — the marketplace for AI agent skills and MCP servers. Every listing is security-reviewed. Creators keep 70%.