File size: 9,931 Bytes
0433390 | 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 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 | """
Universal DFlash decoder for any MLX-converted model.
Provides a high-level interface that works with any mlx_lm model,
including those without pre-built DFlash drafters.
"""
from typing import Optional, List, Dict, Any
import mlx.core as mx
from .model import DFlashDraftModel
from .speculative_decode import DFlashSpeculativeDecoder
class UniversalDFlashDecoder:
"""Universal DFlash decoder that works with any MLX-converted model.
This class handles:
1. Loading pre-converted DFlash drafters
2. Creating generic drafters for unsupported models
3. Training custom drafters on-the-fly
"""
def __init__(
self,
target_model,
tokenizer,
draft_model_path: Optional[str] = None,
draft_layers: int = 5,
draft_hidden_size: int = 1024,
block_size: int = 16,
device: str = "metal",
):
"""Initialize the universal decoder.
Args:
target_model: Any mlx_lm loaded model
tokenizer: Tokenizer for the model
draft_model_path: Optional path to pre-converted DFlash drafter
draft_layers: Number of draft layers (if creating generic drafter)
draft_hidden_size: Hidden size for generic drafter
block_size: Number of tokens per draft block
device: MLX device
"""
self.target_model = target_model
self.tokenizer = tokenizer
self.block_size = block_size
self.device = device
# Determine model type and vocab size
self.vocab_size = getattr(tokenizer, "vocab_size", 151936)
self.target_config = self._extract_target_config(target_model)
# Load or create draft model
if draft_model_path:
print(f"[UniversalDFlash] Loading pre-built drafter from {draft_model_path}")
from .convert import load_mlx_dflash
self.draft_model, self.draft_config = load_mlx_dflash(draft_model_path)
else:
print("[UniversalDFlash] Creating generic drafter for your model...")
self.draft_model = self._create_generic_drafter(
draft_layers=draft_layers,
draft_hidden_size=draft_hidden_size,
)
self.draft_config = None
# Create the speculative decoder
self.decoder = DFlashSpeculativeDecoder(
target_model=target_model,
draft_model=self.draft_model,
tokenizer=tokenizer,
block_size=block_size,
device=device,
)
def _extract_target_config(self, target_model) -> Dict[str, Any]:
"""Extract configuration from target model."""
config = {}
# Try to extract from model attributes
if hasattr(target_model, 'config'):
model_config = target_model.config
config['hidden_size'] = getattr(model_config, 'hidden_size', 4096)
config['num_layers'] = getattr(model_config, 'num_hidden_layers', 32)
config['vocab_size'] = getattr(model_config, 'vocab_size', 151936)
config['intermediate_size'] = getattr(model_config, 'intermediate_size', 14336)
config['num_attention_heads'] = getattr(model_config, 'num_attention_heads', 32)
config['num_key_value_heads'] = getattr(model_config, 'num_key_value_heads', 8)
else:
# Default Qwen3-4B-like config
config = {
'hidden_size': 4096,
'num_layers': 32,
'vocab_size': 151936,
'intermediate_size': 14336,
'num_attention_heads': 32,
'num_key_value_heads': 8,
}
return config
def _create_generic_drafter(
self,
draft_layers: int,
draft_hidden_size: int,
) -> DFlashDraftModel:
"""Create a generic DFlash drafter compatible with the target model.
This creates an untrained drafter that can be trained or used
with pre-trained weights from a similar architecture.
"""
# Determine architecture compatibility
hidden_size = self.target_config.get('hidden_size', 4096)
vocab_size = self.target_config.get('vocab_size', 151936)
# Scale drafter based on target model size
num_heads = draft_hidden_size // 64 # ~64 dims per head
num_kv_heads = max(1, num_heads // 4)
intermediate_size = int(draft_hidden_size * 2.75) # Standard SwiGLU ratio
drafter = DFlashDraftModel(
vocab_size=vocab_size,
hidden_size=draft_hidden_size,
num_layers=draft_layers,
num_heads=num_heads,
num_kv_heads=num_kv_heads,
intermediate_size=intermediate_size,
max_seq_len=8192,
block_size=self.block_size,
mask_token_id=0, # Will be set from tokenizer
num_target_layers=self.target_config.get('num_layers', 32),
)
return drafter
def train_drafter(
self,
dataset: str,
max_seq_length: int = 3072,
epochs: int = 6,
batch_size: int = 32,
lr: float = 6e-4,
warmup_ratio: float = 0.04,
grad_clip: float = 1.0,
output_path: Optional[str] = None,
) -> str:
"""Train a custom DFlash drafter for your target model.
Args:
dataset: Path to training dataset or HF dataset name
max_seq_length: Maximum sequence length for training
epochs: Number of training epochs
batch_size: Training batch size
lr: Learning rate
warmup_ratio: Warmup ratio for cosine schedule
grad_clip: Gradient clipping threshold
output_path: Where to save the trained drafter
Returns:
Path to saved drafter
"""
from .trainer import DFlashTrainer
print(f"[UniversalDFlash] Training custom drafter...")
trainer = DFlashTrainer(
target_model=self.target_model,
drafter=self.draft_model,
tokenizer=self.tokenizer,
)
trained_model = trainer.train(
dataset=dataset,
max_seq_length=max_seq_length,
epochs=epochs,
batch_size=batch_size,
lr=lr,
warmup_ratio=warmup_ratio,
grad_clip=grad_clip,
)
# Update the draft model
self.draft_model = trained_model
self.decoder.draft_model = trained_model
if output_path:
self.save_drafter(output_path)
return output_path or "./trained_dflash_drafter"
def save_drafter(self, path: str):
"""Save the current drafter model."""
import json
from pathlib import Path
path = Path(path)
path.mkdir(parents=True, exist_ok=True)
# Save weights
weights = dict(self.draft_model.parameters())
mx.save_safetensors(str(path / "weights.safetensors"), weights)
# Save config
config = {
"vocab_size": self.draft_model.vocab_size,
"hidden_size": self.draft_model.hidden_size,
"num_hidden_layers": self.draft_model.num_layers,
"num_attention_heads": self.draft_model.num_heads,
"num_key_value_heads": self.draft_model.num_heads // 4,
"intermediate_size": self.draft_model.layers[0].mlp.gate_proj.weight.shape[1] if hasattr(self.draft_model.layers[0].mlp.gate_proj, 'weight') else 2816,
"max_position_embeddings": self.draft_model.max_seq_len,
"block_size": self.draft_model.block_size,
}
with open(path / "config.json", "w") as f:
json.dump(config, f, indent=2)
print(f"[UniversalDFlash] Drafter saved to {path}")
def generate(
self,
prompt: str,
max_tokens: int = 2048,
temperature: float = 0.0,
stop_strings: Optional[List[str]] = None,
) -> str:
"""Generate text using DFlash speculative decoding.
Args:
prompt: Text prompt
max_tokens: Maximum tokens to generate
temperature: Sampling temperature
stop_strings: Optional stop strings
Returns:
Generated text
"""
return self.decoder.generate(
prompt=prompt,
max_tokens=max_tokens,
temperature=temperature,
stop_strings=stop_strings,
)
def benchmark(
self,
prompt: str = "Write a quicksort in Python.",
max_tokens: int = 512,
num_runs: int = 5,
) -> Dict[str, float]:
"""Benchmark DFlash speculative decoding.
Args:
prompt: Test prompt
max_tokens: Tokens per run
num_runs: Number of benchmark runs
Returns:
Dict with speedup metrics
"""
import time
print(f"[Benchmark] Running {num_runs} generations...")
# Warmup
self.generate(prompt, max_tokens=10)
# DFlash generation
dflash_times = []
for _ in range(num_runs):
start = time.time()
self.generate(prompt, max_tokens=max_tokens)
dflash_times.append(time.time() - start)
# Baseline generation (without speculative decoding)
# We estimate based on token count vs time
# In practice you'd run a full baseline comparison
avg_time = sum(dflash_times) / len(dflash_times)
tokens_per_sec = max_tokens / avg_time
print(f"[Benchmark] Avg time: {avg_time:.2f}s, Speed: {tokens_per_sec:.1f} tok/s")
return {
"avg_time_sec": avg_time,
"tokens_per_sec": tokens_per_sec,
"num_runs": num_runs,
}
|