diff --git a/python/sglang/srt/configs/cohere2_moe.py b/python/sglang/srt/configs/cohere2_moe.py index cd470bd69..20c1fd663 100644 --- a/python/sglang/srt/configs/cohere2_moe.py +++ b/python/sglang/srt/configs/cohere2_moe.py @@ -4,57 +4,110 @@ from transformers.configuration_utils import PreTrainedConfig from transformers.models.auto.configuration_auto import CONFIG_MAPPING -try: - from huggingface_hub.dataclasses import strict -except ImportError: # older huggingface_hub - def strict(cls): # type: ignore[misc] - return cls - - -@strict class Cohere2MoeConfig(PreTrainedConfig): model_type = "cohere2_moe" keys_to_ignore_at_inference = ["past_key_values"] - vocab_size: int = 256000 - hidden_size: int = 8192 - intermediate_size: int = 22528 - logit_scale: float = 0.0625 - num_hidden_layers: int = 40 - num_attention_heads: int = 64 - num_key_value_heads: int | None = None - head_dim: int = 128 - hidden_act: str = "silu" - max_position_embeddings: int = 8192 - initializer_range: float = 0.02 - layer_norm_eps: float = 1e-5 - use_cache: bool = True - pad_token_id: int | None = 0 - bos_token_id: int | None = 5 - eos_token_id: int | list[int] | None = 255001 - tie_word_embeddings: bool = True - rope_theta: float | int = 10000.0 - rope_scaling: dict | None = None - attention_bias: bool = False - attention_dropout: float = 0.0 - sliding_window: int | None = 4096 - num_experts_per_tok: int = 2 - num_experts: int = 8 - num_shared_experts: int = 0 - shared_expert_combination_strategy: str = "average" - expert_selection_fn: str = "softmax" - layer_types: list[str] | None = None - first_k_dense_replace: int = 0 - prefix_dense_sliding_window_pattern: int = 1 - norm_topk_prob: bool = True - prefix_dense_intermediate_size: int | None = None - rms_norm_eps: float | None = None - sliding_window_pattern: int = 4 + def __init__( + self, + vocab_size: int = 256000, + hidden_size: int = 8192, + intermediate_size: int = 22528, + logit_scale: float = 0.0625, + num_hidden_layers: int = 40, + num_attention_heads: int = 64, + num_key_value_heads: int | None = None, + head_dim: int = 128, + hidden_act: str = "silu", + max_position_embeddings: int = 8192, + initializer_range: float = 0.02, + layer_norm_eps: float = 1e-5, + use_cache: bool = True, + pad_token_id: int | None = 0, + bos_token_id: int | None = 5, + eos_token_id: int | list[int] | None = 255001, + tie_word_embeddings: bool = True, + rope_theta: float | int = 10000.0, + rope_scaling: dict | None = None, + attention_bias: bool = False, + attention_dropout: float = 0.0, + sliding_window: int | None = 4096, + num_experts_per_tok: int = 2, + num_experts: int = 8, + num_shared_experts: int = 0, + shared_expert_combination_strategy: str = "average", + expert_selection_fn: str = "softmax", + layer_types: list[str] | None = None, + first_k_dense_replace: int = 0, + prefix_dense_sliding_window_pattern: int = 1, + norm_topk_prob: bool = True, + prefix_dense_intermediate_size: int | None = None, + rms_norm_eps: float | None = None, + sliding_window_pattern: int = 4, + **kwargs, + ): + self.vocab_size = vocab_size + self.hidden_size = hidden_size + self.intermediate_size = intermediate_size + self.logit_scale = logit_scale + self.num_hidden_layers = num_hidden_layers + self.num_attention_heads = num_attention_heads + self.num_key_value_heads = ( + num_attention_heads if num_key_value_heads is None else num_key_value_heads + ) + self.head_dim = head_dim + self.hidden_act = hidden_act + self.max_position_embeddings = max_position_embeddings + self.initializer_range = initializer_range + self.layer_norm_eps = layer_norm_eps + self.use_cache = use_cache + self.rope_theta = rope_theta + self.rope_scaling = rope_scaling + self.attention_bias = attention_bias + self.attention_dropout = attention_dropout + self.sliding_window = sliding_window + self.num_experts_per_tok = num_experts_per_tok + self.num_experts = num_experts + self.num_shared_experts = num_shared_experts + self.shared_expert_combination_strategy = shared_expert_combination_strategy + self.expert_selection_fn = expert_selection_fn + self.first_k_dense_replace = first_k_dense_replace + self.prefix_dense_sliding_window_pattern = prefix_dense_sliding_window_pattern + self.norm_topk_prob = norm_topk_prob + self.prefix_dense_intermediate_size = prefix_dense_intermediate_size + self.rms_norm_eps = rms_norm_eps + self.sliding_window_pattern = sliding_window_pattern - def __post_init__(self, **kwargs): - if self.num_key_value_heads is None: - self.num_key_value_heads = self.num_attention_heads + if layer_types is None: + prefix_layers = [ + ( + "sliding_attention" + if ((i + 1) % prefix_dense_sliding_window_pattern) != 0 + else "full_attention" + ) + for i in range(self.first_k_dense_replace) + ] + rest_layers = [ + ( + "sliding_attention" + if ((i + 1) % sliding_window_pattern) != 0 + else "full_attention" + ) + for i in range(self.num_hidden_layers - self.first_k_dense_replace) + ] + self.layer_types = prefix_layers + rest_layers + else: + self.layer_types = layer_types + + super().__init__( + pad_token_id=pad_token_id, + bos_token_id=bos_token_id, + eos_token_id=eos_token_id, + tie_word_embeddings=tie_word_embeddings, + use_cache=use_cache, + **kwargs, + ) if hasattr(self, "standardize_rope_params"): try: @@ -63,27 +116,6 @@ class Cohere2MoeConfig(PreTrainedConfig): except Exception: pass - if self.layer_types is None: - prefix_layers = [ - ( - "sliding_attention" - if ((i + 1) % self.prefix_dense_sliding_window_pattern) != 0 - else "full_attention" - ) - for i in range(self.first_k_dense_replace) - ] - rest_layers = [ - ( - "sliding_attention" - if ((i + 1) % self.sliding_window_pattern) != 0 - else "full_attention" - ) - for i in range(self.num_hidden_layers - self.first_k_dense_replace) - ] - self.layer_types = prefix_layers + rest_layers - - super().__post_init__(**kwargs) - try: CONFIG_MAPPING.register("cohere2_moe", Cohere2MoeConfig) diff --git a/test/registered/unit/configs/test_cohere2_moe_config.py b/test/registered/unit/configs/test_cohere2_moe_config.py new file mode 100644 index 000000000..6aada5d55 --- /dev/null +++ b/test/registered/unit/configs/test_cohere2_moe_config.py @@ -0,0 +1,60 @@ +"""Regression tests for Cohere2MoeConfig import (sgl-project/sglang#28233). + +Before the fix, applying ``huggingface_hub.dataclasses.strict`` to +``Cohere2MoeConfig`` — which is not a stdlib ``@dataclass`` — raised +``StrictDataclassDefinitionError`` at *import* time on +``transformers==5.3.0`` + ``huggingface_hub==1.9.0``, which meant that +``sglang.srt.configs`` could not be imported at all. +""" + +import unittest + +from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import CustomTestCase + +register_cpu_ci(est_time=5, suite="base-a-test-cpu") + + +class TestCohere2MoeConfig(CustomTestCase): + def test_configs_package_imports(self): + """Importing the configs package must not crash at module load.""" + import sglang.srt.configs # noqa: F401 + + def test_derived_defaults_num_key_value_heads(self): + """``num_key_value_heads`` defaults to ``num_attention_heads``.""" + from sglang.srt.configs.cohere2_moe import Cohere2MoeConfig + + cfg = Cohere2MoeConfig() + self.assertEqual(cfg.num_key_value_heads, cfg.num_attention_heads) + + def test_derived_defaults_layer_types(self): + """``layer_types`` is auto-derived to one entry per hidden layer.""" + from sglang.srt.configs.cohere2_moe import Cohere2MoeConfig + + cfg = Cohere2MoeConfig() + self.assertEqual(len(cfg.layer_types), cfg.num_hidden_layers) + + def test_pretrained_config_kwargs_forwarded(self): + """Extra kwargs must flow through ``PreTrainedConfig.__init__``.""" + from sglang.srt.configs.cohere2_moe import Cohere2MoeConfig + + cfg = Cohere2MoeConfig(foo="bar") + cfg_dict = cfg.to_dict() + self.assertEqual(cfg_dict["model_type"], "cohere2_moe") + self.assertEqual(cfg_dict["foo"], "bar") + + def test_use_cache_false_preserved(self): + """``use_cache=False`` must survive ``super().__init__``. + + ``PreTrainedConfig.__init__`` re-assigns ``self.use_cache`` from its + own keyword; if the subclass does not forward the caller's value, + ``use_cache=False`` silently reverts to the base default ``True``. + """ + from sglang.srt.configs.cohere2_moe import Cohere2MoeConfig + + cfg = Cohere2MoeConfig(use_cache=False) + self.assertFalse(cfg.use_cache) + + +if __name__ == "__main__": + unittest.main()