add merge tools
This commit is contained in:
Serge Zaigraeff
Executable
+367
@@ -0,0 +1,367 @@
|
||||
#!/bin/bash
|
||||
|
||||
# Имя окружения conda
|
||||
CONDA_ENV_NAME="whisper-merge"
|
||||
|
||||
# Директории для моделей
|
||||
MODEL_A="/mnt/cloud/llm/whisper/whisper-large-v3-russian"
|
||||
MODEL_B="/mnt/cloud/llm/whisper/whisper-large-v3-ru-podlodka"
|
||||
OUTPUT_DIR="/mnt/cloud/llm/whisper/whisper-large-v3-russian+podlodka"
|
||||
|
||||
# Флаг обновления (по умолчанию false)
|
||||
UPDATE_ENV=false
|
||||
|
||||
# Метод слияния (по умолчанию all)
|
||||
MERGE_METHOD="all"
|
||||
|
||||
# Дополнительные параметры (по умолчанию)
|
||||
SLERP_T=0.5
|
||||
TIES_DENSITY=0.8
|
||||
ENCODER_WEIGHTS="0.6,0.4"
|
||||
DECODER_WEIGHTS="0.4,0.6"
|
||||
DIRECT_ALPHA=0.5
|
||||
|
||||
# Проверка аргументов
|
||||
while [[ "$#" -gt 0 ]]; do
|
||||
case $1 in
|
||||
--update) UPDATE_ENV=true ;;
|
||||
--method) MERGE_METHOD="$2"; shift ;;
|
||||
--slerp-t) SLERP_T="$2"; shift ;;
|
||||
--ties-density) TIES_DENSITY="$2"; shift ;;
|
||||
--encoder-weights) ENCODER_WEIGHTS="$2"; shift ;;
|
||||
--decoder-weights) DECODER_WEIGHTS="$2"; shift ;;
|
||||
--direct-alpha) DIRECT_ALPHA="$2"; shift ;;
|
||||
*) echo "⚠️ Неизвестный параметр: $1"; exit 1 ;;
|
||||
esac
|
||||
shift
|
||||
done
|
||||
|
||||
# Проверка наличия conda
|
||||
if ! command -v conda &> /dev/null; then
|
||||
echo "❌ Conda не установлен. Пожалуйста, установите conda и попробуйте снова."
|
||||
exit 1
|
||||
fi
|
||||
|
||||
# Создание окружения conda, если оно не существует
|
||||
if ! conda env list | grep -q "$CONDA_ENV_NAME"; then
|
||||
echo "🔹 Создание окружения conda: $CONDA_ENV_NAME"
|
||||
conda create -n "$CONDA_ENV_NAME" python=3.12 -y
|
||||
else
|
||||
echo "✅ Окружение conda '$CONDA_ENV_NAME' уже существует."
|
||||
fi
|
||||
|
||||
# Получение пути к conda
|
||||
CONDA_PATH=$(which conda)
|
||||
|
||||
# Проверка, что путь к conda найден
|
||||
if [ -z "$CONDA_PATH" ]; then
|
||||
echo "❌ Не удалось найти путь к conda."
|
||||
exit 1
|
||||
fi
|
||||
|
||||
# Активация окружения conda
|
||||
echo "🔹 Активация окружения conda: $CONDA_ENV_NAME"
|
||||
source $(dirname "$CONDA_PATH")/../etc/profile.d/conda.sh
|
||||
conda activate "$CONDA_ENV_NAME"
|
||||
|
||||
# Если флаг --update установлен, обновляем зависимости
|
||||
if [[ "$UPDATE_ENV" == true ]]; then
|
||||
echo "🔹 Установка и обновление зависимостей"
|
||||
pip install --upgrade pip wheel
|
||||
pip install torch torchaudio
|
||||
pip install transformers datasets evaluate accelerate
|
||||
pip install soundfile librosa tqdm
|
||||
fi
|
||||
|
||||
# Проверка установки torch и transformers
|
||||
if ! python -c "import torch, transformers" &> /dev/null; then
|
||||
echo "❌ PyTorch или Transformers не установлены. Запустите с опцией --update."
|
||||
exit 1
|
||||
fi
|
||||
|
||||
# Создание временного Python-скрипта для слияния
|
||||
cat > merge_models.py << EOL
|
||||
#!/usr/bin/env python3
|
||||
|
||||
import argparse
|
||||
import os
|
||||
import torch
|
||||
from transformers import WhisperForConditionalGeneration, WhisperProcessor
|
||||
from tqdm import tqdm
|
||||
|
||||
def direct_merge(model_a_path, model_b_path, output_path, alpha=0.5):
|
||||
"""Простое прямое слияние весов моделей с коэффициентом alpha"""
|
||||
print(f"🔹 Запуск прямого слияния с alpha={alpha}")
|
||||
|
||||
print(f"📂 Загрузка модели A из {model_a_path}")
|
||||
model_a = WhisperForConditionalGeneration.from_pretrained(model_a_path)
|
||||
|
||||
print(f"📂 Загрузка модели B из {model_b_path}")
|
||||
model_b = WhisperForConditionalGeneration.from_pretrained(model_b_path)
|
||||
|
||||
# Создаем новую модель для хранения объединенных весов
|
||||
print("🔄 Объединение весов моделей...")
|
||||
merged_model = WhisperForConditionalGeneration.from_pretrained(model_a_path)
|
||||
|
||||
# Подсчет общего количества параметров
|
||||
param_count = len(list(merged_model.named_parameters()))
|
||||
|
||||
# Объединение весов
|
||||
with torch.no_grad():
|
||||
for i, (param_name, param) in enumerate(tqdm(merged_model.named_parameters(), total=param_count, desc="Слияние параметров")):
|
||||
if param_name in model_a.state_dict() and param_name in model_b.state_dict():
|
||||
# Объединяем веса с учетом коэффициента альфа
|
||||
merged_param = alpha * model_a.state_dict()[param_name] + (1 - alpha) * model_b.state_dict()[param_name]
|
||||
param.copy_(merged_param)
|
||||
|
||||
# Создаем директорию для выходных данных, если она не существует
|
||||
os.makedirs(output_path, exist_ok=True)
|
||||
|
||||
# Сохраняем объединенную модель
|
||||
print(f"💾 Сохранение объединенной модели в {output_path}")
|
||||
merged_model.save_pretrained(output_path)
|
||||
|
||||
# Сохраняем процессор из модели A
|
||||
print("💾 Сохранение процессора")
|
||||
processor = WhisperProcessor.from_pretrained(model_a_path)
|
||||
processor.save_pretrained(output_path)
|
||||
|
||||
print("✅ Прямое слияние успешно завершено!")
|
||||
return merged_model
|
||||
|
||||
def slerp_weights(model_a, model_b, t):
|
||||
"""Выполняет SLERP (Spherical Linear Interpolation) между весами моделей"""
|
||||
# Получаем словари состояний моделей
|
||||
state_dict_a = model_a.state_dict()
|
||||
state_dict_b = model_b.state_dict()
|
||||
|
||||
# Создаем новый словарь для смешанных весов
|
||||
mixed_weights = {}
|
||||
|
||||
for key in tqdm(state_dict_a.keys(), desc="SLERP слияние"):
|
||||
if key in state_dict_b:
|
||||
# Получаем веса для текущего параметра
|
||||
weight_a = state_dict_a[key].float()
|
||||
weight_b = state_dict_b[key].float()
|
||||
|
||||
# Выполняем SLERP
|
||||
# Нормализуем веса
|
||||
norm_a = torch.norm(weight_a)
|
||||
norm_b = torch.norm(weight_b)
|
||||
|
||||
# Если норма близка к нулю, используем линейную интерполяцию
|
||||
if norm_a < 1e-6 or norm_b < 1e-6:
|
||||
mixed_weights[key] = (1 - t) * weight_a + t * weight_b
|
||||
continue
|
||||
|
||||
weight_a_normalized = weight_a / norm_a
|
||||
weight_b_normalized = weight_b / norm_b
|
||||
|
||||
# Вычисляем косинус угла между векторами
|
||||
cos_theta = torch.sum(weight_a_normalized * weight_b_normalized)
|
||||
cos_theta = torch.clamp(cos_theta, -1.0, 1.0)
|
||||
|
||||
# Вычисляем угол
|
||||
theta = torch.acos(cos_theta)
|
||||
|
||||
# Если угол близок к 0 или π, используем линейную интерполяцию
|
||||
if abs(theta) < 1e-6 or abs(theta - torch.pi) < 1e-6:
|
||||
mixed_weights[key] = (1 - t) * weight_a + t * weight_b
|
||||
else:
|
||||
# SLERP формула: (sin((1-t)*theta) / sin(theta)) * weight_a + (sin(t*theta) / sin(theta)) * weight_b
|
||||
sin_theta = torch.sin(theta)
|
||||
mixed_weights[key] = (torch.sin((1 - t) * theta) / sin_theta) * weight_a + (torch.sin(t * theta) / sin_theta) * weight_b
|
||||
|
||||
# Проверка на NaN
|
||||
if torch.isnan(mixed_weights[key]).any():
|
||||
# Если появились NaN, используем линейную интерполяцию
|
||||
mixed_weights[key] = (1 - t) * weight_a + t * weight_b
|
||||
|
||||
return mixed_weights
|
||||
|
||||
def slerp_merge(model_a_path, model_b_path, output_path, t=0.5):
|
||||
"""Слияние методом SLERP, реализованное напрямую без использования mergekit"""
|
||||
print(f"🔹 Запуск слияния методом SLERP с t={t}")
|
||||
|
||||
print(f"📂 Загрузка модели A из {model_a_path}")
|
||||
model_a = WhisperForConditionalGeneration.from_pretrained(model_a_path)
|
||||
|
||||
print(f"📂 Загрузка модели B из {model_b_path}")
|
||||
model_b = WhisperForConditionalGeneration.from_pretrained(model_b_path)
|
||||
|
||||
# Создаем новую модель для хранения смешанных весов
|
||||
print("🔄 Применение метода SLERP...")
|
||||
merged_model = WhisperForConditionalGeneration.from_pretrained(model_a_path)
|
||||
|
||||
# Получаем смешанные веса
|
||||
mixed_weights = slerp_weights(model_a, model_b, t)
|
||||
|
||||
# Загружаем смешанные веса в модель
|
||||
merged_model.load_state_dict(mixed_weights)
|
||||
|
||||
# Создаем директорию для выходных данных
|
||||
os.makedirs(output_path, exist_ok=True)
|
||||
|
||||
# Сохраняем объединенную модель
|
||||
print(f"💾 Сохранение объединенной модели в {output_path}")
|
||||
merged_model.save_pretrained(output_path)
|
||||
|
||||
# Сохраняем процессор из модели A
|
||||
processor = WhisperProcessor.from_pretrained(model_a_path)
|
||||
processor.save_pretrained(output_path)
|
||||
|
||||
print("✅ Слияние методом SLERP успешно завершено!")
|
||||
return merged_model
|
||||
|
||||
def ties_merge(model_a_path, model_b_path, output_path, density=0.8, encoder_weights=(0.6, 0.4), decoder_weights=(0.4, 0.6)):
|
||||
"""Слияние, реализованное по принципу TIES без использования mergekit"""
|
||||
print(f"🔹 Запуск слияния по принципу TIES с плотностью {density}")
|
||||
print(f" Веса кодировщика: {encoder_weights}, веса декодера: {decoder_weights}")
|
||||
|
||||
print(f"📂 Загрузка модели A из {model_a_path}")
|
||||
model_a = WhisperForConditionalGeneration.from_pretrained(model_a_path)
|
||||
|
||||
print(f"📂 Загрузка модели B из {model_b_path}")
|
||||
model_b = WhisperForConditionalGeneration.from_pretrained(model_b_path)
|
||||
|
||||
# Создаем новую модель для хранения смешанных весов
|
||||
print("🔄 Применение метода TIES...")
|
||||
merged_model = WhisperForConditionalGeneration.from_pretrained(model_a_path)
|
||||
|
||||
# Получаем словари состояний моделей
|
||||
state_dict_a = model_a.state_dict()
|
||||
state_dict_b = model_b.state_dict()
|
||||
|
||||
# Создаем словарь для смешанных весов
|
||||
mixed_weights = {}
|
||||
|
||||
# Для каждого параметра
|
||||
for key in tqdm(state_dict_a.keys(), desc="TIES слияние"):
|
||||
if key in state_dict_b:
|
||||
# Определение типа слоя (кодировщик или декодер)
|
||||
if "encoder" in key:
|
||||
w_a, w_b = encoder_weights
|
||||
elif "decoder" in key:
|
||||
w_a, w_b = decoder_weights
|
||||
else:
|
||||
# Для других параметров используем среднее
|
||||
w_a = (encoder_weights[0] + decoder_weights[0]) / 2
|
||||
w_b = (encoder_weights[1] + decoder_weights[1]) / 2
|
||||
|
||||
# Получаем веса для текущего параметра
|
||||
weight_a = state_dict_a[key]
|
||||
weight_b = state_dict_b[key]
|
||||
|
||||
# Вычисляем абсолютную разницу между весами
|
||||
abs_diff = torch.abs(weight_a - weight_b)
|
||||
|
||||
# Создаем маску для топ-k% элементов с наибольшей разницей
|
||||
# если density=0.8, то мы сохраняем 80% элементов с наибольшей разницей
|
||||
k = int((1 - density) * abs_diff.numel())
|
||||
if k > 0:
|
||||
threshold = torch.kthvalue(abs_diff.flatten(), k).values
|
||||
mask = abs_diff >= threshold
|
||||
else:
|
||||
mask = torch.ones_like(abs_diff, dtype=torch.bool)
|
||||
|
||||
# Создаем смешанные веса:
|
||||
# Для элементов с большой разницей используем веса w_a и w_b
|
||||
# Для остальных используем среднее арифметическое
|
||||
mixed = torch.zeros_like(weight_a)
|
||||
mixed[mask] = w_a * weight_a[mask] + w_b * weight_b[mask]
|
||||
mixed[~mask] = (weight_a[~mask] + weight_b[~mask]) / 2
|
||||
|
||||
mixed_weights[key] = mixed
|
||||
|
||||
# Загружаем смешанные веса в модель
|
||||
merged_model.load_state_dict(mixed_weights)
|
||||
|
||||
# Создаем директорию для выходных данных
|
||||
os.makedirs(output_path, exist_ok=True)
|
||||
|
||||
# Сохраняем объединенную модель
|
||||
print(f"💾 Сохранение объединенной модели в {output_path}")
|
||||
merged_model.save_pretrained(output_path)
|
||||
|
||||
# Сохраняем процессор из модели A
|
||||
processor = WhisperProcessor.from_pretrained(model_a_path)
|
||||
processor.save_pretrained(output_path)
|
||||
|
||||
print("✅ Слияние по принципу TIES успешно завершено!")
|
||||
return merged_model
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description="Объединение моделей Whisper методами SLERP и TIES")
|
||||
parser.add_argument("--model-a", default="$MODEL_A",
|
||||
help="Путь к первой модели")
|
||||
parser.add_argument("--model-b", default="$MODEL_B",
|
||||
help="Путь к второй модели")
|
||||
parser.add_argument("--output-dir", default="$OUTPUT_DIR",
|
||||
help="Директория для сохранения результатов")
|
||||
parser.add_argument("--method", choices=["slerp", "ties", "direct", "all"], default="$MERGE_METHOD",
|
||||
help="Метод слияния: slerp, ties, direct или all")
|
||||
|
||||
# Параметры для SLERP
|
||||
parser.add_argument("--slerp-t", type=float, default=$SLERP_T,
|
||||
help="Коэффициент интерполяции для SLERP (0-1)")
|
||||
|
||||
# Параметры для TIES
|
||||
parser.add_argument("--ties-density", type=float, default=$TIES_DENSITY,
|
||||
help="Плотность смешивания для TIES (0-1)")
|
||||
parser.add_argument("--encoder-weights", type=str, default="$ENCODER_WEIGHTS",
|
||||
help="Веса для кодировщика в формате '0.6,0.4'")
|
||||
parser.add_argument("--decoder-weights", type=str, default="$DECODER_WEIGHTS",
|
||||
help="Веса для декодера в формате '0.4,0.6'")
|
||||
|
||||
# Параметры для прямого слияния
|
||||
parser.add_argument("--direct-alpha", type=float, default=$DIRECT_ALPHA,
|
||||
help="Коэффициент альфа для прямого слияния (0-1)")
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
# Разбор весов для TIES
|
||||
encoder_weights = tuple(map(float, args.encoder_weights.split(',')))
|
||||
decoder_weights = tuple(map(float, args.decoder_weights.split(',')))
|
||||
|
||||
# Создание выходной директории
|
||||
os.makedirs(args.output_dir, exist_ok=True)
|
||||
|
||||
# Выполнение слияния выбранным методом
|
||||
if args.method == "direct" or args.method == "all":
|
||||
direct_output = os.path.join(args.output_dir, "direct")
|
||||
direct_merge(args.model_a, args.model_b, direct_output, args.direct_alpha)
|
||||
|
||||
if args.method == "slerp" or args.method == "all":
|
||||
slerp_output = os.path.join(args.output_dir, "slerp")
|
||||
try:
|
||||
slerp_merge(args.model_a, args.model_b, slerp_output, args.slerp_t)
|
||||
except Exception as e:
|
||||
print(f"❌ Ошибка при слиянии методом SLERP: {e}")
|
||||
print("⚠️ Используется прямое слияние вместо SLERP.")
|
||||
direct_merge(args.model_a, args.model_b, slerp_output, args.slerp_t)
|
||||
|
||||
if args.method == "ties" or args.method == "all":
|
||||
ties_output = os.path.join(args.output_dir, "ties")
|
||||
try:
|
||||
ties_merge(args.model_a, args.model_b, ties_output, args.ties_density, encoder_weights, decoder_weights)
|
||||
except Exception as e:
|
||||
print(f"❌ Ошибка при слиянии методом TIES: {e}")
|
||||
print("⚠️ Используется прямое слияние вместо TIES.")
|
||||
alpha = (encoder_weights[0] + decoder_weights[0]) / 2
|
||||
direct_merge(args.model_a, args.model_b, ties_output, alpha)
|
||||
|
||||
print("\n✅ Все операции слияния успешно завершены!")
|
||||
print(f"📂 Результаты сохранены в {args.output_dir}")
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
EOL
|
||||
|
||||
# Запуск Python-скрипта
|
||||
python merge_models.py
|
||||
|
||||
# Удаление временного скрипта
|
||||
rm merge_models.py
|
||||
|
||||
echo ""
|
||||
echo "✅ Процесс слияния завершен. Результаты сохранены в $OUTPUT_DIR"
|
||||
Reference in New Issue
Block a user