Imported from AMD-AGI/maxtext-slurm (skills/xla-tuning/SKILL.md). Install upstream with npx skills add AMD-AGI/maxtext-slurm --skill xla-tuning. Copyright stays with the author.

XLA / NCCL Flag Tuning for a (Model × Parallelism) Cell

Given a single (MODEL_TAG, PARALLELISM) cell on a fixed nodelist, find the XLA flag / NCCL env-var combination that maximizes steady-state TGS. Typically 1-3 flags capture most of the win; the rest of the sweep produces evidence (HLO + xplane) explaining why the winning flag wins — which is the durable knowledge that survives image/JAX upgrades.

This skill produces or extends one deliverable per model: <MODEL_TAG>-tuning.md in the repo root (and <MODEL_TAG>-tuning.zh.md if a Chinese sibling exists or is requested). It references no other .md files in the repo root.

Common prerequisites

• host-cmd reachable (python3 /maxtext-slurm/.host-cmd/host_cmd.py --ping)
• configs/<MODEL_TAG>.gpu.yml exists and runs at the requested parallelism without OOM
• A working pdbs for the cell (run batch-sweep first if not yet known — XLA tuning is on top of an already-good per_device_batch_size, not a substitute for it)
• Telegram set up (see telegram) for progress updates

Inputs (ask the user; never guess)

The user must provide:

1. MODEL_TAG — must match configs/<MODEL_TAG>.gpu.yml. Examples: deepseek3-671b, qwen3-32b, llama3.1-405b, mixtral-8x22b.

2. CLUSTER_SLOT — one of:

   • explicit nodelist=<list> (recommended; pinned hardware drops TGS noise to ~0.3-1 % which is what you need to detect 1-3 % flag deltas)
   • partition=<name>+nodes=<N> (let slurm pick from the partition pool)
   • nodes=<N> only (agent runs sinfo and picks N idle nodes from the default partition)

3. PARALLELISM — one of:

   • FSDP=N (DCN FSDP only)
   • PP=N (DCN pipeline only)
   • TP=N (ICI tensor parallel)
   • hybrid: PP=N,FSDP=M or FSDP=N,TP=M, etc. (axes multiply to total ranks)

Optional:

• pdbs — per_device_batch_size. If omitted, read from configs/<MODEL_TAG>.gpu.yml. If batch-sweep has been run, use that result.
• CONFIG_TAG — for models with multiple comparable sweep configs (e.g. an MoE model with dense_matmul vs sparse_matmul branches). If omitted, use the model's primary config from gpu.yml.

If any of (1)-(3) are missing, ask the user before starting. Do not assume defaults for these three.

Workflow

- [ ] Step 1 — Pre-flight (host-cmd, image, nodelist, env file, train_env.sh clean, TG)
- [ ] Step 2 — Inventory current state (read gpu.yml + env.sh + container_env.sh + any <MODEL_TAG>-tuning.md if it exists)
- [ ] Step 3 — Wave 1: baseline pair (current env vs raw image; one pair per CONFIG_TAG)
- [ ] Step 4 — Wave 1.5: capture HLO + xplane evidence on baseline-deployed
- [ ] Step 5 — Wave 2-4: flag testing per parallelism catalog (5-10 profiles per wave)
- [ ] Step 6 — Wave 5 (PP only): MaxText pipeline params (microbatches, V chunks)
- [ ] Step 7 — Wave 6: cross-config-tag confirmation of winning recipe (if model has >1 CONFIG_TAG)
- [ ] Step 8 — Write <MODEL_TAG>-tuning.md (and .zh.md mirror if needed)

Step 1 — Pre-flight

Run all of these. Stop and TG-ask if any fail.

# 1. Host-cmd reachable
python3 /maxtext-slurm/.host-cmd/host_cmd.py --timeout 10 "echo HOST_OK; whoami; hostname"

# 2. Resolve CLUSTER_SLOT to a fixed nodelist (write down NODELIST + NUM_NODES + PARTITION)
# (a) nodelist given: NODELIST=<as given>
# (b) partition+N given: pick N idle nodes from `sinfo -p <partition>`
# (c) just N: pick N idle nodes from `sinfo` on the default partition
# Once chosen, NODELIST is FROZEN for the entire sweep. Do not silently swap nodes
# mid-sweep — that invalidates every prior TGS measurement. Recovery options for an
# unhealthy node are in § Autonomous failure recovery below.
python3 /maxtext-slurm/.host-cmd/host_cmd.py --timeout 30 \
  "sinfo -p <PARTITION> -n <NODES> -o '%n %T %f %E' 2>&1"
# Expect every node in idle/mix/alloc. If any are drain/down/fail/maint, follow the
# node-recovery procedure in § Autonomous failure recovery before submitting anything.

# 3. Image tarball accessible (read DOCKER_IMAGE from container_env.sh)
python3 /maxtext-slurm/.host-cmd/host_cmd.py --timeout 10 "ls -la $(grep '^DOCKER_IMAGE' /maxtext-slurm/container_env.sh | cut -d'=' -f2 | tr -d '\"')"

# 4. Other users on the same partition (so submissions don't surprise-conflict)
python3 /maxtext-slurm/.host-cmd/host_cmd.py --timeout 30 "squeue -p <PARTITION> 2>&1 | head"

# 5. Per-model env override (may already contain flags; we'll need to know what's there)
cat /maxtext-slurm/configs/<MODEL_TAG>.env.sh 2>/dev/null || echo "(no per-model env file)"

# 6. train_env.sh state — should be clean (no leftover TUNE_PROFILE block)
grep -c TUNE_PROFILE /maxtext-slurm/train_env.sh   # expect 0; if not, read and decide

# 7. TG works (telegram skill).  Use the repo-relative path; the host-cmd
# runs commands with the maxtext-slurm checkout as cwd.
python3 /maxtext-slurm/.host-cmd/host_cmd.py --timeout 30 \
  "utils/telegram_bot.sh send '<MODEL_TAG> <PARALLELISM> XLA-tuning agent online; pre-flight pass'"

Step 2 — Inventory current state

Read these files and write a one-paragraph "current state" summary that becomes the doc's "Inventory" section:

Do not reference any other .md files in the repo root — they are project-specific and may not exist in other deployments.

Step 3 — Wave 1: baseline pair

The deployed env file may include flags tuned for a different parallelism. Measure both states for each CONFIG_TAG:

restore_default works via XLA's last-wins flag resolution: append --xla_<flag>=<image-default-value> to TUNE_FLAGS for every flag the env file sets. Sign flips between parallelisms are common — never assume an env-file recipe transfers from the parallelism it was tuned on.

Step 4 — Wave 1.5: capture evidence (mandatory)

Submit one baseline-deployed job with profiler + HLO dump. This is the single most valuable run of the sweep — it tells you what mechanism dominates this cell's step time, which prioritizes the candidate flag list.

RAY=1 ./submit.sh <MODEL_TAG>:evidence \
  --partition=<PARTITION> --nodes=<NUM_NODES> --nodelist=<NODELIST> --time=50:00 -- \
  per_device_batch_size=<pdbs> \
  <parallelism overrides; see Submission template> \
  <config-tag overrides; see Submission template> \
  profiler=xplane skip_first_n_steps_for_profiler=10 profiler_steps=3 \
  _env_ENABLE_XLA_DUMP=1 \
  jax_distributed_heartbeat_timeout_seconds=99999

When the job finishes, read:

Use what you find to prioritize the flag catalog below — don't blindly transcribe it. If the dominant cost is something the catalog doesn't address (e.g. grouped-GEMM bubble inside a custom kernel, kernel-launch latency on small messages, NCCL clique creation per microbatch), TG the user with the finding and propose new candidate flags before the canned waves.

Step 5 — Waves 2-4: flag testing

Use the flag catalog below as a starting candidate list, gated by your PARALLELISM. Submit 5-10 profiles per wave (slurm queue serializes them), then pause to review and decide direction. The full sweep is typically 25-50 profiles.

Sign flips between parallelisms are systematic — every flag that helps one parallelism may hurt another. Re-test fresh, even for flags that prior tuning on a different parallelism marked "negative". The only flags to skip are mechanism-irrelevant ones (e.g. pipelined_p2p is meaningless without collective-permute).

Submission template

RAY=1 ./submit.sh <MODEL_TAG>:<short_alias> \
  --partition=<PARTITION> --nodes=<NUM_NODES> --nodelist=<NODELIST> --time=45:00 -- \
  steps=15 dataset_type=synthetic \
  per_device_batch_size=<pdbs> \
  <parallelism overrides> \
  <config-tag overrides> \
  jax_distributed_heartbeat_timeout_seconds=99999 \
  _env_TUNE_PROFILE=<profile_name>

Always pass steps=15 dataset_type=synthetic on every tuning submission (including the Wave 1.5 evidence run and Wave 5/6 follow-ups), regardless of what configs/<MODEL_TAG>.gpu.yml has set. Two reasons:

1. Data-loader noise washes out the signal. TGS deltas of 1-3 % are what you're hunting; grain/c4/HF-tokenizer variance routinely adds 1-2 % per-step jitter that swamps the per-flag effect. Synthetic data is bit-deterministic and produces the lowest-variance TGS measurement available. (See tuning-runs are synthetic-only below for the operational rule.)
2. Tuning is short. 15 steps × ~30 s = ~7-8 min of training per job after compile. Long-step real-data probes belong in the on-demand loss test, not the XLA-flag tuning sweep — totally different question.

If configs/<MODEL_TAG>.gpu.yml has dataset_type: grain (e.g. left over from a loss test), the CLI passthrough dataset_type=synthetic steps=15 overrides it; you do not need to edit the yml. Grain-specific keys (grain_*, tokenizer_*, hf_access_token) are silently ignored when dataset_type=synthetic.

If configs/<MODEL_TAG>.gpu.yml has steps: <large_number> (e.g. steps: 2000 for a loss test), the CLI steps=15 likewise overrides it — no yml edit needed.

<parallelism overrides> derive from PARALLELISM:

<config-tag overrides> come from the model's gpu.yml or model-config-guide. Examples:

• MoE sparse_matmul + DeepEP: sparse_matmul=true use_turbo_grouped_gemm=true use_deepep_dispatch=true
• MoE dense_matmul + capacity-factor dropping: drop the three above (they default to false in gpu.yml)
• Dense (non-MoE): no MoE overrides

Keep the same --nodelist=… for every submission in the sweep. TGS noise from cross-nodelist hardware variance easily exceeds the 1-3 % flag deltas you're hunting. If a node goes unhealthy mid-sweep, follow Step 1 recovery; if it can't be revived, you must drop to N-1 nodes and re-baseline everything before continuing.

train_env.sh editing pattern

Add (at the end of train_env.sh):

# ---- TUNE_PROFILE ----
TUNE_PROFILE="${TUNE_PROFILE:-${EXTRACTED_ENV_MAP[TUNE_PROFILE]:-baseline}}"
TUNE_FLAGS=""
case "${TUNE_PROFILE}" in
    baseline) ;;  # no-op; inherits configs/<MODEL_TAG>.env.sh as-is
    restore_default)
        # Override every flag the env file sets back to its image-default value.
        # Use last-wins: append `--xla_<flag>=<image-default>` here for each one.
        TUNE_FLAGS+=" --xla_<flag>=<image-default-value>"
        ;;
    <profile_name>)
        TUNE_FLAGS+=" --xla_<flag>=<value>"
        ;;
    *)
        echo "[train_env.sh] WARNING: unknown TUNE_PROFILE='$TUNE_PROFILE'" >&2
        ;;
esac
if [[ -n "$TUNE_FLAGS" ]]; then
    XLA_FLAGS="${XLA_FLAGS:+$XLA_FLAGS }${TUNE_FLAGS# }"
    export XLA_FLAGS
fi

For NCCL knobs (and any other env vars), add a SECOND case block at the very end of train_env.sh, after the existing NCCL exports — otherwise they get overwritten:

# ---- TUNE_PROFILE: post-NCCL overrides ----
case "${TUNE_PROFILE:-}" in
    nccl_chan8) export NCCL_NCHANNELS_PER_NET_PEER=8 ;;
    mem95)      export XLA_PYTHON_CLIENT_MEM_FRACTION=.95 ;;
esac

Each submit.sh invocation freezes its own copy of train_env.sh into outputs/.artifacts/artifact_<id>/, so you can edit between submissions to add new profiles without affecting pending jobs (/maxtext-slurm/submit.sh lines 53-69 for the mechanism).

Steady-state TGS computation

Average Tokens/s/device over steps 9-14 (skip warmup steps 0-4, and skip steps 5-8 which are inside the profiler-capture window if profiler is on; step 8 is the profiler-dump step, a massive outlier):

import re, glob

def steady(jobid_glob, lo=9, hi=14):
    paths = sorted(glob.glob(f'/maxtext-slurm/outputs/{jobid_glob}-*.log'))
    if not paths: return None, None, None
    with open(paths[0]) as f: lines = f.readlines()
    steps = {}
    for l in lines:
        m = re.search(r'completed step: (\d+), seconds: ([\d.]+).*Tokens/s/device: ([\d.]+).*loss: ([\d.]+)', l)
        if m and (s := int(m.group(1))) not in steps:
            steps[s] = (float(m.group(2)), float(m.group(3)), float(m.group(4)))
    rng = sorted(s for s in steps if lo <= s <= hi)
    if not rng: return None, None, None
    return (sum(steps[s][0] for s in rng) / len(rng),         # mean step seconds
            sum(steps[s][1] for s in rng) / len(rng),         # mean Tokens/s/device
            steps[max(steps)][2])                              # last loss

Monitoring policy

Apply the Monitoring policy from batch-sweep for hang detection and progressive reporting. The tuning sweep is a 36-50-job batch — every job must be actively monitored (poll every 3-5 min via squeue + log tail). Failure handling has its own dedicated section: see § Autonomous failure recovery for the per-failure-class playbook (RCCL flake / JAX coord timeout / Unknown flag / OOM / cleanup exit=143 / NIC fault / drained node / etc.) and TG-stop triggers.

Two cases that are not in the recovery playbook because they're not failures, just timing:

Step 6 — Wave 5 (PP only): MaxText pipeline params

Skip if PARALLELISM doesn't include PP. Otherwise, after Waves 2-4 plateau, try the pipeline-shape knobs (these change the bubble fraction and per-microbatch overhead):

Bubble fraction = (num_stages − 1) / (num_microbatches × V + num_stages − 1). Use this to predict the trade-off before submitting.

Step 7 — Wave 6: cross-config-tag confirmation

If the model has multiple comparable CONFIG_TAGs, run the winning recipe (top 1-2 from Waves 2-4) on each of the others — one job each. Confirm the recipe doesn't regress them by more than 1 %. If it does, the deployment recommendation has to be config-conditional, not universal.

Step 8 — Output: <MODEL_TAG>-tuning.md

Create or extend <MODEL_TAG>-tuning.md in the repo root. Required sections:

1. Header bullets (date, model, hardware, image, MaxText branch, base config path, sequence length, steps, steady-state metric definition)
2. TL;DR — best recipe per (parallelism × config) cell; comparison vs baseline-deployed; one-sentence "the winning flag wins because "
3. Inventory (from Step 2)
4. Effect of <MODEL_TAG>.env.sh on <PARALLELISM> (Wave 1 result)
5. Tuning leaderboard (top 8 + bottom 5; Profile | Flags | TGS | step (s) | Δ% vs baseline-deployed; mark winner ⭐)
6. Why the winning flag wins (mechanism, backed by Wave 1.5 evidence — cite job IDs and HLO/xplane artifacts)
7. Negative findings (which flags hurt; useful to skip on future sweeps)
8. Sign flips vs other parallelisms (if <MODEL_TAG>-tuning.md already covers another PARALLELISM, add a flag-by-flag delta table; otherwise omit)
9. Recommended deployment (proposed configs/<MODEL_TAG>.env.sh patch; if the winning recipe conflicts with a prior parallelism's recipe, propose a guard like if [[ "${MAXTEXT_DCN_PP:-1}" -le 1 ]])
10. Appendix: data sources (job IDs, profiles, status — full table sorted by job ID)

If a Chinese sibling <MODEL_TAG>-tuning.zh.md exists or is requested, mirror the structure (same headers, same tables, code/flag names stay in English, prose translated).

Per-parallelism flag catalog

These are starting candidates. Re-prioritize and extend based on Wave 1.5 trace evidence.

FSDP-heavy (DCN ring all-gather + reduce-scatter dominate)

PP-heavy (collective-permute + per-stage send/recv dominate)

TP-heavy (ICI all-gather / reduce-scatter on sharded tensors)

Hybrid (e.g. PP×FSDP, FSDP×TP)

Read the Wave 1.5 step time breakdown to identify the dominant axis. Start with that axis's catalog. After 2-3 waves on the dominant axis, add a few flags from the secondary axis at the end. Don't sweep both axes' full catalogs — exponential blow-up.

Mechanism cheatsheet (orientation for "Why the winning flag wins" section)

Pipeline bubble (PP only) is mechanically determined: (num_stages − 1) / (num_microbatches × V + num_stages − 1). No XLA flag changes this — only MaxText pipeline params do (Wave 5).

Autonomous failure recovery

A 36-50-job sweep encounters node failures, RCCL flakes, transient JAX coordination errors, NIC bouncing, and similar issues many times. Recover autonomously when possible; only TG-stop the user when the issue is genuinely beyond the agent's reach. The user is away during the sweep — every interruption costs human attention.

The hard rule is: do not change the nodelist mid-sweep. Cross-nodelist hardware variance (rail layout, IB port wear, ROCm driver skew between nodes) introduces 1-3 % TGS noise that exactly drowns the flag-effect signal you're measuring. Falling back to a different nodelist invalidates the entire leaderboard. If a node from your pinned set cannot be revived after the recovery steps below, TG and stop — do not silently substitute or drop nodes.

Known false-positive log signatures (filter these out before classifying)

These appear in every successful run on this stack. Do NOT treat them as failures, do NOT cancel, do NOT retry:

When triaging a hung or crashed job, grep -v these patterns out first before reading the tail — the actual signal-to-noise on this stack is poor and these warnings can scroll for thousands of lines around the real error.

Per-failure recovery playbook

Recovery primitives via host-cmd

# Resume a drained/down node
python3 /maxtext-slurm/.host-cmd/host_cmd.py --timeout 30 \
  "scontrol update NodeName=<n> State=RESUME Reason=manual; sleep 2; sinfo -n <n>"

# Inspect a node's state in detail
python3 /maxtext-slurm/.host-cmd/host_cmd.py --timeout 30 \
  "scontrol show node <n>"

# Reset ionic NIC stack (safe between jobs; NOT during a running job)
python3 /maxtext-slurm/.host-cmd/host_cmd.py --timeout 60 \
  "ssh <n> 'rmmod ionic_rdma; rmmod ionic; sleep 3; modprobe ionic; modprobe ionic_rdma; sleep 5; ibstat | grep -c \"State: Active\"'"

# Pull recent dmesg around a failure timestamp
python3 /maxtext-slurm/.host-cmd/host_cmd.py --timeout 30 \
  "ssh <n> 'dmesg -T | tail -100 | grep -iE \"oom|kill|fault|reset|amdgpu|hsa|ionic\"'"

# Check GPU temperature / fault state on a node
python3 /maxtext-slurm/.host-cmd/host_cmd.py --timeout 30 \
  "ssh <n> 'rocm-smi --showtemp; rocm-smi --showuse; rocm-smi --showxgmierr 2>/dev/null'"

# Verify all NICs back to Active after a reset
python3 /maxtext-slurm/.host-cmd/host_cmd.py --timeout 30 \
  "ssh <n> 'ibstat | grep -E \"^CA |State:\" | head'"

TG-stop format

When a stop trigger fires, send a single TG message and halt the sweep. The message must include enough context for the user to either fix the cluster issue or decide to resume on a different nodelist:

[xla-tuning STOP] <MODEL_TAG> <PARALLELISM> sweep halted

Reason: <one sentence — node fault / coord timeout / controller down / etc.>
Failing node(s): <hostname[s]>; recovery attempted: <what you tried>; outcome: <what didn't work>

Sweep state at stop:
  - Nodelist: <pinned nodelist>
  - Profiles tested so far: N (top 3 below)
  - Best so far: <profile> @ <TGS> (Δ <%> vs baseline-deployed)
  - <profile> @ <TGS>
  - <profile> @ <TGS>
  - Pending profiles in queue: <list, cancelled by stop>

Artifacts: outputs/<jobid_lo>-* through outputs/<jobid_hi>-*
Awaiting your decision before resuming.

After sending, do not resubmit anything until the user replies. Spend the wait writing partial results into <MODEL_TAG>-tuning.md so the work isn't lost — the leaderboard so far + the stop reason + the artifacts is publishable as a "sweep interrupted at profiles" preliminary report.

Tuning runs are synthetic-only

Every tuning submission must pass steps=15 dataset_type=synthetic, with no exceptions. This is the operational rule that makes the leaderboard interpretable:

The only legitimate reason to use real data in this skill's flow is if the user explicitly asks for a post-tuning loss validation of a winning recipe — that's a separate one-shot run after the sweep, not part of the sweep.

Common pitfalls

• Always pin --nodelist=… for every submission. TGS deltas of 1-3 % are exactly what you're hunting. Hardware variance between nodes (rail layout, IB port wear, ROCm version skew) easily exceeds that. The nodelist is frozen for the entire sweep — recovery options for unhealthy nodes are in § Autonomous failure recovery, not a different nodelist.
• Don't compose XLA_FLAGS from scratch. APPEND your experimental flags via XLA_FLAGS="${XLA_FLAGS:+$XLA_FLAGS }<flag>". The image's compiled-in defaults must be preserved — many old flags trigger Unknown flag aborts in current XLA.
• For every transient failure (RCCL flake, JAX coord timeout, NIC fault, drained node, GitHub 5xx), follow § Autonomous failure recovery — try the recovery sequence, then TG-stop only when the listed stop-trigger fires. The user is away during the sweep; do not interrupt them for routine flakes.
• OOM is never retried. Same flags = same OOM. Drop the profile, record allocation size.
• Cleanup exit=143 with completed step: N-1 present is a benign post-train artifact, especially common on PP. Trust the completed step: lines, not the cleanup status.
• Profiler-on runs perturb steps 5-8. If profiler is on, compute steady-state over steps 9-14 only. If profiler is off, steps 5-14 is fine.
• Sign flips between parallelisms are systematic. Don't assume FSDP findings transfer to PP, or vice versa. Test fresh.
• JOB_NAME ≤ 243 bytes — ext4's 255-byte path-segment limit minus 12 bytes for the jobid prefix. Keep TUNE_PROFILE aliases short (pp8-A, Gco1G, etc.); rely on _env_TUNE_PROFILE=… not the alias to disambiguate.
• profiler= (empty) breaks the YAML enum validator. If you want clean runs, leave the YAML's profiler: "" value and don't pass profiler= as a passthrough flag.

Stop conditions

• A single flag (or 2-3 flag stack) gives ≥ +5 % over baseline-deployed AND doesn't regress other config tags by more than 1 %. → Document, propose deployment patch.
• Or: tested ≥ 25 distinct profiles across waves 2-4 and none beats baseline by > 2 %. → Conclude "structurally bounded at the editable scope". Document the negative finding with HLO + xplane evidence backing the conclusion.
• Either outcome is a successful run. The negative finding is just as valuable — it closes the question with a defensible answer.

Pacing

If still searching at 8 hours of queue time, write up what you have and conclude — the practical ceiling is hit.

Related skills

• batch-sweep — find optimal per_device_batch_size for a (model × parallelism) cell. Run first if pdbs hasn't been tuned for this cell yet — XLA tuning is on top of an already-good batch size.
• profile-drill — kernel-level breakdown of xplane traces; use in Wave 1.5
• job-log-triage — classify hangs, RCCL flakes, OOM signatures
• model-config-guide — for hybrid parallelism choices and <MODEL_TAG>.gpu.yml overrides
• telegram — TG progress at each milestone (kickoff, every positive result, mid-point, final)