core.models.vision.vit_layer_specs

Module Contents

Functions

_te	The Transformer Engine provider, built on first use.
get_vit_layer_with_transformer_engine_spec	Returns ViT layer spec with Transformer Engine layers
get_vit_layer_with_local_spec	Returns ViT layer spec with Mcore local layers
_get_mlp_module_spec

Data

_local

API

core.models.vision.vit_layer_specs._te()

The Transformer Engine provider, built on first use.

Deferred so that importing this module does not require Transformer Engine – only building a Transformer Engine spec does. That is what the HAVE_TE guard used to buy.

core.models.vision.vit_layer_specs._local

‘get_backend(…)’

core.models.vision.vit_layer_specs.get_vit_layer_with_transformer_engine_spec() → megatron.core.transformer.spec_utils.ModuleSpec

Returns ViT layer spec with Transformer Engine layers

core.models.vision.vit_layer_specs.get_vit_layer_with_local_spec() → megatron.core.transformer.spec_utils.ModuleSpec

Returns ViT layer spec with Mcore local layers

core.models.vision.vit_layer_specs._get_mlp_module_spec(use_te: bool = True)