`core.models.backends`#

Module Contents#

Classes#

`BackendSpecProvider`	A protocol for providing the submodules used in Spec building.
`LocalSpecProvider`	Every backend a Megatron-Core-only run uses.
`InferenceSpecProvider`	Every backend an inference-optimized run uses.

Functions#

`require`	Refuse a backend whose optional package is missing or too old.
`select_cross_entropy`	Which cross entropy a config asks for, resolved once rather than on every forward.
`backend_slot`	Ask a provider for a slot added after it may have been written.
`get_backend`	Build the provider for a named backend.
`get_backend_from_config`	Build the provider a TransformerConfig asks for.

Data#

CrossEntropyTarget

loss = target(logits, labels, tp_group), with logits [s, b, vocab/tp] and the target owning every reduction over tp_group. Targets consume logits: they subtract the row max and exponentiate in place, so a caller that needs them afterwards passes a clone.

API#

core.models.backends.CrossEntropyTarget#

None

core.models.backends.require( requirement: str, requested_by: str = 'This backend', instead: str = '', ) → None#

Refuse a backend whose optional package is missing or too old.

Availability comes from the extension module’s own HAVE_* flag, which is set by actually importing the package – a package that is installed but unimportable (a broken CUDA build, a missing shared object) has to fail here, not later.

The requirement is declared once, as REQUIRES on the provider, and checked wherever a caller wants an early, clear refusal. It is deliberately not called when a provider is constructed: a spec may be assembled without the optional package present, and it is instantiating the module that fails.

core.models.backends.select_cross_entropy( cross_entropy_loss_fusion: bool = False, cross_entropy_fusion_impl: str = 'native', cuda_graph_impl: Optional[str] = None, ) → core.models.backends.CrossEntropyTarget#

Which cross entropy a config asks for, resolved once rather than on every forward.

Shared by every provider: the choice depends on the config, not on which backend supplies the rest of the model, which is how it behaved before this was a provider decision.

class core.models.backends.BackendSpecProvider#

Bases: typing.Protocol

A protocol for providing the submodules used in Spec building.

abstractmethod column_parallel_linear() → type#: Which column parallel linear module the backend uses

abstractmethod row_parallel_linear() → type#: Which row parallel linear module the backend uses

abstractmethod fuse_layernorm_and_linear() → bool#: Does the backend support a single module for layernorm and linear

abstractmethod column_parallel_layer_norm_linear() → Optional[type]#: Which module for sequential layernorm and linear

abstractmethod layer_norm( rms_norm: bool = False, for_qk: bool = False, has_residual: bool = False, ) → megatron.core.transformer.torch_norm.LayerNormBuilder#: Which module for layernorm

abstractmethod core_attention() → type#: Which module to use for attention

abstractmethod grouped_mlp_modules( moe_use_grouped_gemm: bool, ) → megatron.core.transformer.moe.moe_layer.ExpertsBuilder#: Which module and submodules to use for grouped mlp

abstractmethod activation_func() → megatron.core.transformer.mlp.TEActivationFunctionBuilder | None#: Which module to use for activation function

mlp_module(grouped: bool = False) → type#: Which module to use for the dense MLP block.

moe_router() → Optional[type]#: Which MoE router to use, or None to keep the MoESubmodules default.

vocab_parallel_cross_entropy() → core.models.backends.CrossEntropyTarget#: Which vocab-parallel cross entropy to use.

class core.models.backends.LocalSpecProvider( cross_entropy_loss_fusion: bool = False, cross_entropy_fusion_impl: str = 'native', cuda_graph_impl: Optional[str] = None, )#

Bases: core.models.backends.BackendSpecProvider

Every backend a Megatron-Core-only run uses.

Initialization

REQUIRES: Optional[str]#: None

column_parallel_linear() → type#: Which column parallel linear module the backend uses

row_parallel_linear() → type[megatron.core.tensor_parallel.layers.RowParallelLinear]#: Which row parallel linear module the backend uses

fuse_layernorm_and_linear() → bool#: Does the backend choose a single module for layernorm and linear

column_parallel_layer_norm_linear() → Optional[type]#: Which module for sequential layernorm and linear

layer_norm( rms_norm: bool = False, for_qk: bool = False, has_residual: bool = False, ) → megatron.core.transformer.torch_norm.LayerNormBuilder#

Which module to use for layer norm.

Apex’s fused LayerNorm when it is available, Torch’s otherwise – the same choice this made before, but decided here rather than by a module-level LNImpl that eight files each kept their own copy of and that this method used to mutate as a side effect.

RMSNorm always comes from Torch: Apex’s fused kernel implements LayerNorm only.

core_attention() → type#: Which module to use for attention

grouped_mlp_modules( moe_use_grouped_gemm: bool, ) → megatron.core.transformer.moe.moe_layer.ExpertsBuilder#: Which module and submodules to use for grouped mlp

activation_func() → megatron.core.transformer.mlp.TEActivationFunctionBuilder | None#: Which module to use for activation function

abstractmethod linear() → type#: Megatron Core has no non-parallel linear; Transformer Engine is the one that does.

mlp_module(grouped: bool = False) → type#: Megatron Core has one MLP block; grouped GEMM is a Transformer Engine feature.

moe_router() → Optional[type]#: Keep the MoESubmodules default.

vocab_parallel_cross_entropy() → core.models.backends.CrossEntropyTarget#: Which vocab-parallel cross entropy to use.

class core.models.backends.InferenceSpecProvider( cross_entropy_loss_fusion: bool = False, cross_entropy_fusion_impl: str = 'native', cuda_graph_impl: Optional[str] = None, )#

Bases: core.models.backends.LocalSpecProvider

Every backend an inference-optimized run uses.

Attention and the grouped linears come from Transformer Engine; the inference-specific gains are in the linear and mixture-of-experts layers.

Initialization

REQUIRES#: ‘transformer_engine’

linear() → type#: Which linear module TE backend uses

column_parallel_linear() → type#: Which column parallel linear module TE backend uses

row_parallel_linear() → type[megatron.core.tensor_parallel.inference_layers.InferenceRowParallelLinear]#: Which row parallel linear module Inference backend uses

fuse_layernorm_and_linear() → bool#: TE backend chooses a single module for layernorm and linear

column_parallel_layer_norm_linear() → type[megatron.core.tensor_parallel.inference_layers.InferenceLayerNormColumnParallelLinear]#: Which module for sequential layernorm and linear

layer_norm( rms_norm: bool = False, for_qk: bool = False, has_residual: bool = False, ) → megatron.core.transformer.torch_norm.LayerNormBuilder#: Which module to use for layer norm

core_attention() → type[megatron.core.extensions.transformer_engine.TEDotProductAttention]#: Which module to use for attention

activation_func() → megatron.core.transformer.mlp.TEActivationFunctionBuilder | None#: Which module to use for activation function

grouped_mlp_modules( moe_use_grouped_gemm: bool, ) → megatron.core.transformer.moe.moe_layer.ExpertsBuilder#: Which module and submodules to use for grouped mlp

moe_router() → Optional[type]#: Inference needs compact [tokens, topk] index routing rather than a dense map.

core.models.backends.backend_slot(

backend: core.models.backends.BackendSpecProvider,

name: str,

default: Callable[[], object],

**kwargs,

)#

Ask a provider for a slot added after it may have been written.

Two kinds of provider predate a newly added slot. One satisfies the protocol structurally and simply has no such attribute. The other subclasses the protocol and so inherits the default defined there – including wrappers that delegate everything else to a fallback provider, which would quietly answer for the backend they wrap. Both are treated as “did not implement this”, so default decides.

default is a callable and is only evaluated when it is needed, so a provider that answers for itself never pays to build an answer it will not use.

core.models.backends.get_backend( transformer_impl: Literal[local, transformer_engine, inference_optimized], *, use_kitchen: bool = False, use_kitchen_attention: bool = False, kitchen_attention_backend: str = 'sdpa', use_te_op_fuser: bool = False, cross_entropy_loss_fusion: bool = False, cross_entropy_fusion_impl: str = 'native', cuda_graph_impl: str | None = None, ) → core.models.backends.BackendSpecProvider#

Build the provider for a named backend.

Kitchen is enabled independently of transformer_impl. It overrides selected operations and delegates the rest to the chosen base provider through fallback. use_te_op_fuser applies only to the Transformer Engine provider.

core.models.backends.get_backend_from_config( config: object, *, transformer_impl: Optional[str] = None, ) → core.models.backends.BackendSpecProvider#

Build the provider a TransformerConfig asks for.

GPT spec builders also accept use_transformer_engine, which can select local layers while config.transformer_impl retains its default of transformer_engine. They pass the selected implementation here so the final norm and MTP use the same backend as the layers. Other callers use config.transformer_impl unchanged.

core.models.backends#