ModelArchitecture
from modal_dojo 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
Disable fused masked softmax. Default: False
multi_latent_attention bool
Use multi-latent attention instead of standard MHA. Default: False
kv_lora_rank int
LoRA rank for compressed key/value projections. Default: 0
qk_head_dim int
Query/key head dimension for multi-latent attention. Default: 0
qk_pos_emb_head_dim int
Query/key positional-embedding head dimension. Default: 0
v_head_dim int
Value head dimension for multi-latent attention. Default: 0
num_experts int
Number of mixture-of-experts experts. Default: 0
moe_layer_freq str
Which layers are MoE, as a Megatron frequency string. 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
Fuse MoE expert GEMMs into a grouped kernel. Default: False
moe_shared_expert_gate bool
Apply a gate to the shared expert. Default: False
moe_router_topk int
Experts selected per token. Default: 0
moe_router_pre_softmax bool
Softmax router scores before top-k. Default: False
moe_router_score_function str
Router scoring function. Default: ""
moe_router_enable_expert_bias bool
Add a learned bias to router scores. Default: False
moe_router_load_balancing_type str
Load-balancing loss used by the router. Default: ""
moe_token_dispatcher_type str
How tokens are dispatched to experts. Default: ""
moe_router_bias_update_rate float | None
Step size for router expert-bias updates.
moe_router_group_topk int
Groups considered before the per-group top-k. Default: 0
moe_router_num_groups int
Expert groups for grouped routing. Default: 0
moe_router_topk_scaling_factor float | None
Scale applied to top-k router scores.
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 transformer-spec import path segments.
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
RoPE frequency scale for long context.
mscale float | None
YaRN mscale applied to RoPE.
mscale_all_dim float | None
YaRN mscale applied across all rotary dimensions.
no_rope_fusion bool
Disable fused RoPE kernels. Default: False