File size: 2,151 Bytes
dfb775d | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 | """Autotune orchestrator — runs three probes and emits an AutotunePlan.
Day 1: returns a hardcoded reference plan (`dry_run=True` always).
Day 2 (on MI300X): wires up the real CK-vs-Triton attention probe and replaces
the hardcoded GEMM/RCCL decisions with documented heuristics.
"""
from __future__ import annotations
from mindxtrain.autotune.attention_probe import probe_attention
from mindxtrain.autotune.gemm_probe import probe_gemm
from mindxtrain.autotune.plan import AutotunePlan, ProbeTiming
from mindxtrain.autotune.rccl_probe import detect_gpu_count, probe_rccl
def run_autotune(gpu_index: int = 0, dry_run: bool = False) -> AutotunePlan:
"""Run the 60-second probe sequence and return an AutotunePlan.
On Day 1 (dry_run=True or no GPU) returns a static reference plan that
the training dispatch can consume to exercise its code path.
"""
if dry_run:
return _reference_plan()
attention_backend, attention_timings = probe_attention(gpu_index=gpu_index)
gemm_heuristic, gemm_timings = probe_gemm(gpu_index=gpu_index)
gpu_count = detect_gpu_count()
rccl_config = probe_rccl(gpu_index=gpu_index, gpu_count=gpu_count)
return AutotunePlan(
attention_backend=attention_backend,
gemm_heuristic=gemm_heuristic,
rccl_config=rccl_config,
fsdp_shard_width=8 if gpu_count == 8 else 1,
probe_timings=[*attention_timings, *gemm_timings],
notes=[
f"real probe: attention + GEMM measured; {gpu_count or 'no'} GPU(s) detected.",
],
)
def _reference_plan() -> AutotunePlan:
"""Hardcoded Day-1 reference plan; values are sane defaults for MI300X gfx942."""
return AutotunePlan(
attention_backend="ck",
gemm_heuristic="hipblaslt_default",
rccl_config="1gpu_noop",
fsdp_shard_width=1,
suggested_lora_rank=16,
suggested_micro_batch_size=4,
probe_timings=[
ProbeTiming(label="dry-run-reference", backend="ck", median_ms=0.0, iterations=1),
],
notes=["dry-run reference plan; replace with real probe output on MI300X via Day-2 bench"],
)
|