368 lines
17 KiB
Bash
Executable File
368 lines
17 KiB
Bash
Executable File
#!/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"
|