File size: 3,402 Bytes
e65937c | 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 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 | """Shared base model components for all speculator types."""
from typing import NamedTuple
from transformers.models.llama.modeling_llama import (
LlamaDecoderLayer,
LlamaRMSNorm,
LlamaRotaryEmbedding,
)
from transformers.models.qwen3.modeling_qwen3 import (
Qwen3DecoderLayer,
Qwen3RMSNorm,
Qwen3RotaryEmbedding,
)
class ModelComponents(NamedTuple):
"""Container for the components of a speculators model.
This groups the building blocks needed to construct a model, enabling
architecture-agnostic code and selective component overriding for
speculative decoding algorithms.
Attributes:
first_layer_class: Class for the first decoder layer. Can be customized
for speculative decoding while keeping other layers standard.
decoder_layer_class: Class for standard decoder layers used throughout
the rest of the model.
norm_class: Normalization layer class (e.g., LlamaRMSNorm, Qwen3RMSNorm).
rotary_emb_class: Rotary positional embedding class for the model.
"""
first_layer_class: type
decoder_layer_class: type
norm_class: type
rotary_emb_class: type
model_classes: dict[str, ModelComponents] = {
"llama": ModelComponents(
LlamaDecoderLayer, # first_layer_class (same as decoder for base models)
LlamaDecoderLayer,
LlamaRMSNorm,
LlamaRotaryEmbedding,
),
"qwen3": ModelComponents(
Qwen3DecoderLayer, # first_layer_class (same as decoder for base models)
Qwen3DecoderLayer,
Qwen3RMSNorm,
Qwen3RotaryEmbedding,
),
}
try:
from transformers.models.qwen3_next.modeling_qwen3_next import (
Qwen3NextDecoderLayer,
Qwen3NextRMSNorm,
Qwen3NextRotaryEmbedding,
)
model_classes["qwen3_next"] = ModelComponents(
Qwen3NextDecoderLayer,
Qwen3NextDecoderLayer,
Qwen3NextRMSNorm,
Qwen3NextRotaryEmbedding,
)
except ImportError:
pass
try:
from transformers.models.qwen3_5.modeling_qwen3_5 import (
Qwen3_5DecoderLayer,
Qwen3_5RMSNorm,
Qwen3_5TextRotaryEmbedding,
)
model_classes["qwen3_5_text"] = ModelComponents(
Qwen3_5DecoderLayer,
Qwen3_5DecoderLayer,
Qwen3_5RMSNorm,
Qwen3_5TextRotaryEmbedding,
)
except ImportError:
pass
try:
from transformers.models.qwen3_5_moe.modeling_qwen3_5_moe import (
Qwen3_5MoeDecoderLayer,
Qwen3_5MoeRMSNorm,
Qwen3_5MoeTextRotaryEmbedding,
)
model_classes["qwen3_5_moe_text"] = ModelComponents(
Qwen3_5MoeDecoderLayer,
Qwen3_5MoeDecoderLayer,
Qwen3_5MoeRMSNorm,
Qwen3_5MoeTextRotaryEmbedding,
)
except ImportError:
pass
def override_components(model_type: str, **overrides) -> ModelComponents:
"""Override specific components from a base model architecture.
Used for speculative decoding to swap custom layers (typically first_layer_class)
while inheriting other components from the base model.
Args:
model_type: Base model type ("llama" or "qwen3").
**overrides: Component fields to override (first_layer_class,
decoder_layer_class, etc).
Returns:
ModelComponents with specified overrides applied.
"""
base = model_classes[model_type]
return base._replace(**overrides)
|