Spaces:
Paused
Paused
File size: 4,355 Bytes
bd89ed8 0b94c41 bd89ed8 0b94c41 bd89ed8 0b94c41 bd89ed8 0b94c41 3b0ae8d 0b94c41 bd89ed8 0b94c41 bd89ed8 0b94c41 bd89ed8 0b94c41 bd89ed8 0b94c41 bd89ed8 0b94c41 bd89ed8 0b94c41 bd89ed8 0b94c41 |
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 |
"""
Holds the interface between the gradio app and the medusa training script
"""
import os
import multiprocessing as mp
from huggingface_hub import HfApi
from huggingface_hub.utils import RepositoryNotFoundError
from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer
import torch
import torch.distributed.run as distributed_run
OUTPUT_DIR = "medusa_heads"
DATASET = "vicuna"
# These can't be changed (e.g. they control the output path)
FIXED_TRAINING_ARGS = \
"""src/medusa_training_script.py
--model_name_or_path {model_id}
--output_dir {output_dir}
--run_name {model_id}-medusa-{dataset}
--dataset {dataset}"""
# These can be freely changed
DEFAULT_TRAINING_ARGS = \
"""--medusa_num_heads 3
--medusa_num_layers 1
--model_max_length 2048
--bf16 True
--num_train_epochs 1
--per_device_train_batch_size 64
--per_device_eval_batch_size 64
--gradient_accumulation_steps 8
--evaluation_strategy no
--save_strategy no
--weight_decay 0.0
--warmup_ratio 0.1
--lr_scheduler_type cosine
--logging_steps 10
--tf32 True
--auto_find_batch_size True
--learning_rate 1e-3"""
def train_medusa_heads(model_id: str, training_args: str, dataset: str):
all_training_args = FIXED_TRAINING_ARGS.format(
model_id=model_id, output_dir=OUTPUT_DIR, dataset=dataset,
) + "\n" + training_args
all_training_arg_list = []
for arg in all_training_args.split("\n"):
all_training_arg_list += arg.split(" ")
print("Full argument list:", all_training_arg_list)
parser = distributed_run.get_args_parser()
args = parser.parse_args(all_training_arg_list)
distributed_run.run(args)
def run(model_id: str, training_args: str, dataset: str) -> str:
print(f"\n\n\nNEW RUN: {model_id}")
api = HfApi()
model_name = model_id.split("/")[-1]
repo_id = f"joaogante/{model_name}-medusa-{dataset}"
# Input validation
if model_id == "":
return """
### Invalid input 🐞
Please fill a model_id.
"""
if api.repo_exists(repo_id):
return f"""
### Invalid input 🐞
{repo_id} already exists, which means that {model_id} has already been used to create medusa heads.
"""
print(f"Valid inputs ✅\nValidating model_id: {model_id}")
# Attempt to load the base model
try:
config = AutoConfig.from_pretrained(model_id)
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=torch.bfloat16)
del config, tokenizer, model
except Exception as e:
return f"""
### {model_id} can't be loaded with AutoClasses 🐞
{e}
"""
print(f"{model_id} can be loaded ✅\nCreating medusa heads (will take a few hours)")
# Run the medusa heads creation
try:
proc = mp.Process(target=train_medusa_heads, args=(model_id, training_args, dataset))
proc.start()
proc.join()
print("Medusa heads training process completed (it might have crashed!)")
except Exception as e:
print("Error ❌\n", e)
return f"""
### Error 😢😢😢
{e}
"""
# Upload the medusa heads to the Hub
try:
# Folder path from https://github.com/FasterDecoding/Medusa/blob/main/medusa/train/train.py#L399
folder_path = (
f"{OUTPUT_DIR}_medusa_{model_name}"
)
if not any([x for x in os.listdir(folder_path) if len(x) >= 3 and x[-3:] == ".pt"]):
raise Exception(
"No model data in the expected model folder, the traning run probably failed. Check the logs for more "
"information."
)
api.create_repo(
repo_id=repo_id,
exist_ok=True,
)
api.upload_folder(
folder_path=folder_path,
repo_id=repo_id,
)
print("Medusa heads upload success ✅\n Uploaded to: ", repo_id)
return f"""
### Success 🔥
Yay! Medusa heads were successfully created and uploaded to the following repo: {repo_id}
"""
except Exception as e:
print("Error ❌\n", e)
try:
api.delete_repo(repo_id)
except RepositoryNotFoundError:
pass
return f"""
### Error 😢😢😢
{e}
"""
|