core.models.hybrid.hybrid_layer_specs

Module Contents

Functions

_get_gated_delta_product_mamba_layer_spec
_get_wide_residual_hybrid_stack_spec	Build a static HybridStack spec whose layers explicitly own wide connections.

Data

moe
moe_inference
_csa_compressor
_csa_indexer
_csa_qk_norm
_hybrid_mtp_block_spec
hybrid_stack_spec
gated_delta_product_stack_spec
hybrid_inference_stack_spec
gated_delta_product_inference_stack_spec
wide_residual_hybrid_stack_spec
wide_residual_gated_delta_product_stack_spec
wide_residual_hybrid_inference_stack_spec
wide_residual_gated_delta_product_inference_stack_spec
mamba_stack_spec
mamba_inference_stack_spec
gdp_stack_spec
gdp_inference_stack_spec
hybrid_dsv4_stack_spec

API

core.models.hybrid.hybrid_layer_specs.moe

‘get_moe_module_spec(…)’

core.models.hybrid.hybrid_layer_specs.moe_inference

‘get_inference_optimized_moe_spec(…)’

core.models.hybrid.hybrid_layer_specs._csa_compressor

‘partial(…)’

core.models.hybrid.hybrid_layer_specs._csa_indexer

‘partial(…)’

core.models.hybrid.hybrid_layer_specs._csa_qk_norm

‘layer_norm(…)’

core.models.hybrid.hybrid_layer_specs._hybrid_mtp_block_spec

‘ModuleSpec(…)’

core.models.hybrid.hybrid_layer_specs._get_gated_delta_product_mamba_layer_spec(in_proj, out_proj)

core.models.hybrid.hybrid_layer_specs.hybrid_stack_spec

‘ModuleSpec(…)’

core.models.hybrid.hybrid_layer_specs.gated_delta_product_stack_spec

‘ModuleSpec(…)’

core.models.hybrid.hybrid_layer_specs.hybrid_inference_stack_spec

‘ModuleSpec(…)’

core.models.hybrid.hybrid_layer_specs.gated_delta_product_inference_stack_spec

‘ModuleSpec(…)’

core.models.hybrid.hybrid_layer_specs._get_wide_residual_hybrid_stack_spec(

stack_spec: megatron.core.transformer.spec_utils.ModuleSpec,

) → megatron.core.transformer.spec_utils.ModuleSpec

Build a static HybridStack spec whose layers explicitly own wide connections.

core.models.hybrid.hybrid_layer_specs.wide_residual_hybrid_stack_spec

‘_get_wide_residual_hybrid_stack_spec(…)’

core.models.hybrid.hybrid_layer_specs.wide_residual_gated_delta_product_stack_spec

‘_get_wide_residual_hybrid_stack_spec(…)’

core.models.hybrid.hybrid_layer_specs.wide_residual_hybrid_inference_stack_spec

‘_get_wide_residual_hybrid_stack_spec(…)’

core.models.hybrid.hybrid_layer_specs.wide_residual_gated_delta_product_inference_stack_spec

‘_get_wide_residual_hybrid_stack_spec(…)’

core.models.hybrid.hybrid_layer_specs.mamba_stack_spec

None

core.models.hybrid.hybrid_layer_specs.mamba_inference_stack_spec

None

core.models.hybrid.hybrid_layer_specs.gdp_stack_spec

None

core.models.hybrid.hybrid_layer_specs.gdp_inference_stack_spec

None

core.models.hybrid.hybrid_layer_specs.hybrid_dsv4_stack_spec

None