runtime/litgpt/config.py

# Copyright Lightning AI. Licensed under the Apache License 2.0, see LICENSE file.

import math
from copy import deepcopy
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any, Literal

import yaml
from typing_extensions import Self


def find_multiple(n: int, k: int) -> int:
    """Utility function for finding the nearest value to n which is a multiple of k.

    NOTE: We define this function in this module rather than `litgpt.utils` so that users can import
    this file to do configuration manipulations in Python environments which do not include all the dependencies
    demanded by `litgpt.utils`.
    """
    assert k > 0
    if n % k == 0:
        return n
    return n + k - (n % k)


@dataclass
class Config:
    name: str = ""
    hf_config: dict = field(default_factory=dict)
    # General size parameters
    block_size: int = 4096
    n_layer: int = 16
    n_embd: int = 4096
    vocab_size: int = 50254
    padding_multiple: int = 512
    padded_vocab_size: int | None = None
    # Transformer block (structure, normalizations)
    norm_class_name: Literal["LayerNorm", "RMSNorm"] = "LayerNorm"
    norm_eps: float = 1e-5
    norm_qk: bool = False
    norm_qk_type: Literal["default", "olmo2"] = "default"
    post_attention_norm: bool = False
    post_mlp_norm: bool = False
    parallel_residual: bool = True
    shared_attention_norm: bool = False
    # Transformer block (self-attention)
    n_head: int = 32
    head_size: int | None = None
    # to use multi-head attention (MHA), set this to `n_head` (default)
    # to use multi-query attention (MQA), set this to 1
    # to use grouped-query attention (GQA), set this to a value in between
    # Example with `n_head=4`
    # ┌───┐┌───┐┌───┐┌───┐     ┌───┐    ┌───┐             ┌───┐
    # │ v ││ v ││ v ││ v │     │ v │    │ v │             │ v │
    # └───┘└───┘└───┘└───┘     └───┘    └───┘             └───┘
    #   │    │    │    │         │        │                 │
    # ┌───┐┌───┐┌───┐┌───┐     ┌───┐    ┌───┐             ┌───┐
    # │ k ││ k ││ k ││ k │     │ k │    │ k │             │ k │
    # └───┘└───┘└───┘└───┘     └───┘    └───┘             └───┘
    #   │    │    │    │      ┌──┴──┐  ┌──┴──┐      ┌────┬──┴─┬────┐
    # ┌───┐┌───┐┌───┐┌───┐  ┌───┐┌───┐┌───┐┌───┐  ┌───┐┌───┐┌───┐┌───┐
    # │ q ││ q ││ q ││ q │  │ q ││ q ││ q ││ q │  │ q ││ q ││ q ││ q │
    # └───┘└───┘└───┘└───┘  └───┘└───┘└───┘└───┘  └───┘└───┘└───┘└───┘
    # ◀──────────────────▶  ◀──────────────────▶  ◀──────────────────▶
    #         MHA                    GQA                   MQA
    #   n_query_groups=4       n_query_groups=2      n_query_groups=1
    #
    # credit https://arxiv.org/pdf/2305.13245.pdf
    n_query_groups: int | None = None
    attn_bias: bool = False
    attention_scores_scalar: int | None = None
    # If `sliding_window_size` is given, sliding window attention with this
    # size is used in layers where `sliding_window_indices` has a 1. The
    # default is all 1, so that sliding window attention is used in all
    # layers. If `len(sliding_window_indices) > n_layer`, we only use the
    # initial part.
    sliding_window_size: int | None = None
    sliding_window_indices: list[int] | None = None
    # Training-only connectivity curriculum. ``qg_global_qg`` begins and ends
    # with the configured sliding pattern while using all-global attention in
    # the middle. Normal inference always uses ``sliding_window_indices``.
    sliding_window_curriculum: Literal["fixed", "qg_global_qg"] = "fixed"
    sliding_window_curriculum_early_fraction: float = 0.25
    sliding_window_curriculum_late_fraction: float = 0.9
    # if `attention_logit_softcapping` is used, cannot use optimized
    # `torch.nn.functional.scaled_dot_product_attention` (which implements
    # Flash attention), may result in higher memory and runtime footprint.
    attention_logit_softcapping: float | None = None
    attention_output_gate: Literal["none", "headwise", "elementwise"] = "none"
    attention_variant: Literal["standard", "shared_diff"] = "standard"
    # Shared-projection local attention plus normalized causal linear memory.
    # The global path reuses MQA Q/K/V and merges before the output projection.
    dual_path_linear_enabled: bool = False
    dual_path_linear_interval: int = 1
    dual_path_linear_rank: int = 8
    dual_path_linear_gate_init: float = 0.1
    dual_path_linear_explicit_backward: bool = False
    shared_diff_rank: int = 8
    # Ofir Press et al. (ACL 2020): s^k (sf)^(L-k) f^k.
    # Zero keeps the ordinary interleaved Transformer.
    sandwich_coefficient: int = 0
    # Token-adaptive residual commits on the final N Sandwich FFN sublayers.
    # Zero disables the router and retains exact official Sandwich behavior.
    # Training-only MTP transformer depth; excluded from normal inference.
    sandwich_adaptive_commit: int = 0
    # Adaptive Depth Memory (ADeM): content-addressed retrieval over earlier
    # residual states. ``control`` instantiates the exact same parameters as
    # ``adaptive`` but bypasses retrieval, enabling strict parameter-matched
    # experiments. ``none`` preserves the original model and checkpoint shape.
    depth_memory_mode: Literal["none", "control", "adaptive"] = "none"
    depth_memory_interval: int = 2
    depth_memory_start: int = 2
    depth_memory_max_sources: int = 4
    depth_memory_max_commit: float = 0.5
    depth_memory_score_scale: Literal["inverse_sqrt", "cosine", "factorized", "recurrent"] = "inverse_sqrt"
    # Fraction of training before a cosine fade of depth-memory commits begins.
    # 1.0 keeps depth memory constant. Values below 1.0 turn it into a
    # training scaffold that is exactly zero at the final optimizer step.
    depth_memory_fade_start_fraction: float = 1.0
    mtp_num_layers: int = 0
    mtp_shared_head: bool = False
    # Kimi K3 post-training: turn the pretrained MTP block into an EAGLE-3
    # drafter by fusing low-, mid-, and high-level AttnRes block outputs.
    mtp_eagle3: bool = False
    # Kimi K3 "Attention Residuals": token-dependent softmax mixing over
    # residual blocks. Zero keeps the ordinary residual stream.
    attn_residual_num_blocks: int = 0
    # Moonshot's released Kimi K3 reference specifies the number of decoder
    # layers per AttnRes block directly (12 in the 93-layer model).  This is
    # the preferred paper-faithful switch; ``attn_residual_num_blocks`` is
    # retained for older experiment checkpoints.
    attn_residual_block_size: int = 0
    # Kimi Linear hybrid: three KDA layers followed by one global MLA layer.
    # KDA itself is provided by the official flash-linear-attention package.
    kda_enabled: bool = False
    kda_mla_ratio: int = 3
    kda_num_heads: int = 12
    kda_head_dim: int = 32
    kda_short_conv_kernel_size: int = 4
    kda_safe_gate: bool = False
    kda_lower_bound: float = -5.0
    kda_full_rank_output_gate: bool = False
    mla_use_output_gate: bool = False
    # Kimi Linear uses NoPE in its global MLA layers.
    mla_use_nope: bool = False
    no_rope_layer_interval: int = 0
    # ResFormer mixes this fraction of layer-0 values into every later layer.
    value_residual_mix: float = 0.0
    xsa_projection: bool = False
    # Rotary position embedding (RoPE)
    rope_base: int = 10000
    rotary_percentage: float = 0.25
    rope_condense_ratio: int = 1
    rope_adjustments: dict | None = None
    rope_interleave: bool = False
    # Transformer block (MLP)
    intermediate_size: int | None = None
    moe_intermediate_size: int | None = None
    # Kimi K3 Stable LatentMoE / SiTU-GLU parameters.  Zero experts selects a
    # dense SiTU-GLU layer; positive values select the paper's normalized
    # latent routed path plus full-width shared experts.
    kimi_situ_beta: float = 4.0
    kimi_situ_linear_beta: float = 25.0
    kimi_latent_moe_num_experts: int = 0
    kimi_latent_moe_top_k: int = 1
    kimi_latent_moe_num_shared_experts: int = 2
    kimi_latent_moe_latent_size: int = 0
    kimi_latent_moe_expert_intermediate_size: int = 0
    kimi_latent_moe_shared_intermediate_size: int = 0
    kimi_latent_moe_quantile_balancing: bool = True
    kimi_latent_moe_qb_bins: int = 1000
    mlp_taper_schedule: Literal["none", "linear", "cosine", "sigmoid"] = "none"
    mlp_taper_start_ratio: float = 1.5
    mlp_taper_end_ratio: float = 0.5
    mlp_taper_multiple: int = 16
    bias: bool = True
    mlp_class_name: Literal[
        "GptNeoxMLP",
        "LLaMAMLP",
        "LLaMAPowLUMLP",
        "KimiSiTUGLUMLP",
        "KimiStableLatentMoEMLP",
        "GemmaMLP",
        "LLaMAMoE",
        "DSwiGLUMLP",
        "AdaptiveDSwiGLUMLP",
        "BlockSparseAdaptiveDSwiGLUMLP",
        "HiddenBlockDSwiGLUMLP",
        "TileRoutedBlockNormDSwiGLUMLP",
        "TileRoutedBlockLayeredNormDSwiGLUMLP",
        "TileRoutedDSwiGLUMLP",
        "TileRoutedGEGLUMLP",
        "TileRoutedReGLUMLP",
        "TileRoutedSquaredReGLUMLP",
        "TileRoutedSigmoidGLUMLP",
        "TileRoutedPowLUGLUMLP",
        "TileRoutedEulerGLUMLP",
        "TileRoutedEulerMishBlendGLUMLP",
        "TileRoutedContextualValueRotationEulerGLUMLP",
        "TileRoutedContextAdaptiveEulerTemperatureGLUMLP",
        "TileRoutedAdaptiveEulerGLUMLP",
        "TileRoutedDiverseEulerGLUMLP",
        "TileRoutedGroupTokenAttentionEulerGLUMLP",
        "TileRoutedHierarchicalEulerGLUMLP",
        "TileRoutedTokenExpertEulerGLUMLP",
        "TileRoutedDynamicGroupMixerEulerGLUMLP",
        "TileRoutedGroupCompetitionEulerGLUMLP",
        "TileRoutedEulerGroupStateMemoryGLUMLP",
        "TileRoutedEulerFusedGroupMixGLUMLP",
        "TileRoutedAdaptiveOperatorFieldEulerGLUMLP",
        "TileRoutedAdaptiveBasisEulerGLUMLP",
        "TileRoutedFusedAdaptiveBasisEulerGLUMLP",
        "TileRoutedAdaptiveDirectionEulerGLUMLP",
        "TileRoutedAdaptiveDirectionEulerFusedGLUMLP",
        "TileRoutedGateLawDSwiGLUMLP",
        "TileRoutedPowRatGLUMLP",
        "TileRoutedSinPowRatGLUMLP",
        "TileRoutedSPONGLUMLP",
        "TileRoutedTrainOnlySPONGLUMLP",
        "TileRoutedSPONPowRatGLUMLP",
        "TileRoutedAttentionMemoryGateDSwiGLUMLP",
        "TileRoutedInnerAttentionFFNMLP",
        "TileRoutedBilinearMemoryStepFFNMLP",
        "TileRoutedRationalGateDSwiGLUMLP",
        "TileRoutedExpGateDSwiGLUMLP",
        "TileRoutedHiddenRMSNormDSwiGLUMLP",
        "TileRoutedLateTokenHiddenRMSPreserveDSwiGLUMLP",
        "TileRoutedRMSMemoryDSwiGLUMLP",
        "TileRoutedFastRMSMemoryDSwiGLUMLP",
        "TileRoutedCoAdaptHiddenControllerDSwiGLUMLP",
        "TileRoutedHiddenDirectionMemoryDSwiGLUMLP",
        "TileRoutedScalarDirectionMemoryDSwiGLUMLP",
        "TileRoutedChannelMemoryDSwiGLUMLP",
        "TileRoutedChannelStateMemoryDSwiGLUMLP",
        "TileRoutedCenteredChannelStateMemoryDSwiGLUMLP",
        "TileRoutedCenteredChannelGroupStateMemoryDSwiGLUMLP",
        "TileRoutedKeyedChannelMemoryDSwiGLUMLP",
        "TileRoutedGroupStateMemoryDSwiGLUMLP",
        "TileRoutedChannelGroupStateMemoryDSwiGLUMLP",
        "TileRoutedGRNDSwiGLUMLP",
        "TileRoutedGroupMixDSwiGLUMLP",
        "TileRoutedTokenGroupGateDSwiGLUMLP",
        "TileRoutedStackedDSwiGLUMLP",
        "TileRoutedChunkedStackedDSwiGLUMLP",
        "TileRoutedGatedStackedDSwiGLUMLP",
        "TileRoutedFullMainGatedStackedDSwiGLUMLP",
        "TileRoutedFullMainRMSBoundedGatedStackedDSwiGLUMLP",
        "TileRoutedFullMainOrthogonalRMSGatedStackedDSwiGLUMLP",
        "TileRoutedFullMainAlignedRMSGatedStackedDSwiGLUMLP",
        "TileRoutedFullMainNormedBranchRMSGatedStackedDSwiGLUMLP",
        "TileRoutedFullMainCompressedRMSGatedStackedDSwiGLUMLP",
        "TileRoutedLateOnlyRMSGatedStackedDSwiGLUMLP",
        "TileRoutedAdditiveLateTokenRMSGatedDSwiGLUMLP",
        "TileRoutedRMSBoundedStackedLayersDSwiGLUMLP",
        "TileRoutedStackedLayersDSwiGLUMLP",
        "TileRoutedLateLayerStackedLayersDSwiGLUMLP",
        "TileRoutedGatedStackedLayersDSwiGLUMLP",
        "TileRoutedAttentionOutputStackedLayersDSwiGLUMLP",
        "TileRoutedStageNormRMSBoundedStackedLayersDSwiGLUMLP",
        "TileRoutedWeightedStageRMSBoundedStackedLayersDSwiGLUMLP",
        "TileRoutedMomentumRMSBoundedStackedLayersDSwiGLUMLP",
        "TileRoutedStaticDSwiGLUMLP",
        "TileRoutedDSwiGLUMLPStaticA2",
        "TileRoutedDSwiGLUMLPStaticGPTS",
    ] = "GptNeoxMLP"
    gelu_approximate: str = "none"
    drelu_target_active: float = 0.1
    drelu_threshold_init: float = 0.5
    drelu_gate_sharpness: float = 10.0
    sparse_mlp_num_groups: int = 8
    sparse_mlp_min_active_groups: int = 1
    sparse_mlp_max_active_groups: int = 2
    grouped_mlp_shared_intermediate_size: int = 0
    grouped_mlp_shared_alpha: float = 1.0
    grouped_mlp_shared_alpha_learnable: bool = False
    grouped_mlp_hidden_rms_late_layer_start: int = 8
    grouped_mlp_hidden_rms_alpha_init: float = 0.1
    grouped_mlp_hidden_rms_alpha_max: float = 0.35
    grouped_mlp_hidden_rms_gate_init: float = 0.1
    grouped_mlp_hidden_rms_output_scale: float = 1.0
    grouped_mlp_hidden_rms_token_gate: bool = True
    rms_budgeted_block: bool = False
    # Nakanishi (2026), "Screening Is Enough". This replaces the complete
    # attention+FFN block stack with parallel gated screening tiles.
    multiscreen_enabled: bool = False
    # Replace every Nth standard Transformer block with a Multiscreen block.
    # Zero disables the hybrid and preserves the existing pure-stack switch.
    multiscreen_layer_interval: int = 0
    # In a periodic hybrid, retain every Nth key/value as a causal landmark.
    # One preserves dense screening; values above one reduce pairwise work.
    multiscreen_landmark_stride: int = 1
    multiscreen_mlp_enabled: bool = False
    multiscreen_num_heads: int = 16
    multiscreen_key_dim: int = 16
    multiscreen_value_dim: int = 64
    multiscreen_window_threshold: int = 256
    multiscreen_query_chunk_size: int = 128
    multiscreen_triton_inference: bool = True
    # Combine the four input projections and use a flattened output GEMM.
    # This preserves parameters and math while reducing launch/read overhead.
    multiscreen_fused_projections: bool = False
    # Recompute complete screening blocks during backward. This retains one
    # residual stream per layer instead of Q/K/V/gate/screening intermediates.
    multiscreen_activation_checkpointing: bool = False
    # Number of consecutive screening blocks per recompute segment. Larger
    # groups save fewer residual boundaries without increasing recompute FLOPs.
    multiscreen_checkpoint_group_size: int = 1
    # Compute the tied language head and CE in bounded token tiles during
    # pretraining instead of materializing B x T x vocab logits.
    multiscreen_chunked_lm_loss: bool = False
    multiscreen_lm_loss_chunk_tokens: int = 2048
    # Paper appendix M: at inference, learned windows larger than the
    # training context can optionally become infinite (SWE).
    multiscreen_swe_context_length: int = 0
    # Zero preserves paper-faithful learned windows. A positive value imposes
    # a strict local-memory ceiling, used by the RWKV hybrid so total decode
    # state cannot grow with context length.
    multiscreen_hard_max_window: int = 0
    # Small active residual calibration used only to close an exact matched
    # parameter budget for pure-Multiscreen comparisons.
    multiscreen_match_adapter_params: int = 0
    # Active final-channel calibration for exact parameter-matched controls.
    model_match_adapter_params: int = 0
    # RWKV-7 global recurrent memory. The hybrid uses a bounded Multiscreen
    # correction after each time-mixing layer and deliberately omits the
    # quadratic Transformer attention and channel-mixing FFN.
    rwkv7_enabled: bool = False
    rwkv7_hybrid_multiscreen: bool = False
    rwkv7_cuda_kernel: bool = True
    rwkv7_multiscreen_interval: int = 1
    # Exact 10,011,392-parameter comparison geometry: remove redundant norm
    # biases and spend the four remaining parameters on active fusion scales.
    rwkv7_exact_10m_norms: bool = False
    rwkv7_match_adapter_params: int = 0
    # Mamba-3 SISO fallback for small GPUs where the RWKV-7 kernel does not
    # reach enough batch x head parallelism.
    mamba3_enabled: bool = False
    mamba3_state_dim: int = 64
    mamba3_expand: int = 2
    mamba3_head_dim: int = 64
    # Number of shared B/C latent-state groups. This is the recurrent analogue
    # of grouped-query attention: value heads share a smaller set of state
    # descriptors instead of each materializing an independent state.
    mamba3_num_groups: int = 1
    mamba3_chunk_size: int = 64
    mamba3_mimo: bool = False
    mamba3_mimo_rank: int = 4
    mamba3_multiscreen_interval: int = 4
    mamba3_recall_attention_interval: int = 0
    mamba3_ffn_interval: int = 0
    mamba3_match_adapter_params: int = 0
    mamba3_activation_checkpointing: bool = False
    mamba3_checkpoint_group_size: int = 1
    # Throughput-first conditional input/output factorization. Each token reads
    # one position-striped embedding bank. The hierarchical head predicts the
    # high and low six-bit vocabulary coordinates independently.
    mamba3_embedding_banks: int = 1
    mamba3_hierarchical_vocab: bool = False
    mamba3_hierarchical_mixtures: int = 1
    mamba3_position_table_size: int = 0
    mamba3_phase_table_size: int = 0
    mamba3_token_scalar_params: int = 0
    rms_budget_dual_stream: bool = False
    rms_budget_attn_mlp: bool = False
    rms_budget_direction_aware: bool = True
    rms_budget_target_init: float = 0.085
    rms_budget_target_min: float = 0.04
    rms_budget_target_max: float = 0.12
    rms_budget_target_learnable: bool = True
    rms_budget_work_scale: float = 1.0
    rms_budget_min_scale: float = 0.0
    rms_budget_exact_target: bool = False
    rms_memory_alpha_max: float = 0.35
    rms_memory_target_init: float = 1.0
    rms_memory_target_min: float = 0.5
    rms_memory_target_max: float = 1.5
    rms_memory_mix_identity_logit: float = 6.0
    rms_memory_mix_offdiag_logit: float = -6.0
    rms_memory_token_gate: bool = False
    rms_memory_gate_init: float = 0.5
    fast_rms_memory_schedule: str = "late8_meanbus"
    fast_rms_memory_alpha_max: float = 0.30
    fast_rms_memory_use_rms_norm: bool = True
    hidden_controller_rank: int = 8
    hidden_controller_rho_frac: float = 0.00035
    hidden_controller_late_layer_start: int = 10
    hidden_direction_memory_slots: int = 4
    hidden_direction_memory_rho_frac: float = 0.00035
    hidden_direction_memory_late_layer_start: int = 10
    hidden_direction_memory_temperature: float = 1.0
    scalar_direction_memory_rho_frac: float = 0.00035
    scalar_direction_memory_late_layer_start: int = 10
    channel_memory_gain_max: float = 0.10
    channel_memory_late_layer_start: int = 0
    group_state_memory_gain_max: float = 0.10
    group_state_memory_late_layer_start: int = 0
    keyed_channel_memory_rank: int = 4
    powlu_m: float = 3.0
    grouped_gate_alpha_init: float = 0.0
    grouped_gate_alpha_max: float = 0.08
    grouped_gate_beta_init: float = 1.0
    grouped_gate_sin_eps: float = 0.0
    grouped_gate_sin_freq: float = 1.0
    grouped_gate_spon_max: float = 0.08
    grouped_gate_layer_schedule: str = "constant"
    grouped_gate_spon_value: bool = False
    attn_gate_rank: int = 8
    attn_gate_slots: int = 4
    attn_gate_alpha_init: float = 0.25
    attn_gate_alpha_max: float = 1.0
    attn_gate_temperature: float = 1.0
    inner_attn_rank: int = 8
    inner_attn_alpha_init: float = 0.25
    inner_attn_alpha_max: float = 0.75
    inner_attn_temperature: float = 1.0
    bilinear_memory_depth: int = 2
    bilinear_memory_alpha_init: float = 0.02
    bilinear_memory_alpha_max: float = 0.35
    bilinear_memory_target_ratio: float = 0.20
    bilinear_memory_bus_mix: float = 0.25
    bilinear_memory_state_update: float = 0.05
    rms_weight_reparam: bool = False
    mlp_matrix_optimizer: str = "adamw"
    grouped_mlp_stack_alpha: float = 0.1
    grouped_mlp_stack_alpha_learnable: bool = False
    grouped_mlp_stack_depth: int = 1
    grouped_mlp_stack_output_scale: float = 1.0
    grouped_mlp_stack_output_scale_learnable: bool = False
    grouped_mlp_stack_init_scale: float = 1.0
    splitstack_late_layer_start: int = 0
    splitstack_initial_gate: float = 0.1
    splitstack_attention_alpha: float = -0.03
    grouped_mlp_contrastive_intermediate_size: int = 0
    grouped_mlp_contrastive_gate: float = 0.0
    grouped_mlp_contrastive_gate_learnable: bool = False
    grouped_mlp_contrastive_gate_max: float = 1.0
    grouped_mlp_branch_start_layer: int = 0
    grouped_mlp_channel_rotation: bool = False
    grouped_mlp_channel_rotation_stride: int | None = None
    grouped_mlp_mix_every_n_layers: int = 0
    grouped_mlp_mix_rank: int = 0
    grouped_mlp_mix_alpha: float = 1.0
    grouped_mlp_mix_alpha_learnable: bool = False
    n_expert: int = 0
    n_shared_expert: int | None = None
    n_expert_groups: int | None = None
    n_topk_groups: int | None = None
    n_topk_scores_per_group: int | None = None
    n_expert_per_token: int = 0
    first_k_dense_replace: int | None = None
    routed_scaling_factor: float = 1.0
    norm_topk_prob: bool = False
    # GPT before/after blocks
    scale_embeddings: bool = False
    lm_head_bias: bool = False
    final_logit_softcapping: float | None = None
    norm_1: bool = True
    norm_2: bool = True
    latent_attention: dict | None = None
    # The base period of the RoPE embeddings for local attention.
    # If not provided, `rope_base` will be used for both local and global attention.
    rope_local_base_freq: float | None = None
    # If provided, must have `>= n_layer` entries, either 0 or 1. For 0,
    # `rope_base` is used, for 1 `rope_local_base_freq` is used. If
    # `len(rope_indices) > n_layer`, we only use the initial part.
    rope_indices: list[int] | None = None
    initializer_range: float | None = None

    def __post_init__(self):
        if not self.name:
            self.name = self.hf_config.get("name", self.name)

        if self.head_size is None:
            assert self.n_embd % self.n_head == 0
            self.head_size = self.n_embd // self.n_head

        # vocab size should be a power of 2 to be optimal on hardware. compute the closest value
        if self.padded_vocab_size is None:
            self.padded_vocab_size = find_multiple(self.vocab_size, self.padding_multiple)
        else:
            # vocab size shouldn't be larger than padded vocab size
            self.vocab_size = min(self.vocab_size, self.padded_vocab_size)

        # compute the number of query groups
        if self.n_query_groups is not None:
            assert self.n_head % self.n_query_groups == 0
        else:
            self.n_query_groups = self.n_head

        assert self.attention_output_gate in {"none", "headwise", "elementwise"}
        assert self.attention_variant in {"standard", "shared_diff"}
        assert 0 <= self.sandwich_coefficient < self.n_layer
        assert 0 <= self.sandwich_adaptive_commit <= self.sandwich_coefficient
        if self.sandwich_adaptive_commit and not self.sandwich_coefficient:
            raise ValueError("Adaptive Sandwich commits require sandwich_coefficient > 0.")
        assert self.depth_memory_mode in {"none", "control", "adaptive"}
        assert self.depth_memory_interval > 0
        assert self.depth_memory_start > 0
        assert self.depth_memory_max_sources > 0
        assert 0.0 < self.depth_memory_max_commit <= 1.0
        assert self.depth_memory_score_scale in {"inverse_sqrt", "cosine", "factorized", "recurrent"}
        assert 0.0 <= self.depth_memory_fade_start_fraction <= 1.0
        if self.mtp_num_layers not in {0, 1}:
            raise ValueError("The paper-faithful path currently supports exactly one MTP module.")
        if self.mtp_shared_head:
            raise ValueError(
                "mtp_shared_head was a legacy shortcut; DeepSeek MTP already shares the real LM head."
            )
        if self.mtp_eagle3:
            if self.mtp_num_layers != 1:
                raise ValueError("EAGLE-3 draft fine-tuning requires mtp_num_layers=1.")
            attnres_blocks = (
                math.ceil(self.n_layer / self.attn_residual_block_size)
                if self.attn_residual_block_size
                else self.attn_residual_num_blocks
            )
            if attnres_blocks < 4:
                raise ValueError(
                    "EAGLE-3 feature fusion requires at least four AttnRes blocks "
                    "for distinct low-, mid-, and high-level features."
                )
        if self.attn_residual_num_blocks < 0:
            raise ValueError("attn_residual_num_blocks cannot be negative.")
        if self.attn_residual_block_size < 0:
            raise ValueError("attn_residual_block_size cannot be negative.")
        if self.attn_residual_num_blocks and self.attn_residual_block_size:
            raise ValueError(
                "Select either legacy attn_residual_num_blocks or "
                "paper-faithful attn_residual_block_size."
            )
        if self.attn_residual_num_blocks:
            num_sublayers = 2 * self.n_layer if self.sandwich_coefficient else self.n_layer
            if num_sublayers % self.attn_residual_num_blocks:
                raise ValueError(
                    "The decoder-layer count must be divisible by "
                    "attn_residual_num_blocks in the legacy path."
                )
        if self.kda_mla_ratio <= 0:
            raise ValueError("kda_mla_ratio must be positive.")
        if self.kda_num_heads <= 0 or self.kda_head_dim <= 0:
            raise ValueError("KDA head count and head dimension must be positive.")
        if self.kda_short_conv_kernel_size <= 0:
            raise ValueError("kda_short_conv_kernel_size must be positive.")
        if self.kda_safe_gate and not -5.0 <= self.kda_lower_bound < 0.0:
            raise ValueError("KDA safe-gate lower bound must be in [-5, 0).")
        if self.kda_enabled and self.latent_attention is None:
            raise ValueError("The KDA hybrid requires latent_attention for its periodic MLA layers.")
        if self.sandwich_coefficient:
            assert not self.parallel_residual
            assert not self.rms_budgeted_block
            assert self.value_residual_mix == 0.0
            assert self.rope_indices is None
        assert self.shared_diff_rank > 0
        if self.attention_variant == "shared_diff":
            assert self.n_head % 2 == 0
            assert self.n_query_groups is not None
            if self.n_query_groups < 2 or self.n_query_groups % 2 != 0:
                raise ValueError(
                    "Shared DIFF requires an even baseline n_query_groups >= 2 because the official "
                    "parameter-matched geometry halves both query and KV head counts; MQA is unsupported."
                )
            assert (self.n_head // 2) % (self.n_query_groups // 2) == 0
            if self.attention_output_gate != "none" or self.norm_qk or self.value_residual_mix > 0.0:
                raise ValueError("Shared DIFF does not support attention gates, QK norm, or value residuals.")
            if self.xsa_projection or self.sliding_window_size is not None:
                raise ValueError("Shared DIFF does not support XSA projection or sliding-window attention.")
        assert self.no_rope_layer_interval >= 0
        assert self.multiscreen_num_heads > 0
        assert self.multiscreen_key_dim >= 2
        assert self.multiscreen_value_dim > 0
        assert self.multiscreen_window_threshold > 1
        assert self.multiscreen_query_chunk_size > 0
        assert self.multiscreen_swe_context_length >= 0
        assert self.multiscreen_hard_max_window >= 0
        if self.rwkv7_hybrid_multiscreen and self.multiscreen_hard_max_window <= 0:
            raise ValueError("The RWKV-7 hybrid requires multiscreen_hard_max_window > 0.")
        if self.rwkv7_hybrid_multiscreen and not self.rwkv7_enabled:
            raise ValueError("rwkv7_hybrid_multiscreen requires rwkv7_enabled.")
        if self.rwkv7_multiscreen_interval <= 0:
            raise ValueError("rwkv7_multiscreen_interval must be positive.")
        if self.rwkv7_match_adapter_params < 0:
            raise ValueError("rwkv7_match_adapter_params cannot be negative.")
        if self.mamba3_enabled and self.rwkv7_enabled:
            raise ValueError("Select either RWKV-7 or Mamba-3 as the global recurrent mixer.")
        if self.mamba3_enabled and self.multiscreen_enabled:
            raise ValueError("Mamba-3 hybrid uses periodic Multiscreen, not the pure Multiscreen block stack.")
        if self.mamba3_state_dim <= 0 or self.mamba3_expand <= 0:
            raise ValueError("Mamba-3 state dimension and expansion must be positive.")
        if self.n_embd * self.mamba3_expand % self.mamba3_head_dim:
            raise ValueError("Mamba-3 expanded width must be divisible by mamba3_head_dim.")
        mamba3_heads = self.n_embd * self.mamba3_expand // self.mamba3_head_dim
        if self.mamba3_num_groups <= 0 or mamba3_heads % self.mamba3_num_groups:
            raise ValueError("Mamba-3 heads must be divisible by mamba3_num_groups.")
        if self.mamba3_chunk_size <= 0 or self.mamba3_multiscreen_interval <= 0:
            raise ValueError("Mamba-3 chunk size and Multiscreen interval must be positive.")
        if self.mamba3_mimo_rank <= 0:
            raise ValueError("Mamba-3 MIMO rank must be positive.")
        if self.mamba3_recall_attention_interval < 0 or self.mamba3_ffn_interval < 0:
            raise ValueError("Mamba-3 recall-attention and FFN intervals cannot be negative.")
        if self.mamba3_match_adapter_params < 0:
            raise ValueError("mamba3_match_adapter_params cannot be negative.")
        if self.mamba3_checkpoint_group_size <= 0:
            raise ValueError("mamba3_checkpoint_group_size must be positive.")
        if self.mamba3_embedding_banks <= 0:
            raise ValueError("mamba3_embedding_banks must be positive.")
        if (
            self.mamba3_position_table_size < 0
            or self.mamba3_phase_table_size < 0
            or self.mamba3_token_scalar_params < 0
        ):
            raise ValueError("Mamba-3 sparse table sizes cannot be negative.")
        if self.mamba3_hierarchical_vocab:
            if not self.mamba3_enabled:
                raise ValueError("The hierarchical vocabulary head requires Mamba-3.")
            if self.padded_vocab_size != 4096:
                raise ValueError("The hierarchical vocabulary head requires a 4096-token vocabulary.")
            if self.mamba3_embedding_banks == 1:
                raise ValueError("The throughput head requires multiple striped embedding banks.")
            if not 1 <= self.mamba3_hierarchical_mixtures <= 8:
                raise ValueError("Hierarchical vocabulary mixtures must be in [1, 8].")
        if self.multiscreen_match_adapter_params < 0:
            raise ValueError("multiscreen_match_adapter_params cannot be negative.")
        if self.dual_path_linear_enabled:
            if self.n_query_groups != 1:
                raise ValueError("dual_path_linear_enabled currently requires MQA.")
            if self.dual_path_linear_interval <= 0:
                raise ValueError("dual_path_linear_interval must be positive.")
            if not 0 < self.dual_path_linear_rank <= self.head_size:
                raise ValueError("dual_path_linear_rank must be in (0, head_size].")
        if self.model_match_adapter_params < 0:
            raise ValueError("model_match_adapter_params cannot be negative.")
        if self.multiscreen_layer_interval < 0:
            raise ValueError("multiscreen_layer_interval cannot be negative.")
        if self.multiscreen_landmark_stride <= 0:
            raise ValueError("multiscreen_landmark_stride must be positive.")
        if self.rwkv7_enabled and self.n_embd % 64:
            raise ValueError("RWKV-7 requires n_embd divisible by its 64-wide heads.")
        if self.rwkv7_enabled and (
            self.sandwich_coefficient or self.depth_memory_mode != "none" or self.mtp_num_layers
        ):
            raise ValueError("RWKV-7 is a complete block replacement and cannot use Sandwich, depth memory, or MTP.")
        if self.rwkv7_enabled and self.multiscreen_enabled:
            raise ValueError(
                "Use rwkv7_hybrid_multiscreen for the hybrid; multiscreen_enabled selects the pure architecture."
            )
        if self.multiscreen_enabled:
            if self.sandwich_coefficient or self.depth_memory_mode != "none":
                raise ValueError(
                    "Multiscreen is a complete block replacement and cannot use "
                    "Sandwich or depth memory."
                )
        elif self.multiscreen_mlp_enabled and self.multiscreen_layer_interval == 0:
            raise ValueError(
                "multiscreen_mlp_enabled requires multiscreen_enabled or "
                "multiscreen_layer_interval."
            )
        assert 0.0 <= self.value_residual_mix <= 1.0
        if self.kimi_situ_beta <= 0 or self.kimi_situ_linear_beta <= 0:
            raise ValueError("SiTU-GLU soft-cap constants must be positive.")
        if self.mlp_class_name == "KimiStableLatentMoEMLP":
            if self.kimi_latent_moe_num_experts <= 0:
                raise ValueError("Stable LatentMoE requires routed experts.")
            if not 0 < self.kimi_latent_moe_top_k <= self.kimi_latent_moe_num_experts:
                raise ValueError("Stable LatentMoE top-k must be in [1, num_experts].")
            if (
                self.kimi_latent_moe_quantile_balancing
                and self.kimi_latent_moe_top_k == self.kimi_latent_moe_num_experts
            ):
                raise ValueError(
                    "Kimi Quantile Balancing requires top-k < num_experts."
                )
            if self.kimi_latent_moe_num_shared_experts != 2:
                raise ValueError("Kimi K3 uses exactly two shared experts.")
            if (
                self.kimi_latent_moe_latent_size <= 0
                or self.kimi_latent_moe_expert_intermediate_size <= 0
                or self.kimi_latent_moe_shared_intermediate_size <= 0
            ):
                raise ValueError("Stable LatentMoE dimensions must be positive.")
            if self.kimi_latent_moe_qb_bins <= 1:
                raise ValueError("Kimi Quantile Balancing requires at least two bins.")
        assert self.mlp_taper_schedule in {"none", "linear", "cosine", "sigmoid"}
        assert self.mlp_taper_multiple > 0

        # compute the intermediate size for MLP if not set
        if self.intermediate_size is None:
            if self.mlp_class_name in {"LLaMAMLP", "LLaMAPowLUMLP"}:
                raise ValueError(f"The config {self.name!r}, needs to set the `intermediate_size`")
            self.intermediate_size = 4 * self.n_embd
        if self.mlp_taper_schedule != "none":
            assert self.n_layer > 1
            assert self.mlp_taper_start_ratio > self.mlp_taper_end_ratio > 0.0
            assert self.intermediate_size % self.mlp_taper_multiple == 0
            if self.first_k_dense_replace is not None:
                raise ValueError("MLP tapering is incompatible with first_k_dense_replace")
            self.tapered_mlp_intermediate_sizes()
        assert 0.0 < self.powlu_m < 10.0
        if self.mlp_class_name in {
            "BlockSparseAdaptiveDSwiGLUMLP",
            "HiddenBlockDSwiGLUMLP",
            "TileRoutedBlockNormDSwiGLUMLP",
            "TileRoutedBlockLayeredNormDSwiGLUMLP",
            "TileRoutedDSwiGLUMLP",
            "TileRoutedGEGLUMLP",
            "TileRoutedReGLUMLP",
            "TileRoutedSquaredReGLUMLP",
            "TileRoutedSigmoidGLUMLP",
            "TileRoutedPowLUGLUMLP",
            "TileRoutedEulerGLUMLP",
        "TileRoutedEulerMishBlendGLUMLP",
        "TileRoutedContextualValueRotationEulerGLUMLP",
        "TileRoutedContextAdaptiveEulerTemperatureGLUMLP",
        "TileRoutedAdaptiveEulerGLUMLP",
        "TileRoutedDiverseEulerGLUMLP",
        "TileRoutedGroupTokenAttentionEulerGLUMLP",
        "TileRoutedHierarchicalEulerGLUMLP",
        "TileRoutedTokenExpertEulerGLUMLP",
        "TileRoutedDynamicGroupMixerEulerGLUMLP",
        "TileRoutedGroupCompetitionEulerGLUMLP",
        "TileRoutedEulerGroupStateMemoryGLUMLP",
        "TileRoutedEulerFusedGroupMixGLUMLP",
        "TileRoutedAdaptiveOperatorFieldEulerGLUMLP",
        "TileRoutedAdaptiveBasisEulerGLUMLP",
        "TileRoutedFusedAdaptiveBasisEulerGLUMLP",
        "TileRoutedAdaptiveDirectionEulerGLUMLP",
        "TileRoutedAdaptiveDirectionEulerFusedGLUMLP",
            "TileRoutedGateLawDSwiGLUMLP",
            "TileRoutedPowRatGLUMLP",
            "TileRoutedSinPowRatGLUMLP",
            "TileRoutedSPONGLUMLP",
            "TileRoutedTrainOnlySPONGLUMLP",
            "TileRoutedSPONPowRatGLUMLP",
            "TileRoutedAttentionMemoryGateDSwiGLUMLP",
            "TileRoutedInnerAttentionFFNMLP",
            "TileRoutedBilinearMemoryStepFFNMLP",
            "TileRoutedRationalGateDSwiGLUMLP",
            "TileRoutedExpGateDSwiGLUMLP",
            "TileRoutedHiddenRMSNormDSwiGLUMLP",
        "TileRoutedLateTokenHiddenRMSPreserveDSwiGLUMLP",
            "TileRoutedRMSMemoryDSwiGLUMLP",
            "TileRoutedFastRMSMemoryDSwiGLUMLP",
        "TileRoutedCoAdaptHiddenControllerDSwiGLUMLP",
        "TileRoutedHiddenDirectionMemoryDSwiGLUMLP",
        "TileRoutedScalarDirectionMemoryDSwiGLUMLP",
        "TileRoutedChannelMemoryDSwiGLUMLP",
        "TileRoutedChannelStateMemoryDSwiGLUMLP",
        "TileRoutedCenteredChannelStateMemoryDSwiGLUMLP",
        "TileRoutedCenteredChannelGroupStateMemoryDSwiGLUMLP",
        "TileRoutedKeyedChannelMemoryDSwiGLUMLP",
        "TileRoutedGroupStateMemoryDSwiGLUMLP",
        "TileRoutedChannelGroupStateMemoryDSwiGLUMLP",
            "TileRoutedGRNDSwiGLUMLP",
            "TileRoutedGroupMixDSwiGLUMLP",
            "TileRoutedTokenGroupGateDSwiGLUMLP",
            "TileRoutedStackedDSwiGLUMLP",
            "TileRoutedChunkedStackedDSwiGLUMLP",
            "TileRoutedGatedStackedDSwiGLUMLP",
            "TileRoutedFullMainGatedStackedDSwiGLUMLP",
            "TileRoutedFullMainRMSBoundedGatedStackedDSwiGLUMLP",
            "TileRoutedFullMainOrthogonalRMSGatedStackedDSwiGLUMLP",
            "TileRoutedFullMainAlignedRMSGatedStackedDSwiGLUMLP",
            "TileRoutedFullMainNormedBranchRMSGatedStackedDSwiGLUMLP",
            "TileRoutedFullMainCompressedRMSGatedStackedDSwiGLUMLP",
            "TileRoutedLateOnlyRMSGatedStackedDSwiGLUMLP",
            "TileRoutedAdditiveLateTokenRMSGatedDSwiGLUMLP",
            "TileRoutedRMSBoundedStackedLayersDSwiGLUMLP",
            "TileRoutedStackedLayersDSwiGLUMLP",
            "TileRoutedLateLayerStackedLayersDSwiGLUMLP",
            "TileRoutedGatedStackedLayersDSwiGLUMLP",
            "TileRoutedAttentionOutputStackedLayersDSwiGLUMLP",
            "TileRoutedStageNormRMSBoundedStackedLayersDSwiGLUMLP",
            "TileRoutedWeightedStageRMSBoundedStackedLayersDSwiGLUMLP",
            "TileRoutedMomentumRMSBoundedStackedLayersDSwiGLUMLP",
            "TileRoutedStaticDSwiGLUMLP",
            "TileRoutedDSwiGLUMLPStaticA2",
            "TileRoutedDSwiGLUMLPStaticGPTS",
        }:
            if self.mlp_class_name in {
                "TileRoutedStackedLayersDSwiGLUMLP",
                "TileRoutedRMSBoundedStackedLayersDSwiGLUMLP",
                "TileRoutedChunkedStackedDSwiGLUMLP",
                "TileRoutedLateLayerStackedLayersDSwiGLUMLP",
                "TileRoutedGatedStackedLayersDSwiGLUMLP",
                "TileRoutedAttentionOutputStackedLayersDSwiGLUMLP",
                "TileRoutedStageNormRMSBoundedStackedLayersDSwiGLUMLP",
                "TileRoutedWeightedStageRMSBoundedStackedLayersDSwiGLUMLP",
                "TileRoutedMomentumRMSBoundedStackedLayersDSwiGLUMLP",
            }:
                assert self.grouped_mlp_stack_depth >= 1
                assert self.intermediate_size % (self.sparse_mlp_num_groups * self.grouped_mlp_stack_depth) == 0
            else:
                assert self.intermediate_size % self.sparse_mlp_num_groups == 0
            assert 1 <= self.sparse_mlp_min_active_groups <= self.sparse_mlp_max_active_groups
            assert self.sparse_mlp_max_active_groups <= self.sparse_mlp_num_groups
            assert self.grouped_mlp_shared_intermediate_size >= 0
            assert self.grouped_mlp_mix_every_n_layers >= 0
            assert self.grouped_mlp_mix_rank >= 0
            assert self.grouped_mlp_hidden_rms_late_layer_start >= 0
            assert self.grouped_mlp_hidden_rms_alpha_max > 0.0
            assert -self.grouped_mlp_hidden_rms_alpha_max < self.grouped_mlp_hidden_rms_alpha_init < self.grouped_mlp_hidden_rms_alpha_max
            assert 0.0 < self.grouped_mlp_hidden_rms_gate_init < 1.0
            assert self.grouped_mlp_hidden_rms_output_scale >= 0.0
            assert self.rms_budget_target_min > 0.0
            assert self.rms_budget_target_max > self.rms_budget_target_min
            assert self.rms_budget_target_min < self.rms_budget_target_init < self.rms_budget_target_max
            assert self.rms_budget_work_scale >= 0.0
            assert self.rms_budget_min_scale >= 0.0
            assert self.rms_memory_alpha_max > 0.0
            assert self.rms_memory_target_min > 0.0
            assert self.rms_memory_target_max > self.rms_memory_target_min
            assert self.rms_memory_target_min < self.rms_memory_target_init < self.rms_memory_target_max
            assert 0.0 < self.rms_memory_gate_init < 1.0
            assert self.mlp_matrix_optimizer in {"adamw", "himuon_tile32"}
            assert self.fast_rms_memory_schedule in {"late8_meanbus", "late9_meanbus", "all_tiny_meanbus"}
            assert self.fast_rms_memory_alpha_max > 0.0
            assert self.hidden_controller_rank > 0
            assert self.hidden_controller_rho_frac >= 0.0
            assert self.hidden_controller_late_layer_start >= 0
            assert self.hidden_direction_memory_slots > 1
            assert self.hidden_direction_memory_rho_frac >= 0.0
            assert self.hidden_direction_memory_late_layer_start >= 0
            assert self.hidden_direction_memory_temperature > 0.0
            assert self.scalar_direction_memory_rho_frac >= 0.0
            assert self.scalar_direction_memory_late_layer_start >= 0
            assert self.channel_memory_gain_max >= 0.0
            assert self.channel_memory_late_layer_start >= 0
            assert self.group_state_memory_gain_max >= 0.0
            assert self.group_state_memory_late_layer_start >= 0
            assert self.keyed_channel_memory_rank > 0
            assert self.grouped_gate_alpha_max > 0.0
            assert -self.grouped_gate_alpha_max < self.grouped_gate_alpha_init < self.grouped_gate_alpha_max
            assert self.grouped_gate_beta_init > 0.0
            assert 0.0 <= self.grouped_gate_sin_eps <= 0.5
            assert self.grouped_gate_sin_freq > 0.0
            assert self.grouped_gate_spon_max >= 0.0
            assert self.grouped_gate_layer_schedule in {"constant", "hiddenrms", "late_positive"}
            assert self.attn_gate_rank > 0
            assert self.attn_gate_slots > 1
            assert self.attn_gate_alpha_max > 0.0
            assert -self.attn_gate_alpha_max < self.attn_gate_alpha_init < self.attn_gate_alpha_max
            assert self.attn_gate_temperature > 0.0
            assert self.inner_attn_rank > 0
            assert self.inner_attn_alpha_max > 0.0
            assert -self.inner_attn_alpha_max < self.inner_attn_alpha_init < self.inner_attn_alpha_max
            assert self.inner_attn_temperature > 0.0
            assert self.bilinear_memory_depth >= 1
            assert self.bilinear_memory_alpha_max > 0.0
            assert -self.bilinear_memory_alpha_max < self.bilinear_memory_alpha_init < self.bilinear_memory_alpha_max
            assert 0.0 < self.bilinear_memory_target_ratio <= 2.0
            assert 0.0 <= self.bilinear_memory_bus_mix <= 1.0
            assert 0.0 <= self.bilinear_memory_state_update <= 1.0
            if self.grouped_mlp_channel_rotation_stride is not None:
                assert self.grouped_mlp_channel_rotation_stride >= 0

        self.rope_n_elem = int(self.rotary_percentage * self.head_size)

        if self.sliding_window_size is not None:
            self.sliding_window_indices = check_indicator_and_length(
                self.sliding_window_indices,
                name="sliding_window_indices",
                required_length=self.n_layer,
            )
        if self.sliding_window_curriculum == "qg_global_qg":
            if self.sliding_window_size is None:
                raise ValueError("qg_global_qg requires sliding_window_size")
            if not (
                0.0 < self.sliding_window_curriculum_early_fraction
                < self.sliding_window_curriculum_late_fraction
                < 1.0
            ):
                raise ValueError("sliding-window curriculum fractions must satisfy 0 < early < late < 1")

        if self.rope_local_base_freq is not None:
            self.rope_indices = check_indicator_and_length(
                self.rope_indices,
                name="rope_indices",
                required_length=self.n_layer,
            )

        if self.latent_attention is not None:
            self.q_lora_rank = self.latent_attention.get("q_lora_rank")
            self.kv_lora_rank = self.latent_attention.get("kv_lora_rank")
            self.qk_rope_head_dim = self.latent_attention.get("qk_rope_head_dim")
            self.qk_nope_head_dim = self.latent_attention.get("qk_nope_head_dim")
            self.v_head_dim = self.latent_attention.get("v_head_dim")
            assert self.q_lora_rank and self.kv_lora_rank
            assert self.qk_nope_head_dim and self.v_head_dim
            if self.mla_use_nope:
                if self.qk_rope_head_dim is None or self.qk_rope_head_dim < 0:
                    raise ValueError("NoPE MLA requires qk_rope_head_dim >= 0.")
            elif not self.qk_rope_head_dim:
                raise ValueError("RoPE MLA requires qk_rope_head_dim > 0.")
            self.qk_head_dim = self.qk_rope_head_dim + self.qk_nope_head_dim
            self.rope_n_elem = self.qk_rope_head_dim
            if (
                self.attention_output_gate != "none"
                or self.no_rope_layer_interval > 0
                or self.value_residual_mix > 0.0
            ):
                raise ValueError("Gated attention, NoPE, and value residuals currently require standard attention")
        if self.first_k_dense_replace is not None:
            assert self.mlp_class_name == "LLaMAMoE"

        if self.n_expert_groups is not None:
            assert self.n_expert % self.n_expert_groups == 0 and self.n_expert_groups > 1
            assert self.n_topk_groups is not None
            experts_per_group = self.n_expert // self.n_expert_groups
            assert self.n_topk_scores_per_group is not None and self.n_topk_scores_per_group <= experts_per_group

    def tapered_mlp_intermediate_sizes(self) -> list[int]:
        """Return the fixed-budget layer widths from Tapered Language Models."""
        if self.intermediate_size is None:
            raise ValueError("intermediate_size must be initialized before computing a taper")
        baseline = self.intermediate_size
        if self.mlp_taper_schedule == "none":
            return [baseline] * self.n_layer

        multiple = self.mlp_taper_multiple
        start = round(baseline * self.mlp_taper_start_ratio / multiple) * multiple
        end = round(baseline * self.mlp_taper_end_ratio / multiple) * multiple
        raw_widths: list[float] = []
        for layer_idx in range(self.n_layer):
            depth = layer_idx / (self.n_layer - 1)
            if self.mlp_taper_schedule == "linear":
                width = start - (start - end) * depth
            elif self.mlp_taper_schedule == "cosine":
                width = end + 0.5 * (start - end) * (1.0 + math.cos(math.pi * depth))
            else:
                width = end + (start - end) / (1.0 + math.exp(10.0 * (depth - 0.5)))
            raw_widths.append(width)

        widths = [round(width / multiple) * multiple for width in raw_widths]
        widths[0] = start
        widths[-1] = end
        target = self.n_layer * baseline
        delta = target - sum(widths)
        if delta % multiple:
            raise ValueError("The tapered MLP budget cannot be represented at the requested multiple")

        # The paper fixes both endpoints, then adjusts interior widths in
        # accelerator-friendly increments while preserving monotonicity.
        while delta:
            step = multiple if delta > 0 else -multiple
            moved = False
            for layer_idx in range(1, self.n_layer - 1):
                candidate = widths[layer_idx] + step
                if widths[layer_idx - 1] >= candidate >= widths[layer_idx + 1]:
                    widths[layer_idx] = candidate
                    delta -= step
                    moved = True
                    if delta == 0:
                        break
            if not moved:
                raise ValueError("Unable to preserve the exact MLP budget and a monotone taper")

        return widths

    @classmethod
    def from_name(cls, name: str, **kwargs: Any) -> Self | None:
        if name not in name_to_config:
            # search through all `config['hf_config']['name']`
            try:
                conf_dict = next(
                    config
                    for config in configs
                    if name == config["hf_config"]["name"]
                    or config["hf_config"]["org"] + "/" + config["hf_config"]["name"] == name
                )
            except StopIteration:
                raise ValueError(f"{name!r} is not a supported config name")
        else:
            conf_dict = name_to_config[name]

        conf_dict = conf_dict.copy()
        conf_dict.update(kwargs)
        return cls(**conf_dict)

    @classmethod
    def from_file(cls, path: str | Path, **kwargs: Any) -> Self:
        with open(path, encoding="utf-8") as fp:
            file_kwargs = yaml.safe_load(fp)
            if file_kwargs is None:
                raise ValueError(f"{path} is empty which is likely unexpected.")
        file_kwargs.update(kwargs)
        return cls(**file_kwargs)

    @classmethod
    def from_checkpoint(cls, path: Path, **kwargs: Any) -> Self:
        """Automatically load `model_config.yaml` and if it doesn't exist - a matching config from `litgpt/config.py`."""
        if (config_path := path / "model_config.yaml").is_file():
            return cls.from_file(config_path, **kwargs)
        if (model_name := path.name) in name_to_config:
            return cls.from_name(model_name, **kwargs)
        raise FileNotFoundError(f"For {str(path)!r} neither 'model_config.yaml' nor matching config exists.")

    @property
    def mlp_class(self) -> type:
        # `self.mlp_class_name` cannot be the type to keep the config serializable
        import litgpt.model

        return getattr(litgpt.model, self.mlp_class_name)

    @property
    def norm_class(self) -> type:
        # `self.norm_class_name` cannot be the type to keep the config serializable

        from functools import partial

        import torch  # Torch import is lazy to make config loading faster

        if self.norm_class_name == "RMSNorm":
            from litgpt.model import RMSNorm

            return partial(RMSNorm, add_unit_offset="Gemma" in self.name)

        if self.norm_class_name == "LayerNorm" and "OLMo" in self.name:
            # this makes it equivalent to `torch.nn.functional.layer_norm`
            # that is used by OLMo
            # Table 5 caption in the OLMo paper shows this - https://aclanthology.org/2024.acl-long.841
            return partial(torch.nn.LayerNorm, elementwise_affine=False)

        return getattr(torch.nn, self.norm_class_name)


def check_indicator_and_length(
    params: list[int] | None,
    name: str,
    required_length: int,
    use_initial_part: bool = True,
    def_val: int = 1,
) -> list[int]:
    if params is None:
        return [def_val] * required_length
    if len(params) != required_length:
        if use_initial_part and len(params) > required_length:
            params = params[:required_length]
        else:
            raise ValueError(f"{name} = {params}, must have length {required_length}")
    if not set(params).issubset({0, 1}):
        raise ValueError(f"{name} = {params}, must only contain 0 and 1")
    return params


########################
# Stability AI StableLM
########################
configs = [
    # https://huggingface.co/stabilityai/stablelm-base-alpha-3b/blob/main/config.json
    dict(name="stablelm-base-alpha-3b", hf_config=dict(org="stabilityai", name="stablelm-base-alpha-3b")),
    # https://huggingface.co/stabilityai/stablelm-base-alpha-7b/blob/main/config.json
    dict(
        name="stablelm-base-alpha-7b",
        hf_config=dict(org="stabilityai", name="stablelm-base-alpha-7b"),
        n_head=48,
        n_embd=6144,
        padding_multiple=256,
    ),
    # https://huggingface.co/stabilityai/stablelm-tuned-alpha-3b/blob/main/config.json
    dict(name="stablelm-tuned-alpha-3b", hf_config=dict(org="stabilityai", name="stablelm-tuned-alpha-3b"), n_head=32),
    # https://huggingface.co/stabilityai/stablelm-tuned-alpha-7b/blob/main/config.json
    dict(
        name="stablelm-tuned-alpha-7b",
        hf_config=dict(org="stabilityai", name="stablelm-tuned-alpha-7b"),
        n_head=48,
        n_embd=6144,
        padding_multiple=256,
    ),
    # https://huggingface.co/stabilityai/stablelm-3b-4e1t/blob/main/config.json
    dict(
        name="stablelm-3b-4e1t",
        hf_config=dict(org="stabilityai", name="stablelm-3b-4e1t"),
        padded_vocab_size=50304,
        n_layer=32,
        n_head=32,
        n_embd=2560,
        parallel_residual=False,
        bias=False,
        mlp_class_name="LLaMAMLP",
        intermediate_size=6912,
    ),
    # https://huggingface.co/stabilityai/stablelm-zephyr-3b/blob/main/config.json
    dict(
        name="stablelm-zephyr-3b",
        hf_config=dict(org="stabilityai", name="stablelm-zephyr-3b"),
        padded_vocab_size=50304,
        n_layer=32,
        n_head=32,
        n_embd=2560,
        parallel_residual=False,
        bias=False,
        mlp_class_name="LLaMAMLP",
        intermediate_size=6912,
    ),
]


##########################
# Stability AI StableCode
##########################
stablecode = [
    # https://huggingface.co/stabilityai/stablecode-completion-alpha-3b/blob/main/config.json
    dict(
        name="stablecode-completion-alpha-3b",
        hf_config=dict(org="stabilityai", name="stablecode-completion-alpha-3b"),
        block_size=16384,
        vocab_size=49152,
        n_layer=32,
        n_embd=2560,
    ),
    # https://huggingface.co/stabilityai/stablecode-completion-alpha-3b-4k/blob/main/config.json
    dict(
        name="stablecode-completion-alpha-3b-4k",
        hf_config=dict(org="stabilityai", name="stablecode-completion-alpha-3b-4k"),
        vocab_size=49152,
        n_layer=32,
        n_embd=2560,
    ),
    # https://huggingface.co/stabilityai/stablecode-instruct-alpha-3b/blob/main/config.json
    dict(
        name="stablecode-instruct-alpha-3b",
        hf_config=dict(org="stabilityai", name="stablecode-instruct-alpha-3b"),
        vocab_size=49152,
        n_layer=32,
        n_embd=2560,
    ),
    # https://huggingface.co/stabilityai/stable-code-3b/blob/main/config.json
    dict(
        name="stable-code-3b",
        hf_config=dict(org="stabilityai", name="stable-code-3b"),
        padded_vocab_size=50304,
        n_layer=32,
        n_embd=2560,
        block_size=16384,
        parallel_residual=False,
        bias=False,
        mlp_class_name="LLaMAMLP",
        intermediate_size=6912,
    ),
]
configs.extend(stablecode)


####################
# EleutherAI Pythia
####################
pythia = [
    # https://huggingface.co/EleutherAI/pythia-14m/blob/main/config.json
    dict(
        name="pythia-14m",
        hf_config=dict(org="EleutherAI", name="pythia-14m"),
        block_size=512,
        n_layer=6,
        n_embd=128,
        n_head=4,
        padding_multiple=128,
    ),
    # https://huggingface.co/EleutherAI/pythia-31m/blob/main/config.json
    dict(
        name="pythia-31m",
        hf_config=dict(org="EleutherAI", name="pythia-31m"),
        block_size=1024,
        n_layer=6,
        n_embd=256,
        n_head=8,
        padding_multiple=128,
    ),
    # https://huggingface.co/EleutherAI/pythia-70m/blob/main/config.json
    dict(
        name="pythia-70m",
        hf_config=dict(org="EleutherAI", name="pythia-70m"),
        block_size=2048,
        n_layer=6,
        n_embd=512,
        n_head=8,
        padding_multiple=128,
    ),
    # https://huggingface.co/EleutherAI/pythia-160m/blob/main/config.json
    dict(
        name="pythia-160m",
        hf_config=dict(org="EleutherAI", name="pythia-160m"),
        block_size=2048,
        n_layer=12,
        n_embd=768,
        n_head=12,
        padding_multiple=128,
    ),
    # https://huggingface.co/EleutherAI/pythia-410m/blob/main/config.json
    dict(
        name="pythia-410m",
        hf_config=dict(org="EleutherAI", name="pythia-410m"),
        block_size=2048,
        n_layer=24,
        n_embd=1024,
        n_head=16,
        padding_multiple=128,
    ),
    # https://huggingface.co/EleutherAI/pythia-1b/blob/main/config.json
    dict(
        name="pythia-1b",
        hf_config=dict(org="EleutherAI", name="pythia-1b"),
        block_size=2048,
        n_embd=2048,
        n_head=8,
        padding_multiple=128,
    ),
    # https://huggingface.co/EleutherAI/pythia-1.4b/blob/main/config.json
    dict(
        name="pythia-1.4b",
        hf_config=dict(org="EleutherAI", name="pythia-1.4b"),
        block_size=2048,
        n_layer=24,
        n_embd=2048,
        n_head=16,
        padding_multiple=128,
    ),
    # https://huggingface.co/EleutherAI/pythia-2.8b/blob/main/config.json
    dict(
        name="pythia-2.8b",
        hf_config=dict(org="EleutherAI", name="pythia-2.8b"),
        block_size=2048,
        n_layer=32,
        n_embd=2560,
        padding_multiple=128,
    ),
    # https://huggingface.co/EleutherAI/pythia-6.9b/blob/main/config.json
    dict(
        name="pythia-6.9b",
        hf_config=dict(org="EleutherAI", name="pythia-6.9b"),
        block_size=2048,
        n_layer=32,
        padding_multiple=256,
    ),
    # https://huggingface.co/EleutherAI/pythia-12b/blob/main/config.json
    dict(
        name="pythia-12b",
        hf_config=dict(org="EleutherAI", name="pythia-12b"),
        block_size=2048,
        n_layer=36,
        n_embd=5120,
        n_head=40,
    ),
]
configs.extend(pythia)
for c in pythia:
    # "pythia-14m" and "pythia-31m" don't have deduped version
    if c["name"] in ("pythia-14m", "pythia-31m"):
        continue
    copy = deepcopy(c)
    copy["name"] = f"{c['name']}-deduped"
    copy["hf_config"]["name"] = f"{c['hf_config']['name']}-deduped"
    configs.append(copy)


#################
# TII UAE Falcon
#################
falcon = [
    # https://huggingface.co/tiiuae/falcon-7b/blob/main/config.json
    dict(
        name="falcon-7b{}",
        hf_config=dict(org="tiiuae", name="falcon-7b{}"),
        block_size=2048,
        vocab_size=65024,
        padded_vocab_size=65024,
        n_layer=32,
        n_head=71,
        n_embd=4544,
        rotary_percentage=1.0,
        n_query_groups=1,
        bias=False,
        # this is not in the config, but in the original model implementation, only for this config
        shared_attention_norm=True,
    ),
    # https://huggingface.co/tiiuae/falcon-40b/blob/main/config.json
    dict(
        name="falcon-40b{}",
        hf_config=dict(org="tiiuae", name="falcon-40b{}"),
        block_size=2048,
        vocab_size=65024,
        padded_vocab_size=65024,
        n_layer=60,
        n_head=128,
        n_embd=8192,
        rotary_percentage=1.0,
        n_query_groups=8,
        bias=False,
    ),
]
for c in falcon:
    for kind in ("", "-instruct"):
        copy = deepcopy(c)
        copy["name"] = c["name"].format(kind)
        copy["hf_config"]["name"] = c["hf_config"]["name"].format(kind)
        configs.append(copy)

# https://huggingface.co/tiiuae/falcon-180b/blob/main/config.json
falcon180b = dict(
    name="falcon-180B{}",
    hf_config=dict(org="tiiuae", name="falcon-180B{}"),
    block_size=2048,
    vocab_size=65024,
    padded_vocab_size=65024,
    n_layer=80,
    n_head=232,
    n_embd=14848,
    rotary_percentage=1.0,
    n_query_groups=8,
    bias=False,
)

for kind in ("", "-chat"):
    copy = deepcopy(falcon180b)
    copy["name"] = falcon180b["name"].format(kind)
    copy["hf_config"]["name"] = falcon180b["hf_config"]["name"].format(kind)
    configs.append(copy)

falcon3 = [
    # https://huggingface.co/tiiuae/Falcon3-1B-Base/blob/main/config.json
    dict(
        name="Falcon3-1B{}",
        hf_config=dict(org="tiiuae", name="Falcon3-1B{}"),
        block_size=4096,
        vocab_size=131072,
        padded_vocab_size=131072,
        n_layer=18,
        n_head=8,
        n_query_groups=4,
        n_embd=2048,
        rotary_percentage=1.0,
        parallel_residual=False,
        rope_base=1000042,
        norm_eps=1e-6,
        bias=False,
        norm_class_name="RMSNorm",
        mlp_class_name="LLaMAMLP",
        intermediate_size=8192,
    ),
    # https://huggingface.co/tiiuae/Falcon3-3B-Base/blob/main/config.json
    dict(
        name="Falcon3-3B{}",
        hf_config=dict(org="tiiuae", name="Falcon3-3B{}"),
        block_size=32768,
        vocab_size=131072,
        padded_vocab_size=131072,
        n_layer=22,
        n_head=12,
        n_query_groups=4,
        n_embd=3072,
        rotary_percentage=1.0,
        parallel_residual=False,
        rope_base=1000042,
        norm_eps=1e-6,
        bias=False,
        norm_class_name="RMSNorm",
        mlp_class_name="LLaMAMLP",
        intermediate_size=9216,
    ),
    # https://huggingface.co/tiiuae/Falcon3-7B-Base/blob/main/config.json
    dict(
        name="Falcon3-7B{}",
        hf_config=dict(org="tiiuae", name="Falcon3-7B{}"),
        block_size=32768,
        vocab_size=131072,
        padded_vocab_size=131072,
        n_layer=28,
        n_head=12,
        n_query_groups=4,
        n_embd=3072,
        rotary_percentage=1.0,
        parallel_residual=False,
        rope_base=1000042,
        norm_eps=1e-6,
        bias=False,
        norm_class_name="RMSNorm",
        mlp_class_name="LLaMAMLP",
        intermediate_size=23040,
    ),
    # https://huggingface.co/tiiuae/Falcon3-10B-Base/blob/main/config.json
    dict(
        name="Falcon3-10B{}",
        hf_config=dict(org="tiiuae", name="Falcon3-10B{}"),
        block_size=32768,
        vocab_size=131072,
        padded_vocab_size=131072,
        n_layer=40,
        n_head=12,
        n_query_groups=4,
        n_embd=3072,
        rotary_percentage=1.0,
        parallel_residual=False,
        rope_base=1000042,
        norm_eps=1e-6,
        bias=False,
        norm_class_name="RMSNorm",
        mlp_class_name="LLaMAMLP",
        intermediate_size=23040,
    ),
]
for c in falcon3:
    for kind in ("-Base", "-Instruct"):
        copy = deepcopy(c)
        copy["name"] = c["name"].format(kind)
        copy["hf_config"]["name"] = c["hf_config"]["name"].format(kind)
        configs.append(copy)


#############################
# OpenLM Research Open LLaMA
#############################
open_LLaMA = [
    # https://huggingface.co/openlm-research/open_llama_3b/blob/main/config.json
    dict(
        name="open_llama_3b",
        hf_config=dict(org="openlm-research", name="open_llama_3b"),
        block_size=2048,
        vocab_size=32000,
        padding_multiple=64,
        n_layer=26,
        n_embd=3200,
        rotary_percentage=1.0,
        parallel_residual=False,
        bias=False,
        norm_class_name="RMSNorm",
        norm_eps=1e-6,
        mlp_class_name="LLaMAMLP",
        intermediate_size=8640,
    ),
    # https://huggingface.co/openlm-research/open_llama_7b/blob/main/config.json
    dict(
        name="open_llama_7b",
        hf_config=dict(org="openlm-research", name="open_llama_7b"),
        block_size=2048,
        vocab_size=32000,
        padding_multiple=64,
        n_layer=32,
        rotary_percentage=1.0,
        parallel_residual=False,
        bias=False,
        norm_class_name="RMSNorm",
        norm_eps=1e-6,
        mlp_class_name="LLaMAMLP",
        intermediate_size=11008,
    ),
    # https://huggingface.co/openlm-research/open_llama_13b/blob/main/config.json
    dict(
        name="open_llama_13b",
        hf_config=dict(org="openlm-research", name="open_llama_13b"),
        block_size=2048,
        vocab_size=32000,
        padding_multiple=64,
        n_layer=40,
        n_head=40,
        n_embd=5120,
        rotary_percentage=1.0,
        parallel_residual=False,
        bias=False,
        norm_class_name="RMSNorm",
        norm_eps=1e-6,
        mlp_class_name="LLaMAMLP",
        intermediate_size=13824,
    ),
]
configs.extend(open_LLaMA)

###############
# Meta LLaMA 2
###############
llama_2 = [
    # https://huggingface.co/meta-llama/Llama-2-7b-hf/blob/main/config.json
    dict(
        name="Llama-2-7b{}-hf",
        hf_config=dict(org="meta-llama", name="Llama-2-7b{}-hf"),
        vocab_size=32000,
        padding_multiple=64,
        n_layer=32,
        rotary_percentage=1.0,
        parallel_residual=False,
        bias=False,
        norm_class_name="RMSNorm",
        mlp_class_name="LLaMAMLP",
        intermediate_size=11008,
    ),
    # https://huggingface.co/meta-llama/Llama-2-13b-hf/blob/main/config.json
    dict(
        name="Llama-2-13b{}-hf",
        hf_config=dict(org="meta-llama", name="Llama-2-13b{}-hf"),
        vocab_size=32000,
        padding_multiple=64,
        n_layer=40,
        n_head=40,
        n_embd=5120,
        rotary_percentage=1.0,
        parallel_residual=False,
        bias=False,
        norm_class_name="RMSNorm",
        mlp_class_name="LLaMAMLP",
        intermediate_size=13824,
    ),
    # https://huggingface.co/meta-llama/Llama-2-70b-hf/blob/main/config.json
    dict(
        name="Llama-2-70b{}-hf",
        hf_config=dict(org="meta-llama", name="Llama-2-70b{}-hf"),
        vocab_size=32000,
        padding_multiple=64,
        n_layer=80,
        n_head=64,
        n_embd=8192,
        n_query_groups=8,
        rotary_percentage=1.0,
        parallel_residual=False,
        bias=False,
        norm_class_name="RMSNorm",
        mlp_class_name="LLaMAMLP",
        intermediate_size=28672,
    ),
]
for c in llama_2:
    for kind in ("", "-chat"):
        copy = deepcopy(c)
        copy["name"] = c["name"].format(kind)
        copy["hf_config"]["name"] = c["hf_config"]["name"].format(kind)
        configs.append(copy)


###############
# Meta LLaMA 3
###############
llama_3 = [
    # https://huggingface.co/meta-llama/Meta-Llama-3-8B/blob/main/config.json
    dict(
        name="Llama-3-8B{}",
        hf_config=dict(org="meta-llama", name="Meta-Llama-3-8B{}"),
        block_size=8192,
        vocab_size=128000,
        padded_vocab_size=128256,
        n_layer=32,
        n_head=32,
        n_query_groups=8,
        rotary_percentage=1.0,
        parallel_residual=False,
        bias=False,
        norm_class_name="RMSNorm",
        mlp_class_name="LLaMAMLP",
        intermediate_size=14336,
        rope_base=500000,
    ),
    # https://huggingface.co/meta-llama/Meta-Llama-3.1-8B/blob/main/config.json
    dict(
        name="Llama-3.1-8B{}",
        hf_config=dict(org="meta-llama", name="Meta-Llama-3.1-8B{}"),
        block_size=131072,
        vocab_size=128000,
        padded_vocab_size=128256,
        n_layer=32,
        n_head=32,
        n_query_groups=8,
        rotary_percentage=1.0,
        parallel_residual=False,
        bias=False,
        norm_class_name="RMSNorm",
        mlp_class_name="LLaMAMLP",
        intermediate_size=14336,
        rope_base=500000,
        rope_adjustments=dict(factor=8.0, low_freq_factor=1.0, high_freq_factor=4.0, original_max_seq_len=8192),
    ),
    # https://huggingface.co/meta-llama/Meta-Llama-3-70B/blob/main/config.json
    dict(
        name="Llama-3-70B{}",
        hf_config=dict(org="meta-llama", name="Meta-Llama-3-70B{}"),
        block_size=8192,
        vocab_size=128000,
        padded_vocab_size=128256,
        n_layer=80,
        n_head=64,
        n_embd=8192,
        n_query_groups=8,
        rotary_percentage=1.0,
        parallel_residual=False,
        bias=False,
        norm_class_name="RMSNorm",
        mlp_class_name="LLaMAMLP",
        intermediate_size=28672,
        rope_base=500000,
    ),
    # https://huggingface.co/meta-llama/Meta-Llama-3.1-70B/blob/main/config.json
    dict(
        name="Llama-3.1-70B{}",
        hf_config=dict(org="meta-llama", name="Meta-Llama-3.1-70B{}"),
        block_size=131072,
        vocab_size=128000,
        padded_vocab_size=128256,
        n_layer=80,
        n_head=64,
        n_embd=8192,
        n_query_groups=8,
        rotary_percentage=1.0,
        parallel_residual=False,
        bias=False,
        norm_class_name="RMSNorm",
        mlp_class_name="LLaMAMLP",
        intermediate_size=28672,
        rope_base=500000,
        rope_adjustments=dict(factor=8.0, low_freq_factor=1.0, high_freq_factor=4.0, original_max_seq_len=8192),
    ),
    # https://huggingface.co/meta-llama/Meta-Llama-3.1-405B/blob/main/config.json
    dict(
        name="Llama-3.1-405B{}",
        hf_config=dict(org="meta-llama", name="Meta-Llama-3.1-405B{}"),
        block_size=131072,
        vocab_size=128000,
        padded_vocab_size=128256,
        n_layer=126,
        n_head=128,
        n_embd=16384,
        n_query_groups=8,
        rotary_percentage=1.0,
        parallel_residual=False,
        bias=False,
        norm_class_name="RMSNorm",
        mlp_class_name="LLaMAMLP",
        intermediate_size=53248,
        rope_base=500000,
        rope_adjustments=dict(factor=8.0, low_freq_factor=1.0, high_freq_factor=4.0, original_max_seq_len=8192),
    ),
    # https://huggingface.co/meta-llama/Llama-3.2-1B/blob/main/config.json
    dict(
        name="Llama-3.2-1B{}",
        hf_config=dict(org="meta-llama", name="Llama-3.2-1B{}"),
        block_size=131072,
        vocab_size=128000,
        padded_vocab_size=128256,
        n_layer=16,
        n_embd=2048,
        n_head=32,
        n_query_groups=8,
        rotary_percentage=1.0,
        parallel_residual=False,
        bias=False,
        norm_class_name="RMSNorm",
        mlp_class_name="LLaMAMLP",
        intermediate_size=8192,
        rope_base=500000,
        rope_adjustments=dict(factor=32.0, low_freq_factor=1.0, high_freq_factor=4.0, original_max_seq_len=8192),
    ),
    # https://huggingface.co/meta-llama/Llama-3.2-3B/blob/main/config.json
    dict(
        name="Llama-3.2-3B{}",
        hf_config=dict(org="meta-llama", name="Llama-3.2-3B{}"),
        block_size=131072,
        vocab_size=128000,
        padded_vocab_size=128256,
        n_layer=28,
        n_embd=3072,
        n_head=24,
        n_query_groups=8,
        rotary_percentage=1.0,
        parallel_residual=False,
        bias=False,
        norm_class_name="RMSNorm",
        mlp_class_name="LLaMAMLP",
        intermediate_size=8192,
        rope_base=500000,
        rope_adjustments=dict(factor=32.0, low_freq_factor=1.0, high_freq_factor=4.0, original_max_seq_len=8192),
    ),
    # https://huggingface.co/meta-llama/Llama-3.3-70B-Instruct/blob/main/config.json
    dict(
        name="Llama-3.3-70B-Instruct",
        hf_config=dict(org="meta-llama", name="Llama-3.3-70B-Instruct"),
        block_size=131072,
        vocab_size=128000,
        padded_vocab_size=128256,
        n_layer=80,
        n_head=64,
        n_embd=8192,
        n_query_groups=8,
        rotary_percentage=1.0,
        parallel_residual=False,
        bias=False,
        norm_class_name="RMSNorm",
        mlp_class_name="LLaMAMLP",
        intermediate_size=28672,
        rope_base=500000,
        rope_adjustments=dict(factor=8.0, low_freq_factor=1.0, high_freq_factor=4.0, original_max_seq_len=8192),
    ),
]
for c in llama_3:
    if c["name"] == "Llama-3.3-70B-Instruct":
        configs.append(c)
        continue
    for kind in ("", "-Instruct"):
        copy = deepcopy(c)
        copy["name"] = c["name"].format(kind)
        copy["hf_config"]["name"] = c["hf_config"]["name"].format(kind)
        configs.append(copy)

#########################
# NVIDIA Llama Nemotron
#########################
configs.append(
    dict(
        name="Llama-3.1-Nemotron-70B-Instruct-HF",
        hf_config=dict(org="nvidia", name="Llama-3.1-Nemotron-70B-Instruct-HF"),
        block_size=131072,
        vocab_size=128000,
        padded_vocab_size=128256,
        n_layer=80,
        n_head=64,
        n_embd=8192,
        n_query_groups=8,
        rotary_percentage=1.0,
        parallel_residual=False,
        bias=False,
        norm_class_name="RMSNorm",
        mlp_class_name="LLaMAMLP",
        intermediate_size=28672,
        rope_base=500000,
        rope_adjustments=dict(factor=8.0, low_freq_factor=1.0, high_freq_factor=4.0, original_max_seq_len=8192),
    ),
)

#################
# Allen AI OLMo
#################
olmo = [
    # https://huggingface.co/allenai/OLMo-1B-hf/blob/main/config.json
    dict(
        name="OLMo-1B-hf",
        hf_config=dict(org="allenai", name="OLMo-1B-hf"),
        vocab_size=50280,
        padded_vocab_size=50304,
        block_size=2048,
        n_embd=2048,
        n_layer=16,
        n_head=16,
        rotary_percentage=1.0,
        parallel_residual=False,
        bias=False,
        norm_class_name="LayerNorm",
        mlp_class_name="LLaMAMLP",
        intermediate_size=8192,
    ),
    # https://huggingface.co/allenai/OLMo-7B-hf/blob/main/config.json
    dict(
        name="OLMo-7B-hf",
        hf_config=dict(org="allenai", name="OLMo-7B-hf"),
        vocab_size=50280,
        padded_vocab_size=50304,
        block_size=2048,
        n_layer=32,
        n_head=32,
        rotary_percentage=1.0,
        parallel_residual=False,
        bias=False,
        norm_class_name="LayerNorm",
        mlp_class_name="LLaMAMLP",
        intermediate_size=11008,
    ),
    # https://huggingface.co/allenai/OLMo-7B-Instruct-hf/blob/main/config.json
    dict(
        name="OLMo-7B-Instruct-hf",
        hf_config=dict(org="allenai", name="OLMo-7B-Instruct-hf"),
        vocab_size=50280,
        padded_vocab_size=50304,
        block_size=2048,
        n_layer=32,
        n_head=32,
        rotary_percentage=1.0,
        parallel_residual=False,
        bias=False,
        norm_class_name="LayerNorm",
        mlp_class_name="LLaMAMLP",
        intermediate_size=11008,
    ),
]

configs.extend(olmo)

olmo2 = [
    # https://huggingface.co/allenai/OLMo-2-1124-7B/blob/main/config.json
    dict(
        name="OLMo-2-1124-7B{}",
        hf_config=dict(org="allenai", name="OLMo-2-1124-7B{}"),
        vocab_size=100278,
        padded_vocab_size=100352,
        block_size=4096,
        n_embd=4096,
        n_layer=32,
        n_head=32,
        n_query_groups=32,
        rotary_percentage=1.0,
        parallel_residual=False,
        bias=False,
        norm_class_name="RMSNorm",
        mlp_class_name="LLaMAMLP",
        norm_eps=1e-06,
        intermediate_size=11008,
        rope_base=500000,
        norm_qk=True,
        post_mlp_norm=True,
        norm_1=False,
        norm_2=False,
        norm_qk_type="olmo2",
        post_attention_norm=True,
    ),
    # https://huggingface.co/allenai/OLMo-2-1124-13B/blob/main/config.json
    dict(
        name="OLMo-2-1124-13B{}",
        hf_config=dict(org="allenai", name="OLMo-2-1124-13B{}"),
        vocab_size=100278,
        padded_vocab_size=100352,
        block_size=4096,
        n_embd=5120,
        n_layer=40,
        n_head=40,
        n_query_groups=40,
        rotary_percentage=1.0,
        parallel_residual=False,
        bias=False,
        norm_class_name="RMSNorm",
        mlp_class_name="LLaMAMLP",
        norm_eps=1e-06,
        intermediate_size=13824,
        rope_base=500000,
        norm_qk=True,
        post_mlp_norm=True,
        norm_1=False,
        norm_2=False,
        norm_qk_type="olmo2",
        post_attention_norm=True,
    ),
]

for c in olmo2:
    for kind in ("", "-SFT", "-DPO", "-Instruct"):
        copy = deepcopy(c)
        copy["name"] = c["name"].format(kind)
        copy["hf_config"]["name"] = c["hf_config"]["name"].format(kind)
        configs.append(copy)

###############
# Google Gemma
###############
gemma = [
    # https://huggingface.co/google/gemma-2b/blob/main/config.json
    dict(
        name="Gemma-2b",
        hf_config=dict(org="google", name="gemma-2b"),
        scale_embeddings=True,
        vocab_size=256000,
        padding_multiple=64,
        n_embd=2048,
        n_layer=18,
        n_head=8,
        n_query_groups=1,
        rotary_percentage=1.0,
        parallel_residual=False,
        bias=False,
        norm_class_name="RMSNorm",
        mlp_class_name="GemmaMLP",
        gelu_approximate="tanh",
        intermediate_size=16384,
    ),
    # https://huggingface.co/google/gemma-7b/blob/main/config.json
    dict(
        name="Gemma-7b",
        hf_config=dict(org="google", name="gemma-7b"),
        scale_embeddings=True,
        vocab_size=256000,
        padding_multiple=64,
        n_embd=3072,
        n_layer=28,
        n_head=16,
        head_size=256,
        rotary_percentage=1.0,
        parallel_residual=False,
        bias=False,
        norm_class_name="RMSNorm",
        mlp_class_name="GemmaMLP",
        gelu_approximate="tanh",
        intermediate_size=24576,
    ),
    # https://huggingface.co/google/gemma-2-2b/blob/main/config.json
    dict(
        name="Gemma-2-2b",
        hf_config=dict(org="google", name="gemma-2-2b"),
        scale_embeddings=True,
        attention_scores_scalar=256,
        vocab_size=256000,
        block_size=8192,
        sliding_window_size=4096,
        # only layer with idx 0, 2, 4, ... have sliding window attention
        sliding_window_indices=[1 if i % 2 == 0 else 0 for i in range(26)],
        intermediate_size=9216,
        n_embd=2304,
        n_layer=26,
        n_head=8,
        n_query_groups=4,
        head_size=256,
        rotary_percentage=1.0,
        parallel_residual=False,
        bias=False,
        norm_class_name="RMSNorm",
        mlp_class_name="GemmaMLP",
        gelu_approximate="tanh",
        post_attention_norm=True,
        post_mlp_norm=True,
        attention_logit_softcapping=50.0,
        final_logit_softcapping=30.0,
    ),
    # https://huggingface.co/google/gemma-2-9b/blob/main/config.json
    dict(
        name="Gemma-2-9b",
        hf_config=dict(org="google", name="gemma-2-9b"),
        scale_embeddings=True,
        attention_scores_scalar=256,
        vocab_size=256000,
        block_size=8192,
        sliding_window_size=4096,
        # only layer with idx 0, 2, 4, ... have sliding window attention
        sliding_window_indices=[1 if i % 2 == 0 else 0 for i in range(42)],
        intermediate_size=14336,
        n_embd=3584,
        n_layer=42,
        n_head=16,
        n_query_groups=8,
        head_size=256,
        rotary_percentage=1.0,
        parallel_residual=False,
        bias=False,
        norm_class_name="RMSNorm",
        mlp_class_name="GemmaMLP",
        gelu_approximate="tanh",
        post_attention_norm=True,
        post_mlp_norm=True,
        attention_logit_softcapping=50.0,
        final_logit_softcapping=30.0,
    ),
    # https://huggingface.co/google/gemma-2-27b/blob/main/config.json
    dict(
        name="Gemma-2-27b",
        hf_config=dict(org="google", name="gemma-2-27b"),
        scale_embeddings=True,
        # In Gemma 2 27B attention scores are scaled not by `sqrt(head_size)` (11.31),
        # but by `sqrt(n_emb // n_head)` = sqrt(4608 // 32) = 12
        attention_scores_scalar=144,
        vocab_size=256000,
        block_size=8192,
        sliding_window_size=4096,
        # only layer with idx 0, 2, 4, ... have sliding window attention
        sliding_window_indices=[1 if i % 2 == 0 else 0 for i in range(46)],
        intermediate_size=36864,
        n_embd=4608,
        n_layer=46,
        n_head=32,
        n_query_groups=16,
        head_size=128,
        rotary_percentage=1.0,
        parallel_residual=False,
        bias=False,
        norm_class_name="RMSNorm",
        mlp_class_name="GemmaMLP",
        gelu_approximate="tanh",
        post_attention_norm=True,
        post_mlp_norm=True,
        attention_logit_softcapping=50.0,
        final_logit_softcapping=30.0,
    ),
]
configs.extend(gemma)
for c in gemma:
    copy = deepcopy(c)
    copy["name"] = f"{c['name']}-it"
    copy["hf_config"]["name"] = f"{c['hf_config']['name']}-it"
    configs.append(copy)

##################
# Google Gemma 3
##################
gemma3 = [
    # https://huggingface.co/google/gemma-3-1b-it/blob/main/config.json
    dict(
        name="Gemma-3-1b-it",
        hf_config=dict(org="google", name="gemma-3-1b-it"),
        scale_embeddings=True,
        attention_scores_scalar=256,
        vocab_size=262144,
        block_size=131072,
        sliding_window_size=512,
        # 5 local layers for every global layer
        sliding_window_indices=[0 if (i + 1) % 6 == 0 else 1 for i in range(26)],
        intermediate_size=6912,
        n_embd=1152,
        n_layer=26,
        n_head=4,
        n_query_groups=1,
        head_size=256,
        rotary_percentage=1.0,
        rope_adjustments=None,
        parallel_residual=False,
        bias=False,
        norm_class_name="RMSNorm",
        mlp_class_name="GemmaMLP",
        gelu_approximate="tanh",
        post_attention_norm=True,
        post_mlp_norm=True,
        norm_qk=True,
        rope_base=1000000,
        rope_local_base_freq=10000,
        # 5 local layers for every global layer
        rope_indices=[0 if (i + 1) % 6 == 0 else 1 for i in range(26)],
    ),
    # https://huggingface.co/google/gemma-3-4b-it/blob/main/config.json
    dict(
        name="Gemma-3-4b-it",
        hf_config=dict(org="google", name="gemma-3-4b-it"),
        scale_embeddings=True,
        attention_scores_scalar=256,
        vocab_size=262144,
        block_size=131072,
        sliding_window_size=1024,
        # 5 local layers for every global layer
        sliding_window_indices=[0 if (i + 1) % 6 == 0 else 1 for i in range(34)],
        intermediate_size=10240,
        n_embd=2560,
        n_layer=34,
        n_head=8,
        n_query_groups=4,
        head_size=256,
        rotary_percentage=1.0,
        rope_adjustments=dict(factor=8.0),
        parallel_residual=False,
        bias=False,
        norm_class_name="RMSNorm",
        mlp_class_name="GemmaMLP",
        gelu_approximate="tanh",
        post_attention_norm=True,
        post_mlp_norm=True,
        norm_qk=True,
        rope_base=1000000,
        rope_local_base_freq=10000,
        # 5 local layers for every global layer
        rope_indices=[0 if (i + 1) % 6 == 0 else 1 for i in range(34)],
    ),
    # https://huggingface.co/google/gemma-3-12b-it/blob/main/config.json
    dict(
        name="Gemma-3-12b-it",
        hf_config=dict(org="google", name="gemma-3-12b-it"),
        scale_embeddings=True,
        attention_scores_scalar=256,
        vocab_size=262144,
        block_size=131072,
        sliding_window_size=1024,
        # 5 local layers for every global layer
        sliding_window_indices=[0 if (i + 1) % 6 == 0 else 1 for i in range(48)],
        intermediate_size=15360,
        n_embd=3840,
        n_layer=48,
        n_head=16,
        n_query_groups=8,
        head_size=256,
        rotary_percentage=1.0,
        rope_adjustments=dict(factor=8.0),
        parallel_residual=False,
        bias=False,
        norm_class_name="RMSNorm",
        mlp_class_name="GemmaMLP",
        gelu_approximate="tanh",
        post_attention_norm=True,
        post_mlp_norm=True,
        norm_qk=True,
        rope_base=1000000,
        rope_local_base_freq=10000,
        # 5 local layers for every global layer
        rope_indices=[0 if (i + 1) % 6 == 0 else 1 for i in range(48)],
    ),
    # https://huggingface.co/google/gemma-3-27b-it/blob/main/config.json
    dict(
        name="Gemma-3-27b-it",
        hf_config=dict(org="google", name="gemma-3-27b-it"),
        scale_embeddings=True,
        attention_scores_scalar=168,
        vocab_size=262144,
        block_size=131072,
        sliding_window_size=1024,
        # 5 local layers for every global layer
        sliding_window_indices=[0 if (i + 1) % 6 == 0 else 1 for i in range(62)],
        intermediate_size=21504,
        n_embd=5376,
        n_layer=62,
        n_head=32,
        n_query_groups=16,
        head_size=128,
        rotary_percentage=1.0,
        rope_adjustments=dict(factor=8.0),
        parallel_residual=False,
        bias=False,
        norm_class_name="RMSNorm",
        mlp_class_name="GemmaMLP",
        gelu_approximate="tanh",
        post_attention_norm=True,
        post_mlp_norm=True,
        norm_qk=True,
        rope_base=1000000,
        rope_local_base_freq=10000,
        # 5 local layers for every global layer
        rope_indices=[0 if (i + 1) % 6 == 0 else 1 for i in range(62)],
    ),
]
configs.extend(gemma3)

##################
# Google CodeGemma
##################
codegemma = [
    # https://huggingface.co/google/codegemma-7b-it/blob/main/config.json
    dict(
        name="CodeGemma-7b-it",
        hf_config=dict(org="google", name="codegemma-7b-it"),
        scale_embeddings=True,
        vocab_size=256000,
        padding_multiple=64,
        n_embd=3072,
        n_layer=28,
        n_head=16,
        head_size=256,
        rotary_percentage=1.0,
        parallel_residual=False,
        bias=False,
        norm_class_name="RMSNorm",
        mlp_class_name="GemmaMLP",
        gelu_approximate="tanh",
        intermediate_size=24576,
    ),
]
configs.extend(codegemma)


##########################
# Stability AI FreeWilly2
##########################
freewilly_2 = [
    # https://huggingface.co/stabilityai/FreeWilly2/blob/main/config.json
    dict(
        name="FreeWilly2",
        hf_config=dict(org="stabilityai", name="FreeWilly2"),
        vocab_size=32000,
        padding_multiple=64,
        n_layer=80,
        n_head=64,
        n_embd=8192,
        n_query_groups=8,
        rotary_percentage=1.0,
        parallel_residual=False,
        bias=False,
        norm_class_name="RMSNorm",
        mlp_class_name="LLaMAMLP",
        intermediate_size=28672,
    )
]
configs.extend(freewilly_2)


##################
# Meta Code Llama
##################
code_llama = [
    # https://huggingface.co/codellama/CodeLlama-7b-hf/blob/main/config.json
    dict(
        name="CodeLlama-7b-hf",
        hf_config=dict(org="codellama", name="CodeLlama-7b-hf"),
        block_size=16384,
        vocab_size=32016,
        padding_multiple=16,
        n_layer=32,
        rotary_percentage=1.0,
        parallel_residual=False,
        bias=False,
        norm_class_name="RMSNorm",
        norm_eps=1e-05,
        mlp_class_name="LLaMAMLP",
        intermediate_size=11008,
        rope_base=1000000,
    ),
    # https://huggingface.co/codellama/CodeLlama-13b-hf/blob/main/config.json
    dict(
        name="CodeLlama-13b-hf",
        hf_config=dict(org="codellama", name="CodeLlama-13b-hf"),
        block_size=16384,
        vocab_size=32016,
        padding_multiple=16,
        n_layer=40,
        n_head=40,
        n_embd=5120,
        rotary_percentage=1.0,
        parallel_residual=False,
        bias=False,
        norm_class_name="RMSNorm",
        norm_eps=1e-05,
        mlp_class_name="LLaMAMLP",
        intermediate_size=13824,
        rope_base=1000000,
    ),
    # https://huggingface.co/codellama/CodeLlama-34b-hf/blob/main/config.json
    dict(
        name="CodeLlama-34b-hf",
        hf_config=dict(org="codellama", name="CodeLlama-34b-hf"),
        block_size=16384,
        vocab_size=32000,
        padded_vocab_size=32000,
        n_layer=48,
        n_head=64,
        n_embd=8192,
        n_query_groups=8,
        rotary_percentage=1.0,
        parallel_residual=False,
        bias=False,
        norm_class_name="RMSNorm",
        norm_eps=1e-05,
        mlp_class_name="LLaMAMLP",
        intermediate_size=22016,
        rope_base=1000000,
    ),
    # https://huggingface.co/codellama/CodeLlama-70b-hf/blob/main/config.json
    dict(
        name="CodeLlama-70b-hf",
        hf_config=dict(org="codellama", name="CodeLlama-70b-hf"),
        block_size=16384,
        vocab_size=32016,
        padding_multiple=16,
        n_layer=80,
        n_head=64,
        n_embd=8192,
        n_query_groups=8,
        rotary_percentage=1.0,
        parallel_residual=False,
        bias=False,
        norm_class_name="RMSNorm",
        norm_eps=1e-05,
        mlp_class_name="LLaMAMLP",
        intermediate_size=28672,
        rope_base=1000000,
    ),
    # https://huggingface.co/codellama/CodeLlama-7b-Python-hf/blob/main/config.json
    dict(
        name="CodeLlama-7b-Python-hf",
        hf_config=dict(org="codellama", name="CodeLlama-7b-Python-hf"),
        block_size=16384,
        vocab_size=32000,
        padded_vocab_size=32000,
        n_layer=32,
        rotary_percentage=1.0,
        parallel_residual=False,
        bias=False,
        norm_class_name="RMSNorm",
        norm_eps=1e-05,
        mlp_class_name="LLaMAMLP",
        intermediate_size=11008,
        rope_base=1000000,
    ),
    # https://huggingface.co/codellama/CodeLlama-13b-Python-hf/blob/main/config.json
    dict(
        name="CodeLlama-13b-Python-hf",
        hf_config=dict(org="codellama", name="CodeLlama-13b-Python-hf"),
        block_size=16384,
        vocab_size=32000,
        padded_vocab_size=32000,
        n_layer=40,
        n_head=40,
        n_embd=5120,
        rotary_percentage=1.0,
        parallel_residual=False,
        bias=False,
        norm_class_name="RMSNorm",
        norm_eps=1e-05,
        mlp_class_name="LLaMAMLP",
        intermediate_size=13824,
        rope_base=1000000,
    ),
    # https://huggingface.co/codellama/CodeLlama-34b-Python-hf/blob/main/config.json
    dict(
        name="CodeLlama-34b-Python-hf",
        hf_config=dict(org="codellama", name="CodeLlama-34b-Python-hf"),
        block_size=16384,
        vocab_size=32000,
        padded_vocab_size=32000,
        n_layer=48,
        n_head=64,
        n_embd=8192,
        n_query_groups=8,
        rotary_percentage=1.0,
        parallel_residual=False,
        bias=False,
        norm_class_name="RMSNorm",
        norm_eps=1e-05,
        mlp_class_name="LLaMAMLP",
        intermediate_size=22016,
        rope_base=1000000,
    ),
    # https://huggingface.co/codellama/CodeLlama-70b-Python-hf/blob/main/config.json
    dict(
        name="CodeLlama-70b-Python-hf",
        hf_config=dict(org="codellama", name="CodeLlama-70b-Python-hf"),
        block_size=16384,
        vocab_size=32016,
        padding_multiple=16,
        n_layer=80,
        n_head=64,
        n_embd=8192,
        n_query_groups=8,
        rotary_percentage=1.0,
        parallel_residual=False,
        bias=False,
        norm_class_name="RMSNorm",
        norm_eps=1e-05,
        mlp_class_name="LLaMAMLP",
        intermediate_size=28672,
        rope_base=1000000,
    ),
    # https://huggingface.co/codellama/CodeLlama-7b-Instruct-hf/blob/main/config.json
    dict(
        name="CodeLlama-7b-Instruct-hf",
        hf_config=dict(org="codellama", name="CodeLlama-7b-Instruct-hf"),
        block_size=16384,
        vocab_size=32016,
        padding_multiple=16,
        n_layer=32,
        rotary_percentage=1.0,
        parallel_residual=False,
        bias=False,
        norm_class_name="RMSNorm",
        norm_eps=1e-05,
        mlp_class_name="LLaMAMLP",
        intermediate_size=11008,
        rope_base=1000000,
    ),
    # https://huggingface.co/codellama/CodeLlama-13b-Instruct-hf/blob/main/config.json
    dict(
        name="CodeLlama-13b-Instruct-hf",
        hf_config=dict(org="codellama", name="CodeLlama-13b-Instruct-hf"),
        block_size=2048,
        vocab_size=32016,
        padding_multiple=16,
        n_layer=40,
        n_head=40,
        n_embd=5120,
        rotary_percentage=1.0,
        parallel_residual=False,
        bias=False,
        norm_class_name="RMSNorm",
        norm_eps=1e-05,
        mlp_class_name="LLaMAMLP",
        intermediate_size=13824,
        rope_base=1000000,
    ),
    # https://huggingface.co/codellama/CodeLlama-34b-Instruct-hf/blob/main/config.json
    dict(
        name="CodeLlama-34b-Instruct-hf",
        hf_config=dict(org="codellama", name="CodeLlama-34b-Instruct-hf"),
        block_size=16384,
        vocab_size=32000,
        padded_vocab_size=32000,
        n_layer=48,
        n_head=64,
        n_embd=8192,
        n_query_groups=8,
        rotary_percentage=1.0,
        parallel_residual=False,
        bias=False,
        norm_class_name="RMSNorm",
        norm_eps=1e-05,
        mlp_class_name="LLaMAMLP",
        intermediate_size=22016,
        rope_base=1000000,
    ),
    # https://huggingface.co/codellama/CodeLlama-70b-Instruct-hf/blob/main/config.json
    dict(
        name="CodeLlama-70b-Instruct-hf",
        hf_config=dict(org="codellama", name="CodeLlama-70b-Instruct-hf"),
        block_size=16384,
        # 32016 is an added token, so not reported in vocab_size
        # https://huggingface.co/codellama/CodeLlama-70b-Instruct-hf/blob/main/tokenizer_config.json
        vocab_size=32015,
        padding_multiple=16,
        n_layer=80,
        n_head=64,
        n_embd=8192,
        n_query_groups=8,
        rotary_percentage=1.0,
        parallel_residual=False,
        bias=False,
        norm_class_name="RMSNorm",
        norm_eps=1e-05,
        mlp_class_name="LLaMAMLP",
        intermediate_size=28672,
        rope_base=1000000,
    ),
]
configs.extend(code_llama)


########################
# garage-bAInd Platypus
########################
platypus = [
    # https://huggingface.co/garage-bAInd/Platypus-30B/blob/main/config.json
    dict(
        name="Platypus-30B",
        hf_config=dict(org="garage-bAInd", name="Platypus-30B"),
        block_size=2048,
        padded_vocab_size=32000,
        n_layer=60,
        n_head=52,
        n_embd=6656,
        rotary_percentage=1.0,
        parallel_residual=False,
        bias=False,
        norm_class_name="RMSNorm",
        norm_eps=1e-06,
        mlp_class_name="LLaMAMLP",
        intermediate_size=17920,
    ),
    # https://huggingface.co/garage-bAInd/Platypus2-7B/blob/main/config.json
    dict(
        name="Platypus2-7B",
        hf_config=dict(org="garage-bAInd", name="Platypus2-7B"),
        padded_vocab_size=32000,
        n_layer=32,
        rotary_percentage=1.0,
        parallel_residual=False,
        bias=False,
        norm_class_name="RMSNorm",
        norm_eps=1e-05,
        mlp_class_name="LLaMAMLP",
        intermediate_size=11008,
    ),
    # https://huggingface.co/garage-bAInd/Platypus2-13B/blob/main/config.json
    dict(
        name="Platypus2-13B",
        hf_config=dict(org="garage-bAInd", name="Platypus2-13B"),
        padded_vocab_size=32000,
        n_layer=40,
        n_head=40,
        n_embd=5120,
        rotary_percentage=1.0,
        parallel_residual=False,
        bias=False,
        norm_class_name="RMSNorm",
        norm_eps=1e-05,
        mlp_class_name="LLaMAMLP",
        intermediate_size=13824,
    ),
    # https://huggingface.co/garage-bAInd/Platypus2-70B/blob/main/config.json
    dict(
        name="Platypus2-70B",
        hf_config=dict(org="garage-bAInd", name="Platypus2-70B"),
        padded_vocab_size=32000,
        n_layer=80,
        n_head=64,
        n_embd=8192,
        rotary_percentage=1.0,
        parallel_residual=False,
        bias=False,
        norm_class_name="RMSNorm",
        mlp_class_name="LLaMAMLP",
        intermediate_size=28672,
    ),
    # https://huggingface.co/garage-bAInd/Camel-Platypus2-13B/blob/main/config.json
    dict(
        name="Camel-Platypus2-13B",
        hf_config=dict(org="garage-bAInd", name="Camel-Platypus2-13B"),
        padded_vocab_size=32000,
        n_layer=40,
        n_head=40,
        n_embd=5120,
        rotary_percentage=1.0,
        parallel_residual=False,
        bias=False,
        norm_class_name="RMSNorm",
        mlp_class_name="LLaMAMLP",
        intermediate_size=13824,
    ),
    # https://huggingface.co/garage-bAInd/Camel-Platypus2-70B/blob/main/config.json
    dict(
        name="Camel-Platypus2-70B",
        hf_config=dict(org="garage-bAInd", name="Camel-Platypus2-70B"),
        padded_vocab_size=32000,
        n_layer=80,
        n_head=64,
        n_embd=8192,
        n_query_groups=8,
        rotary_percentage=1.0,
        parallel_residual=False,
        bias=False,
        norm_class_name="RMSNorm",
        mlp_class_name="LLaMAMLP",
        intermediate_size=28672,
    ),
    # https://huggingface.co/garage-bAInd/Stable-Platypus2-13B/blob/main/config.json
    dict(
        name="Stable-Platypus2-13B",
        hf_config=dict(org="garage-bAInd", name="Stable-Platypus2-13B"),
        padded_vocab_size=32000,
        n_layer=40,
        n_head=40,
        n_embd=5120,
        rotary_percentage=1.0,
        parallel_residual=False,
        bias=False,
        norm_class_name="RMSNorm",
        mlp_class_name="LLaMAMLP",
        intermediate_size=13824,
    ),
    # https://huggingface.co/garage-bAInd/Platypus2-70B-instruct/blob/main/config.json
    dict(
        name="Platypus2-70B-instruct",
        hf_config=dict(org="garage-bAInd", name="Platypus2-70B-instruct"),
        padded_vocab_size=32000,
        n_layer=80,
        n_head=64,
        n_embd=8192,
        n_query_groups=8,
        rotary_percentage=1.0,
        parallel_residual=False,
        bias=False,
        norm_class_name="RMSNorm",
        mlp_class_name="LLaMAMLP",
        intermediate_size=28672,
    ),
]
configs.extend(platypus)


##################################
# togethercomputer LLaMA-2-7B-32K
##################################
together_llama2_32k = [
    # https://huggingface.co/togethercomputer/LLaMA-2-7B-32K/blob/main/config.json
    dict(
        name="LLaMA-2-7B-32K",
        hf_config=dict(org="togethercomputer", name="LLaMA-2-7B-32K"),
        vocab_size=32000,
        padding_multiple=64,
        n_layer=32,
        rotary_percentage=1.0,
        parallel_residual=False,
        bias=False,
        norm_class_name="RMSNorm",
        mlp_class_name="LLaMAMLP",
        intermediate_size=11008,
        rope_condense_ratio=8,
    )
]
configs.extend(together_llama2_32k)


################
# Microsoft Phi
################
phi = [
    # https://huggingface.co/microsoft/phi-1_5/blob/main/config.json
    dict(
        name="phi-1_5",
        hf_config=dict(org="microsoft", name="phi-1_5"),
        vocab_size=50257,
        padded_vocab_size=51200,
        block_size=2048,
        n_embd=2048,
        n_layer=24,
        rotary_percentage=0.5,  # 32 / (n_embd / n_head) = 32 / 64
        shared_attention_norm=True,
        lm_head_bias=True,
        gelu_approximate="tanh",
    ),
    # https://huggingface.co/microsoft/phi-2/blob/main/config.json
    dict(
        name="phi-2",
        hf_config=dict(org="microsoft", name="phi-2"),
        vocab_size=50257,
        padded_vocab_size=51200,
        block_size=2048,
        n_embd=2560,
        n_layer=32,
        rotary_percentage=0.4,  # 32 / (n_embd / n_head) = 32 / 80
        shared_attention_norm=True,
        lm_head_bias=True,
        gelu_approximate="tanh",
    ),
    # https://huggingface.co/microsoft/Phi-3-mini-4k-instruct/blob/main/config.json
    dict(
        name="Phi-3-mini-4k-instruct",
        hf_config=dict(org="microsoft", name="Phi-3-mini-4k-instruct"),
        vocab_size=32000,
        padded_vocab_size=32064,
        block_size=4096,
        n_embd=3072,
        n_layer=32,
        rotary_percentage=1.0,
        bias=False,
        norm_class_name="RMSNorm",
        intermediate_size=8192,
        mlp_class_name="LLaMAMLP",
        parallel_residual=False,
        sliding_window_size=2048,
    ),
    # https://huggingface.co/microsoft/Phi-3-mini-128k-instruct/blob/main/config.json
    dict(
        name="Phi-3-mini-128k-instruct",
        hf_config=dict(org="microsoft", name="Phi-3-mini-128k-instruct"),
        vocab_size=32000,
        padded_vocab_size=32064,
        block_size=131072,
        n_embd=3072,
        n_layer=32,
        rotary_percentage=1.0,
        bias=False,
        norm_class_name="RMSNorm",
        intermediate_size=8192,
        mlp_class_name="LLaMAMLP",
        parallel_residual=False,
        sliding_window_size=262145,
    ),
    # https://huggingface.co/microsoft/Phi-3.5-mini-instruct/blob/main/config.json
    dict(
        name="Phi-3.5-mini-instruct",
        hf_config=dict(org="microsoft", name="Phi-3.5-mini-instruct"),
        vocab_size=32000,
        padded_vocab_size=32064,
        block_size=4096,
        n_embd=3072,
        n_layer=32,
        rotary_percentage=1.0,
        bias=False,
        norm_class_name="RMSNorm",
        intermediate_size=8192,
        mlp_class_name="LLaMAMLP",
        parallel_residual=False,
    ),
    # https://huggingface.co/microsoft/phi-4/blob/main/config.json
    dict(
        name="phi-4",
        hf_config=dict(org="microsoft", name="phi-4"),
        vocab_size=100352,
        padded_vocab_size=100352,
        block_size=16384,
        n_embd=5120,
        n_layer=40,
        n_head=40,
        n_query_groups=10,
        rotary_percentage=1.0,
        bias=False,
        norm_class_name="RMSNorm",
        intermediate_size=17920,
        rope_base=250000,
        mlp_class_name="LLaMAMLP",
        parallel_residual=False,
    ),
    # https://huggingface.co/microsoft/Phi-4-reasoning/blob/main/config.json
    dict(
        name="Phi-4-reasoning",
        hf_config=dict(org="microsoft", name="Phi-4-reasoning"),
        vocab_size=100352,
        padded_vocab_size=100352,
        block_size=32768,
        n_embd=5120,
        n_layer=40,
        n_head=40,
        n_query_groups=10,
        rotary_percentage=1.0,
        bias=False,
        norm_class_name="RMSNorm",
        intermediate_size=17920,
        rope_base=500000,
        mlp_class_name="LLaMAMLP",
        parallel_residual=False,
    ),
    # https://huggingface.co/microsoft/Phi-4-reasoning-plus/blob/main/config.json
    dict(
        name="Phi-4-reasoning-plus",
        hf_config=dict(org="microsoft", name="Phi-4-reasoning-plus"),
        vocab_size=100352,
        padded_vocab_size=100352,
        block_size=32768,
        n_embd=5120,
        n_layer=40,
        n_head=40,
        n_query_groups=10,
        rotary_percentage=1.0,
        bias=False,
        norm_class_name="RMSNorm",
        intermediate_size=17920,
        rope_base=500000,
        mlp_class_name="LLaMAMLP",
        parallel_residual=False,
    ),
    # https://huggingface.co/microsoft/Phi-4-mini-instruct/blob/main/config.json
    dict(
        name="Phi-4-mini-instruct",
        hf_config=dict(org="microsoft", name="Phi-4-mini-instruct"),
        vocab_size=200019,
        padded_vocab_size=200064,
        block_size=131072,
        n_embd=3072,
        n_layer=32,
        n_head=24,
        n_query_groups=8,
        rotary_percentage=0.75,
        bias=False,
        norm_class_name="RMSNorm",
        intermediate_size=8192,
        mlp_class_name="LLaMAMLP",
        parallel_residual=False,
        sliding_window_size=262145,
    ),
    # https://huggingface.co/microsoft/Phi-4-mini-reasoning/blob/main/config.json
    dict(
        name="Phi-4-mini-reasoning",
        hf_config=dict(org="microsoft", name="Phi-4-mini-reasoning"),
        vocab_size=200019,
        padded_vocab_size=200064,
        block_size=131072,
        n_embd=3072,
        n_layer=32,
        n_head=24,
        n_query_groups=8,
        rotary_percentage=0.75,
        bias=False,
        norm_class_name="RMSNorm",
        intermediate_size=8192,
        mlp_class_name="LLaMAMLP",
        parallel_residual=False,
        sliding_window_size=262145,
    ),
]
configs.extend(phi)


#############
# Mistral AI
#############

configs.append(
    # https://huggingface.co/mistralai/mathstral-7B-v0.1/blob/main/config.json
    dict(
        name="Mathstral-7B-v0.1",
        hf_config=dict(org="mistralai", name="mathstral-7B-v0.1"),
        padded_vocab_size=32768,
        block_size=32768,
        n_layer=32,
        n_query_groups=8,
        rotary_percentage=1.0,
        parallel_residual=False,
        bias=False,
        norm_class_name="RMSNorm",
        norm_eps=1e-05,
        mlp_class_name="LLaMAMLP",
        intermediate_size=14336,
        rope_base=1000000,
        sliding_window_size=4096,
    )
)

mistral = [
    # https://huggingface.co/mistralai/Mistral-7B-v0.1/blob/main/config.json
    dict(
        name="Mistral-7B-{}v0.1",
        hf_config=dict(org="mistralai", name="Mistral-7B-{}v0.1"),
        padded_vocab_size=32000,
        block_size=4096,  # should be 32768 but sliding window attention is not implemented
        n_layer=32,
        n_query_groups=8,
        rotary_percentage=1.0,
        parallel_residual=False,
        bias=False,
        norm_class_name="RMSNorm",
        norm_eps=1e-05,
        mlp_class_name="LLaMAMLP",
        intermediate_size=14336,
        sliding_window_size=4096,
    ),
    # https://huggingface.co/mistralai/Mixtral-8x7B-v0.1/blob/main/config.json
    dict(
        name="Mixtral-8x7B-{}v0.1",
        hf_config=dict(org="mistralai", name="Mixtral-8x7B-{}v0.1"),
        padded_vocab_size=32000,
        block_size=32768,
        n_layer=32,
        n_query_groups=8,
        rotary_percentage=1.0,
        parallel_residual=False,
        bias=False,
        norm_class_name="RMSNorm",
        norm_eps=1e-05,
        mlp_class_name="LLaMAMoE",
        intermediate_size=14336,
        rope_base=1000000,
        n_expert=8,
        n_expert_per_token=2,
    ),
    # https://huggingface.co/mistralai/Mixtral-8x22B-Instruct-v0.1/blob/main/config.json
    dict(
        name="Mixtral-8x22B-{}v0.1",
        hf_config=dict(org="mistralai", name="Mixtral-8x22B-{}v0.1"),
        padded_vocab_size=32768,
        block_size=65536,
        n_layer=56,
        n_query_groups=8,
        rotary_percentage=1.0,
        parallel_residual=False,
        bias=False,
        norm_class_name="RMSNorm",
        norm_eps=1e-05,
        mlp_class_name="LLaMAMoE",
        intermediate_size=16384,
        n_head=48,
        n_embd=6144,
        rope_base=1000000,
        n_expert=8,
        n_expert_per_token=2,
    ),
]
for c in mistral:
    for kind in ("", "Instruct-"):
        copy = deepcopy(c)
        copy["name"] = c["name"].format(kind)
        copy["hf_config"]["name"] = c["hf_config"]["name"].format(kind)
        configs.append(copy)
configs.append(
    # https://huggingface.co/unsloth/mistral-7b-v0.2/blob/main/config.json
    dict(
        name="Mistral-7B-v0.2",
        hf_config=dict(org="unsloth", name="Mistral-7B-v0.2"),
        padded_vocab_size=32000,
        block_size=32768,
        n_layer=32,
        n_query_groups=8,
        rotary_percentage=1.0,
        parallel_residual=False,
        bias=False,
        norm_class_name="RMSNorm",
        norm_eps=1e-05,
        mlp_class_name="LLaMAMLP",
        intermediate_size=14336,
        rope_base=1000000,
    )
)
configs.append(
    # https://huggingface.co/mistralai/Mistral-7B-Instruct-v0.2/blob/main/config.json
    dict(
        name="Mistral-7B-Instruct-v0.2",
        hf_config=dict(org="mistralai", name="Mistral-7B-Instruct-v0.2"),
        padded_vocab_size=32000,
        block_size=32768,
        n_layer=32,
        n_query_groups=8,
        rotary_percentage=1.0,
        parallel_residual=False,
        bias=False,
        norm_class_name="RMSNorm",
        norm_eps=1e-05,
        mlp_class_name="LLaMAMLP",
        intermediate_size=14336,
        rope_base=1000000,
    )
)
configs.append(
    # https://huggingface.co/mistralai/Mistral-7B-v0.3/blob/main/config.json
    dict(
        name="Mistral-7B-v0.3",
        hf_config=dict(org="mistralai", name="Mistral-7B-v0.3"),
        padded_vocab_size=32768,
        block_size=32768,
        n_layer=32,
        n_query_groups=8,
        rotary_percentage=1.0,
        parallel_residual=False,
        bias=False,
        norm_class_name="RMSNorm",
        norm_eps=1e-05,
        mlp_class_name="LLaMAMLP",
        intermediate_size=14336,
        rope_base=1000000,
    )
)
configs.append(
    # https://huggingface.co/mistralai/Mistral-7B-Instruct-v0.3/blob/main/config.json
    dict(
        name="Mistral-7B-Instruct-v0.3",
        hf_config=dict(org="mistralai", name="Mistral-7B-Instruct-v0.3"),
        padded_vocab_size=32768,
        block_size=32768,
        n_layer=32,
        n_query_groups=8,
        rotary_percentage=1.0,
        parallel_residual=False,
        bias=False,
        norm_class_name="RMSNorm",
        norm_eps=1e-05,
        mlp_class_name="LLaMAMLP",
        intermediate_size=14336,
        rope_base=1000000,
    )
)
configs.append(
    # https://huggingface.co/mistralai/Mistral-Large-Instruct-2407/blob/main/config.json
    dict(
        name="Mistral-Large-Instruct-2407",
        hf_config=dict(org="mistralai", name="Mistral-Large-Instruct-2407"),
        padded_vocab_size=32768,
        block_size=32768,
        n_layer=88,
        n_head=96,
        n_embd=12288,
        n_query_groups=8,
        rotary_percentage=1.0,
        parallel_residual=False,
        bias=False,
        norm_class_name="RMSNorm",
        norm_eps=1e-05,
        mlp_class_name="LLaMAMLP",
        intermediate_size=28672,
        rope_base=1000000,
    )
)
configs.append(
    # https://huggingface.co/mistralai/Mistral-Large-Instruct-2411/blob/main/config.json
    dict(
        name="Mistral-Large-Instruct-2411",
        hf_config=dict(org="mistralai", name="Mistral-Large-Instruct-2411"),
        padded_vocab_size=32768,
        block_size=32768,
        n_layer=88,
        n_head=96,
        n_embd=12288,
        n_query_groups=8,
        rotary_percentage=1.0,
        parallel_residual=False,
        bias=False,
        norm_class_name="RMSNorm",
        norm_eps=1e-05,
        mlp_class_name="LLaMAMLP",
        intermediate_size=28672,
        rope_base=1000000,
    )
)


############
# TinyLlama
############
tiny_llama = [
    dict(
        name="tiny-llama-1.1b{}",
        hf_config=dict(org="TinyLlama", name="TinyLlama-1.1B{}"),
        block_size=2048,
        vocab_size=32000,
        padding_multiple=64,
        n_layer=22,
        n_head=32,
        n_embd=2048,
        rotary_percentage=1.0,
        parallel_residual=False,
        bias=False,
        norm_class_name="RMSNorm",  # original TinyLlama use FusedRMSNorm
        norm_eps=1e-5,
        mlp_class_name="LLaMAMLP",
        intermediate_size=5632,
        n_query_groups=4,
    )
]
for c in tiny_llama:
    for kind, hf_postfix in (("", "-intermediate-step-1431k-3T"), ("-chat", "-Chat-v1.0")):
        copy = deepcopy(c)
        copy["name"] = c["name"].format(kind)
        copy["hf_config"]["name"] = c["hf_config"]["name"].format(hf_postfix)
        configs.append(copy)


############
# MicroLlama
############
micro_llama = [
    dict(
        name="micro-llama-300M",
        hf_config=dict(org="keeeeenw", name="MicroLlama"),
        block_size=2048,
        vocab_size=32000,
        padding_multiple=64,
        n_layer=12,
        n_head=16,
        n_embd=1024,
        rotary_percentage=1.0,
        parallel_residual=False,
        bias=False,
        norm_class_name="RMSNorm",  # original TinyLlama and MicroLlama use FusedRMSNorm
        norm_eps=1e-5,
        mlp_class_name="LLaMAMLP",
        intermediate_size=5632,
        n_query_groups=4,
    )
]
configs.extend(micro_llama)


##########################
# Trelis Function Calling
##########################
llama_2_function_calling = [
    # https://huggingface.co/Trelis/Llama-2-7b-chat-hf-function-calling-v2/blob/main/config.json
    dict(
        name="Llama-2-7b-chat-hf-function-calling-v2",
        hf_config=dict(org="Trelis", name="Llama-2-7b-chat-hf-function-calling-v2"),
        padding_multiple=64,
        n_layer=32,
        rotary_percentage=1.0,
        parallel_residual=False,
        bias=False,
        norm_class_name="RMSNorm",
        mlp_class_name="LLaMAMLP",
        intermediate_size=11008,
        norm_eps=1e-6,
        block_size=4096,
        vocab_size=32000,
        n_head=32,
        n_embd=4096,
        rope_base=10000,
    )
]

configs.extend(llama_2_function_calling)

##########
# Qwen2.5
##########
qwen_2_5 = [
    # https://huggingface.co/Qwen/Qwen2.5-0.5B/blob/main/config.json
    dict(
        name="Qwen2.5-0.5B{}",
        hf_config=dict(org="Qwen", name="Qwen2.5-0.5B{}"),
        block_size=32768,
        vocab_size=151643,
        padded_vocab_size=151936,
        n_layer=24,
        n_head=14,
        n_embd=896,
        n_query_groups=2,
        rotary_percentage=1.0,
        parallel_residual=False,
        bias=False,
        attn_bias=True,
        norm_class_name="RMSNorm",
        mlp_class_name="LLaMAMLP",
        intermediate_size=4864,
        norm_eps=1e-6,
        rope_base=1000000,
    ),
    # https://huggingface.co/Qwen/Qwen2.5-1.5B/blob/main/config.json
    dict(
        name="Qwen2.5-1.5B{}",
        hf_config=dict(org="Qwen", name="Qwen2.5-1.5B{}"),
        block_size=131072,
        vocab_size=151643,
        padded_vocab_size=151936,
        n_layer=28,
        n_head=12,
        n_embd=1536,
        n_query_groups=2,
        rotary_percentage=1.0,
        parallel_residual=False,
        bias=False,
        attn_bias=True,
        norm_class_name="RMSNorm",
        mlp_class_name="LLaMAMLP",
        intermediate_size=8960,
        norm_eps=1e-6,
        rope_base=1000000,
    ),
    # https://huggingface.co/Qwen/Qwen2.5-3B/blob/main/config.json
    dict(
        name="Qwen2.5-3B{}",
        hf_config=dict(org="Qwen", name="Qwen2.5-3B{}"),
        block_size=32768,
        vocab_size=151643,
        padded_vocab_size=151936,
        n_layer=36,
        n_head=16,
        n_embd=2048,
        n_query_groups=2,
        rotary_percentage=1.0,
        parallel_residual=False,
        bias=False,
        attn_bias=True,
        norm_class_name="RMSNorm",
        mlp_class_name="LLaMAMLP",
        intermediate_size=11008,
        norm_eps=1e-6,
        rope_base=1000000,
    ),
    # https://huggingface.co/Qwen/Qwen2.5-7B/blob/main/config.json
    dict(
        name="Qwen2.5-7B{}",
        hf_config=dict(org="Qwen", name="Qwen2.5-7B{}"),
        block_size=131072,
        vocab_size=151643,
        padded_vocab_size=152064,
        n_layer=28,
        n_head=28,
        n_embd=3584,
        n_query_groups=4,
        rotary_percentage=1.0,
        parallel_residual=False,
        bias=False,
        attn_bias=True,
        norm_class_name="RMSNorm",
        mlp_class_name="LLaMAMLP",
        intermediate_size=18944,
        norm_eps=1e-6,
        rope_base=1000000,
    ),
    # https://huggingface.co/Qwen/Qwen2.5-14B/blob/main/config.json
    dict(
        name="Qwen2.5-14B{}",
        hf_config=dict(org="Qwen", name="Qwen2.5-14B{}"),
        block_size=131072,
        vocab_size=151643,
        padded_vocab_size=152064,
        n_layer=48,
        n_head=40,
        n_embd=5120,
        n_query_groups=8,
        rotary_percentage=1.0,
        parallel_residual=False,
        bias=False,
        attn_bias=True,
        norm_class_name="RMSNorm",
        mlp_class_name="LLaMAMLP",
        intermediate_size=13824,
        norm_eps=1e-5,
        rope_base=1000000,
    ),
    # https://huggingface.co/Qwen/Qwen2.5-32B/blob/main/config.json
    dict(
        name="Qwen2.5-32B{}",
        hf_config=dict(org="Qwen", name="Qwen2.5-32B{}"),
        block_size=131072,
        vocab_size=151643,
        padded_vocab_size=152064,
        n_layer=64,
        n_head=40,
        n_embd=5120,
        n_query_groups=8,
        rotary_percentage=1.0,
        parallel_residual=False,
        bias=False,
        attn_bias=True,
        norm_class_name="RMSNorm",
        mlp_class_name="LLaMAMLP",
        intermediate_size=27648,
        norm_eps=1e-5,
        rope_base=1000000,
    ),
    # https://huggingface.co/Qwen/Qwen2.5-72B/blob/main/config.json
    dict(
        name="Qwen2.5-72B{}",
        hf_config=dict(org="Qwen", name="Qwen2.5-72B{}"),
        block_size=131072,
        vocab_size=151643,
        padded_vocab_size=152064,
        n_layer=80,
        n_head=64,
        n_embd=8192,
        n_query_groups=8,
        rotary_percentage=1.0,
        parallel_residual=False,
        bias=False,
        attn_bias=True,
        norm_class_name="RMSNorm",
        mlp_class_name="LLaMAMLP",
        intermediate_size=29568,
        norm_eps=1e-5,
        rope_base=1000000,
    ),
]

qwen_2_5_coder = [
    # https://huggingface.co/Qwen/Qwen2.5-Coder-0.5B/blob/main/config.json
    dict(
        name="Qwen2.5-Coder-0.5B{}",
        hf_config=dict(org="Qwen", name="Qwen2.5-Coder-0.5B{}"),
        block_size=32768,
        vocab_size=151643,
        padded_vocab_size=151936,
        n_layer=24,
        n_head=14,
        n_embd=896,
        n_query_groups=2,
        rotary_percentage=1.0,
        parallel_residual=False,
        bias=False,
        attn_bias=True,
        norm_class_name="RMSNorm",
        mlp_class_name="LLaMAMLP",
        intermediate_size=4864,
        norm_eps=1e-6,
        rope_base=1000000,
    ),
    # https://huggingface.co/Qwen/Qwen2.5-Coder-1.5B/blob/main/config.json
    dict(
        name="Qwen2.5-Coder-1.5B{}",
        hf_config=dict(org="Qwen", name="Qwen2.5-Coder-1.5B{}"),
        block_size=32768,
        vocab_size=151643,
        padded_vocab_size=151936,
        n_layer=28,
        n_head=12,
        n_embd=1536,
        n_query_groups=2,
        rotary_percentage=1.0,
        parallel_residual=False,
        bias=False,
        attn_bias=True,
        norm_class_name="RMSNorm",
        mlp_class_name="LLaMAMLP",
        intermediate_size=8960,
        norm_eps=1e-6,
        rope_base=1000000,
    ),
    # https://huggingface.co/Qwen/Qwen2.5-Coder-3B/blob/main/config.json
    dict(
        name="Qwen2.5-Coder-3B{}",
        hf_config=dict(org="Qwen", name="Qwen2.5-Coder-3B{}"),
        block_size=32768,
        vocab_size=151643,
        padded_vocab_size=151936,
        n_layer=36,
        n_head=16,
        n_embd=2048,
        n_query_groups=2,
        rotary_percentage=1.0,
        parallel_residual=False,
        bias=False,
        attn_bias=True,
        norm_class_name="RMSNorm",
        mlp_class_name="LLaMAMLP",
        intermediate_size=11008,
        norm_eps=1e-6,
        rope_base=1000000,
    ),
    # https://huggingface.co/Qwen/Qwen2.5-Coder-7B/blob/main/config.json
    dict(
        name="Qwen2.5-Coder-7B{}",
        hf_config=dict(org="Qwen", name="Qwen2.5-Coder-7B{}"),
        block_size=32768,
        vocab_size=151643,
        padded_vocab_size=152064,
        n_layer=28,
        n_head=28,
        n_embd=3584,
        n_query_groups=4,
        rotary_percentage=1.0,
        parallel_residual=False,
        bias=False,
        attn_bias=True,
        norm_class_name="RMSNorm",
        mlp_class_name="LLaMAMLP",
        intermediate_size=18944,
        norm_eps=1e-6,
        rope_base=1000000,
    ),
    # https://huggingface.co/Qwen/Qwen2.5-Coder-14B/blob/main/config.json
    dict(
        name="Qwen2.5-Coder-14B{}",
        hf_config=dict(org="Qwen", name="Qwen2.5-Coder-14B{}"),
        block_size=32768,
        vocab_size=151643,
        padded_vocab_size=152064,
        n_layer=48,
        n_head=40,
        n_embd=5120,
        n_query_groups=8,
        rotary_percentage=1.0,
        parallel_residual=False,
        bias=False,
        attn_bias=True,
        norm_class_name="RMSNorm",
        mlp_class_name="LLaMAMLP",
        intermediate_size=13824,
        norm_eps=1e-5,
        rope_base=1000000,
    ),
    # https://huggingface.co/Qwen/Qwen2.5-Coder-32B/blob/main/config.json
    dict(
        name="Qwen2.5-Coder-32B{}",
        hf_config=dict(org="Qwen", name="Qwen2.5-Coder-32B{}"),
        block_size=32768,
        vocab_size=151643,
        padded_vocab_size=152064,
        n_layer=64,
        n_head=40,
        n_embd=5120,
        n_query_groups=8,
        rotary_percentage=1.0,
        parallel_residual=False,
        bias=False,
        attn_bias=True,
        norm_class_name="RMSNorm",
        mlp_class_name="LLaMAMLP",
        intermediate_size=27648,
        norm_eps=1e-5,
        rope_base=1000000,
    ),
]

qwen_2_5.extend(qwen_2_5_coder)

qwen_2_5_math = [
    # https://huggingface.co/Qwen/Qwen2.5-Math-1.5B/blob/main/config.json
    dict(
        name="Qwen2.5-Math-1.5B{}",
        hf_config=dict(org="Qwen", name="Qwen2.5-Math-1.5B{}"),
        block_size=4096,
        vocab_size=151643,
        padded_vocab_size=151936,
        n_layer=28,
        n_head=12,
        n_embd=1536,
        n_query_groups=2,
        rotary_percentage=1.0,
        parallel_residual=False,
        bias=False,
        attn_bias=True,
        norm_class_name="RMSNorm",
        mlp_class_name="LLaMAMLP",
        intermediate_size=8960,
        norm_eps=1e-6,
        rope_base=10000,
    ),
    # https://huggingface.co/Qwen/Qwen2.5-Math-7B/blob/main/config.json
    dict(
        name="Qwen2.5-Math-7B{}",
        hf_config=dict(org="Qwen", name="Qwen2.5-Math-7B{}"),
        block_size=4096,
        vocab_size=151643,
        padded_vocab_size=152064,
        n_layer=28,
        n_head=28,
        n_embd=3584,
        n_query_groups=4,
        rotary_percentage=1.0,
        parallel_residual=False,
        bias=False,
        attn_bias=True,
        norm_class_name="RMSNorm",
        mlp_class_name="LLaMAMLP",
        intermediate_size=18944,
        norm_eps=1e-6,
        rope_base=10000,
    ),
    # https://huggingface.co/Qwen/Qwen2.5-Math-72B/blob/main/config.json
    dict(
        name="Qwen2.5-Math-72B{}",
        hf_config=dict(org="Qwen", name="Qwen2.5-Math-72B{}"),
        block_size=4096,
        vocab_size=151643,
        padded_vocab_size=152064,
        n_layer=80,
        n_head=64,
        n_embd=8192,
        n_query_groups=8,
        rotary_percentage=1.0,
        parallel_residual=False,
        bias=False,
        attn_bias=True,
        norm_class_name="RMSNorm",
        mlp_class_name="LLaMAMLP",
        intermediate_size=29568,
        norm_eps=1e-5,
        rope_base=10000,
    ),
]

qwen_2_5.extend(qwen_2_5_math)

for c in qwen_2_5:
    for kind in ("", "-Instruct"):
        copy = deepcopy(c)
        copy["name"] = c["name"].format(kind)
        copy["hf_config"]["name"] = c["hf_config"]["name"].format(kind)
        configs.append(copy)

qwen_2_5_1m = [
    # https://huggingface.co/Qwen/Qwen2.5-7B-Instruct-1M/blob/main/config.json
    dict(
        name="Qwen2.5-7B-Instruct-1M",
        hf_config=dict(org="Qwen", name="Qwen2.5-7B-Instruct-1M"),
        block_size=1010000,
        vocab_size=151643,
        padded_vocab_size=152064,
        n_layer=28,
        n_head=28,
        n_embd=3584,
        n_query_groups=4,
        rotary_percentage=1.0,
        parallel_residual=False,
        bias=False,
        attn_bias=True,
        norm_class_name="RMSNorm",
        mlp_class_name="LLaMAMLP",
        intermediate_size=18944,
        norm_eps=1e-5,
        rope_base=10000000,
    ),
    # https://huggingface.co/Qwen/Qwen2.5-14B-Instruct-1M/blob/main/config.json
    dict(
        name="Qwen2.5-14B-Instruct-1M",
        hf_config=dict(org="Qwen", name="Qwen2.5-14B-Instruct-1M"),
        block_size=1010000,
        vocab_size=151643,
        padded_vocab_size=152064,
        n_layer=48,
        n_head=40,
        n_embd=5120,
        n_query_groups=8,
        rotary_percentage=1.0,
        parallel_residual=False,
        bias=False,
        attn_bias=True,
        norm_class_name="RMSNorm",
        mlp_class_name="LLaMAMLP",
        intermediate_size=13824,
        norm_eps=1e-5,
        rope_base=10000000,
    ),
]

configs.extend(qwen_2_5_1m)

##########
# QwQ
##########
qwq = [
    # https://huggingface.co/Qwen/QwQ-32B/blob/main/config.json
    dict(
        name="QwQ-32B",
        hf_config=dict(org="Qwen", name="QwQ-32B"),
        block_size=131072,
        vocab_size=151643,
        padded_vocab_size=152064,
        n_layer=64,
        n_head=40,
        n_embd=5120,
        n_query_groups=8,
        rotary_percentage=1.0,
        parallel_residual=False,
        bias=False,
        attn_bias=True,
        norm_class_name="RMSNorm",
        mlp_class_name="LLaMAMLP",
        intermediate_size=27648,
        norm_eps=1e-5,
        rope_base=1000000,
    ),
    # https://huggingface.co/Qwen/QwQ-32B-Preview/blob/main/config.json
    dict(
        name="QwQ-32B-Preview",
        hf_config=dict(org="Qwen", name="QwQ-32B-Preview"),
        block_size=32768,
        vocab_size=151643,
        padded_vocab_size=152064,
        n_layer=64,
        n_head=40,
        n_embd=5120,
        n_query_groups=8,
        rotary_percentage=1.0,
        parallel_residual=False,
        bias=False,
        attn_bias=True,
        norm_class_name="RMSNorm",
        mlp_class_name="LLaMAMLP",
        intermediate_size=27648,
        norm_eps=1e-5,
        rope_base=1000000,
    ),
]

configs.extend(qwq)

##########
# Qwen3
##########
qwen_3 = [
    # https://huggingface.co/Qwen/Qwen3-0.6B/blob/main/config.json
    dict(
        name="Qwen3-0.6B{}",
        hf_config=dict(org="Qwen", name="Qwen3-0.6B{}"),
        block_size=40960,
        vocab_size=151643,
        padded_vocab_size=151936,
        n_layer=28,
        n_head=16,
        n_embd=1024,
        n_query_groups=8,
        rotary_percentage=1.0,
        parallel_residual=False,
        bias=False,
        norm_class_name="RMSNorm",
        mlp_class_name="LLaMAMLP",
        intermediate_size=3072,
        norm_eps=1e-6,
        rope_base=1000000,
        head_size=128,
        norm_qk=True,
    ),
    # https://huggingface.co/Qwen/Qwen3-1.7B/blob/main/config.json
    dict(
        name="Qwen3-1.7B{}",
        hf_config=dict(org="Qwen", name="Qwen3-1.7B{}"),
        block_size=40960,
        vocab_size=151643,
        padded_vocab_size=151936,
        n_layer=28,
        n_head=16,
        n_embd=2048,
        n_query_groups=8,
        rotary_percentage=1.0,
        parallel_residual=False,
        bias=False,
        norm_class_name="RMSNorm",
        mlp_class_name="LLaMAMLP",
        intermediate_size=6144,
        norm_eps=1e-6,
        rope_base=1000000,
        norm_qk=True,
    ),
    # https://huggingface.co/Qwen/Qwen3-4B/blob/main/config.json
    dict(
        name="Qwen3-4B{}",
        hf_config=dict(org="Qwen", name="Qwen3-4B{}"),
        block_size=40960,
        vocab_size=151643,
        padded_vocab_size=151936,
        n_layer=36,
        n_head=32,
        n_embd=2560,
        n_query_groups=8,
        rotary_percentage=1.0,
        parallel_residual=False,
        bias=False,
        norm_class_name="RMSNorm",
        mlp_class_name="LLaMAMLP",
        intermediate_size=9728,
        norm_eps=1e-6,
        rope_base=1000000,
        head_size=128,
        norm_qk=True,
    ),
    # https://huggingface.co/Qwen/Qwen3-8B/blob/main/config.json
    dict(
        name="Qwen3-8B{}",
        hf_config=dict(org="Qwen", name="Qwen3-8B{}"),
        block_size=40960,
        vocab_size=151643,
        padded_vocab_size=151936,
        n_layer=36,
        n_head=32,
        n_embd=4096,
        n_query_groups=8,
        rotary_percentage=1.0,
        parallel_residual=False,
        bias=False,
        norm_class_name="RMSNorm",
        mlp_class_name="LLaMAMLP",
        intermediate_size=12288,
        norm_eps=1e-6,
        rope_base=1000000,
        norm_qk=True,
    ),
    # https://huggingface.co/Qwen/Qwen3-14B/blob/main/config.json
    dict(
        name="Qwen3-14B{}",
        hf_config=dict(org="Qwen", name="Qwen3-14B{}"),
        block_size=40960,
        vocab_size=151643,
        padded_vocab_size=151936,
        n_layer=40,
        n_head=40,
        n_embd=5120,
        n_query_groups=8,
        rotary_percentage=1.0,
        parallel_residual=False,
        bias=False,
        norm_class_name="RMSNorm",
        mlp_class_name="LLaMAMLP",
        intermediate_size=17408,
        norm_eps=1e-6,
        rope_base=1000000,
        norm_qk=True,
    ),
]
for c in qwen_3:
    for kind in ("", "-Base"):
        copy = deepcopy(c)
        copy["name"] = c["name"].format(kind)
        copy["hf_config"]["name"] = c["hf_config"]["name"].format(kind)
        configs.append(copy)
qwen_3_32b = [
    # https://huggingface.co/Qwen/Qwen3-32B/blob/main/config.json
    dict(
        name="Qwen3-32B",
        hf_config=dict(org="Qwen", name="Qwen3-32B"),
        block_size=40960,
        vocab_size=151643,
        padded_vocab_size=151936,
        n_layer=64,
        n_head=64,
        n_embd=5120,
        n_query_groups=8,
        rotary_percentage=1.0,
        parallel_residual=False,
        bias=False,
        norm_class_name="RMSNorm",
        mlp_class_name="LLaMAMLP",
        intermediate_size=25600,
        norm_eps=1e-6,
        rope_base=1000000,
        head_size=128,
        norm_qk=True,
    ),
]
configs.extend(qwen_3_32b)

qwen_3_moe = [
    # https://huggingface.co/Qwen/Qwen3-30B-A3B/blob/main/config.json
    dict(
        name="Qwen3-30B-A3B",
        hf_config=dict(org="Qwen", name="Qwen3-30B-A3B"),
        block_size=40960,
        head_size=128,
        vocab_size=151643,
        padded_vocab_size=151936,
        n_layer=48,
        n_head=32,
        n_embd=2048,
        n_query_groups=4,
        rotary_percentage=1.0,
        parallel_residual=False,
        bias=False,
        norm_class_name="RMSNorm",
        mlp_class_name="LLaMAMoE",
        intermediate_size=6144,
        moe_intermediate_size=768,
        norm_eps=1e-6,
        rope_base=1000000,
        norm_qk=True,
        n_expert=128,
        n_expert_per_token=8,
    ),
    # https://huggingface.co/Qwen/Qwen3-30B-A3B-Base/blob/main/config.json
    dict(
        name="Qwen3-30B-A3B-Base",
        hf_config=dict(org="Qwen", name="Qwen3-30B-A3B-Base"),
        block_size=40960,
        head_size=128,
        vocab_size=151643,
        padded_vocab_size=151936,
        n_layer=48,
        n_head=32,
        n_embd=2048,
        n_query_groups=4,
        rotary_percentage=1.0,
        parallel_residual=False,
        bias=False,
        norm_class_name="RMSNorm",
        mlp_class_name="LLaMAMoE",
        intermediate_size=6144,
        moe_intermediate_size=768,
        norm_eps=1e-6,
        rope_base=1000000,
        norm_qk=True,
        n_expert=128,
        n_expert_per_token=8,
    ),
    # https://huggingface.co/Qwen/Qwen3-235B-A22B/blob/main/config.json
    dict(
        name="Qwen3-235B-A22B",
        hf_config=dict(org="Qwen", name="Qwen3-235B-A22B"),
        block_size=40960,
        head_size=128,
        vocab_size=151643,
        padded_vocab_size=151936,
        n_layer=94,
        n_head=64,
        n_embd=4096,
        n_query_groups=4,
        rotary_percentage=1.0,
        parallel_residual=False,
        bias=False,
        norm_class_name="RMSNorm",
        mlp_class_name="LLaMAMoE",
        intermediate_size=12288,
        moe_intermediate_size=1536,
        norm_eps=1e-6,
        rope_base=1000000,
        norm_qk=True,
        n_expert=128,
        n_expert_per_token=8,
    ),
]
configs.extend(qwen_3_moe)

qwen_3_2507_thinking_instruct = [
    # https://huggingface.co/Qwen/Qwen3-235B-A22B-Thinking-2507/blob/main/config.json
    dict(
        name="Qwen3-235B-A22B-{}-2507",
        hf_config=dict(org="Qwen", name="Qwen3-235B-A22B-{}-2507"),
        block_size=262144,
        head_size=128,
        vocab_size=151643,
        padded_vocab_size=151936,
        n_layer=94,
        n_head=64,
        n_embd=4096,
        n_query_groups=4,
        rotary_percentage=1.0,
        parallel_residual=False,
        bias=False,
        norm_class_name="RMSNorm",
        mlp_class_name="LLaMAMoE",
        intermediate_size=12288,
        moe_intermediate_size=1536,
        norm_eps=1e-6,
        rope_base=5000000,
        norm_qk=True,
        n_expert=128,
        n_expert_per_token=8,
    ),
    # https://huggingface.co/Qwen/Qwen3-30B-A3B-Thinking-2507/blob/main/config.json
    dict(
        name="Qwen3-30B-A3B-{}-2507",
        hf_config=dict(org="Qwen", name="Qwen3-30B-A3B-{}-2507"),
        block_size=262144,
        head_size=128,
        vocab_size=151643,
        padded_vocab_size=151936,
        n_layer=48,
        n_head=32,
        n_embd=2048,
        n_query_groups=4,
        rotary_percentage=1.0,
        parallel_residual=False,
        bias=False,
        norm_class_name="RMSNorm",
        mlp_class_name="LLaMAMoE",
        intermediate_size=6144,
        moe_intermediate_size=768,
        norm_eps=1e-6,
        rope_base=10000000,
        norm_qk=True,
        n_expert=128,
        n_expert_per_token=8,
    ),
    # https://huggingface.co/Qwen/Qwen3-4B-Thinking-2507/blob/main/config.json
    dict(
        name="Qwen3-4B-{}-2507",
        hf_config=dict(org="Qwen", name="Qwen3-4B-{}-2507"),
        block_size=262144,
        vocab_size=151643,
        padded_vocab_size=151936,
        n_layer=36,
        n_head=32,
        n_embd=2560,
        n_query_groups=8,
        rotary_percentage=1.0,
        parallel_residual=False,
        bias=False,
        norm_class_name="RMSNorm",
        mlp_class_name="LLaMAMLP",
        intermediate_size=9728,
        norm_eps=1e-6,
        rope_base=5000000,
        head_size=128,
        norm_qk=True,
    ),
]

for c in qwen_3_2507_thinking_instruct:
    for kind in ("Thinking", "Instruct"):
        copy = deepcopy(c)
        copy["name"] = c["name"].format(kind)
        copy["hf_config"]["name"] = c["hf_config"]["name"].format(kind)
        configs.append(copy)

#############
# Salamandra
#############
salamandra = [
    # https://huggingface.co/BSC-LT/salamandra-2b-instruct/blob/main/config.json
    dict(
        name="salamandra-2b{}",
        hf_config=dict(org="BSC-LT", name="salamandra-2b{}"),
        block_size=8192,
        vocab_size=256000,
        padded_vocab_size=256000,
        n_layer=24,
        n_head=16,
        n_embd=2048,
        n_query_groups=16,
        rotary_percentage=1.0,
        parallel_residual=False,
        bias=False,
        norm_class_name="RMSNorm",
        mlp_class_name="LLaMAMLP",
        intermediate_size=5440,
        norm_eps=1e-5,
        rope_base=10000,
    ),
    # https://huggingface.co/BSC-LT/salamandra-7b-instruct/blob/main/config.json
    dict(
        name="salamandra-7b{}",
        hf_config=dict(org="BSC-LT", name="salamandra-7b{}"),
        block_size=8192,
        vocab_size=256000,
        padded_vocab_size=256000,
        n_layer=32,
        n_head=32,
        n_embd=4096,
        n_query_groups=8,
        rotary_percentage=1.0,
        parallel_residual=False,
        bias=False,
        norm_class_name="RMSNorm",
        mlp_class_name="LLaMAMLP",
        intermediate_size=11008,
        norm_eps=1e-6,
        rope_base=10000,
    ),
]

for c in salamandra:
    for kind in ("", "-instruct"):
        copy = deepcopy(c)
        copy["name"] = c["name"].format(kind)
        copy["hf_config"]["name"] = c["hf_config"]["name"].format(kind)
        configs.append(copy)


###############
# SmolLM2
###############
smollm2 = [
    # https://huggingface.co/HuggingFaceTB/SmolLM2-135M/blob/main/config.json
    dict(
        name="SmolLM2-135M{}",
        hf_config=dict(org="HuggingFaceTB", name="SmolLM2-135M{}"),
        block_size=8192,
        vocab_size=49152,
        padded_vocab_size=49152,
        n_layer=30,
        n_head=9,
        n_embd=576,
        n_query_groups=3,
        rotary_percentage=1.0,
        parallel_residual=False,
        bias=False,
        norm_class_name="RMSNorm",
        mlp_class_name="LLaMAMLP",
        intermediate_size=1536,
        rope_base=100000,
        norm_eps=1e-5,
    ),
    # https://huggingface.co/HuggingFaceTB/SmolLM2-360M/blob/main/config.json
    dict(
        name="SmolLM2-360M{}",
        hf_config=dict(org="HuggingFaceTB", name="SmolLM2-360M{}"),
        block_size=8192,
        vocab_size=49152,
        padded_vocab_size=49152,
        n_layer=32,
        n_head=15,
        n_embd=960,
        n_query_groups=5,
        rotary_percentage=1.0,
        parallel_residual=False,
        bias=False,
        norm_class_name="RMSNorm",
        mlp_class_name="LLaMAMLP",
        intermediate_size=2560,
        rope_base=100000,
        norm_eps=1e-5,
    ),
    # https://huggingface.co/HuggingFaceTB/SmolLM2-1.7B/blob/main/config.json
    dict(
        name="SmolLM2-1.7B{}",
        hf_config=dict(org="HuggingFaceTB", name="SmolLM2-1.7B{}"),
        block_size=8192,
        vocab_size=49152,
        padded_vocab_size=49152,
        n_layer=24,
        n_head=32,
        n_embd=2048,
        n_query_groups=32,
        rotary_percentage=1.0,
        parallel_residual=False,
        bias=False,
        norm_class_name="RMSNorm",
        mlp_class_name="LLaMAMLP",
        intermediate_size=8192,
        rope_base=130000,
        norm_eps=1e-5,
    ),
]

for c in smollm2:
    for kind in ("", "-Instruct"):
        copy = deepcopy(c)
        copy["name"] = c["name"].format(kind)
        copy["hf_config"]["name"] = c["hf_config"]["name"].format(kind)
        configs.append(copy)

###############
# DeepSeek R1 Distill
###############

r1_distill_llama = [
    # https://huggingface.co/deepseek-ai/DeepSeek-R1-Distill-Llama-8B/blob/main/config.json
    dict(
        name="R1-Distill-Llama-8B",
        hf_config=dict(org="deepseek-ai", name="DeepSeek-R1-Distill-Llama-8B"),
        block_size=131072,
        vocab_size=128000,
        padded_vocab_size=128256,
        n_layer=32,
        n_head=32,
        n_query_groups=8,
        rotary_percentage=1.0,
        parallel_residual=False,
        bias=False,
        norm_class_name="RMSNorm",
        mlp_class_name="LLaMAMLP",
        intermediate_size=14336,
        rope_base=500000,
        rope_adjustments=dict(factor=8.0, low_freq_factor=1.0, high_freq_factor=4.0, original_max_seq_len=8192),
    ),
    # https://huggingface.co/deepseek-ai/DeepSeek-R1-Distill-Llama-70B/blob/main/config.json
    dict(
        name="R1-Distill-Llama-70B",
        hf_config=dict(org="deepseek-ai", name="DeepSeek-R1-Distill-Llama-70B"),
        block_size=131072,
        vocab_size=128000,
        padded_vocab_size=128256,
        n_layer=80,
        n_head=64,
        n_embd=8192,
        n_query_groups=8,
        rotary_percentage=1.0,
        parallel_residual=False,
        bias=False,
        norm_class_name="RMSNorm",
        mlp_class_name="LLaMAMLP",
        intermediate_size=28672,
        rope_base=500000,
        rope_adjustments=dict(factor=8.0, low_freq_factor=1.0, high_freq_factor=4.0, original_max_seq_len=8192),
    ),
]

configs.extend(r1_distill_llama)

name_to_config = {config["name"]: config for config in configs}