File size: 5,414 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 | """
Train a custom DFlash drafter for any MLX-converted model.
This example shows how to:
1. Create a generic DFlash drafter for your model
2. Generate training data using the target model
3. Train the drafter with the DFlash training recipe
4. Save and use the trained drafter
Usage:
python train_custom_drafter.py \
--model mlx-community/Llama-3.1-8B-Instruct-4bit \
--output ./my-dflash-drafter \
--dataset open-web-math \
--samples 10000
"""
import argparse
from pathlib import Path
from mlx_lm import load
from dflash_mlx.universal import UniversalDFlashDecoder
from dflash_mlx.data import generate_training_data, create_mixed_training_data
def main():
parser = argparse.ArgumentParser(description="Train custom DFlash drafter")
parser.add_argument(
"--model",
type=str,
required=True,
help="MLX target model ID (e.g., mlx-community/Llama-3.1-8B-Instruct-4bit)",
)
parser.add_argument(
"--output",
type=str,
required=True,
help="Output directory for trained drafter",
)
parser.add_argument(
"--dataset",
type=str,
default="open-web-math",
help="Dataset name or path for training data",
)
parser.add_argument(
"--samples",
type=int,
default=10000,
help="Number of training samples to generate",
)
parser.add_argument(
"--epochs",
type=int,
default=6,
help="Training epochs",
)
parser.add_argument(
"--batch-size",
type=int,
default=8,
help="Training batch size",
)
parser.add_argument(
"--lr",
type=float,
default=6e-4,
help="Learning rate",
)
parser.add_argument(
"--draft-layers",
type=int,
default=5,
help="Number of draft model layers",
)
parser.add_argument(
"--draft-hidden-size",
type=int,
default=1024,
help="Draft model hidden size",
)
parser.add_argument(
"--block-size",
type=int,
default=16,
help="DFlash block size",
)
parser.add_argument(
"--generate-data",
action="store_true",
help="Generate training data with target model first",
)
args = parser.parse_args()
output_path = Path(args.output)
output_path.mkdir(parents=True, exist_ok=True)
# 1. Load target model
print(f"\n[1] Loading target model: {args.model}")
model, tokenizer = load(args.model)
print(" ✓ Target model loaded")
# 2. Create decoder with generic drafter
print(f"\n[2] Creating DFlash decoder with generic drafter")
print(f" Draft layers: {args.draft_layers}, Hidden size: {args.draft_hidden_size}")
decoder = UniversalDFlashDecoder(
target_model=model,
tokenizer=tokenizer,
draft_layers=args.draft_layers,
draft_hidden_size=args.draft_hidden_size,
block_size=args.block_size,
)
print(" ✓ Decoder initialized")
# 3. Generate or load training data
data_path = output_path / "training_data.jsonl"
if args.generate_data or not data_path.exists():
print(f"\n[3] Generating training data...")
if args.dataset == "mixed":
create_mixed_training_data(
output_path=str(data_path),
total_samples=args.samples,
)
else:
generate_training_data(
target_model=model,
tokenizer=tokenizer,
prompts_dataset=args.dataset,
output_path=str(data_path),
num_samples=args.samples,
temperature=0.0,
)
else:
print(f"\n[3] Using existing training data: {data_path}")
# 4. Train the drafter
print(f"\n[4] Training DFlash drafter...")
print(f" Epochs: {args.epochs}, Batch size: {args.batch_size}, LR: {args.lr}")
trained_drafter = decoder.train_drafter(
dataset=str(data_path),
epochs=args.epochs,
batch_size=args.batch_size,
lr=args.lr,
output_path=str(output_path / "drafter"),
)
# 5. Save final model
print(f"\n[5] Saving trained drafter...")
decoder.save_drafter(str(output_path / "drafter"))
# Save metadata
import json
metadata = {
"target_model": args.model,
"draft_layers": args.draft_layers,
"draft_hidden_size": args.draft_hidden_size,
"block_size": args.block_size,
"training_epochs": args.epochs,
"training_samples": args.samples,
"learning_rate": args.lr,
}
with open(output_path / "metadata.json", "w") as f:
json.dump(metadata, f, indent=2)
print(f"\n{'='*60}")
print("Training complete!")
print(f"{'='*60}")
print(f"\nTo use your trained drafter:")
print(f" from dflash_mlx.universal import UniversalDFlashDecoder")
print(f" from mlx_lm import load")
print(f" model, tokenizer = load('{args.model}')")
print(f" decoder = UniversalDFlashDecoder(")
print(f" target_model=model,")
print(f" tokenizer=tokenizer,")
print(f" draft_model_path='{output_path / 'drafter'}',")
print(f" )")
print(f" output = decoder.generate('Your prompt here')")
if __name__ == "__main__":
main()
|