core.tuning.policy

What the autotune interception should do, as one object.

Triton picks a kernel config by benchmarking its candidates at first call, so the winner depends on wall-clock timings. Two consequences follow, and only one of them is about determinism:

• Two identical runs can select different tile shapes, which changes the floating-point reduction order and therefore the result.

• Every cold process pays for the benchmark, and repeated benchmarks of the same workload disagree with each other, so measurements are noisy.

Pinning the choice to a pure function of the candidate list fixes both. This module holds the intent; :mod:megatron.core.tuning.interception carries it out.

Module Contents

Classes

AutotunePolicy

Immutable configuration inputs for choosing Triton kernel launch parameters.

Functions

use_deterministic_mode	Whether deterministic behaviour is requested for kernel selection.
set_deterministic_mode	Override :func:use_deterministic_mode for the current process.
_string_tuple
coerce_policy	Return value as an :class:AutotunePolicy, converting a mapping.

Data

_deterministic_override
DEFAULT_MODULES
DEFAULT_CONFIG_INVARIANT

API

core.tuning.policy._deterministic_override

None

core.tuning.policy.DEFAULT_MODULES

(‘mamba_ssm’, ‘transformer_engine’, ‘megatron.core’)

core.tuning.policy.DEFAULT_CONFIG_INVARIANT

‘tuple(…)’

core.tuning.policy.use_deterministic_mode() → bool

Whether deterministic behaviour is requested for kernel selection.

core.tuning.policy.set_deterministic_mode(value)

Override :func:use_deterministic_mode for the current process.

core.tuning.policy._string_tuple(name: str, value) → tuple[str, ...]

class core.tuning.policy.AutotunePolicy

Immutable configuration inputs for choosing Triton kernel launch parameters.

This object contains settings only. Installation resolves defaults into a separate policy value; loaded tables, selected configs, and recorded winners belong to the runtime adapter in :mod:megatron.core.tuning.interception. See megatron/core/tuning/README.md for modes, scope, tables and configuration.

mode: Literal[auto, pinned, record] | None

None

None derives the mode from recording and determinism settings. auto leaves Triton alone. pinned replaces the candidate list with one entry chosen without measuring anything. record lets Triton benchmark as usual and captures the winners, so a table can be built; a recording run is deliberately not reproducible.

modules: tuple[str, ...]

None

Module prefixes to act on, matched at package boundaries. Kernels outside them keep Triton’s timed choice.

config_invariant: tuple[str, ...]

‘field(…)’

Qualified kernel names (module.function) whose outputs do not depend on the launch configuration. They keep Triton’s timed choice even in pinned mode, since pinning them costs throughput without changing any result. An empty tuple pins every kernel in scope.

table_path: tuple[str, ...]

()

Directories searched for tuned tables before the packaged defaults. ~ is expanded.

record_path: str | None

None

File prefix for a record run’s per-rank captures, written as <record_path>.rank<N>.json. ~ is expanded.

on_miss: Literal[min_cost, error]

‘min_cost’

What to do when no table entry matches. min_cost is still deterministic, just possibly slower; error refuses to guess.

block_sizes: tuple[tuple[str, int], ...]

()

Explicit BLOCK_* kernel arguments to match when the table misses. Accepts a mapping or pairs; stored as sorted immutable pairs.

verify_every: int

0

Cross-rank agreement check cadence, in training steps. 0 disables. Honoured by the training loop through :func:megatron.core.tuning.maybe_verify_choices.

verify_strict: bool

False

Raise instead of logging a warning when the agreement check finds ranks running different configurations for the same kernel and shape.

enumerate_autotuners: bool

False

Log every multi-config autotuner reached, with whether it is pinned, without changing what is chosen.

chaos: bool

False

Make each rank pick a different config on purpose. A positive control for divergence detectors; never for real runs.

__post_init__()

classmethod from_mapping(

values: collections.abc.Mapping,

) → core.tuning.policy.AutotunePolicy

Build a policy from a mapping, such as a parsed YAML section.

Keys set to None fall back to their defaults, as an omitted key would. Unknown keys raise, so a misspelled option cannot be silently ignored.

resolve(

*,

deterministic: bool = False,

) → core.tuning.policy.AutotunePolicy

Resolve an omitted mode, leaving explicit configuration unchanged.

A recording path implies record. Otherwise, deterministic or PyTorch’s deterministic flag implies pinned; ordinary execution uses auto.

property intercepts: bool

Whether this policy needs the autotuner patched at all.

Enumeration and the agreement check both read what Triton chose, which only the interception records, so either one needs the patch even in auto mode — where it observes the timed choice without changing it.

core.tuning.policy.coerce_policy(value) → core.tuning.policy.AutotunePolicy | None

Return value as an :class:AutotunePolicy, converting a mapping.

None stays None. Anything else is rejected before it can reach the process-wide adapter.