# 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}