Repository navigation
Make Triton autotune configuration selection deterministic - #6076
Conversation
b9a56d1 to
b6f3685
Compare
57eae03 to
217f5c2
Compare
|
/claude strict-review |
Strict Review SummaryFindings — CRITICAL: 0, IMPORTANT: 2, SUGGESTION: 1 This is a carefully engineered and unusually well-documented change. The core mechanism — pinning Triton's candidate list to a single entry chosen by a pure, timing-free rule, restored in a Most impactful findings
Risk assessmentLow risk to existing users. Behavior is fully gated behind deterministic mode / explicit env opt-in; the default |
aa18175 to
fbbed2f
Compare
…l2norm call GatedDeltaNet selects torch_chunk_gated_delta_rule under deterministic_mode "because FLA is not deterministic". Measured on fla 0.4.2 / GB300 with GDN-shaped inputs (q/k L2-normalised, g = -exp(A_log) * softplus(a + dt_bias), B=2, T=4096, H=8, K=V=128, all outputs finite): chunk_gated_delta_rule's forward and five gradients were bit-identical across 4 replays under side-stream contention and across four independent processes on four GPUs with default autotuning, and its call graph has no tl.atomic_* (the only atomics in fla are in path_attn and log_linear_attn). What does move the bits is the Triton autotuner's config choice: pinning every autotuner to its first config changed dk/dv/dg/dbeta (o and dq unchanged), so cross-run agreement requires the autotune configs pinned or the autotune cache shared, which is what NVIDIA#6966 / NVIDIA#6076 provide. The torch fallback costs 110 ms vs 1.1 ms per fwd+bwd (about 100x) and differs from FLA by a relative l2 of 3-4e-3 (bf16 level). Switching deterministic_mode to the FLA kernel once pinning has landed is left to the GDN owners; this commit only makes the docstring say what was measured. torch_chunk_gated_delta_rule(use_qk_l2norm_in_kernel=True) also failed with fla 0.4.2 ("l2norm() got an unexpected keyword argument 'dim'"): FLA's l2norm normalises the last dim and takes eps as its only tuning knob. Call it the way the kernel path in common.py does. Tests replay both implementations under RacingStreams and assert finite, bit-identical forward and backward results, and check the in-kernel l2norm option equals normalising q and k beforehand. Signed-off-by: Yan Xu <yxu1@nvidia.com>
…l2norm call GatedDeltaNet selects torch_chunk_gated_delta_rule under deterministic_mode "because FLA is not deterministic". Measured on fla 0.4.2 / GB300 with GDN-shaped inputs (q/k L2-normalised, g = -exp(A_log) * softplus(a + dt_bias), B=2, T=4096, H=8, K=V=128, all outputs finite): chunk_gated_delta_rule's forward and five gradients were bit-identical across 4 replays under side-stream contention and across four independent processes on four GPUs with default autotuning, and its call graph has no tl.atomic_* (the only atomics in fla are in path_attn and log_linear_attn). What does move the bits is the Triton autotuner's config choice: pinning every autotuner to its first config changed dk/dv/dg/dbeta (o and dq unchanged), so cross-run agreement requires the autotune configs pinned or the autotune cache shared, which is what NVIDIA#6966 / NVIDIA#6076 provide. The torch fallback costs 110 ms vs 1.1 ms per fwd+bwd (about 100x) and differs from FLA by a relative l2 of 3-4e-3 (bf16 level). Switching deterministic_mode to the FLA kernel once pinning has landed is left to the GDN owners; this commit only makes the docstring say what was measured. torch_chunk_gated_delta_rule(use_qk_l2norm_in_kernel=True) also failed with fla 0.4.2 ("l2norm() got an unexpected keyword argument 'dim'"): FLA's l2norm normalises the last dim and takes eps as its only tuning knob. Call it the way the kernel path in common.py does. Tests replay both implementations under RacingStreams and assert finite, bit-identical forward and backward results, and check the in-kernel l2norm option equals normalising q and k beforehand. torch_chunk_gdn2 made the same l2norm(dim=...) call; fixed and covered by the same kind of test. Signed-off-by: Yan Xu <yxu1@nvidia.com>
…l2norm call GatedDeltaNet selects torch_chunk_gated_delta_rule under deterministic_mode "because FLA is not deterministic". Measured on fla 0.4.2 / GB300 with GDN-shaped inputs (q/k L2-normalised, g = -exp(A_log) * softplus(a + dt_bias), B=2, T=4096, H=8, K=V=128, all outputs finite): chunk_gated_delta_rule's forward and five gradients were bit-identical across 4 replays under side-stream contention and across four independent processes on four GPUs with default autotuning, and its call graph has no tl.atomic_* (the only atomics in fla are in path_attn and log_linear_attn). What does move the bits is the Triton autotuner's config choice: pinning every autotuner to its first config changed dk/dv/dg/dbeta (o and dq unchanged), so cross-run agreement requires the autotune configs pinned or the autotune cache shared, which is what NVIDIA#6966 / NVIDIA#6076 provide. The torch fallback costs 110 ms vs 1.1 ms per fwd+bwd (about 100x) and differs from FLA by a relative l2 of 3-4e-3 (bf16 level). Switching deterministic_mode to the FLA kernel once pinning has landed is left to the GDN owners; this commit only makes the docstring say what was measured. torch_chunk_gated_delta_rule(use_qk_l2norm_in_kernel=True) also failed with fla 0.4.2 ("l2norm() got an unexpected keyword argument 'dim'"): FLA's l2norm normalises the last dim and takes eps as its only tuning knob. Call it the way the kernel path in common.py does. Tests replay both implementations under RacingStreams and assert finite, bit-identical forward and backward results, and check the in-kernel l2norm option equals normalising q and k beforehand. torch_chunk_gdn2 made the same l2norm(dim=...) call; fixed and covered by the same kind of test. Signed-off-by: Yan Xu <yxu1@nvidia.com>
…l2norm call GatedDeltaNet selects torch_chunk_gated_delta_rule under deterministic_mode "because FLA is not deterministic". Measured on fla 0.4.2 / GB300 with GDN-shaped inputs (q/k L2-normalised, g = -exp(A_log) * softplus(a + dt_bias), B=2, T=4096, H=8, K=V=128, all outputs finite): chunk_gated_delta_rule's forward and five gradients were bit-identical across 4 replays under side-stream contention and across four independent processes on four GPUs with default autotuning, and its call graph has no tl.atomic_* (the only atomics in fla are in path_attn and log_linear_attn). What does move the bits is the Triton autotuner's config choice: pinning every autotuner to its first config changed dk/dv/dg/dbeta (o and dq unchanged), so cross-run agreement requires the autotune configs pinned or the autotune cache shared, which is what NVIDIA#6966 / NVIDIA#6076 provide. The torch fallback costs 110 ms vs 1.1 ms per fwd+bwd (about 100x) and differs from FLA by a relative l2 of 3-4e-3 (bf16 level). Switching deterministic_mode to the FLA kernel once pinning has landed is left to the GDN owners; this commit only makes the docstring say what was measured. torch_chunk_gated_delta_rule(use_qk_l2norm_in_kernel=True) also failed with fla 0.4.2 ("l2norm() got an unexpected keyword argument 'dim'"): FLA's l2norm normalises the last dim and takes eps as its only tuning knob. Call it the way the kernel path in common.py does. Tests replay both implementations under RacingStreams and assert finite, bit-identical forward and backward results, and check the in-kernel l2norm option equals normalising q and k beforehand. torch_chunk_gdn2 made the same l2norm(dim=...) call; fixed and covered by the same kind of test. Signed-off-by: Yan Xu <yxu1@nvidia.com>
3d753a7 to
a21550f
Compare
The Mamba memory-efficient training path (use_mamba_mem_eff_path, the default) calls mamba_split_conv1d_scan_combined from the external mamba_ssm package. Its kernels carry timing-based @triton.autotune decorators, so the winning config - and with it the floating-point reduction order - can differ per process and per GPU from clock jitter alone. MAMBA_DETERMINISTIC only pinned Megatron's own SSD ops, which the mem-eff path never executes, so two identical deterministic-mode runs could diverge from iteration 1. Verified on a 24-node (96 GPU) Nemotron-3-Ultra pair: all 18 mamba_ssm kernels flipped configs between two identical runs (TRITON_PRINT_AUTOTUNING=1); with this pin the pair is bit-identical across 409824 traced wgrad and collective digests. Also make autotune_configs always pin under deterministic mode: cached autotuning (TRITON_CACHE_AUTOTUNING=1) still derives the cached winner from a first timing benchmark per cache directory, which is not value-deterministic across runs, GPUs, or ephemeral container caches. Signed-off-by: Yan Xu <yxu1@nvidia.com>
Deterministic mode had to choose a kernel config without benchmarking, and picked the cheapest one. That is reproducible but biased toward small tiles and num_stages=1, so it gives up throughput against the config the autotuner would have selected. Add a pre-tuned table, recorded once per GPU architecture by a normal autotuning run, so deterministic mode can pin the fastest config while staying a pure function of its inputs. Preference order is: pre-tuned winner for this kernel and shape, then a TRITON_AUTOTUNE_BLOCK_* override, then the cheapest config with a one-time warning. A missing entry is not fatal - it falls back to the cheapest config and warns, never to timing-based selection - so a drift between our shape key and Triton's only costs throughput, not reproducibility. Selection is applied per call and the candidate list is restored afterwards, so a kernel invoked later with different shapes still sees every candidate. Table hits return the kernel's own Config object rather than a rebuilt one, preserving pre_hook, num_ctas and maxnreg. Record a table with MCORE_DET_TUNE_RECORD=<prefix>; that pass autotunes normally (so it is deliberately not bit-exact) and writes per-rank JSON, merged into the checked-in table by the module's __main__. Signed-off-by: Yan Xu <yxu1@nvidia.com>
Populate the deterministic-mode tuning table from a 96-GPU Nemotron-3-Ultra recording run on GB200 (sm100): 15 mamba_ssm SSD kernels, 16 kernel/shape entries. The recording also quantifies the problem this table solves. Six of the sixteen entries had ranks disagreeing about the winning config within a single run, on identical hardware and identical shapes. _chunk_cumsum_bwd_kernel drew seven different winners across 96 ranks (32/29/18/11/4/1/1 votes), so its candidates are effectively tied and the timing benchmark is close to a coin flip; _chunk_scan_bwd_dc_kernel split 95 to 1. That per-GPU variance is why the failure is probabilistic in GPU count. Because ranks disagree, merging cannot let the last file win: _merge_records now takes a majority vote per kernel and shape, breaking ties on the serialized config so a given set of recordings always yields the same table. Signed-off-by: Yan Xu <yxu1@nvidia.com>
GB300 reports compute capability (10, 3), so the sm100 entries recorded on GB200 do not apply there and those runs fall back to the cheapest config with a warning. Record the sm103 table from a 96-GPU Nemotron-3-Ultra run: 15 mamba_ssm SSD kernels, 16 kernel/shape entries, majority vote over the ranks. Four of the sixteen entries had ranks disagreeing about the winner, against six on GB200. Also document a prerequisite for recording that is easy to get wrong: TRITON_CACHE_AUTOTUNING=1 must be set during a recording pass. With it unset, an installed mamba_ssm that ships its own import-time pin has already reduced every config list to a single entry, so no benchmark runs and nothing is recorded. The first GB300 attempt produced zero entries for exactly this reason. Signed-off-by: Yan Xu <yxu1@nvidia.com>
Widen the deterministic autotune pin from mamba_ssm to a configurable module list, defaulting to "mamba_ssm,transformer_engine" and overridable with DET_AUTOTUNE_PIN_MODULES. Runtime enumeration on the Nemotron-H path finds 18 multi-config Triton autotuners; three are in transformer_engine.common.triton.permutation and are reachable whenever moe_permute_fusion is enabled. Add three diagnostics, all off by default: - DET_AUTOTUNE_ENUMERATE: log every multi-config autotuner invoked, tagged pinned or unpinned, without changing which config is chosen. - verify_autotune_choices(): all-gather each rank's (kernel, shape) to config map and report the kernels ranks disagree on. Call it where every rank arrives, such as a step boundary. DET_AUTOTUNE_VERIFY_STRICT=1 raises instead of warning; DET_AUTOTUNE_CHOICE_LOG writes the per-rank map. - DET_AUTOTUNE_CHAOS: make each rank pick a different config on purpose, reproducibly, as a positive control for the detector. Signed-off-by: Yan Xu <yxu1@nvidia.com>
Kernel-config selection lived under megatron/core/ssm/ops/ and installed from
MambaMixer.__init__, but it is not specific to SSM. The same wall-clock
autotuning affects transformer_engine's MoE permutation kernels, and a model
without Mamba never installed the pin even though those kernels autotune in it
just the same.
Move the machinery to megatron/core/tuning/, split by role: policy.py holds the
intent, selection.py the timing-free choice rules, table.py the tuned configs,
and interception.py the only code that touches Triton internals, so a supported
hook later replaces one file. ssm/ops/determinism.py keeps the SSD tile
workspaces and re-exports the rest, so the SSM kernels' imports are unchanged.
Replace the scattered env reads with one AutotunePolicy. Modes are auto, pinned
and record; deterministic mode implies pinned, and an explicit MCORE_AUTOTUNE_MODE
overrides that so a determinism run can still record. The previous
DET_AUTOTUNE_* and MCORE_DET_TUNE_RECORD names are still accepted.
Store tuned configs as JSON under tuning/tables/ instead of a Python literal,
searched user-path-first, each file carrying the triton and package versions it
was recorded against. Adding an architecture is now a file drop rather than a
source edit, and the merge tool writes that file directly:
MCORE_AUTOTUNE_RECORD=/tmp/rec torchrun ... pretrain.py ...
python -m megatron.core.tuning merge /tmp/rec/*.json -o ~/.mcore/tuning/sm103.json
MCORE_AUTOTUNE_TABLE_PATH=~/.mcore/tuning torchrun ... pretrain.py ...
Add MCORE_AUTOTUNE_ON_MISS=error for callers that would rather fail than accept
the slower deterministic fallback, and a report subcommand that summarizes a
recording and flags the entries ranks disagreed on.
Tests move to tests/unit_tests/tuning/ and cover the new surface: policy
resolution including the legacy names, majority-vote merge, table round-trip,
that a stale entry degrades to a miss rather than an invalid launch, and that
the shipped tables load.
Signed-off-by: Yan Xu <yxu1@nvidia.com>
Activation ran from MambaMixer.__init__, so a model without Mamba never got the policy even though Transformer Engine's MoE permutation kernels autotune by timing in it just the same. Autotuning is not a property of the model. Install from TransformerConfig.__post_init__ instead, which every model construction passes through, and from initialize_megatron for training scripts that reach a kernel before building a config. Both are idempotent. The patch replaces a method on Triton's Autotuner class, so it only has to be installed before the first kernel call, not before the decorators are evaluated. pin_external_mamba_autotuners() stays as a deprecated wrapper for external callers, but nothing in the tree calls it now. Add megatron/core/tuning/README.md covering how pinning works and why any timing-free rule suffices, the three-command pretuning flow, the table format and its staleness handling, where the policy is installed and why not from a layer, how to check what a run did, and the recording gotcha that an installed mamba_ssm's import-time pin leaves nothing to record unless TRITON_CACHE_AUTOTUNING=1 is set. Signed-off-by: Yan Xu <yxu1@nvidia.com>
Two conflicts, both import lists: - megatron/training/training.py: main dropped set_run_config from the global_vars import and added get_train_state. Take both, and keep get_run_config at the sorted position this branch moved it to. The maybe_verify_choices call auto-merged after main's new train_state.iteration increment. - tests/unit_tests/determinism/kernels/test_ssm_kernels.py: keep this branch's megatron.core.tuning imports alongside the mixer-test imports main added. Signed-off-by: Yan Xu <yxu1@nvidia.com> Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Main now builds the MLA RoPE in-place autotuners as separate wrappers
(_autotuned_mla_rope_{fwd,bwd}_inplace_kernel = triton.autotune(...)(kernel)),
so the module attribute a tuned-table entry names is the bare JIT kernel.
test_packaged_entries_name_live_candidates took that for a missing Triton
install and skipped, which also skipped every table entry after it.
When the named attribute is not an autotuner, use the module's autotuner
that runs that kernel. Selection itself is unaffected: it names a tuner by
the kernel it wraps, so the table keys still match.
Signed-off-by: Yan Xu <yxu1@nvidia.com>
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
|
/ok to test a16b075 |
Conflicts were the RNGConfig migration's new _model_rng_config helper and import next to this PR's _triton_autotune_config_from_args in argument_utils.py and yaml_arguments.py; kept both. Signed-off-by: Yan Xu <yxu1@nvidia.com>
|
/ok to test 85a1346 |
initialize_megatron installed args.triton_autotune, which only validate_args sets. Callers that parse arguments and initialize without validating them, such as the RNG ownership test now on main, failed with AttributeError. Build the policy from the parsed arguments with the same helper validate_args uses; it returns an already-built policy unchanged, so validated and YAML arguments install the same policy as before. The initialization test now runs on parsed but unvalidated arguments and registers the run config that initialize_megatron reads since main moved RNG settings into it. Signed-off-by: Yan Xu <yxu1@nvidia.com> Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
|
/ok to test 3144e4c |
|
|
||
| # Before any Triton kernel runs, including the kernel warm-up before model construction. | ||
| # Built from the parsed arguments so callers that skip validate_args still get their policy. | ||
| tuning.install(_triton_autotune_config_from_args(args), deterministic=args.deterministic_mode) |
There was a problem hiding this comment.
What happens if this fails? Should we look at return value and fail on a problem?
There was a problem hiding this comment.
Good catch. The only silent failure was a Triton that is installed but doesn't provide triton.runtime.autotuner.Autotuner. install returned False there, and a run that asked for pinning kept timed autotuning. It raises now instead, and the previous policy stays in place. The other failures (invalid policy values, an unwritable record path) already raised. Since every failure now raises, the return value had nothing left to report, so install returns None and initialize_megatron has nothing to check. If Triton isn't installed at all, no kernel is autotuned, and install doesn't raise.
| InferenceCudaGraphScope, | ||
| ModelType, | ||
| OptimizerConfig, | ||
| # Training arguments, and so checkpoints, carry the validated Triton autotune policy. |
There was a problem hiding this comment.
Unnecessary comment I think, just makes things less clean.
There was a problem hiding this comment.
Removed. The policy is no longer stored on args (see the validate_args thread), so checkpoints don't carry an AutotunePolicy at all. The safe-globals entry is gone too, and the file is back to main.
| apply_determinism_to_args(args) | ||
|
|
||
| # Built here so a bad option fails now; training installs it in initialize_megatron. | ||
| args.triton_autotune = _triton_autotune_config_from_args(args) |
There was a problem hiding this comment.
initialize_megatron doesn't use this, why add it to args?
There was a problem hiding this comment.
There's no reason anymore, so I removed it from both validate_args and validate_yaml. initialize_megatron builds the policy from the arguments right where it installs it, so a bad option now fails there.
| ) | ||
|
|
||
|
|
||
| def _add_triton_autotune_args(parser): |
There was a problem hiding this comment.
Can we add these argument directly from the policy dataclass?
There was a problem hiding this comment.
Yes, done. ArgumentGroupFactory generates them now. It needed two small additions:
- a
prefix, so each field becomes--triton-autotune-<field>/args.triton_autotune_<field>; - support for
tuple[T, ...]fields.
_default_config_from_args takes the same prefix to build the policy back. block_sizes is still added by hand because its NAME=VALUE pairs can't be inferred from the field type. One rename falls out of this: --triton-autotune-enumerate is now --triton-autotune-enumerate-autotuners, matching the field.
The --triton-autotune-* arguments are now generated from AutotunePolicy by ArgumentGroupFactory, which gains a prefix for argument names and destinations and infers variable-length tuple fields. _default_config_from_args takes the same prefix to build the policy back. block_sizes keeps a hand-written argument because its NAME=VALUE pairs cannot be inferred from the field type. --triton-autotune-enumerate is renamed --triton-autotune-enumerate-autotuners to match its field. initialize_megatron is the only consumer of the policy, so validate_args and validate_yaml no longer store it on args, and AutotunePolicy no longer has to be a checkpoint safe global. install raises when a policy needs Triton's autotuner and the installed Triton does not provide it, instead of returning False and leaving a pinned run on timed autotuning. Every failure now raises, so install no longer returns a value. Signed-off-by: Yan Xu <yxu1@nvidia.com> Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
|
/ok to test 1a0f262 |
|
🔄 Merge queue validation started! You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/38078367861 |
Triton picks a kernel configuration by timing its candidates the first time the kernel runs. The winner can differ between ranks and between runs, and a different tiling changes the floating-point reduction order, so two identical runs can produce different numbers. Caching the winners (
TRITON_CACHE_AUTOTUNING=1) doesn't fix this: a rank that misses the cache times the candidates again on its own.This PR adds
megatron.core.tuning. Under--deterministic-mode, every autotuner in scope is pinned to one configuration that is chosen without timing. Runs without--deterministic-modekeep Triton's timed autotuning.How a configuration is chosen
sm100andsm103tables cover the in-tree Mamba, MLA RoPE and mHC kernels.block_sizes, if set.on_miss="error".The default scope is
mamba_ssm,transformer_engineandmegatron.core. Transformer Engine's MoE permutation kernels stay timed: their output doesn't depend on the configuration, and pinning would make them about 5x slower.Usage
--deterministic-modeis all most runs need. Every other option is a--triton-autotune-*argument generated from theAutotunePolicyfields, or atriton_autotunesection with--yaml-cfg.initialize_megatroninstalls the policy before any kernel runs. Code that builds models without the training initializer callsinstall(AutotunePolicy(...)). The README inmegatron/core/tuninglists all options.--triton-autotune-verify-every Ncompares the configurations that ranks chose every N steps.--triton-autotune-enumerate-autotunerslists every autotuner a run reaches and whether it is pinned.--triton-autotune-record-path PREFIX, merge the per-rank files withpython -m megatron.core.tuning merge, and pass the result's directory to--triton-autotune-table-path.Behavior changes
--deterministic-modeinstead ofMAMBA_DETERMINISTIC, which still controls the externalmamba_ssmpackage.block_sizesreplaces theTRITON_AUTOTUNE_BLOCK_*environment variables.--deterministic-mode,TRITON_CACHE_AUTOTUNINGonly matters for autotuners outside the scope, because pinned kernels never time candidates.ArgumentGroupFactorygains aprefixoption and supportstuple[T, ...]fields.Performance
Without a table entry, the cheapest candidate was up to 12.5x (MLA RoPE) and 16x (mHC) slower per kernel than the timed winner. With the packaged tables, total Triton kernel time for these kernels is within 1% of timed autotuning at the measured shapes, and within 6% at a head count that wasn't measured. These numbers are from GB300. The
sm100table reuses thesm103entries and wasn't measured separately.