biondizzle/vllm
1
0
Fork 0
Code Issues Pull Requests Actions 2 Packages Projects Releases Wiki Activity

[Model] Add Grok-2 (#31847)

Browse Source 
Signed-off-by: dangoldbj <dangoldbj23@gmail.com>
This commit is contained in:
Bijaya Dangol
2026-01-08 13:59:48 +01:00
committed by GitHub
parent 18d4e481d0
commit 59d260f5e4
8 changed files with 777 additions and 20 deletions
Download Patch File Download Diff File Expand all files Collapse all files

286
vllm/model_executor/models/grok1.py
View File

@@ -21,8 +21,9 @@
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Inference-only Grok1 model."""
"""Inference-only Grok (Grok1/Grok2) model."""
import math
from collections.abc import Iterable
from itertools import islice
from typing import Any
@@ -35,9 +36,12 @@ from vllm.attention.layer import Attention
from vllm.compilation.decorators import support_torch_compile
from vllm.config import CacheConfig, VllmConfig
from vllm.distributed import get_pp_group, get_tensor_model_parallel_world_size
from vllm.logger import init_logger
from vllm.model_executor.layers.activation import GeluAndMul
from vllm.model_executor.layers.fused_moe import FusedMoE
from vllm.model_executor.layers.layernorm import RMSNorm
from vllm.model_executor.layers.linear import (
MergedColumnParallelLinear,
QKVParallelLinear,
ReplicatedLinear,
RowParallelLinear,
@@ -68,6 +72,100 @@ from .utils import (
DEFAULT_ATTN_OUTPUT_MULTIPLIER = 0.08838834764831845
DEFAULT_OUTPUT_MULTIPLIER_SCALE = 0.5773502691896257
DEFAULT_EMBEDDING_MULTIPLIER_SCALE = 78.38367176906169
DEFAULT_ROUTER_LOGIT_SOFTCAP = 30.0
logger = init_logger(__name__)
def _get_num_experts(config) -> int:
return getattr(config, "num_experts", getattr(config, "num_local_experts", 8))
def _get_moe_intermediate_size(config) -> int:
return getattr(config, "moe_intermediate_size", config.intermediate_size)
def _get_grok_version(config) -> str:
"""Detect Grok version from HF config using multiple heuristics."""
# Check for Grok2-specific attributes (both for robust detection)
has_residual_moe = getattr(config, "residual_moe", False)
has_moe_intermediate_size = hasattr(config, "moe_intermediate_size")
if has_residual_moe or has_moe_intermediate_size:
return "grok2"
return "grok1" # Default to Grok1
def _get_rope_parameters(config) -> dict[str, Any] | None:
rope_parameters = getattr(config, "rope_parameters", None)
if rope_parameters is None:
rope_type = getattr(config, "rope_type", None)
if rope_type is None:
return None
rope_parameters = {"rope_type": rope_type}
rope_theta = getattr(config, "rope_theta", None)
if rope_theta is not None:
rope_parameters["rope_theta"] = rope_theta
scaling_factor = getattr(config, "scaling_factor", None)
if scaling_factor is not None:
rope_parameters["factor"] = scaling_factor
for name in (
"original_max_position_embeddings",
"extrapolation_factor",
"attn_factor",
"beta_fast",
"beta_slow",
):
value = getattr(config, name, None)
if value is not None:
rope_parameters[name] = value
if rope_parameters.get("rope_type") == "original":
rope_parameters = dict(rope_parameters)
rope_parameters["rope_type"] = "default"
return rope_parameters
def _get_moe_renormalize(config) -> bool:
explicit_value = getattr(
config, "moe_router_renormalize", getattr(config, "moe_renormalize", None)
)
if explicit_value is not None:
return bool(explicit_value)
return not getattr(config, "residual_moe", False)
class Grok1MLP(nn.Module):
def __init__(
self,
hidden_size: int,
intermediate_size: int,
quant_config: QuantizationConfig | None = None,
prefix: str = "",
) -> None:
super().__init__()
self.gate_up_proj = MergedColumnParallelLinear(
input_size=hidden_size,
output_sizes=[intermediate_size] * 2,
bias=False,
quant_config=quant_config,
prefix=f"{prefix}.gate_up_proj",
)
self.down_proj = RowParallelLinear(
input_size=intermediate_size,
output_size=hidden_size,
bias=False,
quant_config=quant_config,
prefix=f"{prefix}.down_proj",
)
self.act_fn = GeluAndMul()
def forward(self, x: torch.Tensor) -> torch.Tensor:
x, _ = self.gate_up_proj(x)
x = self.act_fn(x)
x, _ = self.down_proj(x)
return x
class Grok1MoE(nn.Module):
@@ -85,9 +183,11 @@ class Grok1MoE(nn.Module):
top_k: int,
hidden_size: int,
intermediate_size: int,
router_logit_soft_cap: float = 0.0,
params_dtype: torch.dtype | None = None,
quant_config: QuantizationConfig | None = None,
tp_size: int | None = None,
renormalize: bool = False,
prefix: str = "",
):
super().__init__()
@@ -110,12 +210,13 @@ class Grok1MoE(nn.Module):
intermediate_size=intermediate_size,
params_dtype=params_dtype,
reduce_results=True,
renormalize=True,
renormalize=renormalize,
quant_config=quant_config,
tp_size=tp_size,
activation="gelu",
prefix=f"{prefix}.experts",
)
self.router_logit_soft_cap = router_logit_soft_cap
def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
# NOTE: hidden_states can have either 1D or 2D shape.
@@ -123,7 +224,10 @@ class Grok1MoE(nn.Module):
hidden_states = hidden_states.view(-1, self.hidden_size)
# router_logits: (num_tokens, n_experts)
router_logits, _ = self.gate(hidden_states)
router_logits = 30.0 * F.tanh(router_logits / 30.0)
if self.router_logit_soft_cap > 0:
router_logits = self.router_logit_soft_cap * F.tanh(
router_logits / self.router_logit_soft_cap
)
final_hidden_states = self.experts(hidden_states, router_logits)
return final_hidden_states.view(orig_shape)
@@ -187,6 +291,15 @@ class Grok1Attention(nn.Module):
)
attn_logits_soft_cap = max(getattr(config, "attn_logit_softcapping", 30.0), 0.0)
attn_logit_softcapping_method = getattr(
config, "attn_logit_softcapping_method", None
)
if attn_logit_softcapping_method not in (None, "tanh"):
logger.warning_once(
"Grok attention logit softcapping method '%s' is not "
"supported; falling back to default behavior.",
attn_logit_softcapping_method,
)
self.attn = Attention(
self.num_heads,
@@ -238,30 +351,50 @@ class Grok1DecoderLayer(nn.Module):
num_heads=config.num_attention_heads,
max_position=config.max_position_embeddings,
num_kv_heads=config.num_key_value_heads,
rope_parameters=getattr(config, "rope_parameters", None),
rope_parameters=_get_rope_parameters(config),
cache_config=cache_config,
quant_config=quant_config,
prefix=f"{prefix}.attn",
config=config,
) # Pass config to Grok1Attention
# Grok1 uses "num_experts" in its config
num_experts = getattr(config, "num_experts", 8)
num_experts = _get_num_experts(config)
num_experts_per_tok = getattr(config, "num_experts_per_tok", 2)
moe_intermediate_size = _get_moe_intermediate_size(config)
moe_renormalize = _get_moe_renormalize(config)
self.moe_block = Grok1MoE(
num_experts=num_experts,
top_k=num_experts_per_tok,
hidden_size=config.hidden_size,
intermediate_size=config.intermediate_size,
intermediate_size=moe_intermediate_size,
router_logit_soft_cap=max(
getattr(
config,
"router_logit_softcapping",
DEFAULT_ROUTER_LOGIT_SOFTCAP,
),
0.0,
),
quant_config=quant_config,
renormalize=moe_renormalize,
prefix=f"{prefix}.moe_block",
)
self.residual_moe = getattr(config, "residual_moe", False)
self.residual_moe_scale = 1.0 / math.sqrt(2.0)
self.pre_attn_norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
self.post_attn_norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
self.pre_moe_norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
self.post_moe_norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
self.mlp = None
if self.residual_moe:
self.mlp = Grok1MLP(
hidden_size=config.hidden_size,
intermediate_size=config.intermediate_size,
quant_config=quant_config,
prefix=f"{prefix}.mlp",
)
def forward(
self,
@@ -286,7 +419,13 @@ class Grok1DecoderLayer(nn.Module):
# MoE block with normalization
hidden_states, residual = self.pre_moe_norm(hidden_states, residual)
hidden_states = self.moe_block(hidden_states)
if self.residual_moe:
assert self.mlp is not None
hidden_states = (
self.moe_block(hidden_states) + self.mlp(hidden_states)
) * self.residual_moe_scale
else:
hidden_states = self.moe_block(hidden_states)
hidden_states = self.post_moe_norm(hidden_states)
return hidden_states, residual
@@ -294,7 +433,16 @@ class Grok1DecoderLayer(nn.Module):
@support_torch_compile
class Grok1Model(nn.Module):
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
def __init__(
self,
*,
vllm_config: VllmConfig,
prefix: str = "",
ckpt_gate_proj_name: str = "linear",
ckpt_down_proj_name: str = "linear_1",
ckpt_up_proj_name: str = "linear_v",
weight_name_remapping: dict[str, str] | None = None,
):
super().__init__()
config = vllm_config.model_config.hf_config
@@ -305,6 +453,12 @@ class Grok1Model(nn.Module):
self.quant_config = quant_config
self.padding_idx = config.pad_token_id
# Store expert naming for weight loading
self.ckpt_gate_proj_name = ckpt_gate_proj_name
self.ckpt_down_proj_name = ckpt_down_proj_name
self.ckpt_up_proj_name = ckpt_up_proj_name
self.weight_name_remapping = weight_name_remapping or {}
self.vocab_size = config.vocab_size
self.embedding_multiplier_scale = getattr(
@@ -365,14 +519,13 @@ class Grok1Model(nn.Module):
return hidden_states
def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
# Map Grok1's unique expert parameter names to standard names
# Grok1 uses "num_experts" in its config
num_experts = getattr(self.config, "num_experts", 8)
# Map expert parameter names to standard names
num_experts = _get_num_experts(self.config)
return FusedMoE.make_expert_params_mapping(
self,
ckpt_gate_proj_name="linear", # Grok1 specific
ckpt_down_proj_name="linear_1", # Grok1 specific
ckpt_up_proj_name="linear_v", # Grok1 specific
ckpt_gate_proj_name=self.ckpt_gate_proj_name,
ckpt_down_proj_name=self.ckpt_down_proj_name,
ckpt_up_proj_name=self.ckpt_up_proj_name,
num_experts=num_experts,
)
@@ -382,12 +535,18 @@ class Grok1Model(nn.Module):
("qkv_proj", "q_proj", "q"),
("qkv_proj", "k_proj", "k"),
("qkv_proj", "v_proj", "v"),
("mlp.gate_up_proj", "mlp.gate_proj", 0),
("mlp.gate_up_proj", "mlp.up_proj", 1),
]
params_dict = dict(self.named_parameters())
loaded_params: set[str] = set()
expert_params_mapping = self.get_expert_mapping()
for name, loaded_weight in weights:
# Apply version-specific weight name remapping
for old_pattern, new_pattern in self.weight_name_remapping.items():
if old_pattern in name:
name = name.replace(old_pattern, new_pattern)
if self.quant_config is not None and (
scale_name := self.quant_config.get_cache_scale(name)
):
@@ -418,6 +577,8 @@ class Grok1Model(nn.Module):
name = maybe_remap_kv_scale_name(name, params_dict)
if name is None:
continue
if name not in params_dict:
continue
param = params_dict[name]
weight_loader = param.weight_loader
weight_loader(param, loaded_weight, shard_id)
@@ -464,6 +625,8 @@ class Grok1Model(nn.Module):
if "norm.scale" in name:
name = name.replace("scale", "weight")
if name not in params_dict:
continue
param = params_dict[name]
weight_loader = getattr(
param, "weight_loader", default_weight_loader
@@ -473,9 +636,12 @@ class Grok1Model(nn.Module):
return loaded_params
class Grok1ForCausalLM(nn.Module, SupportsLoRA, SupportsPP):
class GrokBaseForCausalLM(nn.Module, SupportsLoRA, SupportsPP):
"""Base class for Grok models with shared logic."""
fall_back_to_pt_during_load = False
# Subclasses should override these
packed_modules_mapping = {
"qkv_proj": [
"q_proj",
@@ -484,6 +650,15 @@ class Grok1ForCausalLM(nn.Module, SupportsLoRA, SupportsPP):
],
}
# Expert weight naming - subclasses override these
ckpt_gate_proj_name: str = "linear"
ckpt_down_proj_name: str = "linear_1"
ckpt_up_proj_name: str = "linear_v"
def get_weight_name_remapping(self) -> dict[str, str]:
"""Return weight name remapping for this version. Override in subclasses."""
return {}
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
super().__init__()
@@ -491,11 +666,15 @@ class Grok1ForCausalLM(nn.Module, SupportsLoRA, SupportsPP):
quant_config = vllm_config.quant_config
self.config = config
self.quant_config = quant_config
self.model = Grok1Model(
vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")
vllm_config=vllm_config,
prefix=maybe_prefix(prefix, "model"),
ckpt_gate_proj_name=self.ckpt_gate_proj_name,
ckpt_down_proj_name=self.ckpt_down_proj_name,
ckpt_up_proj_name=self.ckpt_up_proj_name,
weight_name_remapping=self.get_weight_name_remapping(),
)
self.lm_head = ParallelLMHead(
@@ -512,7 +691,9 @@ class Grok1ForCausalLM(nn.Module, SupportsLoRA, SupportsPP):
config, "output_multiplier_scale", DEFAULT_OUTPUT_MULTIPLIER_SCALE
)
self.logits_processor = LogitsProcessor(
config.vocab_size, scale=self.output_multiplier_scale
config.vocab_size,
scale=self.output_multiplier_scale,
soft_cap=getattr(config, "final_logit_softcapping", None),
)
self.make_empty_intermediate_tensors = (
@@ -553,3 +734,70 @@ class Grok1ForCausalLM(nn.Module, SupportsLoRA, SupportsPP):
def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
return self.model.get_expert_mapping()
class Grok1ForCausalLM(GrokBaseForCausalLM):
"""Grok1-specific implementation."""
# Grok1 expert weight naming
ckpt_gate_proj_name = "linear"
ckpt_down_proj_name = "linear_1"
ckpt_up_proj_name = "linear_v"
def get_weight_name_remapping(self) -> dict[str, str]:
# Grok1 uses standard naming, no remapping needed
return {}
class Grok2ForCausalLM(GrokBaseForCausalLM):
"""Grok2-specific implementation."""
# Grok2 has additional packed modules for MLP
packed_modules_mapping = {
"qkv_proj": [
"q_proj",
"k_proj",
"v_proj",
],
"gate_up_proj": [
"gate_proj",
"up_proj",
],
}
# Grok2 expert weight naming
ckpt_gate_proj_name = "w1"
ckpt_down_proj_name = "w2"
ckpt_up_proj_name = "w3"
def get_weight_name_remapping(self) -> dict[str, str]:
# Grok2 checkpoint uses different naming conventions
return {
".self_attn.": ".attn.",
".block_sparse_moe.": ".moe_block.",
}
# Version dispatch mapping
_GROK_VERSIONS: dict[str, type[GrokBaseForCausalLM]] = {
"grok1": Grok1ForCausalLM,
"grok2": Grok2ForCausalLM,
}
class GrokForCausalLM(GrokBaseForCausalLM):
"""Factory class that dispatches to version-specific implementation."""
def __new__(cls, *, vllm_config: VllmConfig, prefix: str = ""):
config = vllm_config.model_config.hf_config
version = _get_grok_version(config)
instance_cls = _GROK_VERSIONS.get(version)
if instance_cls is None:
raise ValueError(f"Unsupported Grok version: {version}")
# Merge class attributes for LoRA/quantization compatibility
cls.packed_modules_mapping = dict(cls.packed_modules_mapping)
cls.packed_modules_mapping.update(instance_cls.packed_modules_mapping)
return instance_cls(vllm_config=vllm_config, prefix=prefix)

3
vllm/model_executor/models/registry.py
View File

@@ -119,7 +119,8 @@ _TEXT_GENERATION_MODELS = {
"GraniteMoeHybridForCausalLM": ("granitemoehybrid", "GraniteMoeHybridForCausalLM"), # noqa: E501
"GraniteMoeSharedForCausalLM": ("granitemoeshared", "GraniteMoeSharedForCausalLM"), # noqa: E501
"GritLM": ("gritlm", "GritLM"),
"Grok1ModelForCausalLM": ("grok1", "Grok1ForCausalLM"),
"Grok1ModelForCausalLM": ("grok1", "GrokForCausalLM"),
"Grok1ForCausalLM": ("grok1", "GrokForCausalLM"),
"HunYuanMoEV1ForCausalLM": ("hunyuan_v1", "HunYuanMoEV1ForCausalLM"),
"HunYuanDenseV1ForCausalLM": ("hunyuan_v1", "HunYuanDenseV1ForCausalLM"),
"HCXVisionForCausalLM": ("hyperclovax_vision", "HCXVisionForCausalLM"),