---
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.

### 2. Capture The Environment

Run:

```bash
python scripts/musa_env_report.py --output <run-dir>/environment.json
```

Also save the container image digest, source commit, working-tree diff, launch command, and relevant environment switches. Verify physical device visibility from inside the process rather than trusting shell variables alone.

### 3. Establish A Steady Baseline

- Warm up imports, allocator state, compilation, autotuning, dataloader workers, and collectives.
- Exclude compile steps, epoch boundaries, checkpoint saves, validation, and profiler steps.
- Capture enough steady steps to expose variance. Reverse A/B order and repeat when the expected gain is below 3%.
- Keep all ranks on the same shape bucket for a distributed step.

Summarize logs with:

```bash
python scripts/summarize_steps.py train.log --skip-first 2 --json
```

### 4. Quantify Utilization

For one shape:

```bash
python scripts/compute_mfu.py \
  --flops-per-device-step-tflop <F> \
  --step-seconds <T> \
  --peak-tflops-per-device <C> \
  --flops-kind profiler-attributed
```

For mixed buckets, provide a JSON config with per-bucket FLOPs, time, and either step weight or sample count plus global batch. Use the aggregate total-FLOPs / total-time result, not a naive arithmetic mean.

### 5. Descend Through Three Profiling Layers

Use the cheapest layer that answers the current question:

1. Framework layer: separate input pipeline, forward, backward/recompute, optimizer, clipping, and collectives.
2. System layer: use Moore Perf System to inspect CPU/GPU overlap, launches, copies, synchronization, queues, streams, and rank skew.
3. Kernel layer: use Moore Perf Compute only on a small reproducible segment to inspect LaunchStats, MemoryWorkloadAnalysis, SpeedOfLight, occupancy, registers, memory pipelines, and Roofline position.

Do not run full end-to-end training under MCU unless the capture is tightly filtered. MCU replays and serializes kernels to collect counters; its duration is not an end-to-end throughput measurement.

### 6. Form One Evidence-Backed Hypothesis

Examples:

- FA2 is graph-breaking or poorly tiled for the observed head/sequence shape.
- GEMM is using a vendor tensor-core kernel but surrounding cast/reduction traffic dominates.
- repeated static metadata construction creates fill and launch overhead;
- host reads such as `.item()`, `.cpu()`, logging, or metric synchronization serialize the step;
- FSDP queue time reflects rank skew rather than slow collective kernels;
- an activation larger than LLC is repeatedly materialized by cat/split/copy operations;
- short and long buckets need different static compiled policies.

Change one variable at a time. Keep a decision log for both positive and negative experiments.

### 7. Validate And Decide

Require all of the following before retaining a change:

- numerical differences are characterized and acceptable for the target dtype;
- the same checkpoint and data produce finite forward, backward, and optimizer behavior;
- the gain survives repeated non-profiled A/B runs and exceeds normal variance;
- peak memory and all supported buckets remain acceptable;
- the end-to-end result agrees with the microbenchmark direction;
- distributed scaling and checkpoint load/save still work when affected.

Prefer small stable gains that compose, but keep experimental paths disabled by default until their full-training benefit is repeatable.

## Artifact Contract

Create a self-contained run directory with:

```text
run/
  environment.json
  command.txt
  source.txt
  baseline.json
  correctness.json
  profiles/
    framework/
    system/
    compute/
  decisions.md
```

Record exact versions and commands, but sanitize machine-specific and secret values before sharing.