core.models.vision.vit_layer_specs

Module Contents

Functions

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

API

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,

) → megatron.core.transformer.spec_utils.ModuleSpec