`bridge.models.exaone.exaone_moe.exaone_moe_provider`#

Module Contents#

Classes#

`_MTPDenseLayerSpecsList`	Return a dense layer spec when MCore asks which spec to use for MTP.
`ExaoneMoeModelProvider`	Model provider for EXAONE MoE models.

Functions#

_build_exaone_moe_layer_spec

Build EXAONE MoE decoder specs while keeping MTP sub-layers dense.

API#

class bridge.models.exaone.exaone_moe.exaone_moe_provider._MTPDenseLayerSpecsList( data: list[megatron.core.transformer.ModuleSpec], dense_mtp_spec: megatron.core.transformer.ModuleSpec, )#

Bases: list

Return a dense layer spec when MCore asks which spec to use for MTP.

Initialization

Initialize self. See help(type(self)) for accurate signature.

__getitem__(idx)#

bridge.models.exaone.exaone_moe.exaone_moe_provider._build_exaone_moe_layer_spec(

cfg: megatron.bridge.models.gpt_provider.GPTModelProvider,

**kwargs,

) → megatron.core.transformer.ModuleSpec#: Build EXAONE MoE decoder specs while keeping MTP sub-layers dense.

class bridge.models.exaone.exaone_moe.exaone_moe_provider.ExaoneMoeModelProvider#

Bases: megatron.bridge.models.gpt_provider.GPTModelProvider

Model provider for EXAONE MoE models.

transformer_layer_spec: megatron.core.transformer.ModuleSpec | collections.abc.Callable[[megatron.bridge.models.gpt_provider.GPTModelProvider], megatron.core.transformer.ModuleSpec]#: None

normalization: str#: ‘RMSNorm’

activation_func: collections.abc.Callable#: None

gated_linear_unit: bool#: True

position_embedding_type: str#: ‘rope’

add_bias_linear: bool#: False

seq_length: int#: 4096

rotary_base: float#: 1000000.0

rope_scaling: bool#: False

rope_scaling_factor: float#: 8.0

make_vocab_size_divisible_by: int#: 128

mtp_num_layers: int | None#: None

mtp_loss_scaling_factor: float | None#: None

kv_channels: int | None#: 128

attention_dropout: float#: 0.0

hidden_dropout: float#: 0.0

qk_layernorm: bool#: True

moe_grouped_gemm: bool#: True

moe_router_pre_softmax: bool#: True

moe_enable_deepep: bool#: False

moe_token_dispatcher_type: str#: ‘alltoall’

moe_router_load_balancing_type: str#: ‘global_aux_loss’

moe_shared_expert_overlap: bool#: True

moe_expert_capacity_factor: float | None#: None

moe_router_dtype: str#: ‘fp32’

moe_aux_loss_coeff: float#: 0.01

moe_z_loss_coeff: float#: 0.001

moe_permute_fusion: bool#: True

fp8: str | None#: None

fp8_recipe: str | None#: ‘tensorwise’

first_last_layers_bf16: bool#: False

num_layers_at_start_in_bf16: int#: 1

num_layers_at_end_in_bf16: int#: 1

fp8_param: bool#: False

fp8_param_gather: bool#: False

init_method_std: float#: 0.006

layernorm_epsilon: float#: 1e-05

params_dtype: torch.dtype#: None

async_tensor_model_parallel_allreduce: bool#: True

attention_softmax_in_fp32: bool#: True

persist_layer_norm: bool#: True

num_layers_in_first_pipeline_stage: int | None#: None

num_layers_in_last_pipeline_stage: int | None#: None

account_for_embedding_in_pipeline_split: bool#: False

account_for_loss_in_pipeline_split: bool#: False

apply_rope_fusion: bool#: False

bias_activation_fusion: bool#: False

bias_dropout_fusion: bool#: False

masked_softmax_fusion: bool#: False

gradient_accumulation_fusion: bool#: False

moe_router_topk_scaling_factor: float#: 2.5

moe_router_score_function: str#: ‘sigmoid’

moe_router_enable_expert_bias: bool#: True

moe_router_bias_update_rate: float#: 0.001

bridge.models.exaone.exaone_moe.exaone_moe_provider#

Module Contents#

Classes#

Functions#

API#

`bridge.models.exaone.exaone_moe.exaone_moe_provider`#