"```python\nfrom modal_training_gym.common.models.base import ModelArchitecture\n```\n\nTransformer architecture parameters for a specific model.\n\nThese fields map directly to Megatron-LM model-parallel configuration\nflags. Framework launchers read them to generate the correct CLI\narguments for distributed training.\n\n## Model Dimensions\n\n| Field | Type | Default | Description |\n|-------|------|---------|-------------|\n| `num_layers` | `int` | `0` | Number of transformer layers. Default `0`. |\n| `hidden_size` | `int` | `0` | Hidden dimension size. Default `0`. |\n| `ffn_hidden_size` | `int` | `0` | Feed-forward network intermediate size. Default `0`. |\n| `vocab_size` | `int` | `0` | Vocabulary size. Default `0`. |\n\n## Attention\n\n| Field | Type | Default | Description |\n|-------|------|---------|-------------|\n| `num_attention_heads` | `int` | `0` | Number of attention heads. Default `0`. |\n| `group_query_attention` | `bool` | `True` | Enable grouped-query attention (GQA). Default `True`. |\n| `num_query_groups` | `int` | `0` | Number of KV head groups for GQA. Default `0`. |\n| `kv_channels` | `int` | `0` | Per-head key/value channel dimension. Default `0`. |\n\n## Normalization and Activation\n\n| Field | Type | Default | Description |\n|-------|------|---------|-------------|\n| `normalization` | `str` | `\"RMSNorm\"` | Layer normalization type. Default `\"RMSNorm\"`. |\n| `norm_epsilon` | `float` | `1e-06` | Normalization epsilon. Default `1e-6`. |\n| `swiglu` | `bool` | `True` | Use SwiGLU activation in FFN. Default `True`. |\n| `disable_bias_linear` | `bool` | `True` | Disable bias in linear layers. Default `True`. |\n| `qk_layernorm` | `bool` | `True` | Apply layer norm to query and key projections. Default `True`. |\n| `untie_embeddings_and_output_weights` | `bool` | `False` | Use separate output projection weights instead of tying to token embeddings. Default `False`. |\n\n## Mixture of Experts\n\n| Field | Type | Default | Description |\n|-------|------|---------|-------------|\n| `num_experts` | `int` | `0` | Total number of MoE experts. Default `0` (dense model). |\n| `moe_ffn_hidden_size` | `int` | `0` | Per-expert FFN intermediate size. Default `0`. |\n| `moe_shared_expert_intermediate_size` | `int` | `0` | Shared expert FFN intermediate size. Default `0`. |\n\n## MoE Routing\n\n| Field | Type | Default | Description |\n|-------|------|---------|-------------|\n| `moe_router_score_function` | `str` | `\"\"` | Router scoring function (e.g. `\"softmax\"`). Default `\"\"`. |\n| `moe_token_drop_policy` | `str` | `\"\"` | Token drop policy for MoE routing. Default `\"\"`. |\n| `moe_router_dtype` | `str` | `\"\"` | Data type for router computation (e.g. `\"fp32\"`). Default `\"\"`. |\n| `moe_permute_fusion` | `bool` | `False` | Enable permute fusion optimization for MoE. Default `False`. |\n| `moe_aux_loss_coeff` | `float \\| None` | `None` | Auxiliary load-balancing loss coefficient. Default `None`. |\n\n## Checkpoint Conversion\n\n| Field | Type | Default | Description |\n|-------|------|---------|-------------|\n| `megatron_model_type` | `str` | `\"\"` | Slime/Megatron model type string for checkpoint conversion (e.g. `\"qwen3.5-35B-A3B\"`). Used when the training recipe selects a non-bridge conversion mode. Default `\"\"`. |\n\n## Normalization Extras\n\n| Field | Type | Default | Description |\n|-------|------|---------|-------------|\n| `apply_layernorm_1p` | `bool` | `False` | Use zero-centered LayerNorm (add 1 to gamma). Default `False`. |\n\n## Attention Extras\n\n| Field | Type | Default | Description |\n|-------|------|---------|-------------|\n| `use_gated_attention` | `bool` | `False` | Enable gated attention mechanism. Default `False`. |\n| `attention_output_gate` | `bool` | `False` | Enable output gating on attention layers (required by some hybrid architectures such as Qwen 3.6). Default `False`. |\n\n## Position Encoding\n\n| Field | Type | Default | Description |\n|-------|------|---------|-------------|\n| `use_rotary_position_embeddings` | `bool` | `True` | Use RoPE positional encoding. Default `True`. |\n| `rotary_base` | `int` | `10000` | Base frequency for RoPE. Default `10000`. |\n| `rotary_percent` | `float` | `1.0` | Fraction of hidden dims to apply RoPE to. Default `1.0`. |\n\n## Other Fields\n\n| Field | Type | Default | Description |\n|-------|------|---------|-------------|\n| `moe_grouped_gemm` | `bool` | `False` | |\n| `moe_shared_expert_gate` | `bool` | `False` | |\n| `moe_router_topk` | `int` | `0` | |\n| `megatron_spec` | `list[str] \\| None` | `None` | |\n\n**Source:** [`modal_training_gym/common/models/base.py`](https://github.com/modal-projects/training-gym/blob/main/modal_training_gym/common/models/base.py)\n"
ModelArchitecture
Transformer architecture parameters for a specific model.
from modal_training_gym.common.models.base import ModelArchitecture
Transformer architecture parameters for a specific model.
These fields map directly to Megatron-LM model-parallel configuration
flags. Framework launchers read them to generate the correct CLI
arguments for distributed training.
Slime/Megatron model type string for checkpoint conversion (e.g. "qwen3.5-35B-A3B"). Used when the training recipe selects a non-bridge conversion mode. Default "".
