File size: 5,192 Bytes
f9e3832 | 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 | """Build a naive layer-pruned Voxtral draft from Voxtral-Mini-3B-2507.
Loads the full model, keeps a subset of the 30 Llama decoder layers, saves
the pruned model to disk with a corrected config so HF `.from_pretrained()`
reloads it as a smaller `VoxtralForConditionalGeneration`.
Design:
- Full model: 30 decoder layers in `model.language_model.model.layers`.
- Keep layers by 0-indexed positions (default: [0, 6, 12, 18, 24, 29]).
- Preserve everything else unchanged: audio encoder, multi_modal_projector,
embed_tokens, norm, lm_head. This ensures the draft shares the audio path
and vocab with the target (both required for HF `assistant_model=`).
- Save to `--out-dir` so HF can reload with `.from_pretrained()`.
Usage:
python build_draft.py --keep 0,6,12,18,24,29 --out-dir /mnt/drafts/voxtral-mini-3B-draft-6of30
python build_draft.py --keep 0,10,20,29 --out-dir /mnt/drafts/voxtral-mini-3B-draft-4of30
Then load in a benchmark script:
target = VoxtralForConditionalGeneration.from_pretrained("mistralai/Voxtral-Mini-3B-2507", ...)
draft = VoxtralForConditionalGeneration.from_pretrained("/mnt/drafts/voxtral-mini-3B-draft-6of30", ...)
outputs = target.generate(..., assistant_model=draft)
"""
from __future__ import annotations
import argparse
import json
import os
import shutil
from pathlib import Path
import torch
from transformers import (
AutoProcessor,
VoxtralForConditionalGeneration,
)
def parse_args() -> argparse.Namespace:
p = argparse.ArgumentParser()
p.add_argument("--source", default="mistralai/Voxtral-Mini-3B-2507")
p.add_argument("--out-dir", required=True, type=Path)
p.add_argument(
"--keep",
default="0,6,12,18,24,29",
help="Comma-separated 0-indexed layer indices to keep from the 30 "
"Llama decoder layers. Must be sorted ascending, must include "
"at least the first and last for stability.",
)
p.add_argument("--dtype", default="bfloat16", choices=["bfloat16", "float16", "float32"])
p.add_argument("--device", default="cuda")
return p.parse_args()
def main() -> int:
args = parse_args()
keep_layers = [int(x) for x in args.keep.split(",")]
n_keep = len(keep_layers)
assert keep_layers == sorted(keep_layers), "--keep must be ascending"
assert n_keep >= 2, "Keep at least 2 layers"
dtype = {"bfloat16": torch.bfloat16, "float16": torch.float16, "float32": torch.float32}[args.dtype]
print(f"[build_draft] Loading {args.source} on {args.device} dtype={args.dtype}")
model = VoxtralForConditionalGeneration.from_pretrained(
args.source,
torch_dtype=dtype,
low_cpu_mem_usage=True,
)
processor = AutoProcessor.from_pretrained(args.source)
# Introspect layer container. transformers 5.x path:
# VoxtralForConditionalGeneration
# .model (VoxtralModel)
# .language_model (LlamaModel)
# .layers (ModuleList of 30 LlamaDecoderLayer)
# .lm_head
# transformers 4.5x path had .language_model directly under the top-level
# object (no intermediate .model). Support both.
if hasattr(model, "model") and hasattr(model.model, "language_model"):
decoder = model.model.language_model
elif hasattr(model, "language_model"):
lm = model.language_model
decoder = lm.model if hasattr(lm, "model") else lm
else:
raise AttributeError(
"Cannot find decoder. Model has: " + ", ".join(sorted(dict(model.named_children()).keys()))
)
assert hasattr(decoder, "layers"), (
f"Decoder container {type(decoder).__name__} has no .layers. "
f"Children: {list(dict(decoder.named_children()).keys())}"
)
layers = decoder.layers
n_full = len(layers)
print(f"[build_draft] Full model has {n_full} decoder layers")
assert max(keep_layers) < n_full, f"--keep {keep_layers} exceeds n_full={n_full}"
# Rewrite layer_idx on each kept layer so RoPE / KV cache accounting still works
kept = torch.nn.ModuleList([layers[i] for i in keep_layers])
for new_idx, layer in enumerate(kept):
if hasattr(layer, "self_attn") and hasattr(layer.self_attn, "layer_idx"):
layer.self_attn.layer_idx = new_idx
decoder.layers = kept
# Update text_config so config.json reflects the new depth
tcfg = model.config.text_config
tcfg.num_hidden_layers = n_keep
# Some HF paths look for the num_hidden_layers on the top config too.
if hasattr(model.config, "num_hidden_layers"):
model.config.num_hidden_layers = n_keep
# Save
args.out_dir.mkdir(parents=True, exist_ok=True)
print(f"[build_draft] Saving pruned model to {args.out_dir}")
model.save_pretrained(args.out_dir, safe_serialization=True)
processor.save_pretrained(args.out_dir)
# Sanity print
with open(args.out_dir / "config.json") as f:
saved = json.load(f)
print(f"[build_draft] saved text_config.num_hidden_layers = "
f"{saved.get('text_config', {}).get('num_hidden_layers')}")
print(f"[build_draft] Done.")
return 0
if __name__ == "__main__":
raise SystemExit(main())
|