debug followup (#24058)
This commit is contained in:
Xinyuan Tong
@@ -453,6 +453,15 @@ class LlavaBaseForCausalLM(nn.Module):
|
||||
elif forward_batch.forward_mode.is_decode():
|
||||
return self.language_model(input_ids, positions, forward_batch)
|
||||
|
||||
def get_embed_and_head(self):
|
||||
# Spec-decode plumbing: expose the LM's embed/head so the EAGLE draft
|
||||
# can share them with the target. self.language_model is a Llama-family
|
||||
# CausalLM that defines this method.
|
||||
return self.language_model.get_embed_and_head()
|
||||
|
||||
def set_embed_and_head(self, embed, head):
|
||||
self.language_model.set_embed_and_head(embed, head)
|
||||
|
||||
def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]):
|
||||
# Load clip vision model by cfg['mm_vision_tower']:
|
||||
# huggingface_name or path_of_clip_relative_to_llava_model_dir
|
||||
|
||||
@@ -0,0 +1,208 @@
|
||||
# Copyright 2023-2026 SGLang Team
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# 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.
|
||||
# ==============================================================================
|
||||
"""EAGLE draft model for GQA Mistral targets (e.g. Mistral Medium 3.5).
|
||||
|
||||
Reuses ``LlamaForCausalLMEagle`` for the EAGLE machinery (lm_head/embed_tokens
|
||||
construction, optional tied embeddings, capture-aux-hidden-states plumbing) but
|
||||
swaps in a Mistral-specific draft model body that:
|
||||
|
||||
- runs through the standard :class:`LlamaDecoderLayer` (GQA), not the layernorm
|
||||
-less variant ``llama_eagle.LlamaDecoderLayer`` — Mistral's EAGLE checkpoint
|
||||
ships ``layers.0.attention_norm.weight``, so layer 0 expects the input
|
||||
layernorm to be present.
|
||||
- uses ``RowParallelLinear`` for the EAGLE fc fusion layer with a
|
||||
``quant_config``, so the FP8-quantized ``eagle_linear`` weights from the
|
||||
Mistral native checkpoint load via the standard quant pipeline (``LlamaModel``
|
||||
in ``llama_eagle.py`` uses a plain :class:`torch.nn.Linear` which cannot
|
||||
consume FP8 e4m3 tensors).
|
||||
|
||||
The weight name remapping mirrors :class:`MistralForCausalLMMistralFormat` and
|
||||
adds the eagle-specific entries for ``eagle_linear`` → ``model.fc``.
|
||||
"""
|
||||
|
||||
import logging
|
||||
from collections.abc import Iterable
|
||||
from typing import Optional, Tuple
|
||||
|
||||
import regex as re
|
||||
import torch
|
||||
from torch import nn
|
||||
from transformers import PretrainedConfig
|
||||
|
||||
from sglang.srt.distributed import get_pp_group
|
||||
from sglang.srt.layers.layernorm import RMSNorm
|
||||
from sglang.srt.layers.linear import RowParallelLinear
|
||||
from sglang.srt.layers.quantization.base_config import QuantizationConfig
|
||||
from sglang.srt.layers.vocab_parallel_embedding import VocabParallelEmbedding
|
||||
from sglang.srt.model_executor.forward_batch_info import ForwardBatch, PPProxyTensors
|
||||
from sglang.srt.models.llama import LlamaDecoderLayer, LlamaForCausalLM
|
||||
from sglang.srt.models.llama_eagle import LlamaForCausalLMEagle
|
||||
from sglang.srt.utils import add_prefix
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class MistralEagleModel(nn.Module):
|
||||
"""GQA EAGLE draft body with the input-embed ⊕ target-hidden-state fusion."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
config: PretrainedConfig,
|
||||
quant_config: Optional[QuantizationConfig] = None,
|
||||
prefix: str = "",
|
||||
) -> None:
|
||||
super().__init__()
|
||||
self.config = config
|
||||
self.vocab_size = config.vocab_size
|
||||
assert (
|
||||
get_pp_group().world_size == 1
|
||||
), "MistralForCausalLMEagle currently does not support pipeline parallelism"
|
||||
self.pp_group = get_pp_group()
|
||||
self.embed_tokens = VocabParallelEmbedding(
|
||||
config.vocab_size,
|
||||
config.hidden_size,
|
||||
prefix=add_prefix("embed_tokens", prefix),
|
||||
)
|
||||
self.layers = nn.ModuleList(
|
||||
[
|
||||
LlamaDecoderLayer(
|
||||
config=config,
|
||||
layer_id=i,
|
||||
prefix=add_prefix(f"layers.{i}", prefix),
|
||||
quant_config=quant_config,
|
||||
)
|
||||
for i in range(config.num_hidden_layers)
|
||||
]
|
||||
)
|
||||
self.start_layer = 0
|
||||
self.end_layer = config.num_hidden_layers
|
||||
self.fc = RowParallelLinear(
|
||||
config.hidden_size * 2,
|
||||
config.hidden_size,
|
||||
bias=False,
|
||||
quant_config=quant_config,
|
||||
prefix=add_prefix("fc", prefix),
|
||||
input_is_parallel=False,
|
||||
)
|
||||
self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
input_ids: torch.Tensor,
|
||||
positions: torch.Tensor,
|
||||
forward_batch: ForwardBatch,
|
||||
input_embeds: Optional[torch.Tensor] = None,
|
||||
pp_proxy_tensors: Optional[PPProxyTensors] = None,
|
||||
) -> torch.Tensor:
|
||||
if input_embeds is None:
|
||||
hidden_states = self.embed_tokens(input_ids)
|
||||
else:
|
||||
hidden_states = input_embeds
|
||||
|
||||
# EAGLE fusion: concat input embedding with target's previous hidden
|
||||
# state, project back to hidden_size before going through the draft's
|
||||
# transformer layers.
|
||||
hidden_states, _ = self.fc(
|
||||
torch.cat(
|
||||
(hidden_states, forward_batch.spec_info.hidden_states),
|
||||
dim=-1,
|
||||
)
|
||||
)
|
||||
|
||||
residual = None
|
||||
for layer in self.layers:
|
||||
hidden_states, residual = layer(
|
||||
positions, hidden_states, forward_batch, residual
|
||||
)
|
||||
return hidden_states + residual
|
||||
|
||||
|
||||
class MistralForCausalLMEagle(LlamaForCausalLMEagle):
|
||||
"""EAGLE draft for GQA Mistral targets.
|
||||
|
||||
Inherits LlamaForCausalLMEagle for the lm_head/embed_tokens setup and the
|
||||
capture-aux-hidden-state hooks, then overrides ``self.model`` with the
|
||||
quant-aware :class:`MistralEagleModel` and applies Mistral native-format
|
||||
weight remapping during ``load_weights``.
|
||||
"""
|
||||
|
||||
# fmt: off
|
||||
remapping = {
|
||||
r"layers\.(\d+)\.attention_norm\.weight": r"model.layers.\1.input_layernorm.weight",
|
||||
r"layers\.(\d+)\.attention\.wq\.(\w+)": r"model.layers.\1.self_attn.q_proj.\2",
|
||||
r"layers\.(\d+)\.attention\.wk\.(\w+)": r"model.layers.\1.self_attn.k_proj.\2",
|
||||
r"layers\.(\d+)\.attention\.wv\.(\w+)": r"model.layers.\1.self_attn.v_proj.\2",
|
||||
r"layers\.(\d+)\.attention\.wo\.(\w+)": r"model.layers.\1.self_attn.o_proj.\2",
|
||||
r"layers\.(\d+)\.ffn_norm\.weight": r"model.layers.\1.post_attention_layernorm.weight",
|
||||
r"layers\.(\d+)\.feed_forward\.w1\.(\w+)": r"model.layers.\1.mlp.gate_proj.\2",
|
||||
r"layers\.(\d+)\.feed_forward\.w2\.(\w+)": r"model.layers.\1.mlp.down_proj.\2",
|
||||
r"layers\.(\d+)\.feed_forward\.w3\.(\w+)": r"model.layers.\1.mlp.up_proj.\2",
|
||||
r"norm\.weight": "model.norm.weight",
|
||||
# Eagle-specific: the fc layer that fuses input embeds and target
|
||||
# hidden states is named `eagle_linear` in the Mistral checkpoint.
|
||||
# Its FP8 weights live alongside per-tensor activation/weight scales.
|
||||
r"eagle_linear\.weight": r"model.fc.weight",
|
||||
r"eagle_linear\.qscale_act": r"model.fc.input_scale",
|
||||
r"eagle_linear\.qscale_weight": r"model.fc.weight_scale",
|
||||
# tok_embeddings and output are intentionally absent — EAGLE shares
|
||||
# both with the target model and the framework ties them at runtime.
|
||||
}
|
||||
# fmt: on
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
config: PretrainedConfig,
|
||||
quant_config: Optional[QuantizationConfig] = None,
|
||||
prefix: str = "",
|
||||
) -> None:
|
||||
# Run LlamaForCausalLMEagle.__init__ to set up lm_head/embed_tokens/etc.
|
||||
# then replace self.model (which uses a plain torch.nn.Linear for fc and
|
||||
# cannot consume FP8 weights) with our quant-aware draft body.
|
||||
super().__init__(config=config, quant_config=quant_config, prefix=prefix)
|
||||
self.model = MistralEagleModel(
|
||||
config,
|
||||
quant_config=quant_config,
|
||||
prefix=add_prefix("model", prefix),
|
||||
)
|
||||
|
||||
def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]):
|
||||
# Bypass LlamaForCausalLMEagle.load_weights' "prepend model." behaviour
|
||||
# because our remap already emits fully-qualified target names.
|
||||
return LlamaForCausalLM.load_weights(
|
||||
self, self._remap_mistral_to_llama(weights)
|
||||
)
|
||||
|
||||
def _remap_mistral_to_llama(
|
||||
self, weights: Iterable[Tuple[str, torch.Tensor]]
|
||||
) -> Iterable[Tuple[str, torch.Tensor]]:
|
||||
for name, loaded_weight in weights:
|
||||
if name.startswith("model.") or name.startswith("lm_head."):
|
||||
yield name, loaded_weight
|
||||
continue
|
||||
for k, v in self.remapping.items():
|
||||
match = re.fullmatch(k, name)
|
||||
if match:
|
||||
name = match.expand(v)
|
||||
break
|
||||
else:
|
||||
logger.warning(f"Unrecognized weight: {name}. Skipping.")
|
||||
continue
|
||||
if name.endswith(".qscale_act"):
|
||||
name = re.sub(r"\.qscale_act$", ".input_scale", name)
|
||||
elif name.endswith(".qscale_weight"):
|
||||
name = re.sub(r"\.qscale_weight$", ".weight_scale", name)
|
||||
yield name, loaded_weight
|
||||
|
||||
|
||||
EntryClass = [MistralForCausalLMEagle]
|
||||
@@ -27,8 +27,12 @@ def adapt_config_dict(
|
||||
is_moe and (config_dict["moe"].get("num_shared_experts") or 0) > 0
|
||||
)
|
||||
is_eagle = "eagle" in model.lower()
|
||||
if is_eagle and not is_moe:
|
||||
# Dense EAGLE draft model (e.g. Mistral Small 4 EAGLE).
|
||||
is_mla_eagle = is_eagle and any(
|
||||
config_dict.get(k) is not None
|
||||
for k in ("kv_lora_rank", "q_lora_rank", "v_head_dim")
|
||||
)
|
||||
if is_eagle and not is_moe and is_mla_eagle:
|
||||
# Dense MLA EAGLE draft model (e.g. Mistral Small 4 EAGLE).
|
||||
# Uses MLA attention like MistralLarge3 but has no MoE layers.
|
||||
# Set model_type to deepseek_v3 for MLA support, and override
|
||||
# MoE fields so all layers are dense.
|
||||
@@ -47,6 +51,21 @@ def adapt_config_dict(
|
||||
config_dict["topk_method"] = None
|
||||
config_dict["scoring_func"] = "softmax"
|
||||
config_dict["routing_method_type"] = 1
|
||||
elif is_eagle and not is_moe:
|
||||
# Dense GQA EAGLE draft model (e.g. Mistral Medium 3.5 EAGLE).
|
||||
# Routes to a Llama-backbone draft body — no MoE shimming required.
|
||||
config_dict["architectures"] = ["MistralForCausalLMEagle"]
|
||||
config_dict["model_type"] = "mistral"
|
||||
config_dict["rope_is_neox_style"] = False
|
||||
for mla_key in (
|
||||
"q_lora_rank",
|
||||
"qk_rope_head_dim",
|
||||
"qk_nope_head_dim",
|
||||
"kv_lora_rank",
|
||||
"v_head_dim",
|
||||
):
|
||||
if config_dict.get(mla_key) is None:
|
||||
config_dict.pop(mla_key, None)
|
||||
elif is_moe:
|
||||
if is_mistral_large_3:
|
||||
config_dict = _remap_moe_args(config_dict)
|
||||
@@ -328,9 +347,14 @@ class MistralConfigParser:
|
||||
def is_mistral_model(name) -> bool:
|
||||
"""Return True if *name* refers to a Mistral model needing the custom parser."""
|
||||
lower = str(name).lower()
|
||||
return (
|
||||
"mistral-large-3" in lower or "mistral-small-4" in lower or "leanstral" in lower
|
||||
)
|
||||
if "mistral-large-3" in lower or "mistral-small-4" in lower or "leanstral" in lower:
|
||||
return True
|
||||
# EAGLE drafts for Mistral targets ship native-format only (params.json +
|
||||
# consolidated.safetensors, no config.json), so route them through the
|
||||
# custom parser regardless of the base model name.
|
||||
if "eagle" in lower and "mistral" in lower:
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
@lru_cache(maxsize=2)
|
||||
|
||||
Reference in New Issue
Block a user