File size: 3,992 Bytes
f05d2eb | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 | """
Qwen3 Model Configuration.
"""
from dataclasses import dataclass, field
from typing import List, Optional
@dataclass
class Qwen3Config:
"""Configuration class for Qwen3 model."""
vocab_size: int = 151936
hidden_size: int = 1024
intermediate_size: int = 3072
num_hidden_layers: int = 28
num_attention_heads: int = 16
num_key_value_heads: int = 8
head_dim: int = 128
hidden_act: str = "silu"
max_position_embeddings: int = 32768
initializer_range: float = 0.02
rms_norm_eps: float = 1e-6
use_cache: bool = True
tie_word_embeddings: bool = True
rope_theta: float = 1000000.0
rope_scaling: Optional[dict] = None
# ``rope`` (default) preserves legacy behavior; ``alibi`` disables RoPE
# and applies Press et al. head-specific causal distance biases.
positional: str = "rope"
# QK normalisation is part of the published Qwen3-style attention path.
# Keep it on by default so checkpoints/configs written before this flag are
# bit-identical; ``off`` removes both per-head RMSNorm modules.
qk_norm: str = "on"
# Keep Q/K RMSNorm gains and their AdamW state in FP32 during training.
qk_gain_fp32: bool = False
# Training-only RoPE position scales. When non-empty, every attention layer
# independently samples one scale per forward pass. Evaluation always uses
# the standard scale 1.0. This provides a dense positional-augmentation
# control without deleting attention links.
rope_train_position_scale_set: Optional[List[float]] = None
attention_bias: bool = False
attention_dropout: float = 0.0
bos_token_id: int = 151643
eos_token_id: int = 151643
checkpoint_activations: bool = False
attention_type: str = "dense"
sparse_block_size: int = 128
sparse_num_random_blocks: int = 1
sparse_mask_pool_size: int = 1000
sparse_max_seq_length: int = 1024
# When set, only the listed layer indices use sparse attention; all others use dense.
# When None (default), attention_type applies to all layers uniformly.
sparse_layer_indices: Optional[List[int]] = None
# Directory to cache pre-built block masks (one file per layer). If None, masks are rebuilt every run.
sparse_mask_cache_dir: Optional[str] = None
# Seed offset for the frozen (pool_size=1) fixed mask. 0 = legacy (layer-index only).
# Distinct values yield distinct fixed masks (fixed-mask seeds).
sparse_mask_seed: int = 0
# Mixed-sparsity training: if set to a non-empty list, each mask-pool entry draws its
# OWN num_random_blocks uniformly from this set, so the per-step mask (picked at random
# from the pool) also resamples the sparsity LEVEL. None/empty = legacy fixed-nrb pool.
# This is the "Mixed" training arm of the paper; without it the arm silently degenerates
# to a fixed-nrb model.
sparse_mixed_nrb_set: Optional[List[int]] = None
# --- deployment-time baselines (learned MoBA + ASEntmax) ---
# moba_top_k is the per-query block budget used by the top-k deployment ladder.
moba_block_size: int = 128
moba_top_k: int = 32
asentmax_alpha: float = 1.5
asentmax_learn_temp: bool = True
asentmax_init_temp: float = 1.0
def is_layer_sparse(self, layer_idx: int) -> bool:
"""Return True if this layer should use sparse stochastic attention."""
if self.attention_type != "sparse_stochastic":
return False
if self.sparse_layer_indices is None:
return True # all layers sparse (existing behaviour)
return layer_idx in self.sparse_layer_indices
@classmethod
def from_dict(cls, config_dict: dict) -> "Qwen3Config":
return cls(**{k: v for k, v in config_dict.items() if k in cls.__dataclass_fields__})
@classmethod
def from_json(cls, path: str) -> "Qwen3Config":
import json
with open(path, "r") as f:
config_dict = json.load(f)
return cls.from_dict(config_dict)
|