ModelArchitecture
from modal_training_gym.common.models.base import ModelArchitectureMegatron transformer architecture parameters.
Attributes
num_layers int
Number of transformer layers. Default: 0
hidden_size int
Hidden dimension size. Default: 0
ffn_hidden_size int
Feed-forward network intermediate size. Default: 0
num_attention_heads int
Number of attention heads. Default: 0
group_query_attention bool
Whether to enable grouped-query attention. Default: True
num_query_groups int
Number of KV head groups for grouped-query attention. Default: 0
kv_channels int
Per-head key/value channel dimension. Default: 0
vocab_size int
Vocabulary size. Default: 0
normalization str
Layer normalization type. Default: "RMSNorm"
norm_epsilon float
Normalization epsilon. Default: 1e-06
swiglu bool
Whether to use SwiGLU activation in the feed-forward network. Default: True
disable_bias_linear bool
Whether to disable bias in linear layers. Default: True
qk_layernorm bool
Whether to normalize query and key projections. Default: True
untie_embeddings_and_output_weights bool
Whether to use separate output projection weights. Default: False
no_masked_softmax_fusion bool
Default: False
multi_latent_attention bool
Default: False
kv_lora_rank int
Default: 0
qk_head_dim int
Default: 0
qk_pos_emb_head_dim int
Default: 0
v_head_dim int
Default: 0
num_experts int
Number of mixture-of-experts experts. Default: 0
moe_layer_freq str
Default: ""
moe_ffn_hidden_size int
Per-expert feed-forward intermediate size. Default: 0
moe_shared_expert_intermediate_size int
Shared-expert intermediate size. Default: 0
moe_grouped_gemm bool
Default: False
moe_shared_expert_gate bool
Default: False
moe_router_topk int
Default: 0
moe_router_pre_softmax bool
Default: False
moe_router_score_function str
Router scoring function. Default: ""
moe_router_enable_expert_bias bool
Default: False
moe_router_load_balancing_type str
Default: ""
moe_token_dispatcher_type str
Default: ""
moe_router_bias_update_rate float | None
moe_router_group_topk int
Default: 0
moe_router_num_groups int
Default: 0
moe_router_topk_scaling_factor float | None
moe_token_drop_policy str
Token drop policy for mixture-of-experts routing. Default: ""
moe_router_dtype str
Data type for router computation. Default: ""
moe_permute_fusion bool
Whether to enable permute fusion. Default: False
moe_aux_loss_coeff float | None
Auxiliary load-balancing loss coefficient.
megatron_spec list[str] | None
megatron_model_type str
Slime/Megatron model type used for checkpoint conversion outside bridge mode. Default: ""
apply_layernorm_1p bool
Whether to use zero-centered LayerNorm. Default: False
use_gated_attention bool
Whether to enable gated attention. Default: False
attention_output_gate bool
Whether to gate attention outputs. Default: False
use_rotary_position_embeddings bool
Whether to use RoPE. Default: True
rotary_base int
Base frequency for RoPE. Default: 10000
rotary_percent float
Fraction of hidden dimensions that use RoPE. Default: 1.0
rotary_scaling_factor float | None
mscale float | None
mscale_all_dim float | None
no_rope_fusion bool
Default: False