Open Infra Skills

---
name: optimize-musa-training
description: Profile, benchmark, debug, and optimize AI training workloads on Moore Threads MUSA GPUs while preserving numerical behavior. Use for Torch MUSA, MTT GPUs, mthreads-gmi, Moore Perf System (msys), Moore Perf Compute (mcu), MFU/HFU, FlashAttention/FA2, SDPA, torch.compile, FSDP/FSDP2, MCCL, distributed hangs, low-batch utilization, kernel launch or transfer overhead, memory pressure, and CUDA-to-MUSA performance migration.
---

# Optimize MUSA Training

Use a measurement-first workflow to improve MUSA training throughput without changing model semantics. Treat framework timing, system traces, and kernel counters as different layers of evidence.

## Guardrails

- Preserve model architecture, data semantics, optimizer math, precision policy, and checkpoint compatibility unless the user explicitly authorizes a change.
- Establish a versioned baseline before editing code. Compare forward outputs, loss, gradients, memory, and steady-state throughput after every retained change.
- Keep profiler overhead out of the throughput denominator. Measure FLOPs in a profiled run and steady step time in an otherwise equivalent non-profiled run.
- Never label profiler-attributed executed FLOPs as model MFU without stating the FLOP definition and coverage. Distinguish useful model FLOPs, executed hardware FLOPs, and profiler-attributed FLOPs.
- Do not infer MUSA behavior from CUDA behavior. Feature-detect the installed driver, SDK, Torch MUSA, muDNN, muBLAS, MCCL, attention backend, and profiler versions.
- Keep cluster transport separate from profiling logic. Do not require PowerShell, VS Code, a jump host, a specific scheduler, or a particular client operating system.
- Keep credentials, internal hostnames, private image registries, dataset paths, and proprietary reports out of public artifacts.

## Route The Task

1. For environment, import, device-selection, or container failures, read [environment-and-preflight.md](references/environment-and-preflight.md).
2. For MFU/HFU, PyTorch profiling, timeline analysis, transfers, or kernel counters, read [measurement-and-profiling.md](references/measurement-and-profiling.md).
3. For FA2, GEMM, compile, launch, memory, dataloader, FSDP, or MCCL optimization, read [optimization-playbook.md](references/optimization-playbook.md).
4. Before retaining any optimization, read [correctness-and-experiments.md](references/correctness-and-experiments.md).
5. For a real low-batch S5000 case and negative results worth avoiding, read [s5000-case-study.md](references/s5000-case-study.md).

## Workflow

### 1. Define The Metric And Invariants

Record:

- workload, model revision, dataset/sample bucket, precision, sequence shape, batch per device, accumulation, device count, and distributed strategy;
- useful-model FLOPs, executed FLOPs, or profiler-attributed FLOPs;
- peak denominator by SKU and precision;
- model behaviors that must remain unchanged;
- target metric such as samples/s, tokens/s, MFU, HFU, memory, or time-to-train.

Do not optimize a mixed workload as though every sample has the longest shape. Benchmark each meaningful bucket and the actual weighted mixture.