nemo_automodel.components.distributed.blockdiag_cp.state

Runtime knobs and activation-checkpoint-safe state for block-diagonal CP.

The knob normalization (synonym maps and defaults) lives here, in a dependency-free leaf, so every consumer (the kernel-side :func:configure_cp_varlen entry point and any policy-side config parser) derives the accepted values from the same table and can never disagree on synonyms or defaults.

Module Contents

Classes

Name	Description
`_ThreadSharedVar`	A `contextvars.ContextVar`-style get/set/reset holder visible ACROSS threads.

Functions

Name	Description
`_clean`	-
`configure_cp_varlen`	Configure the block-diagonal CP attention path from parsed runtime config.
`cp_attn_fire_count`	Number of block-diagonal CP attention calls since the last reset.
`cp_varlen_runtime_config`	Return the currently configured block-diagonal CP runtime settings.
`normalize_attn_backend`	Canonicalize an attention-backend knob to one of `ATTN_BACKEND_VALUES`.
`normalize_kv_exchange`	Canonicalize a KV-exchange knob to one of `KV_EXCHANGE_VALUES`.
`reset_cp_attn_fire_count`	Zero the CP-attention fire counter (call before each forward).

Data

ATTN_BACKEND_DEFAULT

ATTN_BACKEND_SYNONYMS

ATTN_BACKEND_VALUES

KV_EXCHANGE_DEFAULT

KV_EXCHANGE_SYNONYMS

KV_EXCHANGE_VALUES

_ATTN_BACKEND_REVERSE

_CP_ATTN_BACKEND

_CP_ATTN_FIRE_COUNT

_CP_BLOCKDIAG_STATE

_CP_KV_EXCHANGE

_KV_EXCHANGE_REVERSE

API

class nemo_automodel.components.distributed.blockdiag_cp.state._ThreadSharedVar()

A contextvars.ContextVar-style get/set/reset holder visible ACROSS threads.

The block-diagonal CP state must be readable inside the autograd worker thread that runs activation-checkpointing recompute during backward. A real ContextVar is per-thread and reads its default (None) there, which would silently drop the CP state mid-recompute — the softmax SDPA would fall back to local-only attention. Training steps are sequential, so a single shared slot with token-based restore (supporting nesting) is safe; backward only ever READS the slot.

nemo_automodel.components.distributed.blockdiag_cp.state._ThreadSharedVar.get()

Return the current value (None when unset).

nemo_automodel.components.distributed.blockdiag_cp.state._ThreadSharedVar.reset(
    token
)

Restore the value captured by a previous :meth:set.

nemo_automodel.components.distributed.blockdiag_cp.state._ThreadSharedVar.set(
    value
)

Set the value; returns a token (the previous value) for :meth:reset.

nemo_automodel.components.distributed.blockdiag_cp.state._clean(
    value: typing.Any
) -> str

nemo_automodel.components.distributed.blockdiag_cp.state.configure_cp_varlen(
    attn_backend: str = 'flash',
    kv_exchange: str = 'allgather'
) -> None

Configure the block-diagonal CP attention path from parsed runtime config.

Parameters:

attn_backend

strDefaults to 'flash'

Varlen kernel selection; any synonym accepted by :func:normalize_attn_backend (default "flash").

kv_exchange

strDefaults to 'allgather'

K/V delivery mode; any synonym accepted by :func:normalize_kv_exchange (default "allgather").

nemo_automodel.components.distributed.blockdiag_cp.state.cp_attn_fire_count() -> int

Number of block-diagonal CP attention calls since the last reset.

nemo_automodel.components.distributed.blockdiag_cp.state.cp_varlen_runtime_config() -> dict[str, str]

Return the currently configured block-diagonal CP runtime settings.

Returns: dict[str, str]

A dict with keys "attn_backend" and "kv_exchange" holding the

nemo_automodel.components.distributed.blockdiag_cp.state.normalize_attn_backend(
    value: typing.Any
) -> str

Canonicalize an attention-backend knob to one of ATTN_BACKEND_VALUES.

None/""/auto/default map to the default (flash).

Parameters:

value

Any

The raw user-facing knob value (string-like or None).

Returns: str

The canonical backend name ("flash", "te", or "dense").

Raises:

ValueError: If the value is not a recognized backend or synonym.

nemo_automodel.components.distributed.blockdiag_cp.state.normalize_kv_exchange(
    value: typing.Any
) -> str

Canonicalize a KV-exchange knob to one of KV_EXCHANGE_VALUES.

None/""/auto/default map to the default (allgather).

Parameters:

value

Any

The raw user-facing knob value (string-like or None).

Returns: str

The canonical exchange name ("allgather", "halo", or "a2a").

Raises:

ValueError: If the value is not a recognized exchange mode or synonym.

nemo_automodel.components.distributed.blockdiag_cp.state.reset_cp_attn_fire_count() -> None

Zero the CP-attention fire counter (call before each forward).

nemo_automodel.components.distributed.blockdiag_cp.state.ATTN_BACKEND_DEFAULT = 'flash'

nemo_automodel.components.distributed.blockdiag_cp.state.ATTN_BACKEND_SYNONYMS: dict[str, tuple[str, ...]] = {'flash': ('flash', 'flash_attn', 'flash_attention', 'flash_attention_2', 'fa2')...

nemo_automodel.components.distributed.blockdiag_cp.state.ATTN_BACKEND_VALUES: tuple[str, ...] = tuple(ATTN_BACKEND_SYNONYMS)

nemo_automodel.components.distributed.blockdiag_cp.state.KV_EXCHANGE_DEFAULT = 'allgather'

nemo_automodel.components.distributed.blockdiag_cp.state.KV_EXCHANGE_SYNONYMS: dict[str, tuple[str, ...]] = {'allgather': ('allgather', 'all_gather'), 'halo': ('halo', 'neighbor', 'needed'...

nemo_automodel.components.distributed.blockdiag_cp.state.KV_EXCHANGE_VALUES: tuple[str, ...] = tuple(KV_EXCHANGE_SYNONYMS)

nemo_automodel.components.distributed.blockdiag_cp.state._ATTN_BACKEND_REVERSE = {syn: canon for canon, syns in (ATTN_BACKEND_SYNONYMS.items()) for syn in syns}

nemo_automodel.components.distributed.blockdiag_cp.state._CP_ATTN_BACKEND = ATTN_BACKEND_DEFAULT

nemo_automodel.components.distributed.blockdiag_cp.state._CP_ATTN_FIRE_COUNT: list[int] = [0]

nemo_automodel.components.distributed.blockdiag_cp.state._CP_BLOCKDIAG_STATE = _ThreadSharedVar()

nemo_automodel.components.distributed.blockdiag_cp.state._CP_KV_EXCHANGE = KV_EXCHANGE_DEFAULT

nemo_automodel.components.distributed.blockdiag_cp.state._KV_EXCHANGE_REVERSE = {syn: canon for canon, syns in (KV_EXCHANGE_SYNONYMS.items()) for syn in syns}