| """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"], |
| ) |
|
|
