--- name: xla-tuning description: "Find the XLA flag / NCCL env-var combination that maximizes steady-state TGS for one (model × parallelism) cell. Produces an evidence-backed leaderboard, mechanistic explanation of the winning flag, and a deployment recipe. Use when the user asks to tune XLA flags, tune NCCL, find best collective-permute / all-gather threshold, optimize FSDP/PP/TP, close a parallelism-vs-parallelism throughput gap, or sweep cross-iteration prefetch / overlap-limit / async-stream-priority knobs for a specific model." --- # 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: `-tuning.md` in the repo root (and `-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/.gpu.yml` exists and runs at the requested parallelism without OOM - A working `pdbs` for the cell (run [batch-sweep](../batch-sweep/SKILL.md) 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](../telegram/SKILL.md)) for progress updates ## Inputs (ask the user; never guess) The user must provide: 1. **`MODEL_TAG`** — must match `configs/.gpu.yml`. Examples: `deepseek3-671b`, `qwen3-32b`, `llama3.1-405b`, `mixtral-8x22b`. 2. **`CLUSTER_SLOT`** — one of: - explicit `nodelist=` *(recommended; pinned hardware drops TGS noise to ~0.3-1 % which is what you need to detect 1-3 % flag deltas)* - `partition=+nodes=` *(let slurm pick from the partition pool)* - `nodes=` 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/.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 -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 -tuning.md (and .zh.md mirror if needed) ``` ### Step 1 — Pre-flight Run all of these. Stop and TG-ask if any fail. ```bash # 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= # (b) partition+N given: pick N idle nodes from `sinfo -p ` # (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 -n -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 2>&1 | head" # 5. Per-model env override (may already contain flags; we'll need to know what's there) cat /maxtext-slurm/configs/.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 ' 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: | File | What to extract | |---|---| | `configs/.gpu.yml` | Default parallelism, `per_device_batch_size`, `max_target_length`, `remat_policy`, MoE flags (`sparse_matmul`, `use_turbo_grouped_gemm`, `use_deepep_dispatch`, `capacity_factor`), pipeline params (`pipeline_parallel_layers`, `num_pipeline_microbatches`) | | `configs/.env.sh` (if exists) | Every `XLA_FLAGS` / `NCCL_*` / env override currently deployed for this model — these are *inherited* by all your jobs unless you override with `TUNE_FLAGS` | | `container_env.sh` | `DOCKER_IMAGE`, `MAXTEXT_PATCH_BRANCH`, `MAXTEXT_REPO_DIR` | | `-tuning.md` (if exists) | Prior best recipes per (parallelism × config) cell, prior baseline TGS, prior structural-cost decomposition. If absent, declare "fresh start" — Wave 1.5 will derive it. | **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`: | Profile | TUNE_PROFILE | What it measures | |---|---|---| | `baseline-deployed` | empty (inherits env file as-is) | Production state — *all later Δ% are vs this* | | `baseline-no-env` | `restore_default` (overrides env-file flags back to image defaults) | What the image alone gives — Δ vs `baseline-deployed` tells you whether the deployed flags help, hurt, or are neutral on this parallelism | `restore_default` works via XLA's last-wins flag resolution: append `--xla_=` 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. ```bash RAY=1 ./submit.sh :evidence \ --partition= --nodes= --nodelist= --time=50:00 -- \ per_device_batch_size= \ \ \ 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: | Artifact | What to look for | |---|---| | `outputs/-…//tensorboard/.../plugins/profile/.../*.xplane.pb` | Per-stream timeline. Run [profile-drill](../profile-drill/SKILL.md) for kernel breakdown. **What fraction of step time is in collectives vs compute vs idle bubble?** | | `outputs/-…/xla_dump/module_*.jit_train_step.*_after_optimizations.txt` | Compiled HLO. Search `collective-permute-start`, `all-gather-start`, `reduce-scatter-start` — what's their `frontend_attributes`? Is `is_pipelined` true? Is `is_sync` true? Is the `dependency_set` empty (sync) or non-trivial (async)? | | Same dump file | Search for `while` loops — are paired `*-start` / `*-done` ops inside the same iteration (sync), or paired across iterations (pipelined)? | | Same dump file | Search the inherited XLA_FLAGS line at the top of the log to enumerate the image's compiled-in defaults; many old flags are obsolete and trigger `Unknown flag` aborts | 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](#per-parallelism-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 ```bash RAY=1 ./submit.sh : \ --partition= --nodes= --nodelist= --time=45:00 -- \ steps=15 dataset_type=synthetic \ per_device_batch_size= \ \ \ jax_distributed_heartbeat_timeout_seconds=99999 \ _env_TUNE_PROFILE= ``` **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/.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](#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/.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/.gpu.yml` has `steps: ` (e.g. `steps: 2000` for a loss test), the CLI `steps=15` likewise overrides it — no yml edit needed. `` derive from `PARALLELISM`: | `PARALLELISM` | Passthrough overrides (override gpu.yml's defaults) | |---|---| | `FSDP=N` | none if `N` matches `dcn_fsdp_parallelism` in gpu.yml; else `dcn_fsdp_parallelism= dcn_pipeline_parallelism=1` | | `PP=N` | `dcn_pipeline_parallelism= dcn_fsdp_parallelism=1` plus pipeline params from gpu.yml (`pipeline_parallel_layers`, `num_layers_per_pipeline_stage`, `num_pipeline_microbatches` may need explicit values if the gpu.yml didn't pre-fill them) | | `TP=N` | `ici_tensor_parallelism=` (and likely reduce another ici axis to compensate) | | Hybrid `PP=N,FSDP=M` | both axes; ensure `N × M = num_dcn_replicas` | `` come from the model's gpu.yml or [model-config-guide](../model-config-guide/SKILL.md). 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`): ```bash # ---- TUNE_PROFILE ---- TUNE_PROFILE="${TUNE_PROFILE:-${EXTRACTED_ENV_MAP[TUNE_PROFILE]:-baseline}}" TUNE_FLAGS="" case "${TUNE_PROFILE}" in baseline) ;; # no-op; inherits configs/.env.sh as-is restore_default) # Override every flag the env file sets back to its image-default value. # Use last-wins: append `--xla_=` here for each one. TUNE_FLAGS+=" --xla_=" ;; ) TUNE_FLAGS+=" --xla_=" ;; *) 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: ```bash # ---- 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_/`, 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): ```python 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](../batch-sweep/SKILL.md#monitoring-policy-applies-to-all-sweeps) 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](#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: | Log signature | Diagnosis | Action | |---|---|---| | HLO compilation taking > 15 min on first job after image load | normal for large MoE on cold cache | wait | | HLO compilation > 15 min on later job with same flag set | suspicious — may have hit a compiler bug | cancel, try a slight flag perturbation; if reproducible, mark and skip | ### 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): | Override | Effect | |---|---| | `num_pipeline_microbatches=2N` (e.g. 8 → 16) | halves bubble; doubles per-microbatch comm/MoE round-trip count | | `num_pipeline_microbatches=N/2` | halves per-microbatch overhead; doubles bubble | | `num_layers_per_pipeline_stage=K` (vary K) | changes V chunks per stage; changes bubble divisor | 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_TAG`s, 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: `-tuning.md` Create or extend `-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 `.env.sh` on ``** (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 `-tuning.md` already covers another `PARALLELISM`, add a flag-by-flag delta table; otherwise omit) 9. **Recommended deployment** (proposed `configs/.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 `-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) | Flag | Why it matters here | |---|---| | `--xla_gpu_all_gather_combine_threshold_bytes=` | image default is often 8 GiB which fuses the per-step all-gather into one serial barrier. Sweep 256 MiB / 512 MiB / 1 GiB / 2 GiB; both ag-only AND ag+rs together | | `--xla_gpu_reduce_scatter_combine_threshold_bytes=` | usually leave at default — backward rs is already large per-layer; splitting just adds RCCL launch overhead | | `--xla_gpu_enable_pipelined_all_gather/reduce_scatter/all_reduce=true` | cross-iteration prefetch — usually OOMs on this scale | | `--xla_gpu_enable_while_loop_double_buffering=true` | cross-iteration overlap of `train_step` `while` body — usually negative on FSDP (HBM cost > overlap gain) | | `--xla_gpu_enable_highest_priority_async_stream=true` | image default usually already prioritises async stream sufficiently | | `--xla_gpu_experimental_parallel_collective_overlap_limit=2/4/8` | usually negative on FSDP — ring all-gather already saturates fabric, extra concurrency = contention | | `--xla_gpu_enable_latency_hiding_scheduler=true` | image default may already have it on; check the inherited XLA_FLAGS in the baseline log | | `NCCL_NCHANNELS_PER_NET_PEER=8` | sometimes +1-4 %; sometimes triggers init hangs in combos — see retry policy | | `NCCL_BUFFSIZE=16M` | per-rank workspace — usually neutral or negative | | `NCCL_PROTO=Simple/LL/LL128` | algorithm tweak; usually neutral on large FSDP messages | | `XLA_PYTHON_CLIENT_MEM_FRACTION=.95` | more HBM headroom for prefetch buffers — usually within noise | ### PP-heavy (collective-permute + per-stage send/recv dominate) | Flag | Why it matters here | |---|---| | `--xla_gpu_collective_permute_decomposer_threshold=` | controls when `collective-permute` decomposes into send/recv. Sweep 256 MiB / 1 GiB / 8 GiB; the c-p ops in your HLO have a fixed size (e.g. `bf16[B,V,L,H]`) — find one bracketing threshold and one not-bracketing | | `--xla_gpu_enable_pipelined_p2p=true` | cross-iteration prefetch on `collective-permute`. Usually no-op when the pipeline is `nn.scan`-based with a hard loop carry; verify in HLO | | `--xla_gpu_enable_async_collective_permute=true` (if exists in your XLA build) | makes c-p async-launchable | | `--xla_gpu_experimental_parallel_collective_overlap_limit=2/4/8` | often POSITIVE on PP — multiple independent fabrics (ICI for FSDP-style intra-stage ag, DCN p2p for c-p) enable real concurrency. Sweet spot is usually = number of independent fabrics involved (often 2) | | `--xla_gpu_enable_highest_priority_async_stream=true` | helps when MoE skew creates per-rank stragglers (sparse path) | | `--xla_gpu_enable_while_loop_double_buffering=true` | PP loop carry has different memory headroom; re-test fresh — image-default off-by-default | | `--xla_gpu_enable_pipelined_all_reduce/reduce_scatter/all_gather=true` | per-stage HBM may absorb the prefetch buffers; test (often still OOMs) | | `NCCL_ALGO=Ring/Tree/CollnetDirect` | per-stage 2-rank traffic prefers different algos than ring-based FSDP | | `NCCL_PROTO=LL/LL128/Simple` | low-latency protocol may help small per-stage messages | | `NCCL_NCHANNELS_PER_NET_PEER=2/4/8` | extra channels rarely help 2-rank c-p; test small variations | ### TP-heavy (ICI all-gather / reduce-scatter on sharded tensors) | Flag | Why it matters here | |---|---| | `--xla_gpu_threshold_for_windowed_einsum_mib=` | enables async TP via windowing — primary lever for TP overlap | | `--xla_gpu_enable_async_all_gather=true` (if exists) | TP-style intra-step ag | | `--xla_gpu_collective_inflation_factor=` | controls collective sizing | | `NCCL_NCHANNELS_PER_NET_PEER=` | for in-node hops (most TP traffic is intra-node) | | (most cross-iteration prefetch flags are not relevant to TP — TP doesn't have the iteration-spanning carry that PP has) | | ### 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) | Parallelism | Likely dominant cost | What XLA flags can address | |---|---|---| | **FSDP** | DCN ring all-gather of full per-step weights; backward reduce-scatter of full per-step gradients | break the mega-fused all-gather (combiner threshold); rarely cross-iteration prefetch (OOM) | | **PP** | Per-stage `collective-permute` rendezvous (sync per call); MoE `dispatch+combine` per microbatch; pipeline bubble | concurrent collective execution (overlap_limit on dual-fabric ICI+DCN); decomposer threshold; async-stream priority for skew amelioration | | **TP** | ICI all-gather/reduce-scatter on sharded weight/activation tensors | windowed einsum threshold; async ag/rs on intra-node fabric | | **Hybrid** | Mix; check Wave 1.5 trace for which collective dominates | start from dominant-axis catalog | 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: | Signature | Where | Why it's benign | |---|---|---| | `failed call to cuInit: INTERNAL: CUDA error: Failed call to cuInit: UNKNOWN ERROR (303)` | every rank's worker stderr, dozens of times during Ray actor / JAX backend init | The JAX CUDA backend probes for a CUDA device on every rank even though we run on ROCm. The probe fails (no CUDA), JAX falls back to ROCm, training proceeds normally. Appears bit-identical in successful runs. | | `NCCL WARN MSCCL++: Feature not enabled. ENABLE_MSCCLPP must be defined at compile-time` | every rank during RCCL clique init | The image isn't compiled with MSCCL++; RCCL uses its default algos. No training impact. | | `NCCL WARN LL cutoff points not detected for a supported arch gfx950` (`rccl_wrap.cc`) | every rank during RCCL init | RCCL falls back to default LL thresholds. No training impact. | | `WARNING: AMD GPU device(s) is/are in a low-power state` (rocm-smi during pre-flight) | pre-flight rocm-smi output when GPUs are idle | GPUs leave low-power state automatically when training starts. Not a fault. | 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 | Failure class | Symptom | Recovery sequence | Stop trigger | |---|---|---|---| | **RCCL init flake** | >10 min wall, no `completed step:`, NCCL WARN spam scrolling, `init.cc` / `rccl_wrap.cc` lines | `scancel ` → wait for COMPLETING → resubmit same TUNE_PROFILE on same nodelist | 3 consecutive flakes on the same TUNE_PROFILE → mark `infra-flake⋆`, drop the profile, continue with the rest | | **JAX `DEADLINE_EXCEEDED: GetKeyValue() timed out`** | mid-training `RpcError` or `Coordination timeout`; rank 0 actor disappears with no log signature | Same: `scancel` → resubmit. Often clears on retry. | 2 consecutive timeouts on the same TUNE_PROFILE → cluster coord service may be sick → **TG-stop** with last-known leaderboard | | **`Unknown flag in XLA_FLAGS`** | `parse_flags_from_env.cc:.*Unknown flag: --xla_…` then immediate abort | Open `train_env.sh`, remove the offending flag from the TUNE_PROFILE block, save → resubmit. Pending jobs already have their artifact frozen, so other profiles in the queue are unaffected. | typo on first-ever submission of a profile → fix once and continue; same flag fails after fix → flag is genuinely obsolete in this XLA build → drop the profile | | **OOM** (`RESOURCE_EXHAUSTED: Out of memory while trying to allocate N GiB`) | clean OOM at compile or step 0 | **No retry.** Same flags = same OOM. Record `Total memory: X GiB / Temp: Y GiB` and the requested allocation in the doc. Drop the profile. | n/a — never a stop trigger; just continue | | **Cleanup `exit=143` with `completed step: N-1` present** | training succeeded, post-train Docker cleanup race kills the process tree | **Treat as success.** Extract TGS from `completed step:` lines; ignore the cleanup status. | n/a | | **Single node `drain` / `down` / `fail` state** in `sinfo` | a node from the pinned set is offline | (1) `scontrol update NodeName= State=RESUME Reason=manual` via host-cmd → re-check `sinfo`. (2) If still drained: `scontrol show node ` to read the reason; common recoverable reasons are "Not responding", "Kill task failed", stale-drain-from-previous-job → `RESUME` usually clears them. (3) If `RESUME` doesn't take, ssh into the node and check `dmesg -T \| tail -50` and `rocm-smi` for fault signatures. | node stays drained / down after `RESUME` × 2 attempts → **TG-stop** | | **`ionic_comp_msn cqe with error` / NIC fault on one node** | RDMA completion-queue error; job hangs at first collective | (1) Identify the failing node (rank index from log → look up `JOB_NODELIST` to get hostname). (2) Reset the ionic stack via host-cmd: `ssh 'rmmod ionic_rdma; rmmod ionic; sleep 3; modprobe ionic; modprobe ionic_rdma'` (note: depends on `ionic_rdma` having no in-flight users — safe between jobs). (3) Verify `ssh 'ibstat \| grep -c "State: Active"'` returns 8 (or whatever the per-node NIC count is on this hardware). (4) Resubmit. | NIC reset doesn't restore Active state → **TG-stop** | | **`HSA_STATUS_ERROR` / `rocdevice.cpp: Aborting`** | GPU runtime error mid-training | Cancel; `ssh 'rocm-smi --showtemp \| head'` to check thermal state; `dmesg -T \| grep -iE "amdgpu\|hsa\|gpu hang"` to look for hardware fault signatures. If it's a thermal issue or transient firmware glitch, wait 1-2 min, retry. | hardware fault signature in dmesg (XID-equivalent ECC, GPU reset, fabric link down) → **TG-stop** with the dmesg excerpt | | **`NodeFail` event mid-job** | node drops out during training | Same as "Single node drain/down" above — try `RESUME`. | if multiple nodes fail in the same wave, the cluster has a wider problem → **TG-stop** | | **Slurm `cgroup OOM` killing the entire docker container** with no JAX-side error | host-side memory pressure (rare on training jobs but possible if too many grain workers fork) | Read `journalctl -k --since '