`core.tuning.interception`#

Carrying out an :class:AutotunePolicy against Triton’s autotuner.

This is the only module that touches Triton internals. Everything else works against the policy and the table, so replacing this file is all that is needed when Triton grows a supported hook, or when the upstream packages stop needing the intervention at all.

Module Contents#

Functions#

`active_policy`	The policy currently installed, if any.
`_scope`
`_enumerate`
`_record_choice`
`_record_winner`
`_record_file`
`_check_record_path`
`_dump_records`
`choice_digest`	Hash of every autotune choice this rank has made so far.
`choice_log`	The per-(kernel, shape) choices this rank has made.
`verify_choices`	Compare configs for kernel/shape keys observed on multiple ranks.
`maybe_verify_choices`	Run :func:`verify_choices` if the policy asks for a check at `iteration`.
`install`	Apply `policy` to the whole process, replacing the previous one.
`_reset_runtime_state`
`_autotuner_class`	Triton’s `Autotuner`, or `None` when Triton is not installed.
`_install`

Data#

`logger`
`_installed`
`_policy`
`_tables`
`_selected_configs`
`_scopes`
`_logged_singletons`
`_choice_log`
`_unverified`
`_verified`
`_tune_records`
`_record_rank`
`_enumerated`
`__all__`

API#

core.tuning.interception.logger#: ‘getLogger(…)’

core.tuning.interception._installed#: False

core.tuning.interception._policy: megatron.core.tuning.policy.AutotunePolicy | None#: None

core.tuning.interception._tables: dict#: None

core.tuning.interception._selected_configs: weakref.WeakKeyDictionary#: ‘WeakKeyDictionary(…)’

core.tuning.interception._scopes: weakref.WeakKeyDictionary#: ‘WeakKeyDictionary(…)’

core.tuning.interception._logged_singletons: weakref.WeakSet#: ‘WeakSet(…)’

core.tuning.interception._choice_log: dict#: None

core.tuning.interception._unverified: dict#: None

core.tuning.interception._verified: dict#: None

core.tuning.interception._tune_records: dict#: None

core.tuning.interception._record_rank: int | None#: None

core.tuning.interception._enumerated: set#: ‘set(…)’

core.tuning.interception.active_policy() → megatron.core.tuning.policy.AutotunePolicy | None#: The policy currently installed, if any.

core.tuning.interception._scope(autotuner, policy: megatron.core.tuning.policy.AutotunePolicy)#

core.tuning.interception._enumerate(qualified: str, count: int, state: str) → None#

core.tuning.interception._record_choice( arch: str, qualified: str, key: str, config, pinned: bool, ) → None#

core.tuning.interception._record_winner(arch: str, qualified: str, key: str, config) → None#

core.tuning.interception._record_file(record_path: str, rank: int) → str#

core.tuning.interception._check_record_path(record_path: str) → None#

core.tuning.interception._dump_records() → None#

core.tuning.interception.choice_digest() → str#: Hash of every autotune choice this rank has made so far.

core.tuning.interception.choice_log() → dict#: The per-(kernel, shape) choices this rank has made.

core.tuning.interception.verify_choices(group=None) → bool#

Compare configs for kernel/shape keys observed on multiple ranks.

Call where all ranks arrive, such as a step boundary. Calling it at the moment of choice would deadlock: ranks reach a given kernel at different times, so they would not agree on whether to take part in the collective.

Each check exchanges only the choices made since the previous check and compares them with those already agreed on, so once every kernel and shape has been seen a check moves an empty payload. Pipeline stages and expert ranks can execute different kernels or shapes; an absent key is not a conflicting choice. Agreement only covers observed choices; it does not establish numerical determinism or equal coverage.

core.tuning.interception.maybe_verify_choices(iteration: int, group=None) → bool | None#

Run :func:verify_choices if the policy asks for a check at iteration.

This is the cadence behind AutotunePolicy.verify_every; the training loop calls it once per step and the policy decides whether anything happens. Returns None when no check ran, so a caller can tell “agreed” from “not asked”.

core.tuning.interception.install( policy: megatron.core.tuning.policy.AutotunePolicy | collections.abc.Mapping | None = None, *, deterministic: bool = False, ) → None#

Apply policy to the whole process, replacing the previous one.

Call during initialization, before kernels run. Megatron’s training initializer does this from the --triton-autotune-* arguments; library callers that build models directly call it themselves. None applies the default policy, and a mapping is converted. An omitted mode is derived: a recording path selects record, deterministic or PyTorch’s deterministic flag selects pinned, and anything else auto. Repeated calls reuse the same adapter rather than nesting patches.

Raises if the policy cannot be carried out, such as when its recording path is not writable or when Triton is installed without the autotuner this adapter patches; the previous policy then stays in place. Without Triton no kernel is autotuned, so there is nothing to patch.

core.tuning.interception._reset_runtime_state() → None#

core.tuning.interception._autotuner_class()#: Triton’s Autotuner, or None when Triton is not installed.

core.tuning.interception._install(policy: megatron.core.tuning.policy.AutotunePolicy) → None#

core.tuning.interception.__all__#: [‘active_policy’, ‘choice_digest’, ‘choice_log’, ‘install’, ‘maybe_verify_choices’, ‘verify_choices’…

core.tuning.interception#

Module Contents#

Functions#

Data#

API#

`core.tuning.interception`#