docs/autotune.md

Autotune — the 60-second AOT probe

The single layer that distinguishes mindxtrain from Axolotl, LLaMA-Factory, Unsloth, torchtune, and Primus. Quoted from the frozen design brief in docs/blueprints/:

The single most differentiating angle is the auto-selection layer. No competitor framework — not Axolotl, LLaMA-Factory, Unsloth, torchtune, or Optimum-AMD itself — runs a per-job MI300X micro-benchmark before training to pick CK vs Triton attention backends, hipBLASLt heuristic vs rocBLAS path, AITER vs reference MoE kernels, NCCL_MIN_NCHANNELS, gradient-checkpointing strategy, FSDP shard width, and LoRA rank against the actual (model, dataset shape, sequence length, GPU count) tuple. mindxtrain owns that AOT-only autotune layer.

The AOT-only discipline

JIT autotune (Triton autotune in vLLM cold-start, torch.compile(mode='max-autotune') Inductor, MIOpen find-mode) is forbidden in production training. Reasons:

1. Reproducibility. A run with JIT autotune produces different kernels on different invocations of the same workload, breaking deterministic benchmarks.
2. First-batch latency. Triton autotune on cold start can stall a training step for 5-30 seconds, invisible in the loss curve and very visible in tok/s.
3. Cypherpunk2048 standard. Production paths must be statically declared at deployment. JIT compilation is an in-band runtime decision, which is exactly what the standard prohibits.

The autotune.policy: aot_only field in the YAML is the contract. The training layer reads the AutotunePlan JSON at start, sets env vars + flags, and never re-tunes during the loop. AOTriton (the AOT version of Triton math) is loaded as a precompiled .so; Composable Kernel kernels are pulled from the offline-tuned hipBLASLt cache.

The probe taxonomy

mindxtrain bench runs three probes in sequence inside its 60-second budget. The whole flow is at mindxtrain/autotune/benchmark.py.

1. attention_probe — CK vs Triton SDPA

mindxtrain/autotune/attention_probe.py.

Times torch.nn.functional.scaled_dot_product_attention across four representative shapes (queries × keys × heads × head-dim per the recipe's model.name + data.seq_len) on both backends:

Backend	How
ck	Composable Kernel (default) — hand-tuned ASM/CK kernels via AITER.
triton	AOTriton 0.11.2b0 with TORCH_BLAS_PREFER_HIPBLASLT=0 and PYTORCH_TUNABLEOP_ENABLED=0 toggles.

The probe is real — when torch (--extra ml) and a ROCm-visible GPU are both present, it times the four representative shapes on each backend via torch.nn.attention.sdpa_kernel. Without torch (typical CPU dev box), the probe gracefully returns the canonical ("ck", []) default so bench --dry-run parity holds and the AutotunePlan downstream consumers keep working unchanged.

budget: ~30 s
shapes: 4 representative (qlen, klen, num_heads, head_dim)
output: AttentionBackend ∈ {ck, triton}, list[ProbeTiming]

ProbeTiming is { label, backend, median_ms, iterations } — captured per (shape × backend) so the demo can render a side-by-side timing table in the video.

2. gemm_probe — hipBLASLt heuristic

mindxtrain/autotune/gemm_probe.py.

Per the user-confirmed Day-1 plan ("1 real probe + 2 hardcoded heuristics"), this returns hipblaslt_default for gfx942 based on AMD's documented MI300X tuning guidance. Reference: AMD ROCm 7.2.1 release notes, hipBLASLt 0.10 default heuristics are within 5 % of hand-tuned for the BF16/FP16 GEMMs mindxtrain hits (LoRA rank 16-64, hidden 2048-8192).

Why we don't enumerate. A real hipBLASLt heuristic enumeration is ~1.5 minutes and risks burning the entire 60-second budget. If MMLU eval shows GEMM-bound throughput regression on a specific recipe, revisit later.

Output: hipblaslt_default | hipblaslt_tuned | rocblas_fallback.

3. rccl_probe — collective bandwidth

mindxtrain/autotune/rccl_probe.py.

For 1-GPU runs this is a no-op. For 8-GPU runs it returns 8gpu_xgmi with NCCL_MIN_NCHANNELS=112 set in the plan notes. 2-GPU and 4-GPU groupings raise RuntimeError — MI300X xGMI bandwidth between subsets of 2/4 GPUs is asymmetric, and FSDP shards on those topologies will silently bottleneck.

def probe_rccl(gpu_index: int = 0, gpu_count: int = 1) -> RcclConfig:
    if gpu_count == 1:
        return "1gpu_noop"
    if gpu_count == 8:
        return "8gpu_xgmi"
    raise RuntimeError(f"FSDP on {gpu_count} GPUs is unsafe...")

This is enforced in two places: the rccl_probe raises at probe time, and the XTrainConfig.hardware.gpus field is Literal[1, 8] so the schema rejects bad values at parse time.

The AutotunePlan schema

mindxtrain/autotune/plan.py.

class AutotunePlan(BaseModel):
    schema_version: Literal["1"] = "1"
    gpu_arch: str = "gfx942"
    rocm_version: str = "7.2.1"

    attention_backend: Literal["ck", "triton"] = "ck"
    gemm_heuristic: Literal["hipblaslt_default", "hipblaslt_tuned", "rocblas_fallback"] = "hipblaslt_default"
    rccl_config: Literal["1gpu_noop", "8gpu_xgmi", "unsupported_2_4_gpu"] = "1gpu_noop"

    fsdp_shard_width: Literal[1, 8] = 1
    suggested_lora_rank: int = 16
    suggested_micro_batch_size: int = 4

    probe_timings: list[ProbeTiming] = []
    notes: list[str] = []

Pure data, content-addressed via BLAKE3 in the mindxtrain provenance manifest, fully reproducible across MI300X nodes.

How the training layer consumes the plan

mindxtrain/train/dispatch.py reads the plan and applies it before invoking the backend (real subprocess wrap of accelerate launch -m axolotl.cli.train in mindxtrain/train/sft.py):

def dispatch_training(cfg: XTrainConfig, plan: AutotunePlan, out_dir: Path) -> Path:
    # 1. set env vars: cfg.train.env + plan-driven additions
    #    e.g. plan.rccl_config == "8gpu_xgmi" → set NCCL_MIN_NCHANNELS=112
    #         plan.attention_backend == "ck" → NVTE_CK_USES_BWD_V3=1, etc.
    # 2. compile cfg → Axolotl YAML, override:
    #    train.flash_attention.backend ← plan.attention_backend
    #    train.method.r ← plan.suggested_lora_rank if cfg.train.method.kind == "lora"
    #    train.batch.per_device ← min(cfg, plan.suggested_micro_batch_size)
    # 3. subprocess: accelerate launch -m axolotl.cli.train <yaml>
    # 4. capture stdout/stderr to out_dir/train.log
    # 5. return checkpoint dir

Dry-run / CI path

Every CI pipeline runs mindxtrain bench --dry-run, which skips the GPU probes entirely and emits a hardcoded reference plan. The reference plan has attention_backend: ck, gemm_heuristic: hipblaslt_default, rccl_config: 1gpu_noop, fsdp_shard_width: 1 — sane MI300X 1-GPU defaults that exercise the same code path the real probe writes.

$ uv run mindxtrain bench --dry-run --out plan.json
wrote plan.json (dry_run=True, attention=ck, gemm=hipblaslt_default)

The dry-run path is what makes the GitHub Actions CI matrix CPU-only.

Day 2 implementation budget (target ~30 minutes per probe)

Probe	Target time	Risk
attention_probe	30 min	First-time AOTriton compilation may be slow; warm cache mounted from persistent volume.
gemm_probe	0 (hardcoded)	None — heuristic is documented.
rccl_probe	5 min	1-GPU is no-op; 8-GPU only relevant if we rent the 8× SKU.

Total Day 2 budget: ~35 minutes of MI300X time + writing time. The remaining hours go to verifying the plan flows into Axolotl correctly via Day 3's dispatch wiring.

Where the demo wow-moment lives

Capture autotune_plan.json and the streaming probe output for the 5-minute video. The 60-second autotune dashboard is the single most quotable visual asset in the submission — a measurable kernel selection that no competitor framework ships.

$ mindxtrain bench --gpu 0 --out plan.json
[autotune] CK FA forward, shape=(8, 4096, 32, 128): 12.4 ms (median over 50)
[autotune] Triton FA forward, shape=(8, 4096, 32, 128): 14.7 ms (median over 50)
[autotune] CK FA forward, shape=(8, 4096, 16, 128): 11.0 ms
[autotune] Triton FA forward, shape=(8, 4096, 16, 128): 13.2 ms
[autotune] picked ck (avg 1.2× faster)
[autotune] gemm: hipblaslt_default (gfx942 documented heuristic)
[autotune] rccl: 1gpu_noop
[autotune] wrote plan.json (1.2 KB) in 47 s