File size: 7,862 Bytes
68e5880 | 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 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 | """Actual Q1 no-cache layer-prefix execution for native Qwen3.5 decisions.
Mirrors the mask/position preparation of Transformers Qwen3_5TextModel (5.16.1).
This is a version-audited reference, not a generic model or cache implementation.
"""
from dataclasses import dataclass, replace
import torch
from torch import Tensor, nn
from transformers.masking_utils import (
create_causal_mask,
create_recurrent_attention_mask,
)
from .candidate_projection import (
UnsupportedCandidateHead,
candidate_logits,
)
@dataclass(frozen=True)
class DepthContinuation:
hidden: Tensor # complete sequence residual, BEFORE final norm
position_ids: Tensor
position_embeddings: tuple[Tensor, Tensor]
masks: dict[str, Tensor | None]
depth: int
owner: object
model_signature: tuple
@dataclass
class DepthDecision:
depth: int
logits: Tensor # [1,C]
projection_mode: str
class QwenEarlyExit(nn.Module):
"""Begin once, stop at a real depth, optionally continue without replay.
Q1 means one unpadded complete input sequence. No cache or token generation.
Continuations are ephemeral: do not mutate parameters/train-mode between
begin/advance/readout, or persist them across optimizer steps.
"""
def __init__(self, model: nn.Module):
super().__init__()
self.model = model
if self.base.config.model_type not in {"qwen3_5", "qwen3_5_text"}:
raise ValueError("only Qwen3.5 text/conditional models are supported")
if len(self.backbone.layers) != self.backbone.config.num_hidden_layers:
raise ValueError("layer count/config mismatch")
if not set(self.backbone.config.layer_types) <= {
"linear_attention",
"full_attention",
}:
raise ValueError("unsupported hybrid layer type")
self._owner = object()
@property
def base(self):
return (
self.model.get_base_model()
if hasattr(self.model, "get_base_model")
else self.model
)
@property
def backbone(self):
return (
self.base.model.language_model
if self.base.config.model_type == "qwen3_5"
else self.base.model
)
@property
def full_depth(self):
return len(self.backbone.layers)
def _signature(self):
# Reference guard: optimizer updates and mode/device changes invalidate
# all outstanding continuations. No .data mutation is supported.
return (
tuple((id(module), module.training) for module in self.model.modules()),
tuple(
(id(parameter), parameter._version, parameter.device, parameter.dtype)
for parameter in self.model.parameters()
),
)
def begin(self, input_ids: Tensor) -> DepthContinuation:
if (
input_ids.dtype != torch.long
or input_ids.ndim != 2
or input_ids.shape[0] != 1
or input_ids.shape[1] < 1
):
raise ValueError("Q1 requires nonempty int64 input_ids [1,S], no padding")
for name in (
"image_token_id",
"video_token_id",
"vision_start_token_id",
"vision_end_token_id",
):
token = getattr(self.base.config, name, None)
if token is not None and (input_ids == token).any():
raise ValueError("multimodal placeholders are unsupported")
embedding = self.base.get_input_embeddings()
if ((input_ids < 0) | (input_ids >= embedding.weight.shape[0])).any():
raise ValueError("input token outside vocabulary")
hidden = embedding(input_ids)
positions = (
torch.arange(input_ids.shape[1], device=hidden.device)
.view(1, 1, -1)
.expand(4, 1, -1)
)
text_positions = positions[0]
kwargs = {
"config": self.backbone.config,
"inputs_embeds": hidden,
"attention_mask": None,
"past_key_values": None,
"position_ids": text_positions,
}
masks = {
"full_attention": create_causal_mask(**kwargs),
"linear_attention": create_recurrent_attention_mask(**kwargs),
}
rotary = self.backbone.rotary_emb(hidden, positions[1:])
return DepthContinuation(
hidden, text_positions, rotary, masks, 0, self._owner, self._signature()
)
def _check_state(self, state):
if state.owner is not self._owner:
raise ValueError("continuation belongs to a different executor")
if state.model_signature != self._signature():
raise ValueError(
"stale continuation: model parameters or training mode changed"
)
def advance(self, state: DepthContinuation, target_depth: int) -> DepthContinuation:
self._check_state(state)
if (
type(target_depth) is not int
or not state.depth < target_depth <= self.full_depth
):
raise ValueError("target depth must advance within the model")
hidden = state.hidden
for index in range(state.depth, target_depth):
hidden = self.backbone.layers[index](
hidden,
position_embeddings=state.position_embeddings,
attention_mask=state.masks[self.backbone.config.layer_types[index]],
position_ids=state.position_ids,
past_key_values=None,
use_cache=False,
)
return replace(state, hidden=hidden, depth=target_depth)
def readout(
self, state: DepthContinuation, candidate_token_ids: Tensor
) -> DepthDecision:
self._check_state(state)
if state.depth < 1:
raise ValueError("readout requires at least one executed layer")
head = self.base.get_output_embeddings()
if (
candidate_token_ids.dtype != torch.long
or candidate_token_ids.ndim != 1
or candidate_token_ids.numel() < 2
):
raise ValueError("require at least two int64 candidate tokens [C]")
if candidate_token_ids.device != state.hidden.device:
raise ValueError("candidate IDs must share hidden device")
if candidate_token_ids.unique().numel() != candidate_token_ids.numel():
raise ValueError("duplicate candidate token")
if (
(candidate_token_ids < 0) | (candidate_token_ids >= head.weight.shape[0])
).any():
raise ValueError("candidate token outside vocabulary")
# Never replace the residual used by continuation with normalized hidden.
hidden = self.backbone.norm(state.hidden[:, -1])
try:
logits = candidate_logits(head, hidden, candidate_token_ids)
mode = "candidate_rows"
except UnsupportedCandidateHead:
# Preserve adapters/hooks/parametrizations by executing the real head.
logits = head(hidden).index_select(-1, candidate_token_ids)
mode = "full_head_fallback"
return DepthDecision(state.depth, logits, mode)
def forward(
self, input_ids: Tensor, candidate_token_ids: Tensor, *, depths: tuple[int, ...]
):
if (
not depths
or any(type(d) is not int or not 1 <= d <= self.full_depth for d in depths)
or list(depths) != sorted(set(depths))
):
raise ValueError("depths must be strictly increasing valid layer counts")
state = self.begin(input_ids)
decisions = []
for depth in depths:
state = self.advance(state, depth)
decisions.append(self.readout(state, candidate_token_ids))
return tuple(decisions)
|