@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 173 lines 7.52 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 from huggingface_hub.dataclasses import strict
4 from transformers.configuration_utils import PreTrainedConfig
5 from transformers.modeling_rope_utils import RopeParameters
6
7
8 @strict
9 class HYV4Config(PreTrainedConfig):
10 r"""Configuration class for the HY V4 (`HYV4ForCausalLM`) model.
11
12 Kept field-for-field in sync with
13 `transformers.models.hy_v4.configuration_hy_v4.HYV4Config` so that a
14 checkpoint resolves to the same values through either class. Registering it
15 in vLLM's `_CONFIG_REGISTRY` lets the architecture load without
16 `trust_remote_code` and without an `auto_map` module next to the weights.
17
18 HYV4 is a mixture-of-experts causal language model using Multi-head Latent
19 Attention (MLA), DeepSeek-style sparse attention (DSA), gated MLA,
20 learnable attention sinks, and independent Hyper-Connections (iHC).
21
22 The `layer_types`, `num_nextn_predict_layers` and `mtp_loss_factor` fields
23 below are vLLM-only additions: upstream leaves them to `**kwargs`, but
24 vLLM reads them directly (sparse-attention selection and MTP speculative
25 decoding), so they are declared to give checkpoints that omit them a
26 working default instead of a `None`.
27
28 Args:
29 mlp_layer_types: Per-layer MLP kind, `"dense"` or `"sparse"`. Defaults
30 to one dense layer followed by sparse MoE layers.
31 layer_types: Per-layer attention kind, e.g. `"full_attention"` or
32 `"deepseek_sparse_attention"`.
33 indexer_types: Per-layer DSA indexer kind, `"full"` or `"shared"`. A
34 shared layer reuses the most recent full indexer in the same
35 forward request.
36 index_topk: Maximum number of key positions selected by each DSA query.
37 index_head_dim: Hidden dimension of each DSA indexer head.
38 index_n_heads: Number of DSA indexer heads.
39 enable_lm_head_fp32: Whether the language-model head emits float32
40 logits. Surfaced as ``head_dtype`` so `LogitsProcessor` accumulates
41 the projection into fp32 while the weight stays in the model dtype.
42 enable_ihc: Whether independent Hyper-Connections are enabled.
43 hc_mult: Number of hidden-state channels maintained by iHC.
44 hc_magnitude: Scale applied to the iHC post-gating branch.
45 hc_eps: Numerical epsilon added to iHC sigmoid gates.
46 gated_mla: Whether to gate the MLA output.
47 gating_type: MLA gate granularity, `"elementwise"` or `"headwise"`.
48 learnable_sink: Whether to add a learned per-head attention sink.
49 learnable_sink_init: Initial value of each learned attention-sink
50 logit.
51 swiglu_limit: Magnitude of the routed-expert SwiGLU clamp. Values at or
52 below zero disable the clamp.
53 num_nextn_predict_layers: Number of MTP layers in the checkpoint.
54 mtp_loss_factor: Training-time MTP loss weight (unused at inference).
55
56 """
57
58 model_type = "hy_v4"
59 keys_to_ignore_at_inference = ["past_key_values"]
60 attribute_map = {
61 "num_local_experts": "n_routed_experts",
62 }
63 base_model_tp_plan = {
64 "layers.*.self_attn.q_b_proj": "colwise",
65 "layers.*.self_attn.kv_a_proj_with_mqa": "mla_kv_a_proj",
66 "layers.*.self_attn.kv_b_proj": "colwise",
67 "layers.*.self_attn.o_proj": "rowwise",
68 "layers.*.self_attn.linear_gate": "colwise",
69 "layers.*.mlp.experts.gate_up_proj": "packed_colwise",
70 "layers.*.mlp.experts.down_proj": "rowwise",
71 "layers.*.mlp.experts": "moe_tp_experts",
72 "layers.*.mlp.shared_experts.gate_proj": "colwise",
73 "layers.*.mlp.shared_experts.up_proj": "colwise",
74 "layers.*.mlp.shared_experts.down_proj": "rowwise",
75 "layers.*.mlp.gate_proj": "colwise",
76 "layers.*.mlp.up_proj": "colwise",
77 "layers.*.mlp.down_proj": "rowwise",
78 }
79 base_model_pp_plan = {
80 "embed_tokens": (["input_ids"], ["inputs_embeds"]),
81 "layers": (["hidden_states", "attention_mask"], ["hidden_states"]),
82 "norm": (["hidden_states"], ["hidden_states"]),
83 }
84
85 vocab_size: int = 120832
86 hidden_size: int = 2816
87 intermediate_size: int = 6912
88 moe_intermediate_size: int = 768
89 num_hidden_layers: int = 34
90 num_attention_heads: int = 32
91 num_key_value_heads: int = 32
92 head_dim: int = 256
93 hidden_act: str = "silu"
94 max_position_embeddings: int = 262144
95 initializer_range: float = 0.006
96 rms_norm_eps: float = 1e-5
97 use_cache: bool = True
98 pad_token_id: int | None = 120002
99 bos_token_id: int | None = 120000
100 eos_token_id: int | list[int] | None = 120025
101 tie_word_embeddings: bool = False
102 attention_bias: bool = False
103 attention_dropout: float = 0.0
104 n_routed_experts: int = 256
105 n_shared_experts: int = 1
106 num_experts_per_tok: int = 8
107 routed_scaling_factor: float = 2.827
108 norm_topk_prob: bool = True
109 q_lora_rank: int = 1536
110 kv_lora_rank: int = 512
111 qk_nope_head_dim: int = 192
112 qk_rope_head_dim: int = 64
113 v_head_dim: int = 256
114 mlp_layer_types: list[str] | None = None
115 layer_types: list[str] | None = None
116 index_topk: int = 2048
117 index_head_dim: int = 128
118 index_n_heads: int = 16
119 indexer_types: list[str] | None = None
120 enable_lm_head_fp32: bool = True
121 enable_ihc: bool = True
122 hc_mult: int = 4
123 hc_magnitude: float = 2.0
124 hc_eps: float = 1e-6
125 gated_mla: bool = True
126 gating_type: str = "elementwise"
127 learnable_sink: bool = True
128 learnable_sink_init: float = 0.0
129 swiglu_limit: float = 10.0
130 rope_parameters: RopeParameters | dict | None = None
131 num_nextn_predict_layers: int = 1
132 mtp_loss_factor: float = 0.1
133
134 def __post_init__(self, **kwargs):
135 self.qk_head_dim = self.qk_nope_head_dim + self.qk_rope_head_dim
136 # RoPE applies only to the rope slice, so point `head_dim` at it: the
137 # inherited rotary embedding reads `config.head_dim` and then computes
138 # the right frequencies with no override.
139 self.head_dim = self.qk_rope_head_dim
140
141 if self.mlp_layer_types is None:
142 self.mlp_layer_types = ["dense"] * min(1, self.num_hidden_layers) + [
143 "sparse"
144 ] * max(self.num_hidden_layers - 1, 0)
145 if self.indexer_types is None:
146 self.indexer_types = [
147 "full" if layer_idx == 0 or (layer_idx - 1) % 4 == 0 else "shared"
148 for layer_idx in range(self.num_hidden_layers)
149 ]
150 # `PreTrainedConfig` validates `layer_types` against a fixed allow-list
151 # in which "deepseek_sparse_attention" is the only accepted spelling of
152 # a DSA layer. The model treats "sparse_attention" and "sparse" as
153 # equivalent, so fold them in before the base class validates.
154 if self.layer_types is not None:
155 self.layer_types = [
156 "deepseek_sparse_attention"
157 if lt in ("sparse_attention", "sparse")
158 else lt
159 for lt in self.layer_types
160 ]
161
162 # `ModelConfig.head_dtype` reads this off the HF config, so an fp32 head
163 # needs no model-side dtype juggling: LogitsProcessor picks the
164 # `torch.mm(out_dtype=float32)` path and leaves the weight in the model
165 # dtype. Only fill it in when the checkpoint (or --hf-overrides) has not
166 # already pinned a value.
167 if self.enable_lm_head_fp32 and getattr(self, "head_dtype", None) is None:
168 self.head_dtype = "float32"
169
170 super().__post_init__(**kwargs)
171
172
173 __all__ = ["HYV4Config"]