ML / STRIX HALO

Getting torch.compile Working on Strix Halo (gfx1151)

Andrej Karpathy's autoresearch lets an AI agent run ML experiments autonomously overnight — modify the architecture, train for a fixed time budget, evaluate, keep or discard, repeat. It assumes CUDA. We ported it to an AMD Ryzen AI MAX+ 395 / Radeon 8060S consumer APU on ROCm, and got torch.compile fully working on gfx1151. That last part is the piece people actually search for.

The starting position

There was already an AMD fork of autoresearch — andyluo7/autoresearch — but it targets datacenter MI300X/MI308X parts and runs in eager mode only, no torch.compile. It reports 3.18% MFU. That fork did the important archaeology: it identified the Flash Attention 3 problem and the lerp_() dtype issue. What it didn't do was get the compiler to actually cooperate.

Our target was smaller and stranger: a Strix Halo APU with 64 GB of unified memory shared between CPU and GPU, running TheROCk nightly PyTorch with native gfx1151 support. The specific build in the repo is 2.11.0a0+rocm7.11.0a20260106.

With compile working, the fork reports 24.4% MFU against the datacenter fork's 3.18% — nearly 8x the efficiency, on a chip you can buy in a mini-PC. Here is the comparison table as the repo publishes it:

The repo is explicit about what that val_bpb gap means: it is throughput-bound — fewer optimization steps inside the time budget — not efficiency-bound. The Halo is doing real work per FLOP.

Break #1: Flash Attention 3 is NVIDIA-only

Upstream autoresearch leans on FA3 via the kernels package. That is a hard stop on ROCm — there is no FA3 for AMD, and the package itself is NVIDIA-only. The fix is to route attention through PyTorch's own scaled_dot_product_attention, which on ROCm dispatches to AOTriton.

SDPA expects (B, n_heads, T, head_dim), so the tensors need a transpose out of the (B, T, n_heads, head_dim) layout, and grouped-query attention needs the KV heads expanded manually:

q = q.transpose(1, 2)
k = k.transpose(1, 2)
v = v.transpose(1, 2)

# Expand KV heads if GQA (n_kv_head < n_head)
if self.n_kv_head < self.n_head:
    rep = self.n_head // self.n_kv_head
    k = k.repeat_interleave(rep, dim=1)
    v = v.repeat_interleave(rep, dim=1)

y = F.scaled_dot_product_attention(q, k, v, is_causal=True)

This costs something real, and the port does not hide it: SDPA has no window_size. Upstream's sliding-window (SSSL) pattern degrades to full causal attention on every layer. The window_size argument is simply ignored. That is a live research question the autonomous agent is explicitly invited to poke at — whether it matters, or whether there is another way to get windowing without FA3.

Break #2: the lerp_() dtype casts — the torch.compile fix

This is the hard-won bit. The MuonAdamW optimizer's fused steps are decorated with @torch.compile(dynamic=False, fullgraph=True), and inside them lerp_() is called with scalar-tensor weights. On ROCm, that combination is where compilation falls over — which is why the datacenter fork gave up and ran eager.

The fix is unglamorous and it works: cast the lerp_() weight explicitly to the destination tensor's dtype at every call site. In the AdamW step:

@torch.compile(dynamic=False, fullgraph=True)
def adamw_step_fused(p, grad, exp_avg, exp_avg_sq, step_t, lr_t, beta1_t, beta2_t, eps_t, wd_t):
    p.mul_(1 - lr_t * wd_t)
    exp_avg.lerp_(grad, (1 - beta1_t).to(dtype=exp_avg.dtype))
    exp_avg_sq.lerp_(grad.square(), (1 - beta2_t).to(dtype=exp_avg_sq.dtype))

And the same discipline in the Muon Nesterov-momentum step, where both the buffer update and the interpolated gradient need it:

momentum = momentum_t.to(stacked_grads.dtype)
momentum_buffer.lerp_(stacked_grads, (1 - momentum).to(dtype=momentum_buffer.dtype))
g = stacked_grads.lerp_(momentum_buffer, momentum.to(dtype=stacked_grads.dtype))

The repo notes that PyTorch 2.11 nightly fixes the underlying issue — but the explicit casts are what make the code robust across builds, and the repo's guidance to the agent is to maintain them if it edits the optimizer. If you take one thing from this article for your own gfx1151 port: when fullgraph=True compilation dies inside an optimizer, look for implicit dtype promotion on in-place ops before you look anywhere else.

One related detail that keeps the graph stable: scalar hyperparameters are held as 0-D CPU tensors specifically to avoid torch.compile recompilation when their values change across the schedule. If you pass raw Python floats that vary per step, you will silently re-trace every step and wonder where your throughput went.

Break #3: the ROCm environment

Two env vars decide whether this runs at all, and the repo states them as absolutes:

• NEVER set PYTORCH_HIP_ALLOC_CONF=backend:malloc — it crashes PyTorch. train.py defensively pops both PYTORCH_HIP_ALLOC_CONF and PYTORCH_ALLOC_CONF out of the environment at import time, before torch loads. (expandable_segments is not supported on HIP; use the default allocator.)
• ALWAYS set TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL=1 — this enables the fast AOTriton SDPA kernels. The repo puts the magnitude at a 19x flash-attention speedup. Without it you are on a slow fallback path and your MFU number will be a lie about the hardware.

run.sh handles all of this for you:

#!/bin/bash
# Run autoresearch training on Strix Halo (gfx1151)
cd "$(dirname "$0")"
source venv/bin/activate
unset PYTORCH_HIP_ALLOC_CONF
export TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL=1
export ROCM_HOME=/opt/rocm
export HIP_PATH=/opt/rocm
export PYTORCH_ROCM_ARCH=gfx1151
python train.py "$@"

Tuning for a unified-memory APU

Three more changes, all consequences of the hardware rather than of ROCm:

Batch sizes. TOTAL_BATCH_SIZE=2**17 (~131K tokens per optimizer step) and DEVICE_BATCH_SIZE=16. Baseline peak usage is 6.2 GB out of 64 GB — enormous headroom, and the repo flags scaling up (deeper, wider, or much larger batches) as the most promising first research direction.

Time budget doubled to 10 minutes. The 8060S gets ~127 steps in 5 minutes against an H100's ~953. Ten minutes yields ~250 steps — enough signal for the agent to tell an architecture change from noise, while still permitting roughly 6 experiments per hour, ~50 over a night's sleep.

Auto GPU FLOPS detection. MFU is meaningless if you divide by the wrong denominator, and PyTorch won't tell you a consumer APU's peak. train.py carries a lookup table (H100, H200, A100, B200, MI300X, and "8060s": 49.6e12 for the Radeon 8060S), matches on the device name, and falls back to an GPU_BF16_PEAK_FLOPS env override with a loud warning for unknown parts.

Notably, prepare.py is untouched — same data, same tokenizer, same evaluate_bpb ground-truth harness as upstream. The comparison to H100 numbers is apples-to-apples on the metric.

Running it

git clone https://github.com/bkpaine1/autoresearch-halo.git
cd autoresearch-halo

# Create venv with TheROCk nightly torch
python3 -m venv venv
source venv/bin/activate
pip install --index-url https://rocm.nightlies.amd.com/v2/gfx1151/ --upgrade --no-cache-dir --pre torch
pip install pyarrow rustbpe tiktoken requests numpy matplotlib pandas

# Download data + train tokenizer
python prepare.py --num-shards 10

# Run baseline
./run.sh

Prerequisites: an AMD Strix Halo (or other gfx1151 GPU) with ROCm, TheROCk nightly PyTorch (tested with 2.11.0a0), Python 3.10+. If you already have a TheROCk venv from another project, symlink it in rather than rebuilding:

ln -s /path/to/your/rocm/venv venv
pip install pyarrow rustbpe tiktoken requests
python prepare.py --num-shards 10
./run.sh

To run the autonomous loop, point Claude Code (or any agent) at the repo with program.md as context. It establishes a baseline, edits train.py only, trains for the 10-minute budget, evaluates, keeps improvements and git-resets regressions, and repeats. First run costs about 65 extra seconds for torch.compile warm-up; runs past 15 minutes are killed and treated as failures. Launch a single experiment by hand with:

source venv/bin/activate && unset PYTORCH_HIP_ALLOC_CONF && TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL=1 python train.py

Why this is worth doing

Karpathy's premise is that frontier-style pretraining research can be automated on a single GPU. The corollary this port establishes is that the GPU does not have to be an H100, or even a datacenter part. A consumer APU with 64 GB of unified memory, a nightly ROCm build, an attention swap and a handful of dtype casts gets you a working compiler and a quarter of theoretical peak — and it sits on a desk.

Credit where it is owed: Andrej Karpathy for autoresearch, the nanochat training backbone and the MuonAdamW optimizer; andyluo7 for the first AMD ROCm port and for identifying the FA3→SDPA swap and the lerp_() dtype issue; and TheROCk for nightly PyTorch builds with native gfx1151 support. This port would not exist without the community keeping ROCm alive on consumer hardware. MIT licensed, same as upstream.

Source and full code: github.com/bkpaine1/autoresearch-halo