Spaces:
Sleeping
Sleeping
File size: 7,449 Bytes
7975a75 | 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 | """
src/speculative_decoding.py
Zero-SVD Self-Speculative Decoding for Q-TensorFormer.
Leverages nested Tensor-Train core slicing so that a SINGLE trained model
serves as both its own ultra-fast low-rank draft model and high-capacity target model.
No secondary draft model weights are trained or stored in memory (0 MB parameter overhead).
"""
import time
from typing import Dict, Any, List, Tuple, Optional
import torch
import torch.nn as nn
class ZeroSVDSpeculativeDecoder:
"""
Self-Speculative Decoding engine for Q-TensorFormer.
Operates by slicing nested TT cores on the fly:
- Drafting: Slices weights to draft_rank (e.g. r=1 or r=2) to generate K draft tokens.
- Verification: Slices weights to target_rank (e.g. r=8 or r=16) to verify all K tokens
in a single parallel forward pass.
"""
def __init__(
self,
model: nn.Module,
draft_rank: int = 2,
target_rank: int = 8,
lookahead_k: int = 3,
temperature: float = 0.0,
):
self.model = model
self.draft_rank = draft_rank
self.target_rank = target_rank
self.lookahead_k = lookahead_k
self.temperature = max(0.0, temperature)
def _set_model_rank(self, rank: int):
"""Sets active rank across all TT layers in the model."""
for module in self.model.modules():
if hasattr(module, "set_rank") and callable(module.set_rank):
module.set_rank(rank)
@torch.no_grad()
def generate(
self,
input_ids: torch.Tensor,
max_new_tokens: int = 32,
) -> Tuple[torch.Tensor, Dict[str, Any]]:
"""
Runs self-speculative decoding.
Returns:
generated_ids: Tensor of shape (B, S + generated_tokens)
telemetry: Dict with acceptance_rate, draft_tokens_total, accepted_tokens_total, speedup
"""
device = input_ids.device
curr_ids = input_ids.clone()
total_accepted = 0
total_drafted = 0
speculation_steps = 0
tokens_generated = 0
t_start = time.perf_counter()
while tokens_generated < max_new_tokens:
step_k = min(self.lookahead_k, max_new_tokens - tokens_generated)
if step_k <= 0:
break
# 1. Drafting Phase: Set model to ultra-low draft_rank
self._set_model_rank(self.draft_rank)
draft_tokens = []
draft_probs = []
draft_prefix = curr_ids.clone()
for _ in range(step_k):
outputs = self.model(draft_prefix)
if isinstance(outputs, torch.Tensor):
logits = outputs
elif isinstance(outputs, (tuple, list)):
logits = outputs[0]
elif hasattr(outputs, "logits"):
logits = outputs.logits
else:
logits = outputs
next_logits = logits[:, -1, :]
if self.temperature > 0.0:
probs = torch.softmax(next_logits / self.temperature, dim=-1)
next_tok = torch.multinomial(probs, num_samples=1)
else:
probs = torch.softmax(next_logits, dim=-1)
next_tok = next_logits.argmax(dim=-1, keepdim=True)
tok_id = next_tok[0, 0].item()
draft_tokens.append(tok_id)
draft_probs.append(probs[0, tok_id].item())
draft_prefix = torch.cat([draft_prefix, next_tok], dim=1)
total_drafted += step_k
speculation_steps += 1
# 2. Verification Phase: Set model to target_rank and evaluate in parallel
self._set_model_rank(self.target_rank)
# Form parallel candidate sequence: curr_ids + draft_tokens
candidate_seq = torch.cat([
curr_ids,
torch.tensor([draft_tokens], device=device, dtype=curr_ids.dtype)
], dim=1)
verify_outputs = self.model(candidate_seq)
if isinstance(verify_outputs, torch.Tensor):
verify_logits = verify_outputs
elif isinstance(verify_outputs, (tuple, list)):
verify_logits = verify_outputs[0]
elif hasattr(verify_outputs, "logits"):
verify_logits = verify_outputs.logits
else:
verify_logits = verify_outputs
# Verification indices: positions evaluating the draft tokens
start_pos = curr_ids.shape[1] - 1
n_accepted = 0
for i in range(step_k):
pos = start_pos + i
target_logits_pos = verify_logits[:, pos, :]
target_prob = torch.softmax(target_logits_pos, dim=-1)[0, draft_tokens[i]].item()
draft_prob = max(1e-7, draft_probs[i])
if self.temperature == 0.0:
# Greedy acceptance: exact match on argmax
target_arg = target_logits_pos.argmax(dim=-1).item()
if target_arg == draft_tokens[i]:
n_accepted += 1
else:
# Reject draft and append target's true argmax
break
else:
ratio = target_prob / draft_prob
if ratio >= 1.0 or torch.rand(1).item() < ratio:
n_accepted += 1
else:
break
total_accepted += n_accepted
# Append accepted draft tokens
if n_accepted > 0:
accepted_tensor = torch.tensor([draft_tokens[:n_accepted]], device=device, dtype=curr_ids.dtype)
curr_ids = torch.cat([curr_ids, accepted_tensor], dim=1)
tokens_generated += n_accepted
# Append one correction token from the target model if space remains
if tokens_generated < max_new_tokens:
correction_pos = start_pos + n_accepted
correction_logits = verify_logits[:, correction_pos, :]
if self.temperature > 0.0:
corr_probs = torch.softmax(correction_logits / self.temperature, dim=-1)
corr_tok = torch.multinomial(corr_probs, num_samples=1)
else:
corr_tok = correction_logits.argmax(dim=-1, keepdim=True)
curr_ids = torch.cat([curr_ids, corr_tok], dim=1)
tokens_generated += 1
t_elapsed = time.perf_counter() - t_start
acceptance_rate = total_accepted / max(1, total_drafted)
telemetry = {
"elapsed_seconds": t_elapsed,
"tokens_generated": tokens_generated,
"tokens_per_sec": tokens_generated / max(1e-6, t_elapsed),
"total_drafted": total_drafted,
"total_accepted": total_accepted,
"acceptance_rate": acceptance_rate,
"speculation_steps": speculation_steps,
"draft_rank": self.draft_rank,
"target_rank": self.target_rank,
"lookahead_k": self.lookahead_k,
"zero_extra_vram_overhead": True,
"scientific_classification": "MEASURED",
}
# Restore model to target rank
self._set_model_rank(self.target_rank)
return curr_ids, telemetry
|