Optimize Musa Training

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:

python scripts/musa_env_report.py --output /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:

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

4. Quantify Utilization

For one shape:

python scripts/compute_mfu.py \
  --flops-per-device-step-tflop  \
  --step-seconds  \
  --peak-tflops-per-device  \
  --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:

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.

Source & license

This open-source skill is cataloged on AgentStack and links to its original source — we do not rehost the code.

• Author: open-infra-skills
• Source: open-infra-skills/infra-skills
• License: Apache-2.0

Install and usage instructions live in the source repository linked above.