You signed in with another tab or window. Reload to refresh your session.You signed out in another tab or window. Reload to refresh your session.You switched accounts on another tab or window. Reload to refresh your session.Dismiss alert
{{ message }}
Repository navigation
Commit ebb93b2
Browse filesBrowse the repository at this point in the historyBrowse files
authored
Make Triton autotune configuration selection deterministic (#6076)
Signed-off-by: Yan Xu <yxu1@nvidia.com>
Co-authored-by: Charlie Truong <chtruong@nvidia.com>
Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
| Mamba/SSM Triton ops |`megatron/core/ssm/ops/common/determinism.py`| one fixed autotune config plus a zero-initialized tiled workspace reduced with an ordered `sum`| timing-based autotune, uninitialized workspace |
41
+
| Mamba/SSM Triton ops |`megatron/core/tuning` (config selection), `megatron/core/ssm/ops/common/determinism.py` (workspaces)| one fixed autotune config plus a zero-initialized tiled workspace reduced with an ordered `sum`| timing-based autotune, uninitialized workspace |
42
42
| Mamba/GDP causal conv1d |`megatron/core/ssm/causal_conv1d.py`| causal_conv1d >= 1.6.0 — per-block workspace for the weight and bias gradients, reduced with an ordered `sum`|`atomicAdd` accumulation (order varies per launch) |
43
43
| Transformer Engine attention |`megatron/core/extensions/transformer_engine.py`| requires `NVTE_ALLOW_NONDETERMINISTIC_ALGO=0`, under which TE selects only backends that support deterministic execution (including deterministic FlashAttention backward) | TE picks freely, including atomic-accumulation attention backward |
44
44
| Inference DP scheduling and RL rollout order |`megatron/core/inference/engines/dynamic_engine.py`, `megatron/rl/rl_utils.py`| sort by stable key | completion order |
Copy file name to clipboardExpand all lines: docs/user-guide/deterministic-training.md
+9-12Lines changed: 9 additions & 12 deletions
Display the source diff
Display the rich diff
Original file line number
Diff line number
Diff line change
@@ -23,17 +23,17 @@ When enabled, Megatron applies the env vars and config overrides below via `mega
23
23
24
24
## Environment variables
25
25
26
-
Each variable may be set by the launcher or left unset. If set, the value must be one that has been validated as deterministic — anything else fails hard with an assertion. If unset, `apply_determinism_env` fills the canonical default (except `MAMBA_DETERMINISTIC` and `CAUSAL_CONV1D_DETERMINISTIC`, which their kernels auto-detect from `torch.are_deterministic_algorithms_enabled()`). Must be set before the first cuBLAS / Transformer Engine call — `apply_determinism_to_args` runs early in `validate_args` to guarantee this.
26
+
Each variable may be set by the launcher or left unset. If set, the value must be one that has been validated as deterministic — anything else fails hard with an assertion. If unset, `apply_determinism_env` fills the canonical default (except `MAMBA_DETERMINISTIC` and `CAUSAL_CONV1D_DETERMINISTIC`, which the external `mamba_ssm` and `causal_conv1d` packages auto-detect from `torch.are_deterministic_algorithms_enabled()`). Must be set before the first cuBLAS / Transformer Engine call — `apply_determinism_to_args` runs early in `validate_args` to guarantee this.
|`NCCL_ALGO`| subset of `{Ring, CollnetDirect, CollnetChain, ^NVLS}`|`Ring`| Conservative default — `Ring`'s reduction order is fixed by topology, so it is bit-exact across runs on every supported NCCL version |
31
31
|`NVTE_ALLOW_NONDETERMINISTIC_ALGO`|`0`|`0`| Forces Transformer Engine to use deterministic algorithms |
32
32
|`CUBLAS_WORKSPACE_CONFIG`|`:4096:8` or `:16:8`|`:4096:8`| Deterministic cuBLAS workspace (both sizes are reproducible per NVIDIA docs; `:4096:8` is faster, `:16:8` uses less memory) |
33
-
|`TRITON_CACHE_AUTOTUNING`|`0` or `1`|*(none — opt-in)*| Persists each Triton autotune winner so every rank reuses one choice instead of re-timing it. Left unset, deterministic mode instead pins the cheapest config, which needs no cache — see [Triton autotuning](#triton-autotuning)|
33
+
|`TRITON_CACHE_AUTOTUNING`|`0` or `1`|*(none — opt-in)*| Persists Triton autotune winners for autotuners **outside** the pinned scope, so ranks reuse one choice instead of re-timing it. Kernels inside the scope are pinned without timing either way — see [Triton autotuning](#triton-autotuning)|
34
34
|`TRITON_CACHE_DIR`| any shared-filesystem path |*(none — required only with `TRITON_CACHE_AUTOTUNING=1`)*| No safe default exists: unset, Triton uses a node-local directory and each node autotunes on its own. Required rather than filled in |
35
-
|`TRITON_PRINT_AUTOTUNING`|`1`|*(none — recommended, not set)*| Logs the config each rank selected. Changes no numerics, so it is recommended rather than forced; when `TRITON_CACHE_AUTOTUNING=1` and this is unset, a startup line reminds you. See[Verifying kernel-config agreement](#verifying-kernel-config-agreement)|
36
-
|`MAMBA_DETERMINISTIC`| any string starting with `'1'`|*(none — SSM auto-detects)*|Mamba SSM auto-follows `torch.are_deterministic_algorithms_enabled()` when unset; only an explicit non-deterministic override is rejected |
35
+
|`TRITON_PRINT_AUTOTUNING`|`1`|*(none — not set)*| Logs the config Triton times, which only happens for autotuners outside the pinned scope. For pinned kernels use `--triton-autotune-enumerate-autotuners` and `--triton-autotune-verify-every` — see[Verifying kernel-config agreement](#verifying-kernel-config-agreement)|
36
+
|`MAMBA_DETERMINISTIC`| any string starting with `'1'`|*(none — SSM auto-detects)*|Controls the external `mamba_ssm` package, which auto-follows `torch.are_deterministic_algorithms_enabled()` when unset; only an explicit non-deterministic override is rejected. Megatron's in-tree SSM kernels (`megatron.core.ssm.ops`) do not read it and follow `--deterministic-mode` instead|
37
37
|`CAUSAL_CONV1D_DETERMINISTIC`| any string starting with `'1'`|*(none — the kernel auto-detects)*| causal_conv1d ≥ 1.6.0 auto-follows `torch.are_deterministic_algorithms_enabled()` when unset, reducing the conv weight/bias gradients through a workspace instead of `atomicAdd`; the Mamba and GDP mixers reject a deterministic run without it |
38
38
39
39
If you override `NCCL_ALGO`, the value must be a subset of `{Ring, CollnetDirect, CollnetChain, ^NVLS}`. `Tree` is intentionally excluded: its intra-node chain reduction order is not user-controllable, and the inter-node tree topology can vary across runs without a pinned topology file, so it cannot be vouched for as bit-exact across stacks. `^NVLS` is accepted (banning NVLS is a legitimate user choice on hardware that exposes it); the user is responsible for ensuring whatever NCCL falls back to is deterministic on their environment.
Triton picks a kernel config by timing its candidates, so the winner depends on the machine at that instant and ranks can disagree. Deterministic mode offers two ways to remove that variance:
72
+
Triton picks a kernel config by timing its candidates, so the winner depends on the machine at that instant and ranks can disagree. `--deterministic-mode` pins it: the `megatron.core.tuning` adapter selects one candidate per kernel and shape without timing anything, from a tuned table when one matches (tables ship for `sm100` and `sm103`, covering the `mamba_ssm` SSD kernels and the in-tree MLA RoPE and mHC kernels, or can be recorded with `--triton-autotune-mode record`), otherwise from the cheapest candidate by a pure function of the candidate list. Pinned kernels never consult Triton's autotune cache, so every rank computes the same answer by construction. The cheapest candidate is not necessarily the fastest one; a recorded table recovers the throughput.
73
73
74
-
| Strategy | How to select it | Determinism rests on |
75
-
|---|---|---|
76
-
|**Pinned config** (default) | leave `TRITON_CACHE_AUTOTUNING` unset | Nothing external. `autotune_configs` picks the cheapest config by a pure function of the candidate list, so every rank computes the same answer without timing anything. Slower, since the pinned config is not necessarily the fastest one. |
77
-
|**Cached autotuning**|`TRITON_CACHE_AUTOTUNING=1`**and**`TRITON_CACHE_DIR=<shared path>`| Every rank reading one warm cache. Autotuning still runs and still picks fast configs, but a rank that misses the cache re-times the selection on its own and can pick differently. |
74
+
The default scope is `mamba_ssm`, `transformer_engine` and `megatron.core` (`--triton-autotune-modules`). Kernels whose outputs do not depend on the config — by default only pure data movement, Transformer Engine's MoE permutation kernels, listed in `AutotunePolicy.config_invariant` (`--triton-autotune-config-invariant`) — keep Triton's timed choice, since pinning them would only cost throughput. See `megatron/core/tuning/README.md` for recording tables, precedence rules and the full option list.
78
75
79
-
Cached autotuning is opt-in because its determinism is conditional: the pinned default holds by construction, the cached path holds only while the shared cache does. Setting `TRITON_CACHE_AUTOTUNING=1` without `TRITON_CACHE_DIR` is rejected — unset, Triton falls back to a node-local directory, which is exactly the case the cache is meant to prevent.
76
+
Autotuners outside the scope still time their candidates. `TRITON_CACHE_AUTOTUNING=1` with a shared `TRITON_CACHE_DIR` makes those ranks reuse one cached winner, but a rank that misses the cache re-times the selection on its own and can pick differently, so prefer adding their modules to `--triton-autotune-modules`. Setting `TRITON_CACHE_AUTOTUNING=1` without `TRITON_CACHE_DIR` is rejected — unset, Triton falls back to a node-local directory, which is exactly the case the cache is meant to prevent.
80
77
81
78
## Verifying kernel-config agreement
82
79
83
-
Applies to cached autotuning; the pinned default has nothing to compare. `TRITON_PRINT_AUTOTUNING=1` makes each rank log the config it selects per kernel; group those lines by kernel and key across the per-rank logs, and every group should hold exactly one distinct config.
80
+
`--triton-autotune-verify-every N` compares, every N training steps, the configs that ranks chose for the same kernel and shape, and logs any disagreement (`--triton-autotune-verify-strict` raises instead). Each check exchanges only the choices made since the previous one. `--triton-autotune-enumerate-autotuners` logs, per rank, every multi-config autotuner the run reaches and whether it is pinned, timed as config-invariant, or outside the scope.
84
81
85
-
Note the limit: a rank only logs when it *tunes*, so a run where some ranks hit the cache and others miss cannot be compared this way — the hitting ranks log nothing.
82
+
For autotuners outside the scope, `TRITON_PRINT_AUTOTUNING=1` makes each rank log the config Triton selects; a rank only logs when it tunes, so a run where some ranks hit the autotune cache and others miss cannot be compared this way.
0 commit comments