Skip to content

Make Triton autotune configuration selection deterministic - #6076

Merged
Connor-XY merged 36 commits into
NVIDIA:mainfrom
Connor-XY:yx/det-mamba-autotune-pin
Oct 10, 2026
Merged

Connor-XY merged 36 commits into
NVIDIA:mainfrom
Connor-XY:yx/det-mamba-autotune-pin

Conversation

@Connor-XY

@Connor-XY Connor-XY commented Jul 27, 2026 •

Copy link
Copy Markdown
Contributor

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-mode keep Triton's timed autotuning.

How a configuration is chosen

  1. The packaged table entry for the kernel and shape, if there is one. The sm100 and sm103 tables cover the in-tree Mamba, MLA RoPE and mHC kernels.
  2. Otherwise, a configuration that matches block_sizes, if set.
  3. Otherwise, the cheapest candidate, or an error with on_miss="error".

The default scope is mamba_ssm, transformer_engine and megatron.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-mode is all most runs need. Every other option is a --triton-autotune-* argument generated from the AutotunePolicy fields, or a triton_autotune section with --yaml-cfg. initialize_megatron installs the policy before any kernel runs. Code that builds models without the training initializer calls install(AutotunePolicy(...)). The README in megatron/core/tuning lists all options.

  • --triton-autotune-verify-every N compares the configurations that ranks chose every N steps.
  • --triton-autotune-enumerate-autotuners lists every autotuner a run reaches and whether it is pinned.
  • To build a table for a new GPU, run with --triton-autotune-record-path PREFIX, merge the per-rank files with python -m megatron.core.tuning merge, and pass the result's directory to --triton-autotune-table-path.

Behavior changes

  • Megatron's own SSM kernels follow --deterministic-mode instead of MAMBA_DETERMINISTIC, which still controls the external mamba_ssm package.
  • block_sizes replaces the TRITON_AUTOTUNE_BLOCK_* environment variables.
  • Under --deterministic-mode, TRITON_CACHE_AUTOTUNING only matters for autotuners outside the scope, because pinned kernels never time candidates.
  • ArgumentGroupFactory gains a prefix option and supports tuple[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 sm100 table reuses the sm103 entries and wasn't measured separately.

@copy-pr-bot

copy-pr-bot Bot commented Jul 27, 2026

Copy link
Copy Markdown

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

@Connor-XY
Connor-XY force-pushed the yx/det-mamba-autotune-pin branch from b9a56d1 to b6f3685 Compare July 31, 2026 02:10
@Connor-XY Connor-XY changed the title Pin external mamba_ssm triton autotuners under deterministic mode Pin external Triton autotuners under deterministic mode, with cross-rank choice verification Jul 31, 2026
@Connor-XY Connor-XY changed the title Pin external Triton autotuners under deterministic mode, with cross-rank choice verification Add megatron.core.tuning: pin Triton kernel configs, with pretuning and cross-rank verification Aug 4, 2026
@Connor-XY
Connor-XY force-pushed the yx/det-mamba-autotune-pin branch from 57eae03 to 217f5c2 Compare August 4, 2026 19:38
@Connor-XY

Copy link
Copy Markdown
Contributor Author

/claude strict-review

Comment thread megatron/core/tuning/interception.py Resolved
Comment thread megatron/core/tuning/policy.py Resolved
Comment thread megatron/core/tuning/policy.py Resolved
@claude

claude Bot commented Aug 4, 2026

Copy link
Copy Markdown
Contributor

Strict Review Summary

Findings — 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 finally so per-shape selection still works — is sound and correctly exploits Triton's len(configs) > 1 benchmarking gate. The layering (policy / selection / table / interception) cleanly isolates the one file that touches Triton internals, and the table-lookup-against-live-candidates approach correctly degrades a stale entry to a miss rather than an invalid launch. The CPU-only unit tests cover the meaningful branches. Determinism gating (use_deterministic_mode) means non-deterministic perf runs are genuinely untouched.

Most impactful findings

  1. verify_choices(group=...) sizes gather lists by global world size but passes a subgroup (interception.py:111-117) — the default group=None path works (hence green tests), but any subgroup call, which is the natural use, breaks. Latent trap in a public diagnostic API.

  2. verify_every / MCORE_AUTOTUNE_VERIFY is inert (policy.py:83) — parsed and documented as a step-cadence check, but never read and verify_choices() is never called from the training loop. A user enabling it to catch cross-rank divergence silently gets no protection. Either wire it in or drop it from the surface until the call site lands.

  3. _explicit_mode is a dead field (policy.py:87) — stored, never read; the "explicit mode wins" behavior is already realized in from_env. Minor cleanup.

Risk assessment

Low risk to existing users. Behavior is fully gated behind deterministic mode / explicit env opt-in; the default auto mode leaves Triton untouched, and install() is idempotent and no-ops without an interception. No checkpoint-format, config-default, or public-signature breakage. Legacy DET_AUTOTUNE_* env names are preserved. The megatron/core process-group guidance does not apply (no new parallel_state.get_*_group() reads). The two IMPORTANT items are correctness/UX gaps in the new diagnostic surface rather than risks to the pinning path that is the point of the PR — the pinning itself is validated by the reported 1536-GPU bit-exact runs. Worth resolving the verify_choices group sizing and the inert verify_every before merge so the shipped diagnostics do what they advertise.

@Connor-XY
Connor-XY force-pushed the yx/det-mamba-autotune-pin branch from aa18175 to fbbed2f Compare September 8, 2026 20:30
Connor-XY added a commit to Connor-XY/Megatron-LM that referenced this pull request Sep 8, 2026
…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>
Connor-XY added a commit to Connor-XY/Megatron-LM that referenced this pull request Sep 8, 2026
…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>
Connor-XY added a commit to Connor-XY/Megatron-LM that referenced this pull request Sep 8, 2026
…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>
Connor-XY added a commit to Connor-XY/Megatron-LM that referenced this pull request Sep 9, 2026
…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>
@Connor-XY Connor-XY changed the title Add megatron.core.tuning: pin Triton kernel configs, with pretuning and cross-rank verification Make Triton autotune configuration selection deterministic Sep 10, 2026
@Connor-XY
Connor-XY force-pushed the yx/det-mamba-autotune-pin branch from 3d753a7 to a21550f Compare September 10, 2026 20:27
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>
Connor-XY and others added 2 commits October 7, 2026 12:34
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>
@Connor-XY

Copy link
Copy Markdown
Contributor Author

/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>
@Connor-XY

Copy link
Copy Markdown
Contributor Author

/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>
@Connor-XY

Copy link
Copy Markdown
Contributor Author

/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)

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

What happens if this fails? Should we look at return value and fail on a problem?

@Connor-XY Connor-XY Oct 9, 2026 •

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Comment thread megatron/core/safe_globals.py Outdated
InferenceCudaGraphScope,
ModelType,
OptimizerConfig,
# Training arguments, and so checkpoints, carry the validated Triton autotune policy.

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Unnecessary comment I think, just makes things less clean.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Comment thread megatron/training/arguments.py Outdated
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)

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

initialize_megatron doesn't use this, why add it to args?

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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):

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Can we add these argument directly from the policy dataclass?

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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>
@Connor-XY

Copy link
Copy Markdown
Contributor Author

/ok to test 1a0f262

@nemo-automation-bot

Copy link
Copy Markdown

🔄 Merge queue validation started!

You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/38078367861

This branch was previously deployed

2 inactive deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

26.10 Approved All necessary approvals have been made complexity: high

Projects

None yet

Development

Successfully merging this pull request may close these issues.

7 participants