diff --git a/examples/wizard/.gitignore b/examples/wizard/.gitignore new file mode 100644 index 000000000..e74bf70a2 --- /dev/null +++ b/examples/wizard/.gitignore @@ -0,0 +1,11 @@ +# HuggingFace 令牌（敏感信息，禁止提交） +.hf_token + +# 下载的模型文件与运行产物 +*/models/ +*/output/ +*/offload/ + +# Python 缓存 +__pycache__/ +*.pyc diff --git a/examples/wizard/README.md b/examples/wizard/README.md new file mode 100644 index 000000000..6faac922c --- /dev/null +++ b/examples/wizard/README.md @@ -0,0 +1,300 @@ +# Wizard Merge 示例 + +本目录包含 MindNLP Wizard 模型合并引擎的配方模板与端到端复现案例。 + +## 目录结构 + +``` +examples/wizard/ +├── README.md # 本文档 +├── .hf_token # HuggingFace 令牌文件（需自行填写，已 gitignore） +├── .gitignore +├── linear_merge.yaml # 配方模板：线性加权合并 +├── ties_merge.yaml # 配方模板：TIES 合并 +├── slerp_merge.yaml # 配方模板：球面线性插值 +└── llama3_biomed_dare_ties/ # 端到端示例：Llama3-8B 医学 DARE-TIES 合并 + ├── README.md # 示例详细文档 + ├── requirements.txt # 环境依赖（锁定版本） + ├── dare_ties_biomed.yaml # DARE-TIES 合并配方 + ├── download_models.py # 模型下载脚本 + ├── run_merge.sh # 合并执行脚本 + ├── run_eval.sh # 评测执行脚本 + └── run_lm_eval.py # lm-eval MindSpore 后端入口 +``` + +## 环境要求 + +| 项目 | 版本 | +|------|------| +| Python | >= 3.9, < 3.12（实测 3.11.15） | +| MindSpore | 2.7.1 | +| CANN | >= 8.1.RC1（实测 8.2.RC2） | +| NPU Driver | 25.3.rc1.2 | +| 硬件 | Ascend 910B2（合并可在 CPU 运行，评测需 NPU） | + +**Ascend 环境初始化**：`run_merge.sh` 和 `run_eval.sh` 会自动检测并 source CANN toolkit +与 NPU driver 路径。如果自动检测失败（报 `Unsupported device target Ascend`）， +请根据实际安装位置手动设置，示例： + +```bash +# 设置 ASCEND_HOME 为你机器上 CANN 的实际安装目录（默认 /usr/local/Ascend） +export ASCEND_HOME=/usr/local/Ascend # 按实际路径修改 +source ${ASCEND_HOME}/ascend-toolkit/set_env.sh +export LD_LIBRARY_PATH=${ASCEND_HOME}/driver/lib64:${ASCEND_HOME}/driver/lib64/driver:$LD_LIBRARY_PATH +``` + +> 脚本已支持通过 `ASCEND_HOME` 环境变量指定 CANN 安装路径，未设置时默认 `/usr/local/Ascend/`。 +> 该错误的根因是 `set_env.sh` 未将 `libascend_hal.so` 所在的 driver 目录加入 `LD_LIBRARY_PATH`。 + +安装依赖： + +```bash +cd mindnlp # 项目根目录 + +# 方式 1: 安装 MindNLP 基础依赖 + Wizard 模块兼容依赖 +# 适合开发、调试、运行配方模板 +pip install -r requirements/requirements.txt +pip install -r requirements/wizard-requirements.txt + +# 方式 2: 直接安装端到端示例锁定的完整依赖 +# 适合严格复现 llama3_biomed_dare_ties 示例结果 +pip install -r examples/wizard/llama3_biomed_dare_ties/requirements.txt +``` + +说明： + +- `requirements/wizard-requirements.txt` 提供的是 Wizard 模块的兼容依赖范围，不锁定精确版本。 +- `examples/wizard/llama3_biomed_dare_ties/requirements.txt` 是示例验证通过的锁定环境。 +- 如已按方式 2 安装，一般不需要再单独执行 `pip install lm_eval datasets`。 + +## HuggingFace 令牌配置 + +部分场景（镜像下载、gated 模型）需要 HuggingFace 令牌。两种方式任选其一： + +```bash +# 方式 1: 环境变量 +export HF_TOKEN="hf_xxxxxxxxxxxx" + +# 方式 2: 令牌文件（推荐，所有脚本自动读取） +# 编辑 examples/wizard/.hf_token，将令牌粘贴到非注释行 +``` + +令牌获取地址：https://huggingface.co/settings/tokens + +> `.hf_token` 已被 `.gitignore` 排除，不会提交到版本库。 + +--- + +## 一、配方模板 + +以下为基础合并配方模板，可直接运行或作为自定义配方的起点。 + +### linear — 线性加权合并 + +```yaml +# linear_merge.yaml +models: + - model: Qwen/Qwen2.5-7B + parameters: + weight: 1.0 + - model: Qwen/Qwen2.5-7B-Instruct + parameters: + weight: 0.3 +merge_method: linear +dtype: float16 +``` + +```bash +python -m mindnlp.wizard.merge.scripts.run_yaml \ + examples/wizard/linear_merge.yaml ./output/linear_merged +``` + +### ties — TIES (Trim, Elect Sign & Merge) + +```yaml +# ties_merge.yaml +models: + - model: Qwen/Qwen2.5-7B-Instruct + parameters: + density: 0.5 + weight: 1.0 + - model: Qwen/Qwen2.5-7B-Math + parameters: + density: 0.5 + weight: 0.5 +merge_method: ties +base_model: Qwen/Qwen2.5-7B +parameters: + normalize: true +dtype: bfloat16 +``` + +```bash +python -m mindnlp.wizard.merge.scripts.run_yaml \ + examples/wizard/ties_merge.yaml ./output/ties_merged +``` + +### slerp — 球面线性插值 + +```yaml +# slerp_merge.yaml +models: + - model: meta-llama/Llama-3-8B + - model: meta-llama/Llama-3-8B-Instruct +merge_method: slerp +base_model: meta-llama/Llama-3-8B +parameters: + t: 0.5 +dtype: bfloat16 +``` + +```bash +python -m mindnlp.wizard.merge.scripts.run_yaml \ + examples/wizard/slerp_merge.yaml ./output/slerp_merged +``` + +### 通用选项 + +```bash +# 输出为 MindSpore ckpt 格式（默认 safetensors） +python -m mindnlp.wizard.merge.scripts.run_yaml \ + recipe.yaml ./output --output-format ckpt + +# 复制 tokenizer 到输出目录 +python -m mindnlp.wizard.merge.scripts.run_yaml \ + recipe.yaml ./output --copy-tokenizer + +# 生成模型卡 +python -m mindnlp.wizard.merge.scripts.run_yaml \ + recipe.yaml ./output --write-model-card +``` + +### 自定义配方 + +配方 YAML 基本结构： + +```yaml +models: + - model: + parameters: + weight: <权重值> + density: <密度值，TIES/DARE 专用> +merge_method: <方法名> +base_model: <基准模型，TIES/DARE/SLERP 需要> +parameters: + normalize: true +dtype: +``` + +Wizard 支持 21 种合并方法：linear, slerp, ties, dare_ties, dare_linear, passthrough, +model_stock, nearswap, nuslerp, multislerp, generalized_task_arithmetic, arcee_fusion, +karcher, sce, ram, rectify_embed 等。完整说明见 +[`src/mindnlp/wizard/README.md`](../../src/mindnlp/wizard/README.md)。 + +--- + +## 二、端到端示例：Llama3-8B 医学合并 + +使用 DARE-TIES 方法复现 +[lighteternal/Llama3-merge-biomed-8b](https://huggingface.co/lighteternal/Llama3-merge-biomed-8b)， +覆盖**模型下载 → 合并 → 评测**全流程。 + +**合并后模型**：[chenjingshen/Llama3-8B-merge-biomed-wizard](https://huggingface.co/chenjingshen/Llama3-8B-merge-biomed-wizard) + +### 合并配置 + +| 配置项 | 值 | +|--------|------| +| 合并方法 | DARE-TIES | +| 基座模型 | meta-llama/Meta-Llama-3-8B-Instruct | +| 输出精度 | bfloat16 | +| int8_mask | true | + +### 所需模型 + +| 模型 | HuggingFace 链接 | density | weight | 需要令牌 | +|------|------------------|---------|--------|----------| +| Meta-Llama-3-8B-Instruct (base) | [NousResearch/Meta-Llama-3-8B-Instruct](https://huggingface.co/NousResearch/Meta-Llama-3-8B-Instruct) | 1.0 | 1.0 | 否 | +| Meta-Llama-3-8B-Instruct (delta) | 同上 | 0.60 | 0.5 | 否 | +| Hermes-2-Pro-Llama-3-8B | [NousResearch/Hermes-2-Pro-Llama-3-8B](https://huggingface.co/NousResearch/Hermes-2-Pro-Llama-3-8B) | 0.55 | 0.1 | 否 | +| Llama3-OpenBioLLM-8B | [aaditya/Llama3-OpenBioLLM-8B](https://huggingface.co/aaditya/Llama3-OpenBioLLM-8B) | 0.55 | 0.4 | 否 | + +> 使用 NousResearch 镜像，**不需要**申请 Meta Llama 访问权限。三个模型合计约 48GB。 + +### 评测数据集 + +所有数据集由 `lm-eval-harness` 自动从 HuggingFace 下载，缓存到 `output/datasets/` +（可通过 `HF_DATASETS_CACHE` 自定义）。 + +| 数据集 | HuggingFace 来源 | few-shot | 需要令牌 | +|--------|------------------|----------|----------| +| ARC Challenge | [allenai/ai2_arc](https://huggingface.co/datasets/allenai/ai2_arc) | 25-shot | 否 | +| HellaSwag | [Rowan/hellaswag](https://huggingface.co/datasets/Rowan/hellaswag) | 10-shot | 否 | +| Winogrande | [allenai/winogrande](https://huggingface.co/datasets/allenai/winogrande) | 5-shot | 否 | +| GSM8K | [openai/gsm8k](https://huggingface.co/datasets/openai/gsm8k) | 5-shot | 否 | +| MMLU (6 个医学子集) | [cais/mmlu](https://huggingface.co/datasets/cais/mmlu) | 5-shot | 否 | + +### 快速复现 + +```bash +# 步骤 1: 下载模型（约 48GB） +python examples/wizard/llama3_biomed_dare_ties/download_models.py +# 国内镜像: HF_ENDPOINT=https://hf-mirror.com python ... + +# 步骤 2: 执行合并 +bash examples/wizard/llama3_biomed_dare_ties/run_merge.sh + +# 步骤 3: 评测（依次运行 10 个数据集，约 12 小时） +bash examples/wizard/llama3_biomed_dare_ties/run_eval.sh +``` + +**跳过合并**——直接下载已合并模型后评测： + +```bash +huggingface-cli download chenjingshen/Llama3-8B-merge-biomed-wizard \ + --local-dir ./output/Llama3-merge-biomed-8b + +bash examples/wizard/llama3_biomed_dare_ties/run_eval.sh ./output/Llama3-merge-biomed-8b +``` + +### 评测结果 + +评测设置与 Open LLM Leaderboard v1 一致： + +| 数据集 | 指标 | Wizard 合并 | Llama3-8B-Instruct | OpenBioLLM-8B | +|--------|------|-------------|---------------------|---------------| +| **ARC Challenge** | Accuracy | **59.73%** | 57.17% | 55.38% | +| | Norm. Accuracy | **64.59%** | 60.75% | 58.62% | +| **HellaSwag** | Accuracy | 62.26% | **62.59%** | 61.83% | +| | Norm. Accuracy | 81.35% | **81.53%** | 80.76% | +| **Winogrande** | Accuracy | **76.01%** | 74.51% | 70.88% | +| **GSM8K** | Accuracy | **70.81%** | 68.69% | 10.15% | +| **MMLU-Anatomy** | Accuracy | 71.11% | **72.59%** | 69.62% | +| **MMLU-Clinical Knowledge** | Accuracy | **77.74%** | **77.83%** | 60.38% | +| **MMLU-College Biology** | Accuracy | **80.56%** | **81.94%** | 79.86% | +| **MMLU-College Medicine** | Accuracy | 68.21% | 63.58% | **70.52%** | +| **MMLU-Medical Genetics** | Accuracy | 82.00% | 80.00% | 80.00% | +| **MMLU-Prof. Medicine** | Accuracy | **77.57%** | 71.69% | **77.94%** | + +**综合平均（10 项主指标）**： + +| 模型 | 平均分 | +|------|--------| +| **Wizard 合并** | **75.00%** | +| Llama3-8B-Instruct | 73.31% | +| OpenBioLLM-8B | 65.87% | + +- Wizard 合并以 **75.00%** 的 10 项平均分，超越 Llama3-8B-Instruct（+1.68%）和 OpenBioLLM-8B（+9.12%） +- 4 项通用指标平均（ARC norm / HellaSwag norm / Winogrande / GSM8K）：Wizard **73.19%** vs Instruct 71.37% vs BioLLM 55.10% +- 6 项 MMLU 医学子集平均：Wizard **76.20%** vs Instruct 74.61% vs BioLLM 73.05% +- 10 项中 Wizard 在 6 项胜出 Instruct、8 项胜出 BioLLM +- 验证了 DARE-TIES 合并方法同时提升通用能力与专业领域知识的有效性 + +--- + +## 相关链接 + +- [Wizard Merge 模块文档](../../src/mindnlp/wizard/README.md) +- [MindNLP 主仓库](https://github.com/mindspore-lab/mindnlp) +- [MergeKit 原始仓库](https://github.com/arcee-ai/mergekit) +- [Wizard 复现模型 (HuggingFace)](https://huggingface.co/chenjingshen/Llama3-8B-merge-biomed-wizard) diff --git a/examples/wizard/linear_merge.yaml b/examples/wizard/linear_merge.yaml new file mode 100644 index 000000000..abe6e9d0c --- /dev/null +++ b/examples/wizard/linear_merge.yaml @@ -0,0 +1,16 @@ +# 线性加权合并示例 +# 将两个 Qwen2.5-7B 变体按权重线性合并 +# +# 用法: +# python -m mindnlp.wizard.merge.scripts.run_yaml \ +# examples/wizard/linear_merge.yaml ./output/linear_merged + +models: + - model: Qwen/Qwen2.5-7B + parameters: + weight: 1.0 + - model: Qwen/Qwen2.5-7B-Instruct + parameters: + weight: 0.3 +merge_method: linear +dtype: float16 diff --git a/examples/wizard/llama3_biomed_dare_ties/README.md b/examples/wizard/llama3_biomed_dare_ties/README.md new file mode 100644 index 000000000..91b5fcaab --- /dev/null +++ b/examples/wizard/llama3_biomed_dare_ties/README.md @@ -0,0 +1,230 @@ +# Llama3-8B-merge-biomed — DARE-TIES 医学合并复现 + +使用 MindNLP Wizard 合并引擎，复现 +[lighteternal/Llama3-merge-biomed-8b](https://huggingface.co/lighteternal/Llama3-merge-biomed-8b) +的 DARE-TIES 合并实验。 + +**合并后模型已上传至 HuggingFace**： +[chenjingshen/Llama3-8B-merge-biomed-wizard](https://huggingface.co/chenjingshen/Llama3-8B-merge-biomed-wizard) + +## 目录结构 + +``` +llama3_biomed_dare_ties/ +├── README.md # 本文档 +├── requirements.txt # 环境依赖（锁定版本） +├── dare_ties_biomed.yaml # DARE-TIES 合并配方 +├── download_models.py # 模型下载脚本 +├── run_merge.sh # 合并执行脚本 +├── run_eval.sh # 评测执行脚本 +├── run_lm_eval.py # lm-eval MindSpore 后端入口 +├── models/ # 下载的源模型（自动创建，已 gitignore） +└── output/ # 合并产物与评测结果（自动创建，已 gitignore） + ├── merged/ # 合并后的模型文件 + ├── eval/ # 评测日志与 JSON 结果 + └── datasets/ # HuggingFace 数据集缓存 +``` + +## 合并配置 + +| 配置项 | 值 | +|--------|------| +| 合并方法 | DARE-TIES | +| 基座模型 | meta-llama/Meta-Llama-3-8B-Instruct | +| 输出精度 | bfloat16 | +| int8_mask | true | + +### 合并模型列表 + +| 模型 | HuggingFace 链接 | density | weight | 需要令牌 | +|------|------------------|---------|--------|----------| +| Meta-Llama-3-8B-Instruct (base) | [NousResearch/Meta-Llama-3-8B-Instruct](https://huggingface.co/NousResearch/Meta-Llama-3-8B-Instruct) | 1.0 | 1.0 | 否 | +| Meta-Llama-3-8B-Instruct (delta) | 同上 | 0.60 | 0.5 | 否 | +| Hermes-2-Pro-Llama-3-8B | [NousResearch/Hermes-2-Pro-Llama-3-8B](https://huggingface.co/NousResearch/Hermes-2-Pro-Llama-3-8B) | 0.55 | 0.1 | 否 | +| Llama3-OpenBioLLM-8B | [aaditya/Llama3-OpenBioLLM-8B](https://huggingface.co/aaditya/Llama3-OpenBioLLM-8B) | 0.55 | 0.4 | 否 | + +> 注：本示例使用的模型均为 NousResearch 镜像或开放模型，**不需要**申请 Meta Llama 访问权限。 + +### 评测数据集 + +评测设置与 Open LLM Leaderboard v1 一致。所有数据集由 `lm-eval-harness` 自动从 HuggingFace 下载。 + +| 数据集 | HuggingFace 来源 | few-shot | 需要令牌 | +|--------|------------------|----------|----------| +| ARC Challenge | [allenai/ai2_arc](https://huggingface.co/datasets/allenai/ai2_arc) | 25-shot | 否 | +| HellaSwag | [Rowan/hellaswag](https://huggingface.co/datasets/Rowan/hellaswag) | 10-shot | 否 | +| Winogrande | [allenai/winogrande](https://huggingface.co/datasets/allenai/winogrande) | 5-shot | 否 | +| GSM8K | [openai/gsm8k](https://huggingface.co/datasets/openai/gsm8k) | 5-shot | 否 | +| MMLU (6 个医学子集) | [cais/mmlu](https://huggingface.co/datasets/cais/mmlu) | 5-shot | 否 | + +数据集默认缓存到 `output/datasets/`（可通过 `HF_DATASETS_CACHE` 环境变量自定义）。 + +## 环境要求 + +| 项目 | 版本 | +|------|------| +| Python | >= 3.9, < 3.12（实测 3.11） | +| MindSpore | 2.7.1 | +| CANN | >= 8.1.RC1（实测 8.2.RC2） | +| NPU Driver | 25.3.rc1.2 | +| 硬件 | Ascend 910B2（合并可在 CPU 运行，评测需 NPU） | + +**Ascend 环境**：`run_merge.sh` 和 `run_eval.sh` 会自动检测并 source CANN toolkit 与 driver。 +如遇 `Unsupported device target Ascend` 报错，请根据实际安装位置手动设置： + +```bash +# 设置 ASCEND_HOME 为你机器上 CANN 的实际安装目录（默认 /usr/local/Ascend） +export ASCEND_HOME=/usr/local/Ascend # 按实际路径修改 +source ${ASCEND_HOME}/ascend-toolkit/set_env.sh +export LD_LIBRARY_PATH=${ASCEND_HOME}/driver/lib64:${ASCEND_HOME}/driver/lib64/driver:$LD_LIBRARY_PATH +``` + +## 快速复现 + +### 0. 配置令牌 + +虽然本示例的模型和数据集均不需要令牌，但 HuggingFace 镜像或下载大文件时可能需要。 +两种方式任选其一： + +```bash +# 方式 1: 环境变量 +export HF_TOKEN="hf_xxxxxxxxxxxx" + +# 方式 2: 令牌文件（推荐，脚本会自动读取） +# 编辑 examples/wizard/.hf_token，将令牌粘贴到文件中（非注释行） +``` + +令牌获取地址：https://huggingface.co/settings/tokens + +### 1. 安装依赖 + +```bash +cd mindnlp # 项目根目录 + +# 使用本示例锁定的依赖版本 +pip install -r examples/wizard/llama3_biomed_dare_ties/requirements.txt + +# 或分步安装 +pip install -r requirements/requirements.txt # MindNLP 核心 +pip install -r requirements/wizard-requirements.txt # Wizard 模块 +pip install lm_eval datasets # 评测 +``` + +### 2. 下载模型 + +```bash +# 下载三个源模型到 llama3_biomed_dare_ties/models/ 目录 +python examples/wizard/llama3_biomed_dare_ties/download_models.py + +# 或指定自定义下载目录 +python examples/wizard/llama3_biomed_dare_ties/download_models.py \ + --output-dir /data/hf_models + +# 国内用户可使用镜像加速 +HF_ENDPOINT=https://hf-mirror.com \ + python examples/wizard/llama3_biomed_dare_ties/download_models.py +``` + +> 三个模型合计约 48GB，请确保磁盘空间充足。 + +### 3. 执行合并 + +```bash +# 使用脚本（推荐） +bash examples/wizard/llama3_biomed_dare_ties/run_merge.sh + +# 自定义输出路径 +bash examples/wizard/llama3_biomed_dare_ties/run_merge.sh /path/to/output + +# 或直接调用 Wizard CLI +python -m mindnlp.wizard.merge.scripts.run_yaml \ + examples/wizard/llama3_biomed_dare_ties/dare_ties_biomed.yaml \ + ./output/Llama3-merge-biomed-8b \ + --copy-tokenizer --write-model-card +``` + +合并产物约 16GB（4 个 safetensors 分片 + config.json + tokenizer）。 + +### 4. 评测 + +```bash +# 使用脚本（依次运行 10 个数据集，约 12 小时） +bash examples/wizard/llama3_biomed_dare_ties/run_eval.sh ./output/Llama3-merge-biomed-8b + +# 或单独评测某个数据集 +python examples/wizard/llama3_biomed_dare_ties/run_lm_eval.py \ + --model mindspore \ + --model_args pretrained=./output/Llama3-merge-biomed-8b,dtype=bfloat16 \ + --tasks arc_challenge \ + --num_fewshot 25 \ + --batch_size 1 +``` + +### 跳过合并：直接使用已合并模型 + +如果不想自己执行合并，可以直接从 HuggingFace 下载已合并的模型： + +```bash +# 使用 huggingface-cli +huggingface-cli download chenjingshen/Llama3-8B-merge-biomed-wizard \ + --local-dir ./output/Llama3-merge-biomed-8b + +# 或使用 Python +python -c " +from huggingface_hub import snapshot_download +snapshot_download('chenjingshen/Llama3-8B-merge-biomed-wizard', + local_dir='./output/Llama3-merge-biomed-8b') +" +``` + +然后直接跳到第 4 步执行评测。 + +## 评测结果 + +| 数据集 | 指标 | Wizard 合并 | Llama3-8B-Instruct | OpenBioLLM-8B | +|--------|------|-------------|---------------------|---------------| +| **ARC Challenge** | Accuracy | **59.73%** | 57.17% | 55.38% | +| | Norm. Accuracy | **64.59%** | 60.75% | 58.62% | +| **HellaSwag** | Accuracy | 62.26% | **62.59%** | 61.83% | +| | Norm. Accuracy | 81.35% | **81.53%** | 80.76% | +| **Winogrande** | Accuracy | **76.01%** | 74.51% | 70.88% | +| **GSM8K** | Accuracy | **70.81%** | 68.69% | 10.15% | +| **MMLU-Anatomy** | Accuracy | 71.11% | **72.59%** | 69.62% | +| **MMLU-Clinical Knowledge** | Accuracy | **77.74%** | **77.83%** | 60.38% | +| **MMLU-College Biology** | Accuracy | **80.56%** | **81.94%** | 79.86% | +| **MMLU-College Medicine** | Accuracy | 68.21% | 63.58% | **70.52%** | +| **MMLU-Medical Genetics** | Accuracy | 82.00% | 80.00% | 80.00% | +| **MMLU-Prof. Medicine** | Accuracy | **77.57%** | 71.69% | **77.94%** | + +### 综合平均（10 项主指标） + +| 模型 | 平均分 | +|------|--------| +| **Wizard 合并** | **75.00%** | +| Llama3-8B-Instruct | 73.31% | +| OpenBioLLM-8B | 65.87% | + +- Wizard 合并以 **75.00%** 的 10 项平均分，超越 Llama3-8B-Instruct（+1.68%）和 OpenBioLLM-8B（+9.12%） +- 4 项通用指标平均（ARC norm / HellaSwag norm / Winogrande / GSM8K）：Wizard **73.19%** vs Instruct 71.37% vs BioLLM 55.10% +- 6 项 MMLU 医学子集平均：Wizard **76.20%** vs Instruct 74.61% vs BioLLM 73.05% +- 10 项中 Wizard 在 6 项胜出 Instruct、8 项胜出 BioLLM +- 验证了 DARE-TIES 合并方法同时提升通用能力与专业领域知识的有效性 + +## 实际运行环境 + +| 项目 | 版本 | +|------|------| +| 硬件 | Ascend 910B2 | +| MindSpore | 2.7.1 | +| Python | 3.11.15 | +| CANN | 8.2.RC2 | +| NPU Driver | 25.3.rc1.2 | +| transformers | 4.55.0 | +| lm_eval | 0.4.11 | + +## 参考 + +- [lighteternal/Llama3-merge-biomed-8b](https://huggingface.co/lighteternal/Llama3-merge-biomed-8b) — 原始合并模型 +- [chenjingshen/Llama3-8B-merge-biomed-wizard](https://huggingface.co/chenjingshen/Llama3-8B-merge-biomed-wizard) — Wizard 复现模型 +- [Language Models are Super Mario](https://arxiv.org/abs/2311.03099) — DARE 论文 +- [Resolving Interference When Merging Models](https://arxiv.org/abs/2306.01708) — TIES 论文 diff --git a/examples/wizard/llama3_biomed_dare_ties/dare_ties_biomed.yaml b/examples/wizard/llama3_biomed_dare_ties/dare_ties_biomed.yaml new file mode 100644 index 000000000..ad93229f0 --- /dev/null +++ b/examples/wizard/llama3_biomed_dare_ties/dare_ties_biomed.yaml @@ -0,0 +1,42 @@ +# Reproduces lighteternal/Llama3-merge-biomed-8b +# Original: https://huggingface.co/lighteternal/Llama3-merge-biomed-8b +# +# DARE-TIES merge of Llama-3-8B-Instruct + Hermes-2-Pro + OpenBioLLM +# The first two model entries share the base_model identity (delta=0). +# Only Hermes-2-Pro and OpenBioLLM contribute non-zero task vectors. +# +# 用法: +# python -m mindnlp.wizard.merge.scripts.run_yaml \ +# examples/wizard/llama3_biomed_dare_ties/dare_ties_biomed.yaml \ +# ./output/Llama3-merge-biomed-8b +# +# 如果模型已下载到本地，可将 HF Hub ID 替换为本地路径，例如: +# model: /path/to/NousResearch/Meta-Llama-3-8B-Instruct + +merge_method: dare_ties +base_model: NousResearch/Meta-Llama-3-8B-Instruct + +models: + - model: NousResearch/Meta-Llama-3-8B-Instruct + parameters: + density: 1.0 + weight: 1.0 + + - model: NousResearch/Meta-Llama-3-8B-Instruct + parameters: + density: 0.60 + weight: 0.5 + + - model: NousResearch/Hermes-2-Pro-Llama-3-8B + parameters: + density: 0.55 + weight: 0.1 + + - model: aaditya/Llama3-OpenBioLLM-8B + parameters: + density: 0.55 + weight: 0.4 + +parameters: + int8_mask: true +dtype: bfloat16 diff --git a/examples/wizard/llama3_biomed_dare_ties/download_models.py b/examples/wizard/llama3_biomed_dare_ties/download_models.py new file mode 100644 index 000000000..9fe8d4479 --- /dev/null +++ b/examples/wizard/llama3_biomed_dare_ties/download_models.py @@ -0,0 +1,132 @@ +#!/usr/bin/env python3 +""" +Download the three models needed for the Llama3-merge-biomed-8b DARE-TIES merge. + +Supports retry and mirror fallback (HF_ENDPOINT=https://hf-mirror.com). + +Usage: + # Download to default location (./models/ next to this script) + python download_models.py + + # Download to a custom directory + python download_models.py --output-dir /data/hf_models + + # Use mirror (useful inside mainland China) + HF_ENDPOINT=https://hf-mirror.com python download_models.py + +Environment variables: + HF_TOKEN - HuggingFace access token (required) + HF_ENDPOINT - HuggingFace mirror endpoint (optional) +""" + +import argparse +import os +import sys +import time + +DOWNLOADS = [ + "NousResearch/Meta-Llama-3-8B-Instruct", + "NousResearch/Hermes-2-Pro-Llama-3-8B", + "aaditya/Llama3-OpenBioLLM-8B", +] +MAX_RETRIES = 5 +RETRY_DELAY = 30 + + +def _read_token() -> str: + """Read HF token from env var or .hf_token file.""" + token = os.environ.get("HF_TOKEN", "").strip() + if not token: + token = os.environ.get("HUGGING_FACE_HUB_TOKEN", "").strip() + if token: + return token + # Fallback: read from examples/wizard/.hf_token + token_file = os.path.join( + os.path.dirname(os.path.abspath(__file__)), "..", ".hf_token" + ) + if os.path.isfile(token_file): + with open(token_file, encoding="utf-8") as f: + for line in f: + line = line.strip() + if line and not line.startswith("#"): + return line + return "" + + +def main(): + parser = argparse.ArgumentParser(description="Download models for DARE-TIES merge") + parser.add_argument( + "--output-dir", + default=os.path.join(os.path.dirname(os.path.abspath(__file__)), "models"), + help="Directory to download models into (default: ./models/)", + ) + args = parser.parse_args() + base_dir = args.output_dir + + token = _read_token() + if not token: + print( + "ERROR: Set HF_TOKEN env var, or write your token into " + "examples/wizard/.hf_token", + file=sys.stderr, + ) + sys.exit(1) + + os.makedirs(base_dir, exist_ok=True) + + try: + from huggingface_hub import snapshot_download + except ImportError: + print("ERROR: pip install huggingface_hub", file=sys.stderr) + sys.exit(1) + + endpoint = os.environ.get("HF_ENDPOINT", "") + mirror = "https://hf-mirror.com" + + for repo_id in DOWNLOADS: + local_dir = os.path.join(base_dir, repo_id) + if os.path.isfile(os.path.join(local_dir, "config.json")): + print(f"[skip] {repo_id} (already exists at {local_dir})") + continue + + success = False + for attempt in range(1, MAX_RETRIES + 1): + try: + use_mirror = attempt > MAX_RETRIES // 2 and not endpoint + if use_mirror: + os.environ["HF_ENDPOINT"] = mirror + print(f"[{attempt}/{MAX_RETRIES}] {repo_id} (mirror: {mirror})") + else: + print(f"[{attempt}/{MAX_RETRIES}] {repo_id}") + + snapshot_download( + repo_id=repo_id, + local_dir=local_dir, + token=token, + ) + print(f"[done] {repo_id} -> {local_dir}") + success = True + if use_mirror and "HF_ENDPOINT" in os.environ: + del os.environ["HF_ENDPOINT"] + break + except Exception as e: + print(f"[fail] attempt {attempt}: {e}", file=sys.stderr) + if use_mirror and "HF_ENDPOINT" in os.environ: + del os.environ["HF_ENDPOINT"] + if attempt < MAX_RETRIES: + wait = RETRY_DELAY * attempt + print(f" retry in {wait}s ...", file=sys.stderr) + time.sleep(wait) + + if not success: + print( + f"FATAL: could not download {repo_id} after {MAX_RETRIES} attempts", + file=sys.stderr, + ) + sys.exit(1) + + print(f"\nAll models ready under {base_dir}") + + +if __name__ == "__main__": + main() diff --git a/examples/wizard/llama3_biomed_dare_ties/requirements.txt b/examples/wizard/llama3_biomed_dare_ties/requirements.txt new file mode 100644 index 000000000..ec64092dd --- /dev/null +++ b/examples/wizard/llama3_biomed_dare_ties/requirements.txt @@ -0,0 +1,32 @@ +# ==================================================================== +# Llama3-8B-merge-biomed DARE-TIES 示例 — 环境依赖 +# 基于 Wizard conda 环境实际版本锁定 +# 安装: pip install -r requirements.txt +# ==================================================================== + +# ---- MindSpore 框架 ---- +mindspore==2.7.1 + +# ---- 序列化 / 张量 I/O ---- +safetensors==0.7.0 +ml_dtypes==0.5.4 + +# ---- 配置与数据建模 ---- +PyYAML==6.0.3 +immutables==0.21 +pydantic==2.12.5 +typing_extensions==4.15.0 + +# ---- HuggingFace 生态 ---- +transformers==4.55.0 +tokenizers==0.21.4 +huggingface_hub==0.36.2 + +# ---- 工具库 ---- +numpy==1.26.4 +tqdm==4.67.3 +click==8.3.1 + +# ---- 评测 ---- +lm_eval==0.4.11 +datasets==4.8.2 diff --git a/examples/wizard/llama3_biomed_dare_ties/run_eval.sh b/examples/wizard/llama3_biomed_dare_ties/run_eval.sh new file mode 100755 index 000000000..68f82eaff --- /dev/null +++ b/examples/wizard/llama3_biomed_dare_ties/run_eval.sh @@ -0,0 +1,114 @@ +#!/usr/bin/env bash +set -euo pipefail + +# ============================================================================= +# Evaluate Llama3-merge-biomed-8b on Ascend NPU +# Benchmarks match the Open LLM Leaderboard v1 settings. +# Each dataset produces a separate log + JSON result file. +# +# Usage: +# bash run_eval.sh # 默认模型路径 +# bash run_eval.sh /path/to/merged_model # 自定义模型路径 +# ============================================================================= + +SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)" +MINDNLP_ROOT="$(cd "${SCRIPT_DIR}/../../.." && pwd)" + +MODEL_DIR="${1:-${SCRIPT_DIR}/output/merged}" +EVAL_DTYPE="bfloat16" +BATCH_SIZE="${BATCH_SIZE:-1}" + +EVAL_ROOT="${SCRIPT_DIR}/output/eval" +mkdir -p "$EVAL_ROOT" + +# ---- Environment ---- +export PYTHONUNBUFFERED=1 +export PYTHONPATH="${MINDNLP_ROOT}/src:${PYTHONPATH:-}" + +# Ascend NPU runtime +# Override ASCEND_HOME if CANN is not installed at the default /usr/local/Ascend +ASCEND_HOME="${ASCEND_HOME:-/usr/local/Ascend}" +if [ -f "${ASCEND_HOME}/ascend-toolkit/set_env.sh" ]; then + source "${ASCEND_HOME}/ascend-toolkit/set_env.sh" +fi +if [ -d "${ASCEND_HOME}/driver/lib64" ]; then + export LD_LIBRARY_PATH="${ASCEND_HOME}/driver/lib64:${ASCEND_HOME}/driver/lib64/driver:${LD_LIBRARY_PATH:-}" +fi + +# Activate Wizard conda environment +if command -v conda &>/dev/null; then + source "$(conda info --base)/etc/profile.d/conda.sh" + conda activate Wizard 2>/dev/null || echo "[warn] conda env 'Wizard' not found, using current env" +fi + +# Load HF token from .hf_token file if HF_TOKEN not set +HF_TOKEN_FILE="${SCRIPT_DIR}/../.hf_token" +if [ -z "${HF_TOKEN:-}" ] && [ -f "$HF_TOKEN_FILE" ]; then + _token=$(grep -v '^#' "$HF_TOKEN_FILE" | grep -v '^$' | head -1 | tr -d '[:space:]') + if [ -n "$_token" ]; then + export HF_TOKEN="$_token" + export HUGGING_FACE_HUB_TOKEN="$_token" + export HF_HUB_TOKEN="$_token" + fi +fi +if [ -z "${HF_TOKEN:-}" ]; then + echo "[warn] HF_TOKEN not set and .hf_token empty — gated datasets may fail." +fi +export HF_ENDPOINT="${HF_ENDPOINT:-https://hf-mirror.com}" + +# Dataset cache directory +export HF_DATASETS_CACHE="${HF_DATASETS_CACHE:-${SCRIPT_DIR}/output/datasets}" +mkdir -p "$HF_DATASETS_CACHE" + +LM_EVAL_SCRIPT="${SCRIPT_DIR}/run_lm_eval.py" +MODEL_ARGS="pretrained=${MODEL_DIR},dtype=${EVAL_DTYPE}" + +# ---- Tasks: ordered fastest → slowest ---- +# Format: task_name:num_fewshot +TASKS=( + "mmlu_medical_genetics:5" + "mmlu_anatomy:5" + "mmlu_college_biology:5" + "mmlu_college_medicine:5" + "mmlu_clinical_knowledge:5" + "mmlu_professional_medicine:5" + "gsm8k:5" + "arc_challenge:25" + "winogrande:5" + "hellaswag:10" +) + +SUMMARY_LOG="${SCRIPT_DIR}/output/eval_summary.log" +echo "==== Eval start $(date -u +%Y-%m-%dT%H:%M:%SZ) ====" | tee "$SUMMARY_LOG" +echo "model_dir = ${MODEL_DIR}" | tee -a "$SUMMARY_LOG" +echo "model_args = ${MODEL_ARGS}" | tee -a "$SUMMARY_LOG" +echo "batch_size = ${BATCH_SIZE}" | tee -a "$SUMMARY_LOG" +echo "" | tee -a "$SUMMARY_LOG" + +for entry in "${TASKS[@]}"; do + IFS=':' read -r task nshot <<< "$entry" + out_log="${EVAL_ROOT}/${task}.log" + + echo "---- [${task}] num_fewshot=${nshot} start $(date -u +%H:%M:%S) ----" | tee -a "$SUMMARY_LOG" + + set +e + python "$LM_EVAL_SCRIPT" \ + --model mindspore \ + --model_args "${MODEL_ARGS}" \ + --tasks "${task}" \ + --num_fewshot "${nshot}" \ + --batch_size "${BATCH_SIZE}" \ + --output_path "${EVAL_ROOT}/${task}.json" 2>&1 | tee "$out_log" + exit_code=${PIPESTATUS[0]} + set -e + + if [ $exit_code -eq 0 ]; then + result_line=$(grep -E "acc|exact_match" "$out_log" | tail -1 || echo "parse failed") + echo "[OK] task=${task} exit=${exit_code} ${result_line}" | tee -a "$SUMMARY_LOG" + else + echo "[FAIL] task=${task} exit=${exit_code}" | tee -a "$SUMMARY_LOG" + fi + echo "" | tee -a "$SUMMARY_LOG" +done + +echo "==== Eval done $(date -u +%Y-%m-%dT%H:%M:%SZ) ====" | tee -a "$SUMMARY_LOG" diff --git a/examples/wizard/llama3_biomed_dare_ties/run_lm_eval.py b/examples/wizard/llama3_biomed_dare_ties/run_lm_eval.py new file mode 100644 index 000000000..7a903a442 --- /dev/null +++ b/examples/wizard/llama3_biomed_dare_ties/run_lm_eval.py @@ -0,0 +1,76 @@ +#!/usr/bin/env python3 +""" +Wrapper to run lm-eval-harness with Wizard MindSpore backend. + +Registers the 'mindspore' model type via mindnlp.wizard.merge.eval.mindspore_lm, +then delegates to lm-eval's CLI. + +Usage: + python run_lm_eval.py --model mindspore \ + --model_args pretrained=./output/merged,dtype=bfloat16 \ + --tasks arc_challenge --num_fewshot 25 --batch_size 1 +""" + +import os +import sys + + +def main() -> int: + os.environ.setdefault("DEVICE_TARGET", "Ascend") + + # mindnlp must be imported before torch so mindtorch proxies work. + import mindnlp # pylint: disable=unused-import + import torch + from torch.utils import collect_env + + if not hasattr(collect_env, "get_pretty_env_info"): + collect_env.get_pretty_env_info = ( + lambda: "mindtorch environment info unavailable" + ) + + # Restore DynamicLayer.update for transformers 4.55+ compatibility. + try: + import transformers.cache_utils as _cache_mod + + _DL = getattr(_cache_mod, "DynamicLayer", None) + if _DL is not None: + + def _dynamic_layer_update(self, key_states, value_states, cache_kwargs=None): + if self.keys is None: + self.keys = key_states + self.values = value_states + else: + self.keys = torch.cat([self.keys, key_states], dim=-2) + self.values = torch.cat([self.values, value_states], dim=-2) + return self.keys, self.values + + _DL.update = _dynamic_layer_update + except Exception as _e: + print(f"[warn] DynamicLayer.update restore skipped: {_e}") + + # Force eager attention to avoid mindtorch SDPA reshape bug. + try: + from mindnlp.transformers import AutoModelForCausalLM as _AMCLM # pylint: disable=no-name-in-module + + _orig_from_pretrained = _AMCLM.from_pretrained + + @classmethod + def _from_pretrained_eager(cls, *args, **kwargs): + kwargs.setdefault("attn_implementation", "eager") + return _orig_from_pretrained.__func__(cls, *args, **kwargs) + + _AMCLM.from_pretrained = _from_pretrained_eager + except Exception as _e: + print(f"[warn] eager-attn patch skipped: {_e}") + + # Register lm-eval model type: "mindspore" / "mindnlp". + import mindnlp.wizard.merge.eval.mindspore_lm # pylint: disable=unused-import + from lm_eval.__main__ import cli_evaluate + + sys.argv = ["lm-eval", *sys.argv[1:]] + rc = cli_evaluate() # pylint: disable=assignment-from-no-return + return int(rc) if rc is not None else 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/examples/wizard/llama3_biomed_dare_ties/run_merge.sh b/examples/wizard/llama3_biomed_dare_ties/run_merge.sh new file mode 100755 index 000000000..46f06882f --- /dev/null +++ b/examples/wizard/llama3_biomed_dare_ties/run_merge.sh @@ -0,0 +1,66 @@ +#!/usr/bin/env bash +set -euo pipefail + +# ============================================================================= +# DARE-TIES merge — reproduces lighteternal/Llama3-merge-biomed-8b +# Uses MindNLP Wizard merge backend. +# +# Usage: +# bash run_merge.sh # 使用默认输出路径 +# bash run_merge.sh /path/to/output_dir # 自定义输出路径 +# ============================================================================= + +SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)" +RECIPE="${SCRIPT_DIR}/dare_ties_biomed.yaml" +OUT_DIR="${1:-${SCRIPT_DIR}/output/merged}" +LOG="${SCRIPT_DIR}/output/merge.log" + +# Auto-detect mindnlp src root (examples/wizard/llama3_biomed_dare_ties/ -> src/) +MINDNLP_ROOT="$(cd "${SCRIPT_DIR}/../../.." && pwd)" +export PYTHONPATH="${MINDNLP_ROOT}/src:${PYTHONPATH:-}" + +mkdir -p "$(dirname "$LOG")" + +# Ascend NPU runtime +# Override ASCEND_HOME if CANN is not installed at the default /usr/local/Ascend +ASCEND_HOME="${ASCEND_HOME:-/usr/local/Ascend}" +if [ -f "${ASCEND_HOME}/ascend-toolkit/set_env.sh" ]; then + source "${ASCEND_HOME}/ascend-toolkit/set_env.sh" +fi +if [ -d "${ASCEND_HOME}/driver/lib64" ]; then + export LD_LIBRARY_PATH="${ASCEND_HOME}/driver/lib64:${ASCEND_HOME}/driver/lib64/driver:${LD_LIBRARY_PATH:-}" +fi + +# Activate Wizard conda environment +if command -v conda &>/dev/null; then + source "$(conda info --base)/etc/profile.d/conda.sh" + conda activate Wizard 2>/dev/null || echo "[warn] conda env 'Wizard' not found, using current env" +fi + +# Load HF token from .hf_token file if HF_TOKEN not set +HF_TOKEN_FILE="${SCRIPT_DIR}/../.hf_token" +if [ -z "${HF_TOKEN:-}" ] && [ -f "$HF_TOKEN_FILE" ]; then + _token=$(grep -v '^#' "$HF_TOKEN_FILE" | grep -v '^$' | head -1 | tr -d '[:space:]') + if [ -n "$_token" ]; then + export HF_TOKEN="$_token" + export HUGGING_FACE_HUB_TOKEN="$_token" + export HF_HUB_TOKEN="$_token" + fi +fi +export HF_ENDPOINT="${HF_ENDPOINT:-https://hf-mirror.com}" + +export PYTHONUNBUFFERED=1 + +echo "==== DARE-TIES merge start $(date -u +%Y-%m-%dT%H:%M:%SZ) ====" | tee "$LOG" +echo "Recipe : ${RECIPE}" | tee -a "$LOG" +echo "Output : ${OUT_DIR}" | tee -a "$LOG" +echo "PYTHONPATH: ${PYTHONPATH}" | tee -a "$LOG" + +python -m mindnlp.wizard.merge.scripts.run_yaml \ + "${RECIPE}" \ + "${OUT_DIR}" \ + --copy-tokenizer \ + --write-model-card 2>&1 | tee -a "$LOG" + +echo "==== DARE-TIES merge done $(date -u +%Y-%m-%dT%H:%M:%SZ) ====" | tee -a "$LOG" +echo "Merged model saved to: ${OUT_DIR}" diff --git a/examples/wizard/slerp_merge.yaml b/examples/wizard/slerp_merge.yaml new file mode 100644 index 000000000..53fa32d6f --- /dev/null +++ b/examples/wizard/slerp_merge.yaml @@ -0,0 +1,16 @@ +# SLERP 合并示例 +# 球面线性插值 (Spherical Linear Interpolation)， +# 适用于两个模型之间的平滑权重合并。 +# +# 用法: +# python -m mindnlp.wizard.merge.scripts.run_yaml \ +# examples/wizard/slerp_merge.yaml ./output/slerp_merged + +models: + - model: meta-llama/Llama-3-8B + - model: meta-llama/Llama-3-8B-Instruct +merge_method: slerp +base_model: meta-llama/Llama-3-8B +parameters: + t: 0.5 +dtype: bfloat16 diff --git a/examples/wizard/ties_merge.yaml b/examples/wizard/ties_merge.yaml new file mode 100644 index 000000000..dc4d43226 --- /dev/null +++ b/examples/wizard/ties_merge.yaml @@ -0,0 +1,22 @@ +# TIES 合并示例 +# 使用 TIES (Trim, Elect Sign & Merge) 方法合并， +# 需要指定 base_model 作为参考基准。 +# +# 用法: +# python -m mindnlp.wizard.merge.scripts.run_yaml \ +# examples/wizard/ties_merge.yaml ./output/ties_merged + +models: + - model: Qwen/Qwen2.5-7B-Instruct + parameters: + density: 0.5 + weight: 1.0 + - model: Qwen/Qwen2.5-7B-Math + parameters: + density: 0.5 + weight: 0.5 +merge_method: ties +base_model: Qwen/Qwen2.5-7B +parameters: + normalize: true +dtype: bfloat16 diff --git a/requirements/wizard-requirements.txt b/requirements/wizard-requirements.txt new file mode 100644 index 000000000..fe7f5b5f5 --- /dev/null +++ b/requirements/wizard-requirements.txt @@ -0,0 +1,50 @@ +# ==================================================================== +# Wizard 模块依赖（模型合并引擎） +# +# 用途: +# 1. 提供 Wizard 模块的兼容依赖范围 +# 2. 作为 requirements/requirements.txt 之上的补充依赖 +# +# 说明: +# - MindSpore 由 requirements/requirements.txt 或现有环境提供 +# - 端到端示例的精确复现，请使用: +# examples/wizard/llama3_biomed_dare_ties/requirements.txt +# +# 安装: +# pip install -r requirements/requirements.txt +# pip install -r requirements/wizard-requirements.txt +# ==================================================================== + +# ---- 序列化 / 张量 I/O（必需）---- +safetensors>=0.7,<1.0 +ml_dtypes>=0.5,<1.0 + +# ---- 配置与数据建模（必需）---- +PyYAML>=6.0,<7.0 +immutables>=0.21,<1.0 +pydantic>=2.12,<3.0 +typing_extensions>=4.15,<5.0 + +# ---- HuggingFace 生态（必需）---- +transformers>=4.55,<5.0 +tokenizers>=0.21,<0.22 +huggingface_hub>=0.36,<1.0 + +# ---- LoRA 子命令（可选：wizard-extract-lora / LoRA merge）---- +# 如需使用 LoRA 相关路径，建议额外安装 peft +peft>=0.18,<0.19 + +# ---- 工具库（必需）---- +numpy>=1.26,<2.0 +tqdm>=4.67,<5.0 +click>=8.3,<9.0 + +# ---- 评测后端（推荐）---- +lm_eval>=0.4,<0.5 +datasets>=4.8,<5.0 + +# ---- 进化搜索 / 分布式（可选，仅 evo 子模块使用）---- +# pandas>=2.0,<3.0 +# ray>=2.0,<3.0 +# cma>=3.0,<4.0 +# vllm diff --git a/setup.py b/setup.py index a002a5aea..bbcc36544 100644 --- a/setup.py +++ b/setup.py @@ -171,7 +171,8 @@ def run(self): }, entry_points={ 'console_scripts': [ - 'mtrun=mindtorch.distributed.run:main' + 'mtrun=mindtorch.distributed.run:main', + 'wizard-merge=mindnlp.wizard.merge.scripts.run_yaml:main', ], }, diff --git a/src/mindnlp/wizard/README.md b/src/mindnlp/wizard/README.md new file mode 100644 index 000000000..682794b33 --- /dev/null +++ b/src/mindnlp/wizard/README.md @@ -0,0 +1,284 @@ +# MindNLP Wizard 模型合并模块 + +综合参考 [MergeKit](https://github.com/arcee-ai/mergekit)（Arcee AI）与 PaddleNLP 的相关工程实践，并面向 **MindSpore / Ascend NPU** 适配实现的同架构大模型权重合并引擎。 + +## 📐 设计思路 + +Wizard 的设计同时参考了业界已有 merge 工具的两类经验：一类是以 MergeKit 为代表的配方系统、任务拆解与合并方法设计，另一类是 PaddleNLP 在大模型工程落地中的实现经验。基于此，Wizard 面向 MindSpore / Ascend NPU 重新实现同架构模型合并能力，在尽量保留主流配方语义与执行思路的同时，补齐 `.ckpt`、dtype 安全等本地化工程支持。 + +### 核心设计原则 + +1. **优先兼容上游配方语义，控制实现分歧** + - 在主流同架构合并场景下，尽量保持与 MergeKit 一致的 YAML 配方结构、方法语义与任务拆解方式 + - 对 MindSpore / Ascend 特有能力（如 `.ckpt`、多 NPU、dtype 安全）进行本地化扩展，而不是机械追求接口逐项等同 +2. **dtype 安全优先** + - MindSpore CPU 不原生支持 bfloat16/float16 算术，会静默提升为 float32 + - 引入 `dtype_policy` 模块统一处理：检测 → 提升 → 计算 → 还原 + - 对精度敏感的关键算子通过 `safe_ops` 封装，降低 half 精度在 CPU 路径上的数值风险 +3. **统一懒加载** + - 三种输入格式（safetensors / .bin / .ckpt）全部支持懒加载 + - 仅读取请求的张量字节，不加载整个文件到内存 +4. **DAG 驱动，按需计算** + - 每个张量的合并是一个独立 Task，通过有向无环图（DAG）描述依赖 + - Executor 按拓扑序执行，中间结果用完立即释放，峰值内存远低于全量加载 +5. **格式中立的 I/O 抽象** + - `TensorLoader` / `TensorWriter` 统一读写接口，上层代码不感知底层格式 + - 新增格式只需实现 Loader + Writer，不影响合并逻辑 +6. **原创适配层与上游代码分离** + - MergeKit 移植代码遵循 LGPL-3.0，Wizard 原创文件遵循 Apache 2.0 + - 原创模块集中在：`dtype_policy`、`safe_ops`、`preflight`、`io/_device`、`io/lazy_ckpt`、`eval/` + - 许可证边界清晰，便于合规审查 + +## 📁 目录结构 + +``` +src/mindnlp/wizard/ +├── __init__.py +└── merge/ + ├── merge.py # 合并主入口 (run_merge) + ├── config.py # YAML 配方解析 + ├── common.py # ModelReference 等公共类 + ├── options.py # MergeOptions 运行选项 + ├── graph.py # Task DAG 与 Executor + ├── plan.py # MergePlanner 合并计划生成 + ├── card.py # 模型卡生成 + ├── dtype_policy.py # ★ dtype 安全策略（bf16 保护） + ├── safe_ops.py # ★ 安全张量运算 + ├── preflight.py # ★ 合并前校验 + ├── sparsify.py # 稀疏化（DARE/TIES 用） + ├── multigpu_executor.py # 多设备调度 + ├── architecture/ # 模型架构定义（50+ 种） + ├── merge_methods/ # 20 种合并算法 + ├── io/ # 张量 I/O + │ ├── loader.py # TensorLoader 统一加载 + │ ├── tensor_writer.py # TensorWriter 统一写入 + │ ├── lazy_tensor_loader.py # 懒加载索引 + │ ├── lazy_unpickle.py # .bin 懒加载 + │ ├── lazy_ckpt.py # ★ .ckpt 懒加载 + │ └── _device.py # ★ 设备移动 + ├── tokenizer/ # Tokenizer 合并 + ├── moe/ # MoE 合并与构建 + ├── evo/ # 进化搜索（CMA-ES） + ├── tokensurgeon/ # Token 手术 + ├── eval/ # ★ 评测后端（lm-eval-harness） + ├── scripts/ # CLI 脚本入口 + └── _data/ # 架构 JSON 定义 / Chat 模板 +``` + +> ★ 标记的文件为 Wizard 原创（Apache 2.0），其余移植自 MergeKit（LGPL-3.0）。 + +## 🚀 功能特点 + +- **MergeKit 配方兼容**：尽量保持与 MergeKit 一致的 YAML 配方语义，主流同架构合并配方可直接复用 +- **多格式 I/O**：支持 safetensors / .bin / .ckpt 读取，safetensors / .ckpt 写入 +- **统一懒加载**：全格式按需读取单个张量，降低内存占用 +- **NPU 原生支持**：适配 MindSpore / Ascend NPU 推理与评测 +- **dtype 安全**：bf16/fp16 在 CPU 上的精度保护，避免溢出与失真 +- **20 种合并算法**：linear, slerp, ties, dareties, darelinear, modelstock, nearswap, karcher 等 + +## 📋 支持格式 + + +| 方向 | safetensors | .bin (PyTorch) | .ckpt (MindSpore) | +| --- | ----------- | -------------- | ----------------- | +| 读取 | 单文件 + 分片 | 单文件 + 分片 | 单文件 + 分片 | +| 写入 | 单文件 + 分片 | — | 单文件 + 分片 | + + +## 📊 架构设计 + +下面结合 `src/mindnlp/wizard/merge/` 的真实实现，复盘 Wizard 合并主链路的主要模块与调用关系。 + +### 总体架构 + +主调用链可以概括为： + +`wizard-merge` +→ `scripts/run_yaml.py` +→ `merge.py: run_merge` +→ `preflight.py` +→ `architecture/` +→ `LoaderCache.setup() + loader warmup` +→ `MergePlanner.plan_to_disk()` +→ `TaskUniverse / build_schedule` +→ `Executor` / `MultiDeviceExecutor` +→ 写出权重、配置、tokenizer 与模型卡。 + +```mermaid +flowchart TD + A["CLI / YAML 配方"] + B["入口层: scripts/run_yaml.py / wizard-merge"] + C["主编排: merge.py / preflight.py / LoaderCache"] + D["配置与架构: config.py / options.py / architecture/"] + E["规划层: plan.py / MergePlanner"] + F["任务定义: io/tasks.py / merge_methods/ / tokenizer/"] + G["调度执行: graph.py / multigpu_executor.py"] + H["输出产物: 权重 / config / tokenizer / README / wizard_config.yml"] + + A --> B --> C --> D --> E --> F --> G --> H + + style G fill:#e3f2fd,stroke:#1565c0,stroke-width:2px + style H fill:#fff3e0,stroke:#e65100,stroke-width:2px +``` + + + +### 分层说明 + + +| 层级 | 关键文件 | 真实职责与调用关系 | +| ----------- | ------------------------------------------------------------------------------------------------------- | -------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | +| **CLI 入口** | `scripts/run_yaml.py` | 定义 `wizard-merge` 命令，读取 YAML 文件，构造 `MergeConfiguration` 与 `MergeOptions`，然后调用 `run_merge()`。 | +| **主编排** | `merge.py` `preflight.py` | `run_merge()` 是总入口：先做合并前校验，再识别架构、初始化 `LoaderCache`、预热各模型 loader、调用 `MergePlanner.plan_to_disk()` 生成根任务列表，最后交给执行器运行。 | +| **配置与架构** | `config.py` `options.py` `architecture/` | `config.py` 定义 YAML 对应的数据模型与 `ConfigReader`；`options.py` 负责设备、输出格式、懒加载、多 NPU 等运行期选项；`architecture/` 根据模型配置识别架构并枚举 `WeightInfo`。 | +| **规划层** | `plan.py` | `MergePlanner` 按模块、层、张量展开配方，调用 `merge_methods.make_task()` 组织依赖链，并返回 `SaveTensor`、`FinalizeModel`、可选 `BuildTokenizer` 等根任务。规划层负责“构造任务”，不负责拓扑调度。 | +| **任务定义层** | `io/tasks.py` `merge_methods/` `tokenizer/` | I/O 相关任务定义在 `io/tasks.py`；具体算法 Task 定义在各 `merge_methods/*.py`；tokenizer 相关任务定义在 `tokenizer/`。这些任务共同组成执行器要展开和执行的依赖图。 | +| **调度执行层** | `graph.py` `multigpu_executor.py` | `graph.py` 提供 `Task`、`TaskUniverse`、`build_schedule`、`Executor` 与中间结果回收；开启 `multi_npu` 时，主执行器切换为 `multigpu_executor.py` 中的 `MultiDeviceExecutor`。 | +| **精度与数值支撑** | `dtype_policy.py` `safe_ops.py` | `dtype_policy.py` 负责 bf16/fp16 的工作 dtype 选择与精度策略；`safe_ops.py` 为 CPU half 相关敏感算子提供安全实现。这两者不仅服务算法层，也被预检和部分 I/O/稀疏化路径使用。 | +| **I/O 后端** | `io/loader.py` `io/lazy_tensor_loader.py` `io/lazy_unpickle.py` `io/lazy_ckpt.py` `io/tensor_writer.py` | `io/loader.py` 是统一读取入口；`io/lazy_tensor_loader.py` 负责分片索引与张量路由；`io/lazy_unpickle.py` 处理 `.bin` 懒加载；`io/lazy_ckpt.py` 处理 `.ckpt` 懒加载；`io/tensor_writer.py` 负责 `safetensors` / `ckpt` 写出。 | +| **运行时与输出** | MindSpore / Ascend NPU | 执行完成后，除模型权重外，`run_merge()` 还会保存 `config`、可选 tokenizer、模型卡 `README.md`、`wizard_config.yml`，并在需要时补齐附属文件。 | + + +> `eval/mindspore_lm.py` 提供 lm-eval 的 MindSpore 后端，属于评测集成能力，不在合并主链路内。 +> `moe/`、`evo/`、`tokensurgeon/` 也更适合作为扩展能力理解，而不是这条主合并路径的一部分。 + +### 核心组件 + + +| 组件 | 文件 | 职责 | +| ------------------ | ---------------------- | -------------------------- | +| MergeConfiguration | `config.py` | YAML 配方解析，Pydantic 校验 | +| Architecture | `architecture/` | 自动识别模型架构（50+ 种） | +| MergePlanner | `plan.py` | 将配方展开为根任务及依赖链 | +| Executor | `graph.py` | 展开任务依赖图、拓扑排序、按序执行、内存回收 | +| MergeMethod | `merge_methods/` | 20 种合并算法的统一抽象，并生成具体算法 Task | +| TensorLoader | `io/loader.py` | 多格式统一读取入口，由懒加载路由层调用 | +| TensorWriter | `io/tensor_writer.py` | 多格式统一写出（自动分片） | +| dtype_policy ★ | `dtype_policy.py` | bf16/fp16 精度保护策略 | +| safe_ops ★ | `safe_ops.py` | CPU half 精度安全运算 | +| MindSporeLM ★ | `eval/mindspore_lm.py` | lm-eval MindSpore 推理后端 | + + +## 🛠️ 快速开始 + +### 1. 安装依赖 + +```bash +pip install -r requirements/requirements.txt +pip install -r requirements/wizard-requirements.txt +``` + +### 2. 编写合并配方 + +创建 `recipe.yaml`： + +```yaml +models: + - model: Qwen/Qwen2.5-7B + parameters: + weight: 1.0 + - model: Qwen/Qwen2.5-7B-Instruct + parameters: + weight: 0.5 + density: 0.5 +merge_method: ties +base_model: Qwen/Qwen2.5-7B +parameters: + normalize: true +dtype: bfloat16 +``` + +### 3. 执行合并 + +```bash +# 方式 1: 模块调用 +python -m mindnlp.wizard.merge.scripts.run_yaml recipe.yaml ./merged_model + +# 方式 2: 包 CLI（安装后可用） +wizard-merge recipe.yaml ./merged_model + +# 指定输出格式为 ckpt +python -m mindnlp.wizard.merge.scripts.run_yaml recipe.yaml ./merged_model \ + --output-format ckpt +``` + +### 4. 完整示例 + +`examples/wizard/llama3_biomed_dare_ties/` 提供了端到端的 DARE-TIES 合并复现案例 +（Llama3-8B 医学合并），包含模型下载、合并执行、评测全流程。 + +合并后模型已上传至 +[chenjingshen/Llama3-8B-merge-biomed-wizard](https://huggingface.co/chenjingshen/Llama3-8B-merge-biomed-wizard)。 + +更多配方模板见 `examples/wizard/`。 + +### 5. 合并产物 + +``` +merged_model/ +├── model-00001-of-00003.safetensors +├── model-00002-of-00003.safetensors +├── model-00003-of-00003.safetensors +├── model.safetensors.index.json +├── config.json +├── tokenizer.json +└── wizard_config.yml +``` + +## 🧪 测试 + +```bash +cd mindnlp + +# 运行所有 Wizard 测试 +pytest tests/mindnlp/wizard/ -v + +# 运行单个测试 +pytest tests/mindnlp/wizard/test_ckpt_io.py -v + +# 查看测试覆盖率 +pytest tests/mindnlp/wizard/ --cov=mindnlp.wizard --cov-report=html +``` + +## 📄 许可证 + +Wizard 的大部分代码移植自 [MergeKit](https://github.com/arcee-ai/mergekit)，遵循 **LGPL-3.0**。 + +Wizard 原创文件（`dtype_policy.py`、`safe_ops.py`、`preflight.py`、`io/_device.py`、`io/lazy_ckpt.py`、`eval/`）遵循 **Apache License 2.0**，与 MindNLP 项目许可一致。 + +## 📖 引用 + +如果使用了 Wizard，请引用上游 MergeKit 论文： + +```bibtex +@inproceedings{goddard-etal-2024-arcees, + title = "Arcee{'}s {M}erge{K}it: A Toolkit for Merging Large Language Models", + author = "Goddard, Charles and + Siriwardhana, Shamane and + Ehghaghi, Malikeh and + Meyers, Luke and + Karpukhin, Vladimir and + Benedict, Brian and + McQuade, Mark and + Solawetz, Jacob", + booktitle = "Proceedings of the 2024 Conference on Empirical Methods in Natural Language Processing: Industry Track", + month = nov, + year = "2024", + publisher = "Association for Computational Linguistics", + url = "https://aclanthology.org/2024.emnlp-industry.36", + doi = "10.18653/v1/2024.emnlp-industry.36", + pages = "477--485", +} +``` + +## 🤝 致谢 + +- **MergeKit**（Arcee AI）— 核心架构：Task DAG、YAML 配方系统、20 种合并算法、MoE、Evo、TokenSurgeon。LGPL-3.0 许可。 +- **PaddleNLP**（Baidu）— 工程实践参考：大张量分块策略与 dtype 安全模式。未直接复制代码，仅参考设计思路。 + +## 🔗 相关链接 + +- [MindNLP 主仓库](https://github.com/mindspore-lab/mindnlp) +- [MergeKit 原始仓库](https://github.com/arcee-ai/mergekit) +- [Wizard 合并模型示例 (HuggingFace)](https://huggingface.co/chenjingshen/Llama3-8B-merge-biomed-wizard) + diff --git a/src/mindnlp/wizard/__init__.py b/src/mindnlp/wizard/__init__.py index e69de29bb..325205093 100644 --- a/src/mindnlp/wizard/__init__.py +++ b/src/mindnlp/wizard/__init__.py @@ -0,0 +1,9 @@ +""" +MindNLP Wizard 模块 + +基于 MergeKit 的模型合并引擎，适配 MindSpore / Ascend NPU， +支持同架构大模型权重合并（SLERP、TIES、DARE 等 20 种方法）。 +""" + +__version__ = "0.1.0" +__author__ = "MindNLP Wizard contributors" diff --git a/src/mindnlp/wizard/merge/LICENSE_LGPL-3.0 b/src/mindnlp/wizard/merge/LICENSE_LGPL-3.0 new file mode 100644 index 000000000..0a041280b --- /dev/null +++ b/src/mindnlp/wizard/merge/LICENSE_LGPL-3.0 @@ -0,0 +1,165 @@ + GNU LESSER GENERAL PUBLIC LICENSE + Version 3, 29 June 2007 + + Copyright (C) 2007 Free Software Foundation, Inc. + Everyone is permitted to copy and distribute verbatim copies + of this license document, but changing it is not allowed. + + + This version of the GNU Lesser General Public License incorporates +the terms and conditions of version 3 of the GNU General Public +License, supplemented by the additional permissions listed below. + + 0. Additional Definitions. + + As used herein, "this License" refers to version 3 of the GNU Lesser +General Public License, and the "GNU GPL" refers to version 3 of the GNU +General Public License. + + "The Library" refers to a covered work governed by this License, +other than an Application or a Combined Work as defined below. + + An "Application" is any work that makes use of an interface provided +by the Library, but which is not otherwise based on the Library. +Defining a subclass of a class defined by the Library is deemed a mode +of using an interface provided by the Library. + + A "Combined Work" is a work produced by combining or linking an +Application with the Library. The particular version of the Library +with which the Combined Work was made is also called the "Linked +Version". + + The "Minimal Corresponding Source" for a Combined Work means the +Corresponding Source for the Combined Work, excluding any source code +for portions of the Combined Work that, considered in isolation, are +based on the Application, and not on the Linked Version. + + The "Corresponding Application Code" for a Combined Work means the +object code and/or source code for the Application, including any data +and utility programs needed for reproducing the Combined Work from the +Application, but excluding the System Libraries of the Combined Work. + + 1. Exception to Section 3 of the GNU GPL. + + You may convey a covered work under sections 3 and 4 of this License +without being bound by section 3 of the GNU GPL. + + 2. Conveying Modified Versions. + + If you modify a copy of the Library, and, in your modifications, a +facility refers to a function or data to be supplied by an Application +that uses the facility (other than as an argument passed when the +facility is invoked), then you may convey a copy of the modified +version: + + a) under this License, provided that you make a good faith effort to + ensure that, in the event an Application does not supply the + function or data, the facility still operates, and performs + whatever part of its purpose remains meaningful, or + + b) under the GNU GPL, with none of the additional permissions of + this License applicable to that copy. + + 3. Object Code Incorporating Material from Library Header Files. + + The object code form of an Application may incorporate material from +a header file that is part of the Library. You may convey such object +code under terms of your choice, provided that, if the incorporated +material is not limited to numerical parameters, data structure +layouts and accessors, or small macros, inline functions and templates +(ten or fewer lines in length), you do both of the following: + + a) Give prominent notice with each copy of the object code that the + Library is used in it and that the Library and its use are + covered by this License. + + b) Accompany the object code with a copy of the GNU GPL and this license + document. + + 4. Combined Works. + + You may convey a Combined Work under terms of your choice that, +taken together, effectively do not restrict modification of the +portions of the Library contained in the Combined Work and reverse +engineering for debugging such modifications, if you also do each of +the following: + + a) Give prominent notice with each copy of the Combined Work that + the Library is used in it and that the Library and its use are + covered by this License. + + b) Accompany the Combined Work with a copy of the GNU GPL and this license + document. + + c) For a Combined Work that displays copyright notices during + execution, include the copyright notice for the Library among + these notices, as well as a reference directing the user to the + copies of the GNU GPL and this license document. + + d) Do one of the following: + + 0) Convey the Minimal Corresponding Source under the terms of this + License, and the Corresponding Application Code in a form + suitable for, and under terms that permit, the user to + recombine or relink the Application with a modified version of + the Linked Version to produce a modified Combined Work, in the + manner specified by section 6 of the GNU GPL for conveying + Corresponding Source. + + 1) Use a suitable shared library mechanism for linking with the + Library. A suitable mechanism is one that (a) uses at run time + a copy of the Library already present on the user's computer + system, and (b) will operate properly with a modified version + of the Library that is interface-compatible with the Linked + Version. + + e) Provide Installation Information, but only if you would otherwise + be required to provide such information under section 6 of the + GNU GPL, and only to the extent that such information is + necessary to install and execute a modified version of the + Combined Work produced by recombining or relinking the + Application with a modified version of the Linked Version. (If + you use option 4d0, the Installation Information must accompany + the Minimal Corresponding Source and Corresponding Application + Code. If you use option 4d1, you must provide the Installation + Information in the manner specified by section 6 of the GNU GPL + for conveying Corresponding Source.) + + 5. Combined Libraries. + + You may place library facilities that are a work based on the +Library side by side in a single library together with other library +facilities that are not Applications and are not covered by this +License, and convey such a combined library under terms of your +choice, if you do both of the following: + + a) Accompany the combined library with a copy of the same work based + on the Library, uncombined with any other library facilities, + conveyed under the terms of this License. + + b) Give prominent notice with the combined library that part of it + is a work based on the Library, and explaining where to find the + accompanying uncombined form of the same work. + + 6. Revised Versions of the GNU Lesser General Public License. + + The Free Software Foundation may publish revised and/or new versions +of the GNU Lesser General Public License from time to time. Such new +versions will be similar in spirit to the present version, but may +differ in detail to address new problems or concerns. + + Each version is given a distinguishing version number. If the +Library as you received it specifies that a certain numbered version +of the GNU Lesser General Public License "or any later version" +applies to it, you have the option of following the terms and +conditions either of that published version or of any later version +published by the Free Software Foundation. If the Library as you +received it does not specify a version number of the GNU Lesser +General Public License, you may choose any version of the GNU Lesser +General Public License ever published by the Free Software Foundation. + + If the Library as you received it specifies that a proxy can decide +whether future versions of the GNU Lesser General Public License shall +apply, that proxy's public statement of acceptance of any version is +permanent authorization for you to choose that version for the +Library. diff --git a/src/mindnlp/wizard/merge/__init__.py b/src/mindnlp/wizard/merge/__init__.py index e69de29bb..e3a6eee07 100644 --- a/src/mindnlp/wizard/merge/__init__.py +++ b/src/mindnlp/wizard/merge/__init__.py @@ -0,0 +1,27 @@ +# Originally from MergeKit (https://github.com/arcee-ai/mergekit) +# Copyright (C) 2025 Arcee AI +# SPDX-License-Identifier: LGPL-3.0-only +# Modified for MindSpore/Ascend NPU by MindNLP Wizard contributors. + +"""Public exports for wizard merge package.""" + +from .config import MergeConfiguration, ConfigReader +from .common import ModelReference, ModelPath, dtype_from_name, ImmutableMap +from .options import MergeOptions +from .graph import Task, Executor, TaskUniverse, TaskHandle +from .merge import run_merge + +__all__ = [ + "MergeConfiguration", + "ConfigReader", + "ModelReference", + "ModelPath", + "MergeOptions", + "Task", + "Executor", + "TaskUniverse", + "TaskHandle", + "run_merge", + "dtype_from_name", + "ImmutableMap", +] diff --git a/src/mindnlp/wizard/merge/_data/__init__.py b/src/mindnlp/wizard/merge/_data/__init__.py new file mode 100644 index 000000000..6babb0e71 --- /dev/null +++ b/src/mindnlp/wizard/merge/_data/__init__.py @@ -0,0 +1,4 @@ +# Originally from MergeKit (https://github.com/arcee-ai/mergekit) +# Copyright (C) 2025 Arcee AI +# SPDX-License-Identifier: LGPL-3.0-only +# Modified for MindSpore/Ascend NPU by MindNLP Wizard contributors. diff --git a/src/mindnlp/wizard/merge/_data/architectures/__init__.py b/src/mindnlp/wizard/merge/_data/architectures/__init__.py new file mode 100644 index 000000000..6babb0e71 --- /dev/null +++ b/src/mindnlp/wizard/merge/_data/architectures/__init__.py @@ -0,0 +1,4 @@ +# Originally from MergeKit (https://github.com/arcee-ai/mergekit) +# Copyright (C) 2025 Arcee AI +# SPDX-License-Identifier: LGPL-3.0-only +# Modified for MindSpore/Ascend NPU by MindNLP Wizard contributors. diff --git a/src/mindnlp/wizard/merge/_data/architectures/afmoe_partial.json b/src/mindnlp/wizard/merge/_data/architectures/afmoe_partial.json new file mode 100644 index 000000000..e8a418b5d --- /dev/null +++ b/src/mindnlp/wizard/merge/_data/architectures/afmoe_partial.json @@ -0,0 +1,95 @@ +{ + "model_type": "_afmoe_partial", + "architectures": [ + "_AfmoePartialForCausalLM" + ], + "pre_weights": [ + { + "name": "model.embed_tokens.weight", + "is_embed": true + } + ], + "post_weights": [ + { + "name": "model.norm.weight" + }, + { + "name": "lm_head.weight", + "is_embed": true, + "optional": true, + "tied_names": [ + "model.embed_tokens.weight" + ] + } + ], + "num_layers_config_key": "num_hidden_layers", + "layer_templates": { + "weights": [ + { + "name": "model.layers.${layer_index}.input_layernorm.weight" + }, + { + "name": "model.layers.${layer_index}.pre_mlp_layernorm.weight" + }, + { + "name": "model.layers.${layer_index}.mlp.down_proj.weight", + "optional": true + }, + { + "name": "model.layers.${layer_index}.mlp.gate_proj.weight", + "optional": true + }, + { + "name": "model.layers.${layer_index}.mlp.up_proj.weight", + "optional": true + }, + { + "name": "model.layers.${layer_index}.mlp.shared_experts.down_proj.weight", + "optional": true + }, + { + "name": "model.layers.${layer_index}.mlp.shared_experts.gate_proj.weight", + "optional": true + }, + { + "name": "model.layers.${layer_index}.mlp.shared_experts.up_proj.weight", + "optional": true + }, + { + "name": "model.layers.${layer_index}.mlp.expert_bias", + "optional": true + }, + { + "name": "model.layers.${layer_index}.mlp.router.gate.weight", + "optional": true + }, + { + "name": "model.layers.${layer_index}.post_mlp_layernorm.weight" + }, + { + "name": "model.layers.${layer_index}.self_attn.q_norm.weight" + }, + { + "name": "model.layers.${layer_index}.self_attn.q_proj.weight" + }, + { + "name": "model.layers.${layer_index}.self_attn.k_norm.weight" + }, + { + "name": "model.layers.${layer_index}.self_attn.k_proj.weight" + }, + { + "name": "model.layers.${layer_index}.self_attn.o_proj.weight" + }, + { + "name": "model.layers.${layer_index}.self_attn.v_proj.weight" + }, + { + "name": "model.layers.${layer_index}.self_attn.gate_proj.weight" + }, + { + "name": "model.layers.${layer_index}.post_attention_layernorm.weight" + } + ] + } +} diff --git a/src/mindnlp/wizard/merge/_data/architectures/apertus.json b/src/mindnlp/wizard/merge/_data/architectures/apertus.json new file mode 100644 index 000000000..4b6d90aef --- /dev/null +++ b/src/mindnlp/wizard/merge/_data/architectures/apertus.json @@ -0,0 +1,72 @@ +{ + "model_type": "apertus", + "architectures": [ + "ApertusForCausalLM" + ], + "pre_weights": [ + { + "name": "model.embed_tokens.weight", + "is_embed": true + } + ], + "post_weights": [ + { + "name": "model.norm.weight" + }, + { + "name": "lm_head.weight", + "is_embed": true, + "optional": true, + "tied_names": [ + "model.embed_tokens.weight" + ] + } + ], + "num_layers_config_key": "num_hidden_layers", + "layer_templates": { + "weights": [ + { + "name": "model.layers.${layer_index}.attention_layernorm.weight" + }, + { + "name": "model.layers.${layer_index}.feedforward_layernorm.weight" + }, + { + "name": "model.layers.${layer_index}.mlp.act_fn.alpha_n" + }, + { + "name": "model.layers.${layer_index}.mlp.act_fn.alpha_p" + }, + { + "name": "model.layers.${layer_index}.mlp.act_fn.beta" + }, + { + "name": "model.layers.${layer_index}.mlp.act_fn.eps" + }, + { + "name": "model.layers.${layer_index}.mlp.down_proj.weight" + }, + { + "name": "model.layers.${layer_index}.mlp.up_proj.weight" + }, + { + "name": "model.layers.${layer_index}.self_attn.k_norm.weight" + }, + { + "name": "model.layers.${layer_index}.self_attn.k_proj.weight" + }, + { + "name": "model.layers.${layer_index}.self_attn.q_norm.weight" + }, + { + "name": "model.layers.${layer_index}.self_attn.q_proj.weight" + }, + { + "name": "model.layers.${layer_index}.self_attn.v_proj.weight" + }, + { + "name": "model.layers.${layer_index}.self_attn.o_proj.weight" + } + ] + } +} diff --git a/src/mindnlp/wizard/merge/_data/architectures/arcee.json b/src/mindnlp/wizard/merge/_data/architectures/arcee.json new file mode 100644 index 000000000..2a1088a35 --- /dev/null +++ b/src/mindnlp/wizard/merge/_data/architectures/arcee.json @@ -0,0 +1,54 @@ +{ + "model_type": "arcee", + "architectures": [ + "ArceeForCausalLM" + ], + "pre_weights": [ + { + "name": "model.embed_tokens.weight", + "is_embed": true + } + ], + "post_weights": [ + { + "name": "model.norm.weight" + }, + { + "name": "lm_head.weight", + "is_embed": true, + "optional": true, + "tied_names": [ + "model.embed_tokens.weight" + ] + } + ], + "num_layers_config_key": "num_hidden_layers", + "layer_templates": { + "weights": [ + { + "name": "model.layers.${layer_index}.input_layernorm.weight" + }, + { + "name": "model.layers.${layer_index}.mlp.down_proj.weight" + }, + { + "name": "model.layers.${layer_index}.mlp.up_proj.weight" + }, + { + "name": "model.layers.${layer_index}.post_attention_layernorm.weight" + }, + { + "name": "model.layers.${layer_index}.self_attn.k_proj.weight" + }, + { + "name": "model.layers.${layer_index}.self_attn.o_proj.weight" + }, + { + "name": "model.layers.${layer_index}.self_attn.q_proj.weight" + }, + { + "name": "model.layers.${layer_index}.self_attn.v_proj.weight" + } + ] + } +} diff --git a/src/mindnlp/wizard/merge/_data/architectures/baichuan.json b/src/mindnlp/wizard/merge/_data/architectures/baichuan.json new file mode 100644 index 000000000..3d28020c6 --- /dev/null +++ b/src/mindnlp/wizard/merge/_data/architectures/baichuan.json @@ -0,0 +1,47 @@ +{ + "model_type": "baichuan", + "architectures": [ + "BaichuanForCausalLM" + ], + "pre_weights": [ + { + "name": "model.embed_tokens.weight", + "is_embed": true + } + ], + "post_weights": [ + { + "name": "model.norm.weight" + }, + { + "name": "lm_head.weight", + "is_embed": true + } + ], + "num_layers_config_key": "num_hidden_layers", + "layer_templates": { + "weights": [ + { + "name": "model.layers.${layer_index}.input_layernorm.weight" + }, + { + "name": "model.layers.${layer_index}.self_attn.W_pack.weight" + }, + { + "name": "model.layers.${layer_index}.self_attn.o_proj.weight" + }, + { + "name": "model.layers.${layer_index}.post_attention_layernorm.weight" + }, + { + "name": "model.layers.${layer_index}.mlp.gate_proj.weight" + }, + { + "name": "model.layers.${layer_index}.mlp.down_proj.weight" + }, + { + "name": "model.layers.${layer_index}.mlp.up_proj.weight" + } + ] + } +} diff --git a/src/mindnlp/wizard/merge/_data/architectures/bert-masked-lm.json b/src/mindnlp/wizard/merge/_data/architectures/bert-masked-lm.json new file mode 100644 index 000000000..d6430e402 --- /dev/null +++ b/src/mindnlp/wizard/merge/_data/architectures/bert-masked-lm.json @@ -0,0 +1,119 @@ +{ + "model_type": "bert", + "architectures": [ + "BertForMaskedLM" + ], + "pre_weights": [ + { + "name": "bert.embeddings.position_embeddings.weight" + }, + { + "name": "bert.embeddings.token_type_embeddings.weight" + }, + { + "name": "bert.embeddings.word_embeddings.weight", + "is_embed": true + }, + { + "name": "bert.embeddings.LayerNorm.bias", + "aliases": [ + "bert.embeddings.LayerNorm.beta" + ] + }, + { + "name": "bert.embeddings.LayerNorm.weight", + "aliases": [ + "bert.embeddings.LayerNorm.gamma" + ] + }, + { + "name": "bert.embeddings.position_ids", + "optional": true, + "force_dtype": "int64" + } + ], + "post_weights": [ + { + "name": "bert.pooler.dense.weight" + }, + { + "name": "bert.pooler.dense.bias" + }, + { + "name": "cls.predictions.bias" + }, + { + "name": "cls.predictions.decoder.weight", + "optional": true, + "tied_names": [ + "bert.embeddings.word_embeddings.weight" + ], + "is_embed": true + } + ], + "num_layers_config_key": "num_hidden_layers", + "layer_templates": { + "weights": [ + { + "name": "bert.encoder.layer.${layer_index}.attention.self.query.weight" + }, + { + "name": "bert.encoder.layer.${layer_index}.attention.self.query.bias" + }, + { + "name": "bert.encoder.layer.${layer_index}.attention.self.key.weight" + }, + { + "name": "bert.encoder.layer.${layer_index}.attention.self.key.bias" + }, + { + "name": "bert.encoder.layer.${layer_index}.attention.self.value.weight" + }, + { + "name": "bert.encoder.layer.${layer_index}.attention.self.value.bias" + }, + { + "name": "bert.encoder.layer.${layer_index}.attention.output.dense.weight" + }, + { + "name": "bert.encoder.layer.${layer_index}.attention.output.dense.bias" + }, + { + "name": "bert.encoder.layer.${layer_index}.attention.output.LayerNorm.bias", + "aliases": [ + "bert.encoder.layer.${layer_index}.attention.output.LayerNorm.beta" + ] + }, + { + "name": "bert.encoder.layer.${layer_index}.attention.output.LayerNorm.weight", + "aliases": [ + "bert.encoder.layer.${layer_index}.attention.output.LayerNorm.gamma" + ] + }, + { + "name": "bert.encoder.layer.${layer_index}.intermediate.dense.weight" + }, + { + "name": "bert.encoder.layer.${layer_index}.intermediate.dense.bias" + }, + { + "name": "bert.encoder.layer.${layer_index}.output.dense.weight" + }, + { + "name": "bert.encoder.layer.${layer_index}.output.dense.bias" + }, + { + "name": "bert.encoder.layer.${layer_index}.output.LayerNorm.bias", + "aliases": [ + "bert.encoder.layer.${layer_index}.output.LayerNorm.beta" + ] + }, + { + "name": "bert.encoder.layer.${layer_index}.output.LayerNorm.weight", + "aliases": [ + "bert.encoder.layer.${layer_index}.output.LayerNorm.gamma" + ] + } + ] + } +} diff --git a/src/mindnlp/wizard/merge/_data/architectures/bert-sequence-classification.json b/src/mindnlp/wizard/merge/_data/architectures/bert-sequence-classification.json new file mode 100644 index 000000000..81a61ff7d --- /dev/null +++ b/src/mindnlp/wizard/merge/_data/architectures/bert-sequence-classification.json @@ -0,0 +1,118 @@ +{ + "model_type": "bert", + "architectures": [ + "BertForSequenceClassification", + "BertForMultipleChoice", + "BertForTokenClassification" + ], + "pre_weights": [ + { + "name": "bert.embeddings.position_embeddings.weight" + }, + { + "name": "bert.embeddings.token_type_embeddings.weight" + }, + { + "name": "bert.embeddings.word_embeddings.weight", + "is_embed": true + }, + { + "name": "bert.embeddings.LayerNorm.bias", + "aliases": [ + "bert.embeddings.LayerNorm.beta" + ] + }, + { + "name": "bert.embeddings.LayerNorm.weight", + "aliases": [ + "bert.embeddings.LayerNorm.gamma" + ] + }, + { + "name": "bert.embeddings.position_ids", + "optional": true, + "force_dtype": "int64" + } + ], + "post_weights": [ + { + "name": "bert.pooler.dense.weight", + "optional": true + }, + { + "name": "bert.pooler.dense.bias", + "optional": true + }, + { + "name": "classifier.bias" + }, + { + "name": "classifier.weight" + } + ], + "num_layers_config_key": "num_hidden_layers", + "layer_templates": { + "weights": [ + { + "name": "bert.encoder.layer.${layer_index}.attention.self.query.weight" + }, + { + "name": "bert.encoder.layer.${layer_index}.attention.self.query.bias" + }, + { + "name": "bert.encoder.layer.${layer_index}.attention.self.key.weight" + }, + { + "name": "bert.encoder.layer.${layer_index}.attention.self.key.bias" + }, + { + "name": "bert.encoder.layer.${layer_index}.attention.self.value.weight" + }, + { + "name": "bert.encoder.layer.${layer_index}.attention.self.value.bias" + }, + { + "name": "bert.encoder.layer.${layer_index}.attention.output.dense.weight" + }, + { + "name": "bert.encoder.layer.${layer_index}.attention.output.dense.bias" + }, + { + "name": "bert.encoder.layer.${layer_index}.attention.output.LayerNorm.bias", + "aliases": [ + "bert.encoder.layer.${layer_index}.attention.output.LayerNorm.beta" + ] + }, + { + "name": "bert.encoder.layer.${layer_index}.attention.output.LayerNorm.weight", + "aliases": [ + "bert.encoder.layer.${layer_index}.attention.output.LayerNorm.gamma" + ] + }, + { + "name": "bert.encoder.layer.${layer_index}.intermediate.dense.weight" + }, + { + "name": "bert.encoder.layer.${layer_index}.intermediate.dense.bias" + }, + { + "name": "bert.encoder.layer.${layer_index}.output.dense.weight" + }, + { + "name": "bert.encoder.layer.${layer_index}.output.dense.bias" + }, + { + "name": "bert.encoder.layer.${layer_index}.output.LayerNorm.bias", + "aliases": [ + "bert.encoder.layer.${layer_index}.output.LayerNorm.beta" + ] + }, + { + "name": "bert.encoder.layer.${layer_index}.output.LayerNorm.weight", + "aliases": [ + "bert.encoder.layer.${layer_index}.output.LayerNorm.gamma" + ] + } + ] + } +} diff --git a/src/mindnlp/wizard/merge/_data/architectures/bert.json b/src/mindnlp/wizard/merge/_data/architectures/bert.json new file mode 100644 index 000000000..5de5f5b02 --- /dev/null +++ b/src/mindnlp/wizard/merge/_data/architectures/bert.json @@ -0,0 +1,175 @@ +{ + "model_type": "bert", + "architectures": [ + "BertModel" + ], + "pre_weights": [ + { + "name": "embeddings.position_embeddings.weight", + "aliases": [ + "bert.embeddings.position_embeddings.weight" + ] + }, + { + "name": "embeddings.token_type_embeddings.weight", + "aliases": [ + "bert.embeddings.token_type_embeddings.weight" + ] + }, + { + "name": "embeddings.word_embeddings.weight", + "is_embed": true, + "aliases": [ + "bert.embeddings.word_embeddings.weight" + ] + }, + { + "name": "embeddings.LayerNorm.bias", + "aliases": [ + "embeddings.LayerNorm.beta", + "bert.embeddings.LayerNorm.bias", + "bert.embeddings.LayerNorm.beta" + ] + }, + { + "name": "embeddings.LayerNorm.weight", + "aliases": [ + "embeddings.LayerNorm.gamma", + "bert.embeddings.LayerNorm.weight", + "bert.embeddings.LayerNorm.gamma", + "bert.embeddings.LayerNorm.weight" + ] + }, + { + "name": "embeddings.position_ids", + "optional": true, + "force_dtype": "int64", + "aliases": [ + "bert.embeddings.position_ids" + ] + } + ], + "post_weights": [ + { + "name": "pooler.dense.weight", + "aliases": [ + "bert.pooler.dense.weight" + ] + }, + { + "name": "pooler.dense.bias", + "aliases": [ + "bert.pooler.dense.bias" + ] + } + ], + "num_layers_config_key": "num_hidden_layers", + "layer_templates": { + "weights": [ + { + "name": "encoder.layer.${layer_index}.attention.self.query.weight", + "aliases": [ + "bert.encoder.layer.${layer_index}.attention.self.query.weight" + ] + }, + { + "name": "encoder.layer.${layer_index}.attention.self.query.bias", + "aliases": [ + "bert.encoder.layer.${layer_index}.attention.self.query.bias" + ] + }, + { + "name": "encoder.layer.${layer_index}.attention.self.key.weight", + "aliases": [ + "bert.encoder.layer.${layer_index}.attention.self.key.weight" + ] + }, + { + "name": "encoder.layer.${layer_index}.attention.self.key.bias", + "aliases": [ + "bert.encoder.layer.${layer_index}.attention.self.key.bias" + ] + }, + { + "name": "encoder.layer.${layer_index}.attention.self.value.weight", + "aliases": [ + "bert.encoder.layer.${layer_index}.attention.self.value.weight" + ] + }, + { + "name": "encoder.layer.${layer_index}.attention.self.value.bias", + "aliases": [ + "bert.encoder.layer.${layer_index}.attention.self.value.bias" + ] + }, + { + "name": "encoder.layer.${layer_index}.attention.output.dense.weight", + "aliases": [ + "bert.encoder.layer.${layer_index}.attention.output.dense.weight" + ] + }, + { + "name": "encoder.layer.${layer_index}.attention.output.dense.bias", + "aliases": [ + "bert.encoder.layer.${layer_index}.attention.output.dense.bias" + ] + }, + { + "name": "encoder.layer.${layer_index}.attention.output.LayerNorm.bias", + "aliases": [ + "encoder.layer.${layer_index}.attention.output.LayerNorm.beta", + "bert.encoder.layer.${layer_index}.attention.output.LayerNorm.bias", + "bert.encoder.layer.${layer_index}.attention.output.LayerNorm.beta" + ] + }, + { + "name": "encoder.layer.${layer_index}.attention.output.LayerNorm.weight", + "aliases": [ + "encoder.layer.${layer_index}.attention.output.LayerNorm.gamma", + "bert.encoder.layer.${layer_index}.attention.output.LayerNorm.weight", + "bert.encoder.layer.${layer_index}.attention.output.LayerNorm.gamma" + ] + }, + { + "name": "encoder.layer.${layer_index}.intermediate.dense.weight", + "aliases": [ + "bert.encoder.layer.${layer_index}.intermediate.dense.weight" + ] + }, + { + "name": "encoder.layer.${layer_index}.intermediate.dense.bias", + "aliases": [ + "bert.encoder.layer.${layer_index}.intermediate.dense.bias" + ] + }, + { + "name": "encoder.layer.${layer_index}.output.dense.weight", + "aliases": [ + "bert.encoder.layer.${layer_index}.output.dense.weight" + ] + }, + { + "name": "encoder.layer.${layer_index}.output.dense.bias", + "aliases": [ + "bert.encoder.layer.${layer_index}.output.dense.bias" + ] + }, + { + "name": "encoder.layer.${layer_index}.output.LayerNorm.bias", + "aliases": [ + "encoder.layer.${layer_index}.output.LayerNorm.beta", + "bert.encoder.layer.${layer_index}.output.LayerNorm.bias", + "bert.encoder.layer.${layer_index}.output.LayerNorm.beta" + ] + }, + { + "name": "encoder.layer.${layer_index}.output.LayerNorm.weight", + "aliases": [ + "encoder.layer.${layer_index}.output.LayerNorm.gamma", + "bert.encoder.layer.${layer_index}.output.LayerNorm.weight", + "bert.encoder.layer.${layer_index}.output.LayerNorm.gamma" + ] + } + ] + } +} diff --git a/src/mindnlp/wizard/merge/_data/architectures/chatglm.json b/src/mindnlp/wizard/merge/_data/architectures/chatglm.json new file mode 100644 index 000000000..b094685f8 --- /dev/null +++ b/src/mindnlp/wizard/merge/_data/architectures/chatglm.json @@ -0,0 +1,50 @@ +{ + "model_type": "chatglm", + "architectures": [ + "ChatGLMModel" + ], + "pre_weights": [ + { + "name": "transformer.embedding.word_embeddings.weight", + "is_embed": true + }, + { + "name": "transformer.rotary_pos_emb.inv_freq" + } + ], + "post_weights": [ + { + "name": "transformer.encoder.final_layernorm.weight" + }, + { + "name": "transformer.output_layer.weight", + "is_embed": true + } + ], + "num_layers_config_key": "num_hidden_layers", + "layer_templates": { + "weights": [ + { + "name": "transformer.encoder.layers.${layer_index}.input_layernorm.weight" + }, + { + "name": "transformer.encoder.layers.${layer_index}.mlp.dense_4h_to_h.weight" + }, + { + "name": "transformer.encoder.layers.${layer_index}.mlp.dense_h_to_4h.weight" + }, + { + "name": "transformer.encoder.layers.${layer_index}.post_attention_layernorm.weight" + }, + { + "name": "transformer.encoder.layers.${layer_index}.self_attention.dense.weight" + }, + { + "name": "transformer.encoder.layers.${layer_index}.self_attention.query_key_value.bias" + }, + { + "name": "transformer.encoder.layers.${layer_index}.self_attention.query_key_value.weight" + } + ] + } +} diff --git a/src/mindnlp/wizard/merge/_data/architectures/cohere.json b/src/mindnlp/wizard/merge/_data/architectures/cohere.json new file mode 100644 index 000000000..c80331032 --- /dev/null +++ b/src/mindnlp/wizard/merge/_data/architectures/cohere.json @@ -0,0 +1,51 @@ +{ + "model_type": "cohere", + "architectures": [ + "CohereForCausalLM" + ], + "pre_weights": [ + { + "name": "model.embed_tokens.weight", + "is_embed": true + } + ], + "post_weights": [ + { + "name": "model.norm.weight" + }, + { + "name": "lm_head.weight", + "is_embed": true, + "optional": true + } + ], + "num_layers_config_key": "num_hidden_layers", + "layer_templates": { + "weights": [ + { + "name": "model.layers.${layer_index}.input_layernorm.weight" + }, + { + "name": "model.layers.${layer_index}.mlp.down_proj.weight" + }, + { + "name": "model.layers.${layer_index}.mlp.gate_proj.weight" + }, + { + "name": "model.layers.${layer_index}.mlp.up_proj.weight" + }, + { + "name": "model.layers.${layer_index}.self_attn.q_proj.weight" + }, + { + "name": "model.layers.${layer_index}.self_attn.k_proj.weight" + }, + { + "name": "model.layers.${layer_index}.self_attn.v_proj.weight" + }, + { + "name": "model.layers.${layer_index}.self_attn.o_proj.weight" + } + ] + } +} diff --git a/src/mindnlp/wizard/merge/_data/architectures/cohere2.json b/src/mindnlp/wizard/merge/_data/architectures/cohere2.json new file mode 100644 index 000000000..755e7db24 --- /dev/null +++ b/src/mindnlp/wizard/merge/_data/architectures/cohere2.json @@ -0,0 +1,51 @@ +{ + "model_type": "cohere2", + "architectures": [ + "Cohere2ForCausalLM" + ], + "pre_weights": [ + { + "name": "model.embed_tokens.weight", + "is_embed": true + } + ], + "post_weights": [ + { + "name": "model.norm.weight" + }, + { + "name": "lm_head.weight", + "is_embed": true, + "optional": true + } + ], + "num_layers_config_key": "num_hidden_layers", + "layer_templates": { + "weights": [ + { + "name": "model.layers.${layer_index}.self_attn.q_proj.weight" + }, + { + "name": "model.layers.${layer_index}.self_attn.k_proj.weight" + }, + { + "name": "model.layers.${layer_index}.self_attn.v_proj.weight" + }, + { + "name": "model.layers.${layer_index}.self_attn.o_proj.weight" + }, + { + "name": "model.layers.${layer_index}.mlp.gate_proj.weight" + }, + { + "name": "model.layers.${layer_index}.mlp.up_proj.weight" + }, + { + "name": "model.layers.${layer_index}.mlp.down_proj.weight" + }, + { + "name": "model.layers.${layer_index}.input_layernorm.weight" + } + ] + } +} diff --git a/src/mindnlp/wizard/merge/_data/architectures/distilbert-masked-lm.json b/src/mindnlp/wizard/merge/_data/architectures/distilbert-masked-lm.json new file mode 100644 index 000000000..1a079811c --- /dev/null +++ b/src/mindnlp/wizard/merge/_data/architectures/distilbert-masked-lm.json @@ -0,0 +1,105 @@ +{ + "model_type": "distilbert", + "architectures": [ + "DistilBertForMaskedLM" + ], + "pre_weights": [ + { + "name": "distilbert.embeddings.position_embeddings.weight" + }, + { + "name": "distilbert.embeddings.word_embeddings.weight", + "is_embed": true + }, + { + "name": "distilbert.embeddings.LayerNorm.bias", + "aliases": [ + "distilbert.embeddings.LayerNorm.beta" + ] + }, + { + "name": "distilbert.embeddings.LayerNorm.weight", + "aliases": [ + "distilbert.embeddings.LayerNorm.gamma" + ] + } + ], + "post_weights": [ + { + "name": "vocab_transform.weight" + }, + { + "name": "vocab_transform.bias" + }, + { + "name": "vocab_layer_norm.bias" + }, + { + "name": "vocab_layer_norm.weight" + }, + { + "name": "vocab_projector.weight", + "is_embed": true, + "optional": true, + "tied_names": [ + "distilbert.embeddings.word_embeddings.weight" + ] + }, + { + "name": "vocab_projector.bias" + } + ], + "num_layers_config_key": "num_hidden_layers", + "layer_templates": { + "weights": [ + { + "name": "distilbert.transformer.layer.${layer_index}.attention.k_lin.weight" + }, + { + "name": "distilbert.transformer.layer.${layer_index}.attention.k_lin.bias" + }, + { + "name": "distilbert.transformer.layer.${layer_index}.attention.q_lin.weight" + }, + { + "name": "distilbert.transformer.layer.${layer_index}.attention.q_lin.bias" + }, + { + "name": "distilbert.transformer.layer.${layer_index}.attention.v_lin.weight" + }, + { + "name": "distilbert.transformer.layer.${layer_index}.attention.v_lin.bias" + }, + { + "name": "distilbert.transformer.layer.${layer_index}.attention.out_lin.weight" + }, + { + "name": "distilbert.transformer.layer.${layer_index}.attention.out_lin.bias" + }, + { + "name": "distilbert.transformer.layer.${layer_index}.sa_layer_norm.bias" + }, + { + "name": "distilbert.transformer.layer.${layer_index}.sa_layer_norm.weight" + }, + { + "name": "distilbert.transformer.layer.${layer_index}.ffn.lin1.weight" + }, + { + "name": "distilbert.transformer.layer.${layer_index}.ffn.lin1.bias" + }, + { + "name": "distilbert.transformer.layer.${layer_index}.ffn.lin2.weight" + }, + { + "name": "distilbert.transformer.layer.${layer_index}.ffn.lin2.bias" + }, + { + "name": "distilbert.transformer.layer.${layer_index}.output_layer_norm.bias" + }, + { + "name": "distilbert.transformer.layer.${layer_index}.output_layer_norm.weight" + } + ] + } +} diff --git a/src/mindnlp/wizard/merge/_data/architectures/distilbert-sequence-classification.json b/src/mindnlp/wizard/merge/_data/architectures/distilbert-sequence-classification.json new file mode 100644 index 000000000..00c727639 --- /dev/null +++ b/src/mindnlp/wizard/merge/_data/architectures/distilbert-sequence-classification.json @@ -0,0 +1,94 @@ +{ + "model_type": "distilbert", + "architectures": [ + "DistilBertForSequenceClassification" + ], + "pre_weights": [ + { + "name": "distilbert.embeddings.position_embeddings.weight" + }, + { + "name": "distilbert.embeddings.word_embeddings.weight", + "is_embed": true + }, + { + "name": "distilbert.embeddings.LayerNorm.bias", + "aliases": [ + "distilbert.embeddings.LayerNorm.beta" + ] + }, + { + "name": "distilbert.embeddings.LayerNorm.weight", + "aliases": [ + "distilbert.embeddings.LayerNorm.gamma" + ] + } + ], + "post_weights": [ + { + "name": "classifier.bias" + }, + { + "name": "classifier.weight" + }, + { + "name": "pre_classifier.bias" + }, + { + "name": "pre_classifier.weight" + } + ], + "num_layers_config_key": "num_hidden_layers", + "layer_templates": { + "weights": [ + { + "name": "distilbert.transformer.layer.${layer_index}.attention.k_lin.weight" + }, + { + "name": "distilbert.transformer.layer.${layer_index}.attention.k_lin.bias" + }, + { + "name": "distilbert.transformer.layer.${layer_index}.attention.q_lin.weight" + }, + { + "name": "distilbert.transformer.layer.${layer_index}.attention.q_lin.bias" + }, + { + "name": "distilbert.transformer.layer.${layer_index}.attention.v_lin.weight" + }, + { + "name": "distilbert.transformer.layer.${layer_index}.attention.v_lin.bias" + }, + { + "name": "distilbert.transformer.layer.${layer_index}.attention.out_lin.weight" + }, + { + "name": "distilbert.transformer.layer.${layer_index}.attention.out_lin.bias" + }, + { + "name": "distilbert.transformer.layer.${layer_index}.sa_layer_norm.bias" + }, + { + "name": "distilbert.transformer.layer.${layer_index}.sa_layer_norm.weight" + }, + { + "name": "distilbert.transformer.layer.${layer_index}.ffn.lin1.weight" + }, + { + "name": "distilbert.transformer.layer.${layer_index}.ffn.lin1.bias" + }, + { + "name": "distilbert.transformer.layer.${layer_index}.ffn.lin2.weight" + }, + { + "name": "distilbert.transformer.layer.${layer_index}.ffn.lin2.bias" + }, + { + "name": "distilbert.transformer.layer.${layer_index}.output_layer_norm.bias" + }, + { + "name": "distilbert.transformer.layer.${layer_index}.output_layer_norm.weight" + } + ] + } +} diff --git a/src/mindnlp/wizard/merge/_data/architectures/distilbert-token-classification.json b/src/mindnlp/wizard/merge/_data/architectures/distilbert-token-classification.json new file mode 100644 index 000000000..456c512f1 --- /dev/null +++ b/src/mindnlp/wizard/merge/_data/architectures/distilbert-token-classification.json @@ -0,0 +1,88 @@ +{ + "model_type": "distilbert", + "architectures": [ + "DistilBertForTokenClassification" + ], + "pre_weights": [ + { + "name": "distilbert.embeddings.position_embeddings.weight" + }, + { + "name": "distilbert.embeddings.word_embeddings.weight", + "is_embed": true + }, + { + "name": "distilbert.embeddings.LayerNorm.bias", + "aliases": [ + "distilbert.embeddings.LayerNorm.beta" + ] + }, + { + "name": "distilbert.embeddings.LayerNorm.weight", + "aliases": [ + "distilbert.embeddings.LayerNorm.gamma" + ] + } + ], + "post_weights": [ + { + "name": "classifier.bias" + }, + { + "name": "classifier.weight" + } + ], + "num_layers_config_key": "num_hidden_layers", + "layer_templates": { + "weights": [ + { + "name": "distilbert.transformer.layer.${layer_index}.attention.k_lin.weight" + }, + { + "name": "distilbert.transformer.layer.${layer_index}.attention.k_lin.bias" + }, + { + "name": "distilbert.transformer.layer.${layer_index}.attention.q_lin.weight" + }, + { + "name": "distilbert.transformer.layer.${layer_index}.attention.q_lin.bias" + }, + { + "name": "distilbert.transformer.layer.${layer_index}.attention.v_lin.weight" + }, + { + "name": "distilbert.transformer.layer.${layer_index}.attention.v_lin.bias" + }, + { + "name": "distilbert.transformer.layer.${layer_index}.attention.out_lin.weight" + }, + { + "name": "distilbert.transformer.layer.${layer_index}.attention.out_lin.bias" + }, + { + "name": "distilbert.transformer.layer.${layer_index}.sa_layer_norm.bias" + }, + { + "name": "distilbert.transformer.layer.${layer_index}.sa_layer_norm.weight" + }, + { + "name": "distilbert.transformer.layer.${layer_index}.ffn.lin1.weight" + }, + { + "name": "distilbert.transformer.layer.${layer_index}.ffn.lin1.bias" + }, + { + "name": "distilbert.transformer.layer.${layer_index}.ffn.lin2.weight" + }, + { + "name": "distilbert.transformer.layer.${layer_index}.ffn.lin2.bias" + }, + { + "name": "distilbert.transformer.layer.${layer_index}.output_layer_norm.bias" + }, + { + "name": "distilbert.transformer.layer.${layer_index}.output_layer_norm.weight" + } + ] + } +} diff --git a/src/mindnlp/wizard/merge/_data/architectures/distilbert.json b/src/mindnlp/wizard/merge/_data/architectures/distilbert.json new file mode 100644 index 000000000..bd9971623 --- /dev/null +++ b/src/mindnlp/wizard/merge/_data/architectures/distilbert.json @@ -0,0 +1,81 @@ +{ + "model_type": "distilbert", + "architectures": [ + "DistilBertModel" + ], + "pre_weights": [ + { + "name": "distilbert.embeddings.position_embeddings.weight" + }, + { + "name": "distilbert.embeddings.word_embeddings.weight", + "is_embed": true + }, + { + "name": "distilbert.embeddings.LayerNorm.bias", + "aliases": [ + "distilbert.embeddings.LayerNorm.beta" + ] + }, + { + "name": "distilbert.embeddings.LayerNorm.weight", + "aliases": [ + "distilbert.embeddings.LayerNorm.gamma" + ] + } + ], + "post_weights": [], + "num_layers_config_key": "num_hidden_layers", + "layer_templates": { + "weights": [ + { + "name": "distilbert.transformer.layer.${layer_index}.attention.k_lin.weight" + }, + { + "name": "distilbert.transformer.layer.${layer_index}.attention.k_lin.bias" + }, + { + "name": "distilbert.transformer.layer.${layer_index}.attention.q_lin.weight" + }, + { + "name": "distilbert.transformer.layer.${layer_index}.attention.q_lin.bias" + }, + { + "name": "distilbert.transformer.layer.${layer_index}.attention.v_lin.weight" + }, + { + "name": "distilbert.transformer.layer.${layer_index}.attention.v_lin.bias" + }, + { + "name": "distilbert.transformer.layer.${layer_index}.attention.out_lin.weight" + }, + { + "name": "distilbert.transformer.layer.${layer_index}.attention.out_lin.bias" + }, + { + "name": "distilbert.transformer.layer.${layer_index}.sa_layer_norm.bias" + }, + { + "name": "distilbert.transformer.layer.${layer_index}.sa_layer_norm.weight" + }, + { + "name": "distilbert.transformer.layer.${layer_index}.ffn.lin1.weight" + }, + { + "name": "distilbert.transformer.layer.${layer_index}.ffn.lin1.bias" + }, + { + "name": "distilbert.transformer.layer.${layer_index}.ffn.lin2.weight" + }, + { + "name": "distilbert.transformer.layer.${layer_index}.ffn.lin2.bias" + }, + { + "name": "distilbert.transformer.layer.${layer_index}.output_layer_norm.bias" + }, + { + "name": "distilbert.transformer.layer.${layer_index}.output_layer_norm.weight" + } + ] + } +} diff --git a/src/mindnlp/wizard/merge/_data/architectures/exaone.json b/src/mindnlp/wizard/merge/_data/architectures/exaone.json new file mode 100644 index 000000000..e9024473f --- /dev/null +++ b/src/mindnlp/wizard/merge/_data/architectures/exaone.json @@ -0,0 +1,78 @@ +{ + "model_type": "exaone", + "architectures": [ + "ExaoneForCausalLM" + ], + "pre_weights": [ + { + "name": "transformer.wte.weight", + "is_embed": true, + "output_space": "running_residual" + } + ], + "num_layers_config_key": "num_hidden_layers", + "layer_templates": { + "weights": [ + { + "name": "transformer.h.${layer_index}.ln_1.weight", + "input_space": "running_residual" + }, + { + "name": "transformer.h.${layer_index}.attn.attention.q_proj.weight", + "input_space": "running_residual", + "output_space": "attn_qk_${layer_index}", + "head_split": "output", + "is_kq": true + }, + { + "name": "transformer.h.${layer_index}.attn.attention.k_proj.weight", + "input_space": "running_residual", + "output_space": "attn_qk_${layer_index}", + "head_split": "output", + "is_kq": true + }, + { + "name": "transformer.h.${layer_index}.attn.attention.v_proj.weight", + "input_space": "running_residual", + "output_space": "attn_v_${layer_index}", + "head_split": "output" + }, + { + "name": "transformer.h.${layer_index}.attn.attention.out_proj.weight", + "input_space": "attn_v_${layer_index}", + "output_space": "running_residual", + "head_split": "input" + }, + { + "name": "transformer.h.${layer_index}.ln_2.weight", + "input_space": "running_residual" + }, + { + "name": "transformer.h.${layer_index}.mlp.c_fc_0.weight", + "input_space": "running_residual", + "output_space": "up_${layer_index}" + }, + { + "name": "transformer.h.${layer_index}.mlp.c_fc_1.weight", + "input_space": "running_residual", + "output_space": "up_${layer_index}" + }, + { + "name": "transformer.h.${layer_index}.mlp.c_proj.weight", + "input_space": "up_${layer_index}", + "output_space": "running_residual" + } + ] + }, + "post_weights": [ + { + "name": "transformer.ln_f.weight", + "input_space": "running_residual" + }, + { + "name": "lm_head.weight", + "input_space": "running_residual", + "is_embed": true + } + ] +} diff --git a/src/mindnlp/wizard/merge/_data/architectures/falcon.json b/src/mindnlp/wizard/merge/_data/architectures/falcon.json new file mode 100644 index 000000000..648dab062 --- /dev/null +++ b/src/mindnlp/wizard/merge/_data/architectures/falcon.json @@ -0,0 +1,53 @@ +{ + "model_type": "falcon", + "architectures": [ + "FalconForCausalLM" + ], + "pre_weights": [ + { + "name": "transformer.word_embeddings.weight", + "is_embed": true + } + ], + "post_weights": [ + { + "name": "transformer.ln_f.weight" + }, + { + "name": "transformer.ln_f.bias" + }, + { + "name": "lm_head.weight", + "is_embed": true + } + ], + "num_layers_config_key": "num_hidden_layers", + "layer_templates": { + "weights": [ + { + "name": "transformer.h.${layer_index}.ln_attn.bias" + }, + { + "name": "transformer.h.${layer_index}.ln_attn.weight" + }, + { + "name": "transformer.h.${layer_index}.ln_mlp.bias" + }, + { + "name": "transformer.h.${layer_index}.ln_mlp.weight" + }, + { + "name": "transformer.h.${layer_index}.mlp.dense_4h_to_h.weight" + }, + { + "name": "transformer.h.${layer_index}.mlp.dense_h_to_4h.weight" + }, + { + "name": "transformer.h.${layer_index}.self_attention.dense.weight" + }, + { + "name": "transformer.h.${layer_index}.self_attention.query_key_value.weight" + } + ] + } +} diff --git a/src/mindnlp/wizard/merge/_data/architectures/glm4.json b/src/mindnlp/wizard/merge/_data/architectures/glm4.json new file mode 100644 index 000000000..3ff0d6953 --- /dev/null +++ b/src/mindnlp/wizard/merge/_data/architectures/glm4.json @@ -0,0 +1,72 @@ +{ + "model_type": "glm4", + "architectures": [ + "Glm4ForCausalLM" + ], + "pre_weights": [ + { + "name": "model.embed_tokens.weight", + "is_embed": true + } + ], + "num_layers_config_key": "num_hidden_layers", + "layer_templates": { + "weights": [ + { + "name": "model.layers.${layer_index}.input_layernorm.weight" + }, + { + "name": "model.layers.${layer_index}.mlp.down_proj.weight" + }, + { + "name": "model.layers.${layer_index}.mlp.gate_up_proj.weight" + }, + { + "name": "model.layers.${layer_index}.post_attention_layernorm.weight" + }, + { + "name": "model.layers.${layer_index}.post_mlp_layernorm.weight" + }, + { + "name": "model.layers.${layer_index}.post_self_attn_layernorm.weight" + }, + { + "name": "model.layers.${layer_index}.self_attn.q_proj.weight" + }, + { + "name": "model.layers.${layer_index}.self_attn.q_proj.bias", + "optional": true + }, + { + "name": "model.layers.${layer_index}.self_attn.k_proj.weight" + }, + { + "name": "model.layers.${layer_index}.self_attn.k_proj.bias", + "optional": true + }, + { + "name": "model.layers.${layer_index}.self_attn.v_proj.weight" + }, + { + "name": "model.layers.${layer_index}.self_attn.v_proj.bias", + "optional": true + }, + { + "name": "model.layers.${layer_index}.self_attn.o_proj.weight" + } + ] + }, + "post_weights": [ + { + "name": "model.norm.weight" + }, + { + "name": "lm_head.weight", + "is_embed": true, + "optional": true, + "tied_names": [ + "model.embed_tokens.weight" + ] + } + ] +} diff --git a/src/mindnlp/wizard/merge/_data/architectures/glm4moe.json b/src/mindnlp/wizard/merge/_data/architectures/glm4moe.json new file mode 100644 index 000000000..64e048889 --- /dev/null +++ b/src/mindnlp/wizard/merge/_data/architectures/glm4moe.json @@ -0,0 +1,77 @@ +{ + "model_type": "glm4_moe", + "architectures": [ + "Glm4MoeForCausalLM" + ], + "pre_weights": [ + { + "name": "model.embed_tokens.weight", + "is_embed": true + } + ], + "num_layers_config_key": "num_hidden_layers", + "layer_templates": { + "weights": [ + { + "name": "model.layers.${layer_index}.input_layernorm.weight" + }, + { + "name": "model.layers.${layer_index}.mlp.down_proj.weight" + }, + { + "name": "model.layers.${layer_index}.mlp.gate_proj.weight" + }, + { + "name": "model.layers.${layer_index}.mlp.up_proj.weight" + }, + { + "name": "model.layers.${layer_index}.post_attention_layernorm.weight" + }, + { + "name": "model.layers.${layer_index}.self_attn.q_norm.weight", + "optional": true + }, + { + "name": "model.layers.${layer_index}.self_attn.k_norm.weight", + "optional": true + }, + { + "name": "model.layers.${layer_index}.self_attn.q_proj.weight" + }, + { + "name": "model.layers.${layer_index}.self_attn.q_proj.bias", + "optional": true + }, + { + "name": "model.layers.${layer_index}.self_attn.k_proj.weight" + }, + { + "name": "model.layers.${layer_index}.self_attn.k_proj.bias", + "optional": true + }, + { + "name": "model.layers.${layer_index}.self_attn.v_proj.weight" + }, + { + "name": "model.layers.${layer_index}.self_attn.v_proj.bias", + "optional": true + }, + { + "name": "model.layers.${layer_index}.self_attn.o_proj.weight" + } + ] + }, + "post_weights": [ + { + "name": "model.norm.weight" + }, + { + "name": "lm_head.weight", + "is_embed": true, + "optional": true, + "tied_names": [ + "model.embed_tokens.weight" + ] + } + ] +} diff --git a/src/mindnlp/wizard/merge/_data/architectures/gpt-neox.json b/src/mindnlp/wizard/merge/_data/architectures/gpt-neox.json new file mode 100644 index 000000000..e8bf29a3b --- /dev/null +++ b/src/mindnlp/wizard/merge/_data/architectures/gpt-neox.json @@ -0,0 +1,74 @@ +{ + "model_type": "gpt_neox", + "architectures": [ + "GPTNeoXForCausalLM" + ], + "pre_weights": [ + { + "name": "gpt_neox.embed_in.weight", + "is_embed": true + } + ], + "post_weights": [ + { + "name": "gpt_neox.final_layer_norm.bias" + }, + { + "name": "gpt_neox.final_layer_norm.weight" + }, + { + "name": "embed_out.weight", + "is_embed": true + } + ], + "num_layers_config_key": "num_hidden_layers", + "layer_templates": { + "weights": [ + { + "name": "gpt_neox.layers.${layer_index}.attention.dense.weight" + }, + { + "name": "gpt_neox.layers.${layer_index}.attention.dense.bias" + }, + { + "name": "gpt_neox.layers.${layer_index}.attention.query_key_value.weight" + }, + { + "name": "gpt_neox.layers.${layer_index}.attention.query_key_value.bias" + }, + { + "name": "gpt_neox.layers.${layer_index}.input_layernorm.weight" + }, + { + "name": "gpt_neox.layers.${layer_index}.input_layernorm.bias" + }, + { + "name": "gpt_neox.layers.${layer_index}.mlp.dense_4h_to_h.weight" + }, + { + "name": "gpt_neox.layers.${layer_index}.mlp.dense_4h_to_h.bias" + }, + { + "name": "gpt_neox.layers.${layer_index}.mlp.dense_h_to_4h.weight" + }, + { + "name": "gpt_neox.layers.${layer_index}.mlp.dense_h_to_4h.bias" + }, + { + "name": "gpt_neox.layers.${layer_index}.post_attention_layernorm.weight" + }, + { + "name": "gpt_neox.layers.${layer_index}.post_attention_layernorm.bias" + }, + { + "name": "gpt_neox.layers.${layer_index}.attention.bias" + }, + { + "name": "gpt_neox.layers.${layer_index}.attention.masked_bias" + }, + { + "name": "gpt_neox.layers.${layer_index}.attention.rotary_emb.inv_freq" + } + ] + } +} diff --git a/src/mindnlp/wizard/merge/_data/architectures/gpt2-sequence-classification.json b/src/mindnlp/wizard/merge/_data/architectures/gpt2-sequence-classification.json new file mode 100644 index 000000000..54cf31f61 --- /dev/null +++ b/src/mindnlp/wizard/merge/_data/architectures/gpt2-sequence-classification.json @@ -0,0 +1,66 @@ +{ + "model_type": "gpt2", + "architectures": [ + "GPT2ForSequenceClassification" + ], + "pre_weights": [ + { + "name": "transformer.wte.weight" + }, + { + "name": "transformer.wpe.weight" + } + ], + "post_weights": [ + { + "name": "transformer.ln_f.weight" + }, + { + "name": "transformer.ln_f.bias" + }, + { + "name": "score.weight" + } + ], + "num_layers_config_key": "n_layer", + "layer_templates": { + "weights": [ + { + "name": "transformer.h.${layer_index}.attn.c_attn.weight" + }, + { + "name": "transformer.h.${layer_index}.attn.c_attn.bias" + }, + { + "name": "transformer.h.${layer_index}.attn.c_proj.weight" + }, + { + "name": "transformer.h.${layer_index}.attn.c_proj.bias" + }, + { + "name": "transformer.h.${layer_index}.ln_1.weight" + }, + { + "name": "transformer.h.${layer_index}.ln_1.bias" + }, + { + "name": "transformer.h.${layer_index}.ln_2.weight" + }, + { + "name": "transformer.h.${layer_index}.ln_2.bias" + }, + { + "name": "transformer.h.${layer_index}.mlp.c_proj.weight" + }, + { + "name": "transformer.h.${layer_index}.mlp.c_proj.bias" + }, + { + "name": "transformer.h.${layer_index}.mlp.c_fc.weight" + }, + { + "name": "transformer.h.${layer_index}.mlp.c_fc.bias" + } + ] + } +} diff --git a/src/mindnlp/wizard/merge/_data/architectures/gpt2.json b/src/mindnlp/wizard/merge/_data/architectures/gpt2.json new file mode 100644 index 000000000..fc7a32014 --- /dev/null +++ b/src/mindnlp/wizard/merge/_data/architectures/gpt2.json @@ -0,0 +1,121 @@ +{ + "model_type": "gpt2", + "architectures": [ + "GPT2LMHeadModel" + ], + "pre_weights": [ + { + "name": "transformer.wte.weight", + "is_embed": true, + "aliases": [ + "wte.weight" + ] + }, + { + "name": "transformer.wpe.weight", + "aliases": [ + "wpe.weight" + ] + } + ], + "post_weights": [ + { + "name": "transformer.ln_f.weight", + "aliases": [ + "ln_f.weight" + ] + }, + { + "name": "transformer.ln_f.bias", + "aliases": [ + "ln_f.bias" + ] + }, + { + "name": "lm_head.weight", + "is_embed": true, + "optional": true, + "tied_names": [ + "transformer.wte.weight", + "wte.weight" + ] + } + ], + "num_layers_config_key": "n_layer", + "layer_templates": { + "weights": [ + { + "name": "transformer.h.${layer_index}.attn.c_attn.weight", + "aliases": [ + "h.${layer_index}.attn.c_attn.weight" + ] + }, + { + "name": "transformer.h.${layer_index}.attn.c_attn.bias", + "aliases": [ + "h.${layer_index}.attn.c_attn.bias" + ] + }, + { + "name": "transformer.h.${layer_index}.attn.c_proj.weight", + "aliases": [ + "h.${layer_index}.attn.c_proj.weight" + ] + }, + { + "name": "transformer.h.${layer_index}.attn.c_proj.bias", + "aliases": [ + "h.${layer_index}.attn.c_proj.bias" + ] + }, + { + "name": "transformer.h.${layer_index}.ln_1.weight", + "aliases": [ + "h.${layer_index}.ln_1.weight" + ] + }, + { + "name": "transformer.h.${layer_index}.ln_1.bias", + "aliases": [ + "h.${layer_index}.ln_1.bias" + ] + }, + { + "name": "transformer.h.${layer_index}.ln_2.weight", + "aliases": [ + "h.${layer_index}.ln_2.weight" + ] + }, + { + "name": "transformer.h.${layer_index}.ln_2.bias", + "aliases": [ + "h.${layer_index}.ln_2.bias" + ] + }, + { + "name": "transformer.h.${layer_index}.mlp.c_proj.weight", + "aliases": [ + "h.${layer_index}.mlp.c_proj.weight" + ] + }, + { + "name": "transformer.h.${layer_index}.mlp.c_proj.bias", + "aliases": [ + "h.${layer_index}.mlp.c_proj.bias" + ] + }, + { + "name": "transformer.h.${layer_index}.mlp.c_fc.weight", + "aliases": [ + "h.${layer_index}.mlp.c_fc.weight" + ] + }, + { + "name": "transformer.h.${layer_index}.mlp.c_fc.bias", + "aliases": [ + "h.${layer_index}.mlp.c_fc.bias" + ] + } + ] + } +} diff --git a/src/mindnlp/wizard/merge/_data/architectures/gptbigcode.json b/src/mindnlp/wizard/merge/_data/architectures/gptbigcode.json new file mode 100644 index 000000000..c12bac5c7 --- /dev/null +++ b/src/mindnlp/wizard/merge/_data/architectures/gptbigcode.json @@ -0,0 +1,72 @@ +{ + "model_type": "gpt_bigcode", + "architectures": [ + "GPTBigCodeForCausalLM" + ], + "pre_weights": [ + { + "name": "transformer.wte.weight", + "is_embed": true + }, + { + "name": "transformer.wpe.weight" + } + ], + "post_weights": [ + { + "name": "transformer.ln_f.weight" + }, + { + "name": "transformer.ln_f.bias" + }, + { + "name": "lm_head.weight", + "is_embed": true, + "optional": true, + "tied_names": [ + "transformer.wte.weight" + ] + } + ], + "num_layers_config_key": "n_layer", + "layer_templates": { + "weights": [ + { + "name": "transformer.h.${layer_index}.attn.c_attn.weight" + }, + { + "name": "transformer.h.${layer_index}.attn.c_attn.bias" + }, + { + "name": "transformer.h.${layer_index}.attn.c_proj.weight" + }, + { + "name": "transformer.h.${layer_index}.attn.c_proj.bias" + }, + { + "name": "transformer.h.${layer_index}.ln_1.weight" + }, + { + "name": "transformer.h.${layer_index}.ln_1.bias" + }, + { + "name": "transformer.h.${layer_index}.ln_2.weight" + }, + { + "name": "transformer.h.${layer_index}.ln_2.bias" + }, + { + "name": "transformer.h.${layer_index}.mlp.c_proj.weight" + }, + { + "name": "transformer.h.${layer_index}.mlp.c_proj.bias" + }, + { + "name": "transformer.h.${layer_index}.mlp.c_fc.weight" + }, + { + "name": "transformer.h.${layer_index}.mlp.c_fc.bias" + } + ] + } +} diff --git a/src/mindnlp/wizard/merge/_data/architectures/granite4moehybrid.json b/src/mindnlp/wizard/merge/_data/architectures/granite4moehybrid.json new file mode 100644 index 000000000..1ef6c64b0 --- /dev/null +++ b/src/mindnlp/wizard/merge/_data/architectures/granite4moehybrid.json @@ -0,0 +1,99 @@ +{ + "model_type": "granitemoehybrid", + "architectures": [ + "GraniteMoeHybridForCausalLM" + ], + "pre_weights": [ + { + "name": "model.embed_tokens.weight", + "is_embed": true + } + ], + "post_weights": [ + { + "name": "model.norm.weight" + }, + { + "name": "lm_head.weight", + "is_embed": true, + "optional": true, + "tied_names": [ + "model.embed_tokens.weight" + ] + } + ], + "num_layers_config_key": "num_hidden_layers", + "layer_templates": { + "weights": [ + { + "name": "model.layers.${layer_index}.input_layernorm.weight" + }, + { + "name": "model.layers.${layer_index}.shared_mlp.input_linear.weight" + }, + { + "name": "model.layers.${layer_index}.shared_mlp.output_linear.weight" + }, + { + "name": "model.layers.${layer_index}.block_sparse_moe.input_linear.weight" + }, + { + "name": "model.layers.${layer_index}.block_sparse_moe.router.layer.weight" + }, + { + "name": "model.layers.${layer_index}.block_sparse_moe.output_linear.weight" + }, + { + "name": "model.layers.${layer_index}.post_attention_layernorm.weight" + }, + { + "name": "model.layers.${layer_index}.self_attn.k_proj.weight", + "optional": true + }, + { + "name": "model.layers.${layer_index}.self_attn.o_proj.weight", + "optional": true + }, + { + "name": "model.layers.${layer_index}.self_attn.q_proj.weight", + "optional": true + }, + { + "name": "model.layers.${layer_index}.self_attn.v_proj.weight", + "optional": true + }, + { + "name": "model.layers.${layer_index}.mamba.A_log", + "optional": true + }, + { + "name": "model.layers.${layer_index}.mamba.conv1d.bias", + "optional": true + }, + { + "name": "model.layers.${layer_index}.mamba.conv1d.weight", + "optional": true + }, + { + "name": "model.layers.${layer_index}.mamba.D", + "optional": true + }, + { + "name": "model.layers.${layer_index}.mamba.dt_bias", + "optional": true + }, + { + "name": "model.layers.${layer_index}.mamba.in_proj.weight", + "optional": true + }, + { + "name": "model.layers.${layer_index}.mamba.norm.weight", + "optional": true + }, + { + "name": "model.layers.${layer_index}.mamba.out_proj.weight", + "optional": true + } + ] + } +} diff --git a/src/mindnlp/wizard/merge/_data/architectures/internlm2.json b/src/mindnlp/wizard/merge/_data/architectures/internlm2.json new file mode 100644 index 000000000..888faa484 --- /dev/null +++ b/src/mindnlp/wizard/merge/_data/architectures/internlm2.json @@ -0,0 +1,51 @@ +{ + "model_type": "internlm2", + "architectures": [ + "InternLM2ForCausalLM" + ], + "pre_weights": [ + { + "name": "model.tok_embeddings.weight", + "is_embed": true + } + ], + "post_weights": [ + { + "name": "model.norm.weight" + }, + { + "name": "output.weight", + "is_embed": true, + "optional": true, + "tied_names": [ + "model.tok_embeddings.weight" + ] + } + ], + "num_layers_config_key": "num_hidden_layers", + "layer_templates": { + "weights": [ + { + "name": "model.layers.${layer_index}.attention_norm.weight" + }, + { + "name": "model.layers.${layer_index}.ffn_norm.weight" + }, + { + "name": "model.layers.${layer_index}.attention.wqkv.weight" + }, + { + "name": "model.layers.${layer_index}.attention.wo.weight" + }, + { + "name": "model.layers.${layer_index}.feed_forward.w1.weight" + }, + { + "name": "model.layers.${layer_index}.feed_forward.w2.weight" + }, + { + "name": "model.layers.${layer_index}.feed_forward.w3.weight" + } + ] + } +} diff --git a/src/mindnlp/wizard/merge/_data/architectures/jais.json b/src/mindnlp/wizard/merge/_data/architectures/jais.json new file mode 100644 index 000000000..c08cca95f --- /dev/null +++ b/src/mindnlp/wizard/merge/_data/architectures/jais.json @@ -0,0 +1,70 @@ +{ + "model_type": "jais", + "architectures": [ + "JAISLMHeadModel" + ], + "pre_weights": [ + { + "name": "transformer.wte.weight", + "is_embed": true + }, + { + "name": "transformer.relative_pe.slopes" + } + ], + "post_weights": [ + { + "name": "transformer.ln_f.weight" + }, + { + "name": "transformer.ln_f.bias" + } + ], + "num_layers_config_key": "n_layer", + "layer_templates": { + "weights": [ + { + "name": "transformer.h.${layer_index}.attn.c_attn.weight" + }, + { + "name": "transformer.h.${layer_index}.attn.c_attn.bias" + }, + { + "name": "transformer.h.${layer_index}.attn.c_proj.weight" + }, + { + "name": "transformer.h.${layer_index}.attn.c_proj.bias" + }, + { + "name": "transformer.h.${layer_index}.ln_1.weight" + }, + { + "name": "transformer.h.${layer_index}.ln_1.bias" + }, + { + "name": "transformer.h.${layer_index}.ln_2.weight" + }, + { + "name": "transformer.h.${layer_index}.ln_2.bias" + }, + { + "name": "transformer.h.${layer_index}.mlp.c_fc.weight" + }, + { + "name": "transformer.h.${layer_index}.mlp.c_fc.bias" + }, + { + "name": "transformer.h.${layer_index}.mlp.c_fc2.weight" + }, + { + "name": "transformer.h.${layer_index}.mlp.c_fc2.bias" + }, + { + "name": "transformer.h.${layer_index}.mlp.c_proj.weight" + }, + { + "name": "transformer.h.${layer_index}.mlp.c_proj.bias" + } + ] + } +} diff --git a/src/mindnlp/wizard/merge/_data/architectures/lfm2.json b/src/mindnlp/wizard/merge/_data/architectures/lfm2.json new file mode 100644 index 000000000..d383d93c1 --- /dev/null +++ b/src/mindnlp/wizard/merge/_data/architectures/lfm2.json @@ -0,0 +1,80 @@ +{ + "model_type": "lfm2", + "architectures": [ + "Lfm2ForCausalLM" + ], + "pre_weights": [ + { + "name": "model.embed_tokens.weight", + "is_embed": true + }, + { + "name": "model.embedding_norm.weight", + "is_embed": true + } + + ], + "post_weights": [ + { + "name": "lm_head.weight", + "is_embed": true, + "optional": true + } + ], + "num_layers_config_key": "num_hidden_layers", + "layer_templates": { + "weights": [ + { + "name": "model.layers.${layer_index}.feed_forward.w1.weight" + }, + { + "name": "model.layers.${layer_index}.feed_forward.w2.weight" + }, + { + "name": "model.layers.${layer_index}.feed_forward.w3.weight" + }, + { + "name": "model.layers.${layer_index}.ffn_norm.weight" + }, + { + "name": "model.layers.${layer_index}.operator_norm.weight" + }, + { + "name": "model.layers.${layer_index}.conv.conv.weight", + "optional": true + }, + { + "name": "model.layers.${layer_index}.conv.in_proj.weight", + "optional": true + }, + { + "name": "model.layers.${layer_index}.conv.out_proj.weight", + "optional": true + }, + { + "name": "model.layers.${layer_index}.self_attn.k_layernorm.weight", + "optional": true + }, + { + "name": "model.layers.${layer_index}.self_attn.k_proj.weight", + "optional": true + }, + { + "name": "model.layers.${layer_index}.self_attn.out_proj.weight", + "optional": true + }, + { + "name": "model.layers.${layer_index}.self_attn.q_layernorm.weight", + "optional": true + }, + { + "name": "model.layers.${layer_index}.self_attn.q_proj.weight", + "optional": true + }, + { + "name": "model.layers.${layer_index}.self_attn.v_proj.weight", + "optional": true + } + ] + } +} diff --git a/src/mindnlp/wizard/merge/_data/architectures/llama-sequence-classification.json b/src/mindnlp/wizard/merge/_data/architectures/llama-sequence-classification.json new file mode 100644 index 000000000..7c73d12e1 --- /dev/null +++ b/src/mindnlp/wizard/merge/_data/architectures/llama-sequence-classification.json @@ -0,0 +1,82 @@ +{ + "model_type": "llama", + "architectures": [ + "LlamaForSequenceClassification" + ], + "pre_weights": [ + { + "name": "model.embed_tokens.weight", + "is_embed": true, + "output_space": "running_residual" + } + ], + "num_layers_config_key": "num_hidden_layers", + "layer_templates": { + "weights": [ + { + "name": "model.layers.${layer_index}.input_layernorm.weight", + "input_space": "running_residual" + }, + { + "name": "model.layers.${layer_index}.self_attn.q_proj.weight", + "input_space": "running_residual", + "output_space": "attn_qk_${layer_index}", + "head_split": "output", + "is_kq": true + }, + { + "name": "model.layers.${layer_index}.self_attn.k_proj.weight", + "input_space": "running_residual", + "output_space": "attn_qk_${layer_index}", + "head_split": "output", + "is_kq": true + }, + { + "name": "model.layers.${layer_index}.self_attn.v_proj.weight", + "input_space": "running_residual", + "output_space": "attn_v_${layer_index}", + "head_split": "output" + }, + { + "name": "model.layers.${layer_index}.self_attn.o_proj.weight", + "input_space": "attn_v_${layer_index}", + "output_space": "running_residual", + "head_split": "input" + }, + { + "name": "model.layers.${layer_index}.post_attention_layernorm.weight", + "input_space": "running_residual" + }, + { + "name": "model.layers.${layer_index}.mlp.up_proj.weight", + "input_space": "running_residual", + "output_space": "up_${layer_index}" + }, + { + "name": "model.layers.${layer_index}.mlp.gate_proj.weight", + "input_space": "running_residual", + "output_space": "up_${layer_index}" + }, + { + "name": "model.layers.${layer_index}.mlp.down_proj.weight", + "input_space": "up_${layer_index}", + "output_space": "running_residual" + } + ] + }, + "post_weights": [ + { + "name": "model.norm.weight", + "input_space": "running_residual" + }, + { + "name": "score.weight", + "input_space": "running_residual" + }, + { + "name": "score.bias", + "input_space": "running_residual", + "optional": true + } + ] +} diff --git a/src/mindnlp/wizard/merge/_data/architectures/llama.json b/src/mindnlp/wizard/merge/_data/architectures/llama.json new file mode 100644 index 000000000..00918a2c0 --- /dev/null +++ b/src/mindnlp/wizard/merge/_data/architectures/llama.json @@ -0,0 +1,83 @@ +{ + "model_type": "llama", + "architectures": [ + "LlamaForCausalLM", + "LLaMaForCausalLM" + ], + "pre_weights": [ + { + "name": "model.embed_tokens.weight", + "is_embed": true, + "output_space": "running_residual" + } + ], + "num_layers_config_key": "num_hidden_layers", + "layer_templates": { + "weights": [ + { + "name": "model.layers.${layer_index}.input_layernorm.weight", + "input_space": "running_residual" + }, + { + "name": "model.layers.${layer_index}.self_attn.q_proj.weight", + "input_space": "running_residual", + "output_space": "attn_qk_${layer_index}", + "head_split": "output", + "is_kq": true + }, + { + "name": "model.layers.${layer_index}.self_attn.k_proj.weight", + "input_space": "running_residual", + "output_space": "attn_qk_${layer_index}", + "head_split": "output", + "is_kq": true + }, + { + "name": "model.layers.${layer_index}.self_attn.v_proj.weight", + "input_space": "running_residual", + "output_space": "attn_v_${layer_index}", + "head_split": "output" + }, + { + "name": "model.layers.${layer_index}.self_attn.o_proj.weight", + "input_space": "attn_v_${layer_index}", + "output_space": "running_residual", + "head_split": "input" + }, + { + "name": "model.layers.${layer_index}.post_attention_layernorm.weight", + "input_space": "running_residual" + }, + { + "name": "model.layers.${layer_index}.mlp.up_proj.weight", + "input_space": "running_residual", + "output_space": "up_${layer_index}" + }, + { + "name": "model.layers.${layer_index}.mlp.gate_proj.weight", + "input_space": "running_residual", + "output_space": "up_${layer_index}" + }, + { + "name": "model.layers.${layer_index}.mlp.down_proj.weight", + "input_space": "up_${layer_index}", + "output_space": "running_residual" + } + ] + }, + "post_weights": [ + { + "name": "model.norm.weight", + "input_space": "running_residual" + }, + { + "name": "lm_head.weight", + "input_space": "running_residual", + "is_embed": true, + "optional": true, + "tied_names": [ + "model.embed_tokens.weight" + ] + } + ] +} diff --git a/src/mindnlp/wizard/merge/_data/architectures/llama4.json b/src/mindnlp/wizard/merge/_data/architectures/llama4.json new file mode 100644 index 000000000..0edd8749f --- /dev/null +++ b/src/mindnlp/wizard/merge/_data/architectures/llama4.json @@ -0,0 +1,207 @@ +{ + "kind": "modular", + "architectures": [ + "Llama4ForConditionalGeneration" + ], + "model_type": "llama4", + "tagalong_files": [ + "preprocessor_config.json", + "processor_config.json" + ], + "vocab_size_config_key": "text_config.vocab_size", + "modules": { + "text_decoder": { + "weight_prefix": "language_model.", + "architecture": { + "model_type": "llama4_text", + "architectures": [], + "pre_weights": [ + { + "name": "model.embed_tokens.weight", + "is_embed": true + } + ], + "num_layers_config_key": "text_config.num_hidden_layers", + "layer_templates": { + "weights": [ + { + "name": "model.layers.${layer_index}.input_layernorm.weight" + }, + { + "name": "model.layers.${layer_index}.self_attn.q_proj.weight" + }, + { + "name": "model.layers.${layer_index}.self_attn.k_proj.weight" + }, + { + "name": "model.layers.${layer_index}.self_attn.v_proj.weight" + }, + { + "name": "model.layers.${layer_index}.self_attn.o_proj.weight" + }, + { + "name": "model.layers.${layer_index}.post_attention_layernorm.weight" + }, + { + "name": "model.layers.${layer_index}.feed_forward.shared_expert.down_proj.weight", + "optional": true + }, + { + "name": "model.layers.${layer_index}.feed_forward.shared_expert.gate_proj.weight", + "optional": true + }, + { + "name": "model.layers.${layer_index}.feed_forward.shared_expert.up_proj.weight", + "optional": true + }, + { + "name": "model.layers.${layer_index}.feed_forward.experts.down_proj", + "optional": true + }, + { + "name": "model.layers.${layer_index}.feed_forward.experts.gate_up_proj", + "optional": true + }, + { + "name": "model.layers.${layer_index}.feed_forward.router.weight", + "optional": true + }, + { + "name": "model.layers.${layer_index}.feed_forward.down_proj.weight", + "optional": true + }, + { + "name": "model.layers.${layer_index}.feed_forward.gate_proj.weight", + "optional": true + }, + { + "name": "model.layers.${layer_index}.feed_forward.up_proj.weight", + "optional": true + } + ] + }, + "post_weights": [ + { + "name": "model.norm.weight" + }, + { + "name": "lm_head.weight", + "is_embed": true, + "optional": true, + "tied_names": [ + "model.embed_tokens.weight" + ] + } + ] + } + }, + "multi_modal_projector": { + "weight_prefix": "multi_modal_projector.", + "architecture": { + "model_type": "", + "architectures": [], + "pre_weights": [ + { + "name": "linear_1.weight" + } + ], + "post_weights": [], + "layer_templates": { + "weights": [] + }, + "override_num_layers": 0 + } + }, + "vision_model": { + "weight_prefix": "vision_model.", + "architecture": { + "model_type": "llama4_vision_model", + "architectures": [], + "pre_weights": [ + { + "name": "class_embedding" + }, + { + "name": "patch_embedding.linear.weight" + }, + { + "name": "positional_embedding_vlm" + }, + { + "name": "layernorm_pre.bias" + }, + { + "name": "layernorm_pre.weight" + } + ], + "post_weights": [ + { + "name": "layernorm_post.weight" + }, + { + "name": "layernorm_post.bias" + }, + { + "name": "vision_adapter.mlp.fc1.weight" + }, + { + "name": "vision_adapter.mlp.fc2.weight" + } + ], + "layer_templates": { + "weights": [ + { + "name": "model.layers.${layer_index}.input_layernorm.weight" + }, + { + "name": "model.layers.${layer_index}.input_layernorm.bias" + }, + { + "name": "model.layers.${layer_index}.mlp.fc1.weight" + }, + { + "name": "model.layers.${layer_index}.mlp.fc1.bias" + }, + { + "name": "model.layers.${layer_index}.mlp.fc2.weight" + }, + { + "name": "model.layers.${layer_index}.mlp.fc2.bias" + }, + { + "name": "model.layers.${layer_index}.post_attention_layernorm.weight" + }, + { + "name": "model.layers.${layer_index}.post_attention_layernorm.bias" + }, + { + "name": "model.layers.${layer_index}.self_attn.k_proj.weight" + }, + { + "name": "model.layers.${layer_index}.self_attn.k_proj.bias" + }, + { + "name": "model.layers.${layer_index}.self_attn.o_proj.weight" + }, + { + "name": "model.layers.${layer_index}.self_attn.o_proj.bias" + }, + { + "name": "model.layers.${layer_index}.self_attn.q_proj.weight" + }, + { + "name": "model.layers.${layer_index}.self_attn.q_proj.bias" + }, + { + "name": "model.layers.${layer_index}.self_attn.v_proj.weight" + }, + { + "name": "model.layers.${layer_index}.self_attn.v_proj.bias" + } + ] + }, + "num_layers_config_key": "vision_config.num_hidden_layers" + } + } + } +} diff --git a/src/mindnlp/wizard/merge/_data/architectures/mamba.json b/src/mindnlp/wizard/merge/_data/architectures/mamba.json new file mode 100644 index 000000000..1c4735325 --- /dev/null +++ b/src/mindnlp/wizard/merge/_data/architectures/mamba.json @@ -0,0 +1,60 @@ +{ + "model_type": "mamba", + "architectures": [ + "MambaForCausalLM" + ], + "pre_weights": [ + { + "name": "backbone.embeddings.weight", + "is_embed": true + } + ], + "post_weights": [ + { + "name": "backbone.norm_f.weight" + }, + { + "name": "lm_head.weight", + "is_embed": true, + "optional": true, + "tied_names": [ + "backbone.embeddings.weight" + ] + } + ], + "num_layers_config_key": "num_hidden_layers", + "layer_templates": { + "weights": [ + { + "name": "backbone.layers.${layer_index}.mixer.A_log" + }, + { + "name": "backbone.layers.${layer_index}.mixer.conv1d.bias" + }, + { + "name": "backbone.layers.${layer_index}.mixer.conv1d.weight" + }, + { + "name": "backbone.layers.${layer_index}.mixer.D" + }, + { + "name": "backbone.layers.${layer_index}.mixer.dt_proj.bias" + }, + { + "name": "backbone.layers.${layer_index}.mixer.dt_proj.weight" + }, + { + "name": "backbone.layers.${layer_index}.mixer.in_proj.weight" + }, + { + "name": "backbone.layers.${layer_index}.mixer.out_proj.weight" + }, + { + "name": "backbone.layers.${layer_index}.mixer.x_proj.weight" + }, + { + "name": "backbone.layers.${layer_index}.norm.weight" + } + ] + } +} diff --git a/src/mindnlp/wizard/merge/_data/architectures/mistral.json b/src/mindnlp/wizard/merge/_data/architectures/mistral.json new file mode 100644 index 000000000..0c729fda4 --- /dev/null +++ b/src/mindnlp/wizard/merge/_data/architectures/mistral.json @@ -0,0 +1,57 @@ +{ + "model_type": "mistral", + "architectures": [ + "MistralForCausalLM" + ], + "pre_weights": [ + { + "name": "model.embed_tokens.weight", + "is_embed": true + } + ], + "num_layers_config_key": "num_hidden_layers", + "layer_templates": { + "weights": [ + { + "name": "model.layers.${layer_index}.input_layernorm.weight" + }, + { + "name": "model.layers.${layer_index}.self_attn.q_proj.weight" + }, + { + "name": "model.layers.${layer_index}.self_attn.k_proj.weight" + }, + { + "name": "model.layers.${layer_index}.self_attn.v_proj.weight" + }, + { + "name": "model.layers.${layer_index}.self_attn.o_proj.weight" + }, + { + "name": "model.layers.${layer_index}.post_attention_layernorm.weight" + }, + { + "name": "model.layers.${layer_index}.mlp.up_proj.weight" + }, + { + "name": "model.layers.${layer_index}.mlp.gate_proj.weight" + }, + { + "name": "model.layers.${layer_index}.mlp.down_proj.weight" + } + ] + }, + "post_weights": [ + { + "name": "model.norm.weight" + }, + { + "name": "lm_head.weight", + "is_embed": true, + "optional": true, + "tied_names": [ + "model.embed_tokens.weight" + ] + } + ] +} diff --git a/src/mindnlp/wizard/merge/_data/architectures/mistral3.json b/src/mindnlp/wizard/merge/_data/architectures/mistral3.json new file mode 100644 index 000000000..f39fc7c5b --- /dev/null +++ b/src/mindnlp/wizard/merge/_data/architectures/mistral3.json @@ -0,0 +1,146 @@ +{ + "kind": "modular", + "architectures": [ + "Mistral3ForConditionalGeneration" + ], + "model_type": "mistral3", + "tagalong_files": [ + "preprocessor_config.json", + "processor_config.json" + ], + "vocab_size_config_key": "text_config.vocab_size", + "modules": { + "text_decoder": { + "weight_prefix": "language_model.", + "architecture": { + "model_type": "mistral", + "architectures": [], + "pre_weights": [ + { + "name": "model.embed_tokens.weight", + "is_embed": true + } + ], + "num_layers_config_key": "text_config.num_hidden_layers", + "layer_templates": { + "weights": [ + { + "name": "model.layers.${layer_index}.input_layernorm.weight" + }, + { + "name": "model.layers.${layer_index}.self_attn.q_proj.weight" + }, + { + "name": "model.layers.${layer_index}.self_attn.k_proj.weight" + }, + { + "name": "model.layers.${layer_index}.self_attn.v_proj.weight" + }, + { + "name": "model.layers.${layer_index}.self_attn.o_proj.weight" + }, + { + "name": "model.layers.${layer_index}.post_attention_layernorm.weight" + }, + { + "name": "model.layers.${layer_index}.mlp.up_proj.weight" + }, + { + "name": "model.layers.${layer_index}.mlp.gate_proj.weight" + }, + { + "name": "model.layers.${layer_index}.mlp.down_proj.weight" + } + ] + }, + "post_weights": [ + { + "name": "model.norm.weight" + }, + { + "name": "lm_head.weight", + "is_embed": true, + "optional": true, + "tied_names": [ + "model.embed_tokens.weight" + ] + } + ] + } + }, + "multi_modal_projector": { + "weight_prefix": "multi_modal_projector.", + "architecture": { + "model_type": "", + "architectures": [], + "pre_weights": [ + { + "name": "linear_1.weight" + }, + { + "name": "linear_2.weight" + }, + { + "name": "norm.weight" + }, + { + "name": "patch_merger.merging_layer.weight" + } + ], + "post_weights": [], + "layer_templates": { + "weights": [] + }, + "override_num_layers": 0 + } + }, + "vision_tower": { + "weight_prefix": "vision_tower.", + "architecture": { + "model_type": "pixtral", + "architectures": [], + "pre_weights": [ + { + "name": "patch_conv.weight" + }, + { + "name": "ln_pre.weight" + } + ], + "post_weights": [], + "layer_templates": { + "weights": [ + { + "name": "transformer.layers.${layer_index}.ffn_norm.weight" + }, + { + "name": "transformer.layers.${layer_index}.feed_forward.down_proj.weight" + }, + { + "name": "transformer.layers.${layer_index}.feed_forward.gate_proj.weight" + }, + { + "name": "transformer.layers.${layer_index}.feed_forward.up_proj.weight" + }, + { + "name": "transformer.layers.${layer_index}.attention_norm.weight" + }, + { + "name": "transformer.layers.${layer_index}.attention.k_proj.weight" + }, + { + "name": "transformer.layers.${layer_index}.attention.o_proj.weight" + }, + { + "name": "transformer.layers.${layer_index}.attention.q_proj.weight" + }, + { + "name": "transformer.layers.${layer_index}.attention.v_proj.weight" + } + ] + }, + "num_layers_config_key": "vision_config.num_hidden_layers" + } + } + } +} diff --git a/src/mindnlp/wizard/merge/_data/architectures/olmo2.json b/src/mindnlp/wizard/merge/_data/architectures/olmo2.json new file mode 100644 index 000000000..13604967f --- /dev/null +++ b/src/mindnlp/wizard/merge/_data/architectures/olmo2.json @@ -0,0 +1,57 @@ +{ + "model_type": "olmo2", + "architectures": [ + "Olmo2ForCausalLM" + ], + "pre_weights": [ + { + "name": "model.embed_tokens.weight", + "is_embed": true + } + ], + "num_layers_config_key": "num_hidden_layers", + "layer_templates": { + "weights": [ + { + "name": "model.layers.${layer_index}.self_attn.q_proj.weight" + }, + { + "name": "model.layers.${layer_index}.self_attn.k_proj.weight" + }, + { + "name": "model.layers.${layer_index}.self_attn.v_proj.weight" + }, + { + "name": "model.layers.${layer_index}.self_attn.o_proj.weight" + }, + { + "name": "model.layers.${layer_index}.post_attention_layernorm.weight" + }, + { + "name": "model.layers.${layer_index}.mlp.up_proj.weight" + }, + { + "name": "model.layers.${layer_index}.mlp.gate_proj.weight" + }, + { + "name": "model.layers.${layer_index}.mlp.down_proj.weight" + }, + { + "name": "model.layers.${layer_index}.post_feedforward_layernorm.weight" + } + ] + }, + "post_weights": [ + { + "name": "model.norm.weight" + }, + { + "name": "lm_head.weight", + "is_embed": true, + "optional": true, + "tied_names": [ + "model.embed_tokens.weight" + ] + } + ] +} diff --git a/src/mindnlp/wizard/merge/_data/architectures/phi-1.json b/src/mindnlp/wizard/merge/_data/architectures/phi-1.json new file mode 100644 index 000000000..612bd7dc4 --- /dev/null +++ b/src/mindnlp/wizard/merge/_data/architectures/phi-1.json @@ -0,0 +1,66 @@ +{ + "model_type": "mixformer-sequential", + "architectures": [ + "MixFormerSequentialForCausalLM" + ], + "pre_weights": [ + { + "name": "layers.0.wte.weight", + "is_embed": true + } + ], + "num_layers_config_key": "n_layer", + "layer_templates": { + "weights": [ + { + "name": "layers.${layer_index}.ln.bias" + }, + { + "name": "layers.${layer_index}.ln.weight" + }, + { + "name": "layers.${layer_index}.mixer.Wqkv.bias" + }, + { + "name": "layers.${layer_index}.mixer.Wqkv.weight" + }, + { + "name": "layers.${layer_index}.mixer.out_proj.bias" + }, + { + "name": "layers.${layer_index}.mixer.out_proj.weight" + }, + { + "name": "layers.${layer_index}.mixer.rotary_emb.inv_freq" + }, + { + "name": "layers.${layer_index}.mlp.fc1.bias" + }, + { + "name": "layers.${layer_index}.mlp.fc1.weight" + }, + { + "name": "layers.${layer_index}.mlp.fc2.bias" + }, + { + "name": "layers.${layer_index}.mlp.fc2.weight" + } + ] + }, + "post_weights": [ + { + "name": "layers.${num_layers}.linear.bias", + "is_embed": true + }, + { + "name": "layers.${num_layers}.linear.weight", + "is_embed": true + }, + { + "name": "layers.${num_layers}.ln.bias" + }, + { + "name": "layers.${num_layers}.ln.weight" + } + ] +} diff --git a/src/mindnlp/wizard/merge/_data/architectures/phi2-old.json b/src/mindnlp/wizard/merge/_data/architectures/phi2-old.json new file mode 100644 index 000000000..a735ecc47 --- /dev/null +++ b/src/mindnlp/wizard/merge/_data/architectures/phi2-old.json @@ -0,0 +1,62 @@ +{ + "model_type": "phi-msft", + "architectures": [ + "PhiForCausalLM" + ], + "pre_weights": [ + { + "name": "transformer.embd.wte.weight", + "is_embed": true + } + ], + "post_weights": [ + { + "name": "lm_head.linear.bias" + }, + { + "name": "lm_head.linear.weight", + "is_embed": true + }, + { + "name": "lm_head.ln.bias" + }, + { + "name": "lm_head.ln.weight" + } + ], + "num_layers_config_key": "n_layer", + "layer_templates": { + "weights": [ + { + "name": "transformer.h.${layer_index}.ln.bias" + }, + { + "name": "transformer.h.${layer_index}.ln.weight" + }, + { + "name": "transformer.h.${layer_index}.mixer.out_proj.bias" + }, + { + "name": "transformer.h.${layer_index}.mixer.out_proj.weight" + }, + { + "name": "transformer.h.${layer_index}.mixer.Wqkv.bias" + }, + { + "name": "transformer.h.${layer_index}.mixer.Wqkv.weight" + }, + { + "name": "transformer.h.${layer_index}.mlp.fc1.bias" + }, + { + "name": "transformer.h.${layer_index}.mlp.fc1.weight" + }, + { + "name": "transformer.h.${layer_index}.mlp.fc2.bias" + }, + { + "name": "transformer.h.${layer_index}.mlp.fc2.weight" + } + ] + } +} diff --git a/src/mindnlp/wizard/merge/_data/architectures/phi2.json b/src/mindnlp/wizard/merge/_data/architectures/phi2.json new file mode 100644 index 000000000..f509e51cf --- /dev/null +++ b/src/mindnlp/wizard/merge/_data/architectures/phi2.json @@ -0,0 +1,74 @@ +{ + "model_type": "phi", + "architectures": [ + "PhiForCausalLM" + ], + "pre_weights": [ + { + "name": "model.embed_tokens.weight", + "is_embed": true + } + ], + "post_weights": [ + { + "name": "lm_head.bias" + }, + { + "name": "lm_head.weight", + "is_embed": true + }, + { + "name": "model.final_layernorm.bias" + }, + { + "name": "model.final_layernorm.weight" + } + ], + "num_layers_config_key": "num_hidden_layers", + "layer_templates": { + "weights": [ + { + "name": "model.layers.${layer_index}.input_layernorm.bias" + }, + { + "name": "model.layers.${layer_index}.input_layernorm.weight" + }, + { + "name": "model.layers.${layer_index}.self_attn.dense.bias" + }, + { + "name": "model.layers.${layer_index}.self_attn.dense.weight" + }, + { + "name": "model.layers.${layer_index}.self_attn.q_proj.bias" + }, + { + "name": "model.layers.${layer_index}.self_attn.q_proj.weight" + }, + { + "name": "model.layers.${layer_index}.self_attn.k_proj.bias" + }, + { + "name": "model.layers.${layer_index}.self_attn.k_proj.weight" + }, + { + "name": "model.layers.${layer_index}.self_attn.v_proj.bias" + }, + { + "name": "model.layers.${layer_index}.self_attn.v_proj.weight" + }, + { + "name": "model.layers.${layer_index}.mlp.fc1.bias" + }, + { + "name": "model.layers.${layer_index}.mlp.fc1.weight" + }, + { + "name": "model.layers.${layer_index}.mlp.fc2.bias" + }, + { + "name": "model.layers.${layer_index}.mlp.fc2.weight" + } + ] + } +} diff --git a/src/mindnlp/wizard/merge/_data/architectures/phi3-small.json b/src/mindnlp/wizard/merge/_data/architectures/phi3-small.json new file mode 100644 index 000000000..f27dfac4d --- /dev/null +++ b/src/mindnlp/wizard/merge/_data/architectures/phi3-small.json @@ -0,0 +1,69 @@ +{ + "model_type": "phi3small", + "architectures": [ + "Phi3SmallForCausalLM" + ], + "pre_weights": [ + { + "name": "model.embed_tokens.weight", + "is_embed": true + } + ], + "post_weights": [ + { + "name": "lm_head.weight", + "is_embed": true, + "optional": true, + "tied_names": [ + "model.embed_tokens.weight" + ] + }, + { + "name": "model.final_layernorm.weight" + }, + { + "name": "model.final_layernorm.bias" + } + ], + "num_layers_config_key": "num_hidden_layers", + "layer_templates": { + "weights": [ + { + "name": "model.layers.${layer_index}.input_layernorm.weight" + }, + { + "name": "model.layers.${layer_index}.input_layernorm.bias" + }, + { + "name": "model.layers.${layer_index}.post_attention_layernorm.weight" + }, + { + "name": "model.layers.${layer_index}.post_attention_layernorm.bias" + }, + { + "name": "model.layers.${layer_index}.self_attn.dense.weight" + }, + { + "name": "model.layers.${layer_index}.self_attn.dense.bias" + }, + { + "name": "model.layers.${layer_index}.self_attn.query_key_value.weight" + }, + { + "name": "model.layers.${layer_index}.self_attn.query_key_value.bias" + }, + { + "name": "model.layers.${layer_index}.mlp.up_proj.weight" + }, + { + "name": "model.layers.${layer_index}.mlp.up_proj.bias" + }, + { + "name": "model.layers.${layer_index}.mlp.down_proj.weight" + }, + { + "name": "model.layers.${layer_index}.mlp.down_proj.bias" + } + ] + } +} diff --git a/src/mindnlp/wizard/merge/_data/architectures/phi3.json b/src/mindnlp/wizard/merge/_data/architectures/phi3.json new file mode 100644 index 000000000..6c606b842 --- /dev/null +++ b/src/mindnlp/wizard/merge/_data/architectures/phi3.json @@ -0,0 +1,44 @@ +{ + "model_type": "phi3", + "architectures": [ + "Phi3ForCausalLM" + ], + "pre_weights": [ + { + "name": "model.embed_tokens.weight", + "is_embed": true + } + ], + "post_weights": [ + { + "name": "lm_head.weight", + "is_embed": true + }, + { + "name": "model.norm.weight" + } + ], + "num_layers_config_key": "num_hidden_layers", + "layer_templates": { + "weights": [ + { + "name": "model.layers.${layer_index}.input_layernorm.weight" + }, + { + "name": "model.layers.${layer_index}.post_attention_layernorm.weight" + }, + { + "name": "model.layers.${layer_index}.self_attn.o_proj.weight" + }, + { + "name": "model.layers.${layer_index}.self_attn.qkv_proj.weight" + }, + { + "name": "model.layers.${layer_index}.mlp.gate_up_proj.weight" + }, + { + "name": "model.layers.${layer_index}.mlp.down_proj.weight" + } + ] + } +} diff --git a/src/mindnlp/wizard/merge/_data/architectures/qwen.json b/src/mindnlp/wizard/merge/_data/architectures/qwen.json new file mode 100644 index 000000000..33bf88dcc --- /dev/null +++ b/src/mindnlp/wizard/merge/_data/architectures/qwen.json @@ -0,0 +1,50 @@ +{ + "model_type": "qwen", + "architectures": [ + "QWenLMHeadModel" + ], + "pre_weights": [ + { + "name": "transformer.wte.weight", + "is_embed": true + } + ], + "post_weights": [ + { + "name": "transformer.ln_f.weight" + }, + { + "name": "lm_head.weight", + "is_embed": true + } + ], + "num_layers_config_key": "num_hidden_layers", + "layer_templates": { + "weights": [ + { + "name": "transformer.h.${layer_index}.attn.c_attn.bias" + }, + { + "name": "transformer.h.${layer_index}.attn.c_attn.weight" + }, + { + "name": "transformer.h.${layer_index}.attn.c_proj.weight" + }, + { + "name": "transformer.h.${layer_index}.ln_1.weight" + }, + { + "name": "transformer.h.${layer_index}.ln_2.weight" + }, + { + "name": "transformer.h.${layer_index}.mlp.c_proj.weight" + }, + { + "name": "transformer.h.${layer_index}.mlp.w1.weight" + }, + { + "name": "transformer.h.${layer_index}.mlp.w2.weight" + } + ] + } +} diff --git a/src/mindnlp/wizard/merge/_data/architectures/qwen2.json b/src/mindnlp/wizard/merge/_data/architectures/qwen2.json new file mode 100644 index 000000000..c71315231 --- /dev/null +++ b/src/mindnlp/wizard/merge/_data/architectures/qwen2.json @@ -0,0 +1,66 @@ +{ + "model_type": "qwen2", + "architectures": [ + "Qwen2ForCausalLM" + ], + "pre_weights": [ + { + "name": "model.embed_tokens.weight", + "is_embed": true + } + ], + "post_weights": [ + { + "name": "model.norm.weight" + }, + { + "name": "lm_head.weight", + "is_embed": true, + "optional": true, + "tied_names": [ + "model.embed_tokens.weight" + ] + } + ], + "num_layers_config_key": "num_hidden_layers", + "layer_templates": { + "weights": [ + { + "name": "model.layers.${layer_index}.input_layernorm.weight" + }, + { + "name": "model.layers.${layer_index}.mlp.down_proj.weight" + }, + { + "name": "model.layers.${layer_index}.mlp.gate_proj.weight" + }, + { + "name": "model.layers.${layer_index}.mlp.up_proj.weight" + }, + { + "name": "model.layers.${layer_index}.post_attention_layernorm.weight" + }, + { + "name": "model.layers.${layer_index}.self_attn.k_proj.bias" + }, + { + "name": "model.layers.${layer_index}.self_attn.k_proj.weight" + }, + { + "name": "model.layers.${layer_index}.self_attn.o_proj.weight" + }, + { + "name": "model.layers.${layer_index}.self_attn.q_proj.bias" + }, + { + "name": "model.layers.${layer_index}.self_attn.q_proj.weight" + }, + { + "name": "model.layers.${layer_index}.self_attn.v_proj.bias" + }, + { + "name": "model.layers.${layer_index}.self_attn.v_proj.weight" + } + ] + } +} diff --git a/src/mindnlp/wizard/merge/_data/architectures/qwen2_5_vl.json b/src/mindnlp/wizard/merge/_data/architectures/qwen2_5_vl.json new file mode 100644 index 000000000..379cd5d89 --- /dev/null +++ b/src/mindnlp/wizard/merge/_data/architectures/qwen2_5_vl.json @@ -0,0 +1,170 @@ +{ + "kind": "modular", + "model_type": "qwen2_5_vl", + "architectures": [ + "Qwen2_5_VLForConditionalGeneration" + ], + "num_vision_layers_config_key": "vision_config.depth", + "vocab_size_config_key": "vocab_size", + "tagalong_files": [ + "preprocessor_config.json", + "vocab.json" + ], + "modules": { + "text_decoder": { + "architecture": { + "architectures": [ + "Qwen2_5ForCausalLM" + ], + "model_type": "qwen2_5", + "num_layers_config_key": "num_hidden_layers", + "pre_weights": [ + { + "name": "model.embed_tokens.weight", + "is_embed": true + } + ], + "post_weights": [ + { + "name": "model.norm.weight" + }, + { + "name": "lm_head.weight", + "is_embed": true, + "optional": true, + "tied_names": [ + "model.embed_tokens.weight" + ] + } + ], + "layer_templates": { + "weights": [ + { + "name": "model.layers.${layer_index}.input_layernorm.weight" + }, + { + "name": "model.layers.${layer_index}.mlp.down_proj.weight" + }, + { + "name": "model.layers.${layer_index}.mlp.gate_proj.weight" + }, + { + "name": "model.layers.${layer_index}.mlp.up_proj.weight" + }, + { + "name": "model.layers.${layer_index}.post_attention_layernorm.weight" + }, + { + "name": "model.layers.${layer_index}.self_attn.k_proj.bias" + }, + { + "name": "model.layers.${layer_index}.self_attn.k_proj.weight" + }, + { + "name": "model.layers.${layer_index}.self_attn.o_proj.weight" + }, + { + "name": "model.layers.${layer_index}.self_attn.q_proj.bias" + }, + { + "name": "model.layers.${layer_index}.self_attn.q_proj.weight" + }, + { + "name": "model.layers.${layer_index}.self_attn.v_proj.bias" + }, + { + "name": "model.layers.${layer_index}.self_attn.v_proj.weight" + } + ] + } + } + }, + "multi_modal_projector": { + "model_type": "", + "architecture": { + "model_type": "", + "architectures": [], + "override_num_layers": 2, + "pre_weights": [ + { + "name": "visual.merger.ln_q.weight" + } + ], + "post_weights": [], + "layer_templates": { + "weights": [ + { + "name": "visual.merger.mlp.0.weight" + }, + { + "name": "visual.merger.mlp.0.bias" + }, + { + "name": "visual.merger.mlp.2.weight" + }, + { + "name": "visual.merger.mlp.2.bias" + } + ] + } + } + }, + "vision_tower": + { + "architecture": { + "model_type": "", + "architectures": [], + "num_layers_config_key": "vision_config.depth", + "pre_weights": [ + { + "name": "visual.patch_embed.proj.weight", + "is_embed": true + } + ], + "post_weights": [], + "layer_templates": { + "weights": [ + { + "name": "visual.blocks.${layer_index}.attn.proj.bias" + }, + { + "name": "visual.blocks.${layer_index}.attn.proj.weight" + }, + { + "name": "visual.blocks.${layer_index}.attn.qkv.bias" + }, + { + "name": "visual.blocks.${layer_index}.attn.qkv.weight" + }, + { + "name": "visual.blocks.${layer_index}.mlp.down_proj.bias" + }, + { + "name": "visual.blocks.${layer_index}.mlp.down_proj.weight" + }, + { + "name": "visual.blocks.${layer_index}.mlp.gate_proj.bias" + }, + { + "name": "visual.blocks.${layer_index}.mlp.gate_proj.weight" + }, + { + "name": "visual.blocks.${layer_index}.mlp.up_proj.bias" + }, + { + "name": "visual.blocks.${layer_index}.mlp.up_proj.weight" + }, + { + "name": "visual.blocks.${layer_index}.norm1.weight" + }, + { + "name": "visual.blocks.${layer_index}.norm2.weight" + } + + + ] + } + } + } + } +} diff --git a/src/mindnlp/wizard/merge/_data/architectures/qwen2_vl.json b/src/mindnlp/wizard/merge/_data/architectures/qwen2_vl.json new file mode 100644 index 000000000..7e5ee7422 --- /dev/null +++ b/src/mindnlp/wizard/merge/_data/architectures/qwen2_vl.json @@ -0,0 +1,171 @@ +{ + "kind": "modular", + "model_type": "qwen2_vl", + "architectures": [ + "Qwen2VLForConditionalGeneration" + ], + "num_vision_layers_config_key": "vision_config.depth", + "vocab_size_config_key": "vocab_size", + "tagalong_files": [ + "preprocessor_config.json", + "vocab.json" + ], + "modules": { + "text_decoder": { + "architecture":{ + "architectures": [ + "Qwen2ForCausalLM" + ], + "model_type": "qwen2", + "num_layers_config_key": "num_hidden_layers", + "pre_weights": [ + { + "name": "model.embed_tokens.weight", + "is_embed": true + } + ], + "post_weights": [ + { + "name": "model.norm.weight" + }, + { + "name": "lm_head.weight", + "is_embed": true, + "optional": true, + "tied_names": [ + "model.embed_tokens.weight" + ] + } + ], + "layer_templates": { + "weights": [ + { + "name": "model.layers.${layer_index}.input_layernorm.weight" + }, + { + "name": "model.layers.${layer_index}.mlp.down_proj.weight" + }, + { + "name": "model.layers.${layer_index}.mlp.gate_proj.weight" + }, + { + "name": "model.layers.${layer_index}.mlp.up_proj.weight" + }, + { + "name": "model.layers.${layer_index}.post_attention_layernorm.weight" + }, + { + "name": "model.layers.${layer_index}.self_attn.k_proj.bias" + }, + { + "name": "model.layers.${layer_index}.self_attn.k_proj.weight" + }, + { + "name": "model.layers.${layer_index}.self_attn.o_proj.weight" + }, + { + "name": "model.layers.${layer_index}.self_attn.q_proj.bias" + }, + { + "name": "model.layers.${layer_index}.self_attn.q_proj.weight" + }, + { + "name": "model.layers.${layer_index}.self_attn.v_proj.bias" + }, + { + "name": "model.layers.${layer_index}.self_attn.v_proj.weight" + } + ] + } + } + }, + "multi_modal_projector": { + "model_type": "", + "architecture": { + "model_type": "", + "architectures": [], + "override_num_layers": 2, + "pre_weights": [ + { + "name": "visual.merger.ln_q.weight" + }, + { + "name": "visual.merger.ln_q.bias" + } + ], + "post_weights": [], + "layer_templates": { + "weights": [ + { + "name": "visual.merger.mlp.0.weight" + }, + { + "name": "visual.merger.mlp.0.bias" + }, + { + "name": "visual.merger.mlp.2.weight" + }, + { + "name": "visual.merger.mlp.2.bias" + } + ] + } + } + }, + "vision_tower": + { + "architecture": { + "model_type": "", + "architectures": [], + "num_layers_config_key": "vision_config.depth", + "pre_weights": [ + { + "name": "visual.patch_embed.proj.weight", + "is_embed": true + } + ], + "post_weights": [], + "layer_templates": { + "weights": [ + { + "name": "visual.blocks.${layer_index}.norm1.weight" + }, + { + "name": "visual.blocks.${layer_index}.norm1.bias" + }, + { + "name": "visual.blocks.${layer_index}.norm2.weight" + }, + { + "name": "visual.blocks.${layer_index}.norm2.bias" + }, + { + "name": "visual.blocks.${layer_index}.attn.qkv.weight" + }, + { + "name": "visual.blocks.${layer_index}.attn.qkv.bias" + }, + { + "name": "visual.blocks.${layer_index}.attn.proj.weight" + }, + { + "name": "visual.blocks.${layer_index}.attn.proj.bias" + }, + { + "name": "visual.blocks.${layer_index}.mlp.fc1.weight" + }, + { + "name": "visual.blocks.${layer_index}.mlp.fc1.bias" + }, + { + "name": "visual.blocks.${layer_index}.mlp.fc2.weight" + }, + { + "name": "visual.blocks.${layer_index}.mlp.fc2.bias" + } + ] + } + } + } + } +} diff --git a/src/mindnlp/wizard/merge/_data/architectures/qwen3.json b/src/mindnlp/wizard/merge/_data/architectures/qwen3.json new file mode 100644 index 000000000..917f287c2 --- /dev/null +++ b/src/mindnlp/wizard/merge/_data/architectures/qwen3.json @@ -0,0 +1,63 @@ +{ + "model_type": "qwen3", + "architectures": [ + "Qwen3ForCausalLM" + ], + "pre_weights": [ + { + "name": "model.embed_tokens.weight", + "is_embed": true + } + ], + "post_weights": [ + { + "name": "model.norm.weight" + }, + { + "name": "lm_head.weight", + "is_embed": true, + "optional": true, + "tied_names": [ + "model.embed_tokens.weight" + ] + } + ], + "num_layers_config_key": "num_hidden_layers", + "layer_templates": { + "weights": [ + { + "name": "model.layers.${layer_index}.input_layernorm.weight" + }, + { + "name": "model.layers.${layer_index}.mlp.down_proj.weight" + }, + { + "name": "model.layers.${layer_index}.mlp.gate_proj.weight" + }, + { + "name": "model.layers.${layer_index}.mlp.up_proj.weight" + }, + { + "name": "model.layers.${layer_index}.post_attention_layernorm.weight" + }, + { + "name": "model.layers.${layer_index}.self_attn.k_norm.weight" + }, + { + "name": "model.layers.${layer_index}.self_attn.k_proj.weight" + }, + { + "name": "model.layers.${layer_index}.self_attn.q_norm.weight" + }, + { + "name": "model.layers.${layer_index}.self_attn.q_proj.weight" + }, + { + "name": "model.layers.${layer_index}.self_attn.v_proj.weight" + }, + { + "name": "model.layers.${layer_index}.self_attn.o_proj.weight" + } + ] + } +} diff --git a/src/mindnlp/wizard/merge/_data/architectures/roberta-masked-lm.json b/src/mindnlp/wizard/merge/_data/architectures/roberta-masked-lm.json new file mode 100644 index 000000000..1aae76a1b --- /dev/null +++ b/src/mindnlp/wizard/merge/_data/architectures/roberta-masked-lm.json @@ -0,0 +1,107 @@ +{ + "model_type": "roberta", + "architectures": [ + "RobertaForMaskedLM" + ], + "pre_weights": [ + { + "name": "roberta.embeddings.position_embeddings.weight" + }, + { + "name": "roberta.embeddings.word_embeddings.weight", + "is_embed": true + }, + { + "name": "roberta.embeddings.token_type_embeddings.weight" + }, + { + "name": "roberta.embeddings.LayerNorm.weight" + }, + { + "name": "roberta.embeddings.LayerNorm.bias" + }, + { + "name": "roberta.embeddings.position_ids", + "optional": true, + "force_dtype": "int64" + } + ], + "post_weights": [ + { + "name": "lm_head.bias" + }, + { + "name": "lm_head.dense.weight" + }, + { + "name": "lm_head.dense.bias" + }, + { + "name": "lm_head.layer_norm.weight" + }, + { + "name": "lm_head.layer_norm.bias" + }, + { + "name": "lm_head.decoder.weight", + "is_embed": true, + "optional": true, + "tied_names": [ + "roberta.embeddings.word_embeddings.weight" + ] + } + ], + "num_layers_config_key": "num_hidden_layers", + "layer_templates": { + "weights": [ + { + "name": "roberta.encoder.layer.${layer_index}.attention.output.dense.weight" + }, + { + "name": "roberta.encoder.layer.${layer_index}.attention.output.dense.bias" + }, + { + "name": "roberta.encoder.layer.${layer_index}.attention.output.LayerNorm.weight" + }, + { + "name": "roberta.encoder.layer.${layer_index}.attention.output.LayerNorm.bias" + }, + { + "name": "roberta.encoder.layer.${layer_index}.attention.self.query.weight" + }, + { + "name": "roberta.encoder.layer.${layer_index}.attention.self.query.bias" + }, + { + "name": "roberta.encoder.layer.${layer_index}.attention.self.key.weight" + }, + { + "name": "roberta.encoder.layer.${layer_index}.attention.self.key.bias" + }, + { + "name": "roberta.encoder.layer.${layer_index}.attention.self.value.weight" + }, + { + "name": "roberta.encoder.layer.${layer_index}.attention.self.value.bias" + }, + { + "name": "roberta.encoder.layer.${layer_index}.intermediate.dense.weight" + }, + { + "name": "roberta.encoder.layer.${layer_index}.intermediate.dense.bias" + }, + { + "name": "roberta.encoder.layer.${layer_index}.output.dense.weight" + }, + { + "name": "roberta.encoder.layer.${layer_index}.output.dense.bias" + }, + { + "name": "roberta.encoder.layer.${layer_index}.output.LayerNorm.weight" + }, + { + "name": "roberta.encoder.layer.${layer_index}.output.LayerNorm.bias" + } + ] + } +} diff --git a/src/mindnlp/wizard/merge/_data/architectures/roberta-sequence-classification.json b/src/mindnlp/wizard/merge/_data/architectures/roberta-sequence-classification.json new file mode 100644 index 000000000..6e137e6d8 --- /dev/null +++ b/src/mindnlp/wizard/merge/_data/architectures/roberta-sequence-classification.json @@ -0,0 +1,95 @@ +{ + "model_type": "roberta", + "architectures": [ + "RobertaForSequenceClassification" + ], + "pre_weights": [ + { + "name": "roberta.embeddings.position_embeddings.weight" + }, + { + "name": "roberta.embeddings.word_embeddings.weight" + }, + { + "name": "roberta.embeddings.token_type_embeddings.weight" + }, + { + "name": "roberta.embeddings.LayerNorm.weight" + }, + { + "name": "roberta.embeddings.LayerNorm.bias" + }, + { + "name": "roberta.embeddings.position_ids", + "optional": true, + "force_dtype": "int64" + } + ], + "post_weights": [ + { + "name": "classifier.dense.weight" + }, + { + "name": "classifier.dense.bias" + }, + { + "name": "classifier.out_proj.weight" + }, + { + "name": "classifier.out_proj.bias" + } + ], + "num_layers_config_key": "num_hidden_layers", + "layer_templates": { + "weights": [ + { + "name": "roberta.encoder.layer.${layer_index}.attention.output.dense.weight" + }, + { + "name": "roberta.encoder.layer.${layer_index}.attention.output.dense.bias" + }, + { + "name": "roberta.encoder.layer.${layer_index}.attention.output.LayerNorm.weight" + }, + { + "name": "roberta.encoder.layer.${layer_index}.attention.output.LayerNorm.bias" + }, + { + "name": "roberta.encoder.layer.${layer_index}.attention.self.query.weight" + }, + { + "name": "roberta.encoder.layer.${layer_index}.attention.self.query.bias" + }, + { + "name": "roberta.encoder.layer.${layer_index}.attention.self.key.weight" + }, + { + "name": "roberta.encoder.layer.${layer_index}.attention.self.key.bias" + }, + { + "name": "roberta.encoder.layer.${layer_index}.attention.self.value.weight" + }, + { + "name": "roberta.encoder.layer.${layer_index}.attention.self.value.bias" + }, + { + "name": "roberta.encoder.layer.${layer_index}.intermediate.dense.weight" + }, + { + "name": "roberta.encoder.layer.${layer_index}.intermediate.dense.bias" + }, + { + "name": "roberta.encoder.layer.${layer_index}.output.dense.weight" + }, + { + "name": "roberta.encoder.layer.${layer_index}.output.dense.bias" + }, + { + "name": "roberta.encoder.layer.${layer_index}.output.LayerNorm.weight" + }, + { + "name": "roberta.encoder.layer.${layer_index}.output.LayerNorm.bias" + } + ] + } +} diff --git a/src/mindnlp/wizard/merge/_data/architectures/roberta-token-classification.json b/src/mindnlp/wizard/merge/_data/architectures/roberta-token-classification.json new file mode 100644 index 000000000..2b9cae9f6 --- /dev/null +++ b/src/mindnlp/wizard/merge/_data/architectures/roberta-token-classification.json @@ -0,0 +1,89 @@ +{ + "model_type": "roberta", + "architectures": [ + "RobertaForTokenClassification" + ], + "pre_weights": [ + { + "name": "roberta.embeddings.position_embeddings.weight" + }, + { + "name": "roberta.embeddings.word_embeddings.weight" + }, + { + "name": "roberta.embeddings.token_type_embeddings.weight" + }, + { + "name": "roberta.embeddings.LayerNorm.weight" + }, + { + "name": "roberta.embeddings.LayerNorm.bias" + }, + { + "name": "roberta.embeddings.position_ids", + "optional": true, + "force_dtype": "int64" + } + ], + "post_weights": [ + { + "name": "classifier.weight" + }, + { + "name": "classifier.bias" + } + ], + "num_layers_config_key": "num_hidden_layers", + "layer_templates": { + "weights": [ + { + "name": "roberta.encoder.layer.${layer_index}.attention.output.dense.weight" + }, + { + "name": "roberta.encoder.layer.${layer_index}.attention.output.dense.bias" + }, + { + "name": "roberta.encoder.layer.${layer_index}.attention.output.LayerNorm.weight" + }, + { + "name": "roberta.encoder.layer.${layer_index}.attention.output.LayerNorm.bias" + }, + { + "name": "roberta.encoder.layer.${layer_index}.attention.self.query.weight" + }, + { + "name": "roberta.encoder.layer.${layer_index}.attention.self.query.bias" + }, + { + "name": "roberta.encoder.layer.${layer_index}.attention.self.key.weight" + }, + { + "name": "roberta.encoder.layer.${layer_index}.attention.self.key.bias" + }, + { + "name": "roberta.encoder.layer.${layer_index}.attention.self.value.weight" + }, + { + "name": "roberta.encoder.layer.${layer_index}.attention.self.value.bias" + }, + { + "name": "roberta.encoder.layer.${layer_index}.intermediate.dense.weight" + }, + { + "name": "roberta.encoder.layer.${layer_index}.intermediate.dense.bias" + }, + { + "name": "roberta.encoder.layer.${layer_index}.output.dense.weight" + }, + { + "name": "roberta.encoder.layer.${layer_index}.output.dense.bias" + }, + { + "name": "roberta.encoder.layer.${layer_index}.output.LayerNorm.weight" + }, + { + "name": "roberta.encoder.layer.${layer_index}.output.LayerNorm.bias" + } + ] + } +} diff --git a/src/mindnlp/wizard/merge/_data/architectures/roberta.json b/src/mindnlp/wizard/merge/_data/architectures/roberta.json new file mode 100644 index 000000000..cab186c0d --- /dev/null +++ b/src/mindnlp/wizard/merge/_data/architectures/roberta.json @@ -0,0 +1,89 @@ +{ + "model_type": "roberta", + "architectures": [ + "RobertaModel" + ], + "pre_weights": [ + { + "name": "roberta.embeddings.position_embeddings.weight" + }, + { + "name": "roberta.embeddings.word_embeddings.weight" + }, + { + "name": "roberta.embeddings.token_type_embeddings.weight" + }, + { + "name": "roberta.embeddings.LayerNorm.weight" + }, + { + "name": "roberta.embeddings.LayerNorm.bias" + }, + { + "name": "roberta.embeddings.position_ids", + "optional": true, + "force_dtype": "int64" + } + ], + "post_weights": [ + { + "name": "pooler.dense.weight" + }, + { + "name": "pooler.dense.bias" + } + ], + "num_layers_config_key": "num_hidden_layers", + "layer_templates": { + "weights": [ + { + "name": "roberta.encoder.layer.${layer_index}.attention.output.dense.weight" + }, + { + "name": "roberta.encoder.layer.${layer_index}.attention.output.dense.bias" + }, + { + "name": "roberta.encoder.layer.${layer_index}.attention.output.LayerNorm.weight" + }, + { + "name": "roberta.encoder.layer.${layer_index}.attention.output.LayerNorm.bias" + }, + { + "name": "roberta.encoder.layer.${layer_index}.attention.self.query.weight" + }, + { + "name": "roberta.encoder.layer.${layer_index}.attention.self.query.bias" + }, + { + "name": "roberta.encoder.layer.${layer_index}.attention.self.key.weight" + }, + { + "name": "roberta.encoder.layer.${layer_index}.attention.self.key.bias" + }, + { + "name": "roberta.encoder.layer.${layer_index}.attention.self.value.weight" + }, + { + "name": "roberta.encoder.layer.${layer_index}.attention.self.value.bias" + }, + { + "name": "roberta.encoder.layer.${layer_index}.intermediate.dense.weight" + }, + { + "name": "roberta.encoder.layer.${layer_index}.intermediate.dense.bias" + }, + { + "name": "roberta.encoder.layer.${layer_index}.output.dense.weight" + }, + { + "name": "roberta.encoder.layer.${layer_index}.output.dense.bias" + }, + { + "name": "roberta.encoder.layer.${layer_index}.output.LayerNorm.weight" + }, + { + "name": "roberta.encoder.layer.${layer_index}.output.LayerNorm.bias" + } + ] + } +} diff --git a/src/mindnlp/wizard/merge/_data/architectures/solar.json b/src/mindnlp/wizard/merge/_data/architectures/solar.json new file mode 100644 index 000000000..78fd5998f --- /dev/null +++ b/src/mindnlp/wizard/merge/_data/architectures/solar.json @@ -0,0 +1,82 @@ +{ + "model_type": "solar", + "architectures": [ + "SolarForCausalLM" + ], + "pre_weights": [ + { + "name": "model.embed_tokens.weight", + "is_embed": true, + "output_space": "running_residual" + } + ], + "num_layers_config_key": "num_hidden_layers", + "layer_templates": { + "weights": [ + { + "name": "model.layers.${layer_index}.input_layernorm.weight", + "input_space": "running_residual" + }, + { + "name": "model.layers.${layer_index}.self_attn.k_proj.weight", + "input_space": "running_residual", + "output_space": "attn_qk_${layer_index}", + "head_split": "output", + "is_kq": true + }, + { + "name": "model.layers.${layer_index}.self_attn.q_proj.weight", + "input_space": "running_residual", + "output_space": "attn_qk_${layer_index}", + "head_split": "output", + "is_kq": true + }, + { + "name": "model.layers.${layer_index}.self_attn.v_proj.weight", + "input_space": "running_residual", + "output_space": "attn_v_${layer_index}", + "head_split": "output" + }, + { + "name": "model.layers.${layer_index}.self_attn.o_proj.weight", + "input_space": "attn_v_${layer_index}", + "output_space": "running_residual", + "head_split": "input" + }, + { + "name": "model.layers.${layer_index}.post_attention_layernorm.weight", + "input_space": "running_residual" + }, + { + "name": "model.layers.${layer_index}.mlp.gate_proj.weight", + "input_space": "running_residual", + "output_space": "up_${layer_index}" + }, + { + "name": "model.layers.${layer_index}.mlp.up_proj.weight", + "input_space": "running_residual", + "output_space": "up_${layer_index}" + }, + { + "name": "model.layers.${layer_index}.mlp.down_proj.weight", + "input_space": "up_${layer_index}", + "output_space": "running_residual" + } + ] + }, + "post_weights": [ + { + "name": "model.norm.weight", + "input_space": "running_residual" + }, + { + "name": "lm_head.weight", + "input_space": "running_residual", + "is_embed": true, + "optional": true, + "tied_names": [ + "model.lm_head.weight" + ] + } + ] +} diff --git a/src/mindnlp/wizard/merge/_data/architectures/stablelm.json b/src/mindnlp/wizard/merge/_data/architectures/stablelm.json new file mode 100644 index 000000000..827d523a4 --- /dev/null +++ b/src/mindnlp/wizard/merge/_data/architectures/stablelm.json @@ -0,0 +1,98 @@ +{ + "model_type": "stablelm_epoch", + "architectures": [ + "StableLMEpochForCausalLM" + ], + "pre_weights": [ + { + "name": "model.embed_tokens.weight", + "is_embed": true, + "output_space": "h_0" + } + ], + "num_layers_config_key": "num_hidden_layers", + "layer_templates": { + "weights": [ + { + "name": "model.layers.${layer_index}.input_layernorm.weight", + "input_space": "h_${layer_index}" + }, + { + "name": "model.layers.${layer_index}.input_layernorm.bias", + "input_space": "h_${layer_index}" + }, + { + "name": "model.layers.${layer_index}.self_attn.q_proj.weight", + "input_space": "h_${layer_index}", + "output_space": "attn_qk_${layer_index}" + }, + { + "name": "model.layers.${layer_index}.self_attn.k_proj.weight", + "input_space": "h_${layer_index}", + "output_space": "attn_qk_${layer_index}" + }, + { + "name": "model.layers.${layer_index}.self_attn.v_proj.weight", + "input_space": "h_${layer_index}", + "output_space": "attn_v_${layer_index}" + }, + { + "name": "model.layers.${layer_index}.self_attn.o_proj.weight", + "input_space": "attn_v_${layer_index}", + "output_space": "post_attn_${layer_index}" + }, + { + "name": "model.layers.${layer_index}.post_attention_layernorm.weight", + "input_space": "h_a_${layer_index}" + }, + { + "name": "model.layers.${layer_index}.post_attention_layernorm.bias", + "input_space": "h_a_${layer_index}" + }, + { + "name": "model.layers.${layer_index}.mlp.up_proj.weight", + "input_space": "h_a_${layer_index}", + "output_space": "up_${layer_index}" + }, + { + "name": "model.layers.${layer_index}.mlp.gate_proj.weight", + "input_space": "h_a_${layer_index}", + "output_space": "up_${layer_index}" + }, + { + "name": "model.layers.${layer_index}.mlp.down_proj.weight", + "input_space": "up_${layer_index}", + "output_space": "post_mlp_${layer_index}" + } + ], + "procedural_spaces": [ + { + "name": "h_a_${layer_index}", + "type": "residual", + "inputs": [ + "h_${layer_index}", + "post_attn_${layer_index}" + ] + }, + { + "name": "h_${layer_index+1}", + "type": "residual", + "inputs": [ + "h_a_${layer_index}", + "post_mlp_${layer_index}" + ] + } + ] + }, + "post_weights": [ + { + "name": "model.norm.weight", + "input_space": "h_${num_layers}" + }, + { + "name": "lm_head.weight", + "input_space": "h_${num_layers}", + "is_embed": true + } + ] +} diff --git a/src/mindnlp/wizard/merge/_data/architectures/stablelm2.json b/src/mindnlp/wizard/merge/_data/architectures/stablelm2.json new file mode 100644 index 000000000..d2d2f461a --- /dev/null +++ b/src/mindnlp/wizard/merge/_data/architectures/stablelm2.json @@ -0,0 +1,74 @@ +{ + "model_type": "stablelm", + "architectures": [ + "StableLmForCausalLM" + ], + "pre_weights": [ + { + "name": "model.embed_tokens.weight", + "is_embed": true + } + ], + "post_weights": [ + { + "name": "model.norm.weight" + }, + { + "name": "model.norm.bias" + }, + { + "name": "lm_head.weight", + "is_embed": true + } + ], + "num_layers_config_key": "num_hidden_layers", + "layer_templates": { + "weights": [ + { + "name": "model.layers.${layer_index}.input_layernorm.weight" + }, + { + "name": "model.layers.${layer_index}.input_layernorm.bias" + }, + { + "name": "model.layers.${layer_index}.mlp.down_proj.weight" + }, + { + "name": "model.layers.${layer_index}.mlp.gate_proj.weight" + }, + { + "name": "model.layers.${layer_index}.mlp.up_proj.weight" + }, + { + "name": "model.layers.${layer_index}.post_attention_layernorm.weight" + }, + { + "name": "model.layers.${layer_index}.post_attention_layernorm.bias" + }, + { + "name": "model.layers.${layer_index}.self_attn.q_proj.weight" + }, + { + "name": "model.layers.${layer_index}.self_attn.q_proj.bias", + "optional": true + }, + { + "name": "model.layers.${layer_index}.self_attn.k_proj.weight" + }, + { + "name": "model.layers.${layer_index}.self_attn.k_proj.bias", + "optional": true + }, + { + "name": "model.layers.${layer_index}.self_attn.v_proj.weight" + }, + { + "name": "model.layers.${layer_index}.self_attn.v_proj.bias", + "optional": true + }, + { + "name": "model.layers.${layer_index}.self_attn.o_proj.weight" + } + ] + } +} diff --git a/src/mindnlp/wizard/merge/_data/architectures/starcoder2.json b/src/mindnlp/wizard/merge/_data/architectures/starcoder2.json new file mode 100644 index 000000000..c22668999 --- /dev/null +++ b/src/mindnlp/wizard/merge/_data/architectures/starcoder2.json @@ -0,0 +1,81 @@ +{ + "model_type": "starcoder2", + "architectures": [ + "Starcoder2ForCausalLM" + ], + "pre_weights": [ + { + "name": "model.embed_tokens.weight", + "is_embed": true + } + ], + "post_weights": [ + { + "name": "lm_head.weight", + "is_embed": true, + "optional": true, + "tied_names": [ + "model.embed_tokens.weight" + ] + }, + { + "name": "model.norm.bias" + }, + { + "name": "model.norm.weight" + } + ], + "num_layers_config_key": "num_hidden_layers", + "layer_templates": { + "weights": [ + { + "name": "model.layers.${layer_index}.input_layernorm.bias" + }, + { + "name": "model.layers.${layer_index}.input_layernorm.weight" + }, + { + "name": "model.layers.${layer_index}.self_attn.q_proj.bias" + }, + { + "name": "model.layers.${layer_index}.self_attn.q_proj.weight" + }, + { + "name": "model.layers.${layer_index}.self_attn.k_proj.bias" + }, + { + "name": "model.layers.${layer_index}.self_attn.k_proj.weight" + }, + { + "name": "model.layers.${layer_index}.self_attn.v_proj.bias" + }, + { + "name": "model.layers.${layer_index}.self_attn.v_proj.weight" + }, + { + "name": "model.layers.${layer_index}.self_attn.o_proj.bias" + }, + { + "name": "model.layers.${layer_index}.self_attn.o_proj.weight" + }, + { + "name": "model.layers.${layer_index}.post_attention_layernorm.bias" + }, + { + "name": "model.layers.${layer_index}.post_attention_layernorm.weight" + }, + { + "name": "model.layers.${layer_index}.mlp.c_fc.bias" + }, + { + "name": "model.layers.${layer_index}.mlp.c_fc.weight" + }, + { + "name": "model.layers.${layer_index}.mlp.c_proj.bias" + }, + { + "name": "model.layers.${layer_index}.mlp.c_proj.weight" + } + ] + } +} diff --git a/src/mindnlp/wizard/merge/_data/architectures/t5.json b/src/mindnlp/wizard/merge/_data/architectures/t5.json new file mode 100644 index 000000000..9ed8a7b80 --- /dev/null +++ b/src/mindnlp/wizard/merge/_data/architectures/t5.json @@ -0,0 +1,170 @@ +{ + "kind": "modular", + "architectures": [ + "T5ForConditionalGeneration" + ], + "model_type": "t5", + "modules": { + "decoder": { + "architecture": { + "model_type": "", + "architectures": [], + "pre_weights": [ + { + "name": "decoder.embed_tokens.weight", + "is_embed": true, + "optional": true, + "tied_names": [ + "shared.weight", + "lm_head.weight", + "encoder.embed_tokens.weight" + ] + } + ], + "num_layers_config_key": "num_decoder_layers", + "layer_templates": { + "weights": [ + { + "name": "decoder.block.${layer_index}.layer.0.layer_norm.weight" + }, + { + "name": "decoder.block.${layer_index}.layer.0.SelfAttention.q.weight" + }, + { + "name": "decoder.block.${layer_index}.layer.0.SelfAttention.k.weight" + }, + { + "name": "decoder.block.${layer_index}.layer.0.SelfAttention.v.weight" + }, + { + "name": "decoder.block.${layer_index}.layer.0.SelfAttention.o.weight" + }, + { + "name": "decoder.block.${layer_index}.layer.0.SelfAttention.relative_attention_bias.weight", + "optional": true + }, + { + "name": "decoder.block.${layer_index}.layer.1.EncDecAttention.q.weight" + }, + { + "name": "decoder.block.${layer_index}.layer.1.EncDecAttention.k.weight" + }, + { + "name": "decoder.block.${layer_index}.layer.1.EncDecAttention.v.weight" + }, + { + "name": "decoder.block.${layer_index}.layer.1.EncDecAttention.o.weight" + }, + { + "name": "decoder.block.${layer_index}.layer.1.layer_norm.weight" + }, + { + "name": "decoder.block.${layer_index}.layer.2.DenseReluDense.wi_0.weight" + }, + { + "name": "decoder.block.${layer_index}.layer.2.DenseReluDense.wi_1.weight" + }, + { + "name": "decoder.block.${layer_index}.layer.2.DenseReluDense.wo.weight" + }, + { + "name": "decoder.block.${layer_index}.layer.2.layer_norm.weight" + } + ] + }, + "post_weights": [ + { + "name": "decoder.final_layer_norm.weight" + } + ] + } + }, + "encoder": { + "architecture": { + "model_type": "", + "architectures": [], + "pre_weights": [ + { + "name": "encoder.embed_tokens.weight", + "is_embed": true, + "optional": true, + "tied_names": [ + "shared.weight", + "lm_head.weight", + "decoder.embed_tokens.weight" + ] + } + ], + "num_layers_config_key": "num_hidden_layers", + "layer_templates": { + "weights": [ + { + "name": "encoder.block.${layer_index}.layer.0.layer_norm.weight" + }, + { + "name": "encoder.block.${layer_index}.layer.0.SelfAttention.q.weight" + }, + { + "name": "encoder.block.${layer_index}.layer.0.SelfAttention.k.weight" + }, + { + "name": "encoder.block.${layer_index}.layer.0.SelfAttention.v.weight" + }, + { + "name": "encoder.block.${layer_index}.layer.0.SelfAttention.o.weight" + }, + { + "name": "encoder.block.${layer_index}.layer.0.SelfAttention.relative_attention_bias.weight", + "optional": true + }, + { + "name": "encoder.block.${layer_index}.layer.1.DenseReluDense.wi_0.weight" + }, + { + "name": "encoder.block.${layer_index}.layer.1.DenseReluDense.wi_1.weight" + }, + { + "name": "encoder.block.${layer_index}.layer.1.DenseReluDense.wo.weight" + }, + { + "name": "encoder.block.${layer_index}.layer.1.layer_norm.weight" + } + ] + }, + "post_weights": [ + { + "name": "encoder.final_layer_norm.weight" + } + ] + } + }, + "shared": { + "architecture": { + "model_type": "", + "architectures": [], + "pre_weights": [ + { + "name": "shared.weight", + "is_embed": true + } + ], + "layer_templates": { + "weights": [] + }, + "post_weights": [ + { + "name": "lm_head.weight", + "is_embed": true, + "optional": true, + "tied_names": [ + "shared.weight", + "encoder.embed_tokens.weight", + "decoder.embed_tokens.weight" + ] + } + ], + "override_num_layers": 0 + } + } + } +} diff --git a/src/mindnlp/wizard/merge/_data/architectures/whisper.json b/src/mindnlp/wizard/merge/_data/architectures/whisper.json new file mode 100644 index 000000000..82ac444b1 --- /dev/null +++ b/src/mindnlp/wizard/merge/_data/architectures/whisper.json @@ -0,0 +1,196 @@ +{ + "kind": "modular", + "architectures": [ + "WhisperForConditionalGeneration" + ], + "model_type": "whisper", + "tagalong_files": [ + "preprocessor_config.json", + "normalizer.json" + ], + "modules": { + "decoder": { + "weight_prefix": "model.decoder", + "architecture": { + "model_type": "", + "architectures": [], + "pre_weights": [ + { + "name": "embed_tokens.weight", + "is_embed": true + }, + { + "name": "embed_positions.weight" + } + ], + "num_layers_config_key": "decoder_layers", + "layer_templates": { + "weights": [ + { + "name": "layers.${layer_index}.encoder_attn.k_proj.weight" + }, + { + "name": "layers.${layer_index}.encoder_attn.out_proj.bias" + }, + { + "name": "layers.${layer_index}.encoder_attn.out_proj.weight" + }, + { + "name": "layers.${layer_index}.encoder_attn.q_proj.bias" + }, + { + "name": "layers.${layer_index}.encoder_attn.q_proj.weight" + }, + { + "name": "layers.${layer_index}.encoder_attn.v_proj.bias" + }, + { + "name": "layers.${layer_index}.encoder_attn.v_proj.weight" + }, + { + "name": "layers.${layer_index}.encoder_attn_layer_norm.bias" + }, + { + "name": "layers.${layer_index}.encoder_attn_layer_norm.weight" + }, + { + "name": "layers.${layer_index}.fc1.bias" + }, + { + "name": "layers.${layer_index}.fc1.weight" + }, + { + "name": "layers.${layer_index}.fc2.bias" + }, + { + "name": "layers.${layer_index}.fc2.weight" + }, + { + "name": "layers.${layer_index}.final_layer_norm.bias" + }, + { + "name": "layers.${layer_index}.final_layer_norm.weight" + }, + { + "name": "layers.${layer_index}.self_attn.k_proj.weight" + }, + { + "name": "layers.${layer_index}.self_attn.out_proj.bias" + }, + { + "name": "layers.${layer_index}.self_attn.out_proj.weight" + }, + { + "name": "layers.${layer_index}.self_attn.q_proj.bias" + }, + { + "name": "layers.${layer_index}.self_attn.q_proj.weight" + }, + { + "name": "layers.${layer_index}.self_attn.v_proj.bias" + }, + { + "name": "layers.${layer_index}.self_attn.v_proj.weight" + }, + { + "name": "layers.${layer_index}.self_attn_layer_norm.bias" + }, + { + "name": "layers.${layer_index}.self_attn_layer_norm.weight" + } + ] + }, + "post_weights": [ + { + "name": "layer_norm.bias" + }, + { + "name": "layer_norm.weight" + } + ] + } + }, + "encoder": { + "weight_prefix": "model.encoder.", + "architecture": { + "model_type": "", + "architectures": [], + "pre_weights": [ + { + "name": "embed_positions.weight" + }, + { + "name": "conv1.bias" + }, + { + "name": "conv1.weight" + }, + { + "name": "conv2.bias" + }, + { + "name": "conv2.weight" + } + ], + "post_weights": [ + { + "name": "layer_norm.bias" + }, + { + "name": "layer_norm.weight" + } + ], + "layer_templates": { + "weights": [ + { + "name": "layers.${layer_index}.fc1.bias" + }, + { + "name": "layers.${layer_index}.fc1.weight" + }, + { + "name": "layers.${layer_index}.fc2.bias" + }, + { + "name": "layers.${layer_index}.fc2.weight" + }, + { + "name": "layers.${layer_index}.final_layer_norm.bias" + }, + { + "name": "layers.${layer_index}.final_layer_norm.weight" + }, + { + "name": "layers.${layer_index}.self_attn.k_proj.weight" + }, + { + "name": "layers.${layer_index}.self_attn.out_proj.bias" + }, + { + "name": "layers.${layer_index}.self_attn.out_proj.weight" + }, + { + "name": "layers.${layer_index}.self_attn.q_proj.bias" + }, + { + "name": "layers.${layer_index}.self_attn.q_proj.weight" + }, + { + "name": "layers.${layer_index}.self_attn.v_proj.bias" + }, + { + "name": "layers.${layer_index}.self_attn.v_proj.weight" + }, + { + "name": "layers.${layer_index}.self_attn_layer_norm.bias" + }, + { + "name": "layers.${layer_index}.self_attn_layer_norm.weight" + } + ] + }, + "num_layers_config_key": "encoder_layers" + } + } + } +} diff --git a/src/mindnlp/wizard/merge/_data/chat_templates/__init__.py b/src/mindnlp/wizard/merge/_data/chat_templates/__init__.py new file mode 100644 index 000000000..6babb0e71 --- /dev/null +++ b/src/mindnlp/wizard/merge/_data/chat_templates/__init__.py @@ -0,0 +1,4 @@ +# Originally from MergeKit (https://github.com/arcee-ai/mergekit) +# Copyright (C) 2025 Arcee AI +# SPDX-License-Identifier: LGPL-3.0-only +# Modified for MindSpore/Ascend NPU by MindNLP Wizard contributors. diff --git a/src/mindnlp/wizard/merge/_data/chat_templates/alpaca.jinja b/src/mindnlp/wizard/merge/_data/chat_templates/alpaca.jinja new file mode 100644 index 000000000..45837b0af --- /dev/null +++ b/src/mindnlp/wizard/merge/_data/chat_templates/alpaca.jinja @@ -0,0 +1,29 @@ +{{ (messages|selectattr('role', 'equalto', 'system')|list|last).content|trim if (messages|selectattr('role', 'equalto', 'system')|list) else '' }} + +{% for message in messages %} +{% if message['role'] == 'user' %} +### Instruction: +{{ message['content']|trim -}} +{% if not loop.last %} + + +{% endif %} +{% elif message['role'] == 'assistant' %} +### Response: +{{ message['content']|trim -}} +{% if not loop.last %} + + +{% endif %} +{% elif message['role'] == 'user_context' %} +### Input: +{{ message['content']|trim -}} +{% if not loop.last %} + + +{% endif %} +{% endif %} +{% endfor %} +{% if add_generation_prompt and messages[-1]['role'] != 'assistant' %} +### Response: +{% endif %} diff --git a/src/mindnlp/wizard/merge/_data/chat_templates/chatml.jinja b/src/mindnlp/wizard/merge/_data/chat_templates/chatml.jinja new file mode 100644 index 000000000..4f3444551 --- /dev/null +++ b/src/mindnlp/wizard/merge/_data/chat_templates/chatml.jinja @@ -0,0 +1,2 @@ +{% for message in messages %}{{'<|im_start|>' + message['role'] + '\n' + message['content'] + '<|im_end|>' + '\n'}}{% endfor %} +{% if add_generation_prompt %}{{ '<|im_start|>assistant\n' }}{% endif %} diff --git a/src/mindnlp/wizard/merge/_data/chat_templates/exaone.jinja b/src/mindnlp/wizard/merge/_data/chat_templates/exaone.jinja new file mode 100644 index 000000000..3a4d07ae3 --- /dev/null +++ b/src/mindnlp/wizard/merge/_data/chat_templates/exaone.jinja @@ -0,0 +1,14 @@ +{% for message in messages %} + {% if loop.first and message['role'] != 'system' %} + {{ '[|system|][|endofturn|]\n' }} + {% endif %} + {{ '[|' + message['role'] + '|]' + message['content'] }} + {% if message['role'] == 'user' %} + {{ '\n' }} + {% else %} + {{ '[|endofturn|]\n' }} + {% endif %} +{% endfor %} +{% if add_generation_prompt %} + {{ '[|assistant|]' }} +{% endif %} diff --git a/src/mindnlp/wizard/merge/_data/chat_templates/llama3.jinja b/src/mindnlp/wizard/merge/_data/chat_templates/llama3.jinja new file mode 100644 index 000000000..0fcec78aa --- /dev/null +++ b/src/mindnlp/wizard/merge/_data/chat_templates/llama3.jinja @@ -0,0 +1,7 @@ +{% set loop_messages = messages %} +{% for message in loop_messages %} +{% set content = '<|start_header_id|>' + message['role'] + '<|end_header_id|>\n\n'+ message['content'] | trim + '<|eot_id|>' %} +{% if loop.index0 == 0 %}{% set content = bos_token + content %}{% endif %} +{{ content }} +{% endfor %} +{% if add_generation_prompt %}{{ '<|start_header_id|>assistant<|end_header_id|>\n\n' }}{% endif %} diff --git a/src/mindnlp/wizard/merge/_data/chat_templates/mistral.jinja b/src/mindnlp/wizard/merge/_data/chat_templates/mistral.jinja new file mode 100644 index 000000000..40b37ad7f --- /dev/null +++ b/src/mindnlp/wizard/merge/_data/chat_templates/mistral.jinja @@ -0,0 +1,24 @@ +{%- if messages[0]['role'] == 'system' %} + {%- set system_message = messages[0]['content'] %} + {%- set loop_messages = messages[1:] %} +{%- else %} + {%- set loop_messages = messages %} +{%- endif %} + +{{- bos_token }} +{%- for message in loop_messages %} + {%- if (message['role'] == 'user') != (loop.index0 % 2 == 0) %} + {{- raise_exception('After the optional system message, conversation roles must alternate user/assistant/user/assistant/...') }} + {%- endif %} + {%- if message['role'] == 'user' %} + {%- if loop.first and system_message is defined %} + {{- ' [INST] ' + system_message + '\n\n' + message['content'] + ' [/INST]' }} + {%- else %} + {{- ' [INST] ' + message['content'] + ' [/INST]' }} + {%- endif %} + {%- elif message['role'] == 'assistant' %} + {{- ' ' + message['content'] + eos_token}} + {%- else %} + {{- raise_exception('Only user and assistant roles are supported, with the exception of an initial optional system message!') }} + {%- endif %} +{%- endfor %} diff --git a/src/mindnlp/wizard/merge/architecture/__init__.py b/src/mindnlp/wizard/merge/architecture/__init__.py new file mode 100644 index 000000000..0fb9d58bf --- /dev/null +++ b/src/mindnlp/wizard/merge/architecture/__init__.py @@ -0,0 +1,121 @@ +# Copyright (C) 2025 Arcee AI +# SPDX-License-Identifier: LGPL-3.0-only +# Modified for MindSpore/Ascend NPU by MindNLP Wizard contributors. +# + +import logging +from functools import lru_cache +from typing import TYPE_CHECKING, Optional + +from transformers import PretrainedConfig + +from .auto import infer_architecture_info +from .base import ( + ConfiguredModelArchitecture, + ConfiguredModuleArchitecture, + ModelArchitecture, + ModuleArchitecture, + ModuleDefinition, + WeightInfo, +) +from .json_definitions import NAME_TO_ARCH +from .moe_defs import ( + AfmoeModuleArchitecture, + Glm4MoeModuleArchitecture, + MixtralModuleArchitecture, + Qwen3MoeModuleArchitecture, +) +from ..options import MergeOptions + +if TYPE_CHECKING: + from ..config import MergeConfiguration + +LOG = logging.getLogger(__name__) + +WARNED_ARCHITECTURE_NAMES = set() + + +def arch_info_for_config(config: PretrainedConfig) -> Optional[ModelArchitecture]: + if len(config.architectures) != 1: + raise RuntimeError("More than one architecture in config?") + arch_name = config.architectures[0] + + if arch_name == MixtralModuleArchitecture.ARCHITECTURE_NAME: + module = MixtralModuleArchitecture.from_config(config) + return ModelArchitecture( + modules={"default": ModuleDefinition(architecture=module)}, + architectures=[arch_name], + model_type="mixtral", + ) + elif arch_name == Qwen3MoeModuleArchitecture.ARCHITECTURE_NAME: + module = Qwen3MoeModuleArchitecture.from_config(config) + return ModelArchitecture( + modules={"default": ModuleDefinition(architecture=module)}, + architectures=[arch_name], + model_type="qwen3_moe", + ) + elif arch_name == AfmoeModuleArchitecture.ARCHITECTURE_NAME: + module = AfmoeModuleArchitecture.from_config(config) + return ModelArchitecture( + modules={"default": ModuleDefinition(architecture=module)}, + architectures=[arch_name], + model_type="afmoe", + ) + elif arch_name == Glm4MoeModuleArchitecture.ARCHITECTURE_NAME: + module = Glm4MoeModuleArchitecture.from_config(config) + return ModelArchitecture( + modules={"default": ModuleDefinition(architecture=module)}, + architectures=[arch_name], + model_type="glm4_moe", + ) + elif arch_name in NAME_TO_ARCH: + candidates = list(NAME_TO_ARCH[arch_name]) + if len(candidates) == 1: + return candidates[0] + + for c in candidates: + if c.expected_model_type == config.model_type: + return c + LOG.warning( + f"Multiple architectures for {arch_name}, none match model type {config.model_type}" + ) + + if arch_name not in WARNED_ARCHITECTURE_NAMES: + LOG.warning(f"No JSON architecture found for {arch_name}") + WARNED_ARCHITECTURE_NAMES.add(arch_name) + return None + + +def get_architecture_info( + config: "MergeConfiguration", options: MergeOptions +) -> ModelArchitecture: + models = config.referenced_models() + if not models: + raise ValueError("No models referenced in config") + + model_arch_info = [ + arch_info_for_config(m.config(trust_remote_code=options.trust_remote_code)) + for m in models + ] + if all(arch is not None for arch in model_arch_info): + if not options.allow_crimes and any( + arch != model_arch_info[0] for arch in model_arch_info + ): + raise RuntimeError( + "Must specify --allow-crimes to attempt to mix different architectures" + ) + return model_arch_info[0] + + return infer_architecture_info(tuple(models), config.base_model, options) + + +__all__ = [ + "ModelArchitecture", + "ModuleArchitecture", + "ModuleDefinition", + "ConfiguredModuleArchitecture", + "ConfiguredModelArchitecture", + "WeightInfo", + "get_architecture_info", + "arch_info_for_config", +] diff --git a/src/mindnlp/wizard/merge/architecture/auto.py b/src/mindnlp/wizard/merge/architecture/auto.py new file mode 100644 index 000000000..323e227f8 --- /dev/null +++ b/src/mindnlp/wizard/merge/architecture/auto.py @@ -0,0 +1,203 @@ +# Copyright (C) 2025 Arcee AI +# SPDX-License-Identifier: LGPL-3.0-only +# Modified for MindSpore/Ascend NPU by MindNLP Wizard contributors. +# +# A local torch import is retained inside get_transformers_info() because +# that function introspects HuggingFace transformers models (torch-based). + +import logging +import re +from collections import defaultdict +from functools import lru_cache +from typing import List, Optional, Tuple + +from .base import ( + ModelArchitecture, + ModuleDefinition, + WeightInfo, +) +from .json_definitions import ( + JsonLayerTemplates, + JsonModuleArchDef, + JsonModuleArchitecture, +) +from ..common import ModelReference, get_auto_cls +from ..options import MergeOptions + +try: + from transformers.modeling_utils import _get_tied_weight_keys +except ImportError: + _get_tied_weight_keys = None + +RE_LAYER_INDEX = re.compile(r"\.(\d+)\.") + +LOG = logging.getLogger(__name__) + + +def get_model_tensor_names(model: ModelReference, options: MergeOptions) -> List[str]: + loader = model.lazy_loader( + cache_dir=options.transformers_cache, lazy_loader=options.lazy_loader + ) + return list(loader.index.tensor_paths.keys()) + + +def get_transformers_info(model: ModelReference, options: MergeOptions) -> tuple: + import torch # needed for HuggingFace transformers model introspection + + try: + cfg = model.config( + trust_remote_code=options.trust_remote_code, + ) + auto_cls = get_auto_cls(cfg.architectures[0]) + except Exception as e: + LOG.warning( + f"Unable to load config for {model.model} - tied/ignored weights will not be detected", + exc_info=e, + ) + return None, None, set() + try: + with torch.device("meta"): + model = auto_cls.from_pretrained( + model.model.path, + revision=model.model.revision, + trust_remote_code=options.trust_remote_code, + device_map="meta", + ) + except Exception as e: + LOG.warning( + f"Unable to load model {model.model} with transformers - tied/ignored weights will not be detected", + exc_info=e, + ) + return None, None, set() + + ignore_on_save = getattr(model, "_keys_to_ignore_on_save", None) + if _get_tied_weight_keys is None: + LOG.warning( + "Unable to get tied weights - incompatible transformers version", + ) + tied_keys = None + else: + tied_keys = _get_tied_weight_keys(model) + if ignore_on_save is not None: + ignore_on_save = set(ignore_on_save) + + embed_names = set() + _embed_out = model.get_output_embeddings() + _embed_in = model.get_input_embeddings() + for name, module in model.named_modules(): + if ( + isinstance(module, torch.nn.Embedding) + or module == _embed_out + or module == _embed_in + ): + embed_names.add(name + ".weight") + return ignore_on_save, tied_keys, embed_names + + +@lru_cache(maxsize=128) +def infer_architecture_info( + models: Tuple[ModelReference, ...], + base_model: Optional[ModelReference], + options: MergeOptions, +) -> ModelArchitecture: + model_tensor_names = { + model: set(get_model_tensor_names(model, options)) + for model in (set(models).union({base_model} if base_model else {})) + } + models = list(models) + if base_model is None: + base_model = models.pop(0) + all_tensor_names = set().union(*model_tensor_names.values()) + in_all_models = all_tensor_names.intersection(*model_tensor_names.values()) + + ignore_on_save, tied_keys, embed_names = get_transformers_info(base_model, options) + + module_prefixes = set() + module_layer_counts = defaultdict(int) + module_templates = defaultdict(set) + module_loose_weights = defaultdict(set) + + for tensor_name in all_tensor_names: + if ignore_on_save and tensor_name in ignore_on_save: + continue + if match := RE_LAYER_INDEX.search(tensor_name): + prefix = tensor_name[: match.start()] + module_prefixes.add(prefix) + layer_idx = int(match.group(1)) + module_layer_counts[prefix] = max( + module_layer_counts[prefix], layer_idx + 1 + ) + module_templates[prefix] = module_templates[prefix].union( + set([RE_LAYER_INDEX.sub(".${layer_index}.", tensor_name, count=1)]) + ) + + if len(module_prefixes) == 1: + prefix = module_prefixes.pop() + module_templates = {"": module_templates[prefix]} + module_layer_counts = {"": module_layer_counts[prefix]} + module_loose_weights = {"": module_loose_weights[prefix]} + module_prefixes = {""} + else: + module_prefixes.add("") + + for tensor_name in all_tensor_names: + if RE_LAYER_INDEX.search(tensor_name): + continue + for prefix in module_prefixes: + if tensor_name.startswith(prefix): + module_loose_weights[prefix].add(tensor_name[len(prefix) :]) + break + + if "" in module_prefixes and not (module_loose_weights[""] or module_templates[""]): + module_prefixes.remove("") + if not module_prefixes: + raise ValueError("No modules found in models") + + logging.warning(f"Inferred {len(module_prefixes)} modules:") + for prefix in module_prefixes: + logging.warning( + f" {repr(prefix or 'default')} with {module_layer_counts[prefix]} layers, {len(module_templates[prefix])} templates, and {len(module_loose_weights[prefix])} loose weights" + ) + + def _wi(template: str, prefix: str) -> WeightInfo: + full_name = prefix + template + optional = (full_name.replace("${layer_index}", "0") not in in_all_models) or ( + tied_keys is not None + and any(re.search(pat, full_name) for pat in tied_keys) + ) + is_embed = (full_name in embed_names) or ( + tied_keys is not None + and any(re.search(pat, full_name) for pat in tied_keys) + ) + return WeightInfo( + name=template, + optional=optional, + is_embed=is_embed, + ) + + module_archs = {} + for prefix in module_prefixes: + num_layers = module_layer_counts[prefix] + module_archs[prefix or "default"] = JsonModuleArchitecture( + definition=JsonModuleArchDef( + model_type="", + architectures=[], + pre_weights=[_wi(t, "") for t in module_loose_weights[prefix]], + layer_templates=JsonLayerTemplates( + weights=[_wi(t, "") for t in module_templates[prefix]] + ), + post_weights=[], + num_layers_config_key=None, + override_num_layers=num_layers, + ), + ) + + res = ModelArchitecture( + modules={ + key: ModuleDefinition(architecture=value) + for key, value in module_archs.items() + }, + architectures=[], + model_type="", + ) + return res diff --git a/src/mindnlp/wizard/merge/architecture/base.py b/src/mindnlp/wizard/merge/architecture/base.py new file mode 100644 index 000000000..f766a5abf --- /dev/null +++ b/src/mindnlp/wizard/merge/architecture/base.py @@ -0,0 +1,152 @@ +# Copyright (C) 2025 Arcee AI +# SPDX-License-Identifier: LGPL-3.0-only +# Modified for MindSpore/Ascend NPU by MindNLP Wizard contributors. +# + +from abc import ABC, abstractmethod +from typing import Dict, List, Optional, Tuple + +from pydantic import BaseModel, Field +from transformers import PretrainedConfig + +from ..common import get_config_value + + +class WeightInfo(BaseModel, frozen=True): + """Information about an individual weight tensor in a model. + + Attributes: + name (str): + The name of the tensor representing the weight. + is_embed (bool): + Indicates whether the weight is for an embedding or language model head. + optional (bool): + Indicates whether the weight can be omitted from a model. + aliases (Optional[List[str]]): + List of alternative names for the weight, if applicable. + force_dtype (Optional[str]): + Mandatory dtype for the weight, if applicable. + """ + + name: str + is_embed: bool = False + optional: bool = False + aliases: Optional[Tuple[str, ...]] = None + force_dtype: Optional[str] = None + tied_names: Optional[Tuple[str, ...]] = None + + +def _prefix_weight(weight: WeightInfo, prefix: Optional[str] = None) -> WeightInfo: + if prefix is None: + return weight + return WeightInfo( + name=prefix + weight.name, + aliases=tuple(prefix + alias for alias in weight.aliases or ()) or None, + tied_names=tuple(prefix + tied_name for tied_name in weight.tied_names or ()) + or None, + **weight.model_dump(exclude={"name", "aliases", "tied_names"}), + ) + + +class ModuleArchitecture(ABC): + @abstractmethod + def pre_weights(self, config: PretrainedConfig) -> List[WeightInfo]: + """Return a list of all weights preceding the first layer.""" + + @abstractmethod + def post_weights(self, config: PretrainedConfig) -> List[WeightInfo]: + """Return a list of all weights following the final layer.""" + + @abstractmethod + def layer_weights( + self, index: int, config: PretrainedConfig + ) -> Optional[List[WeightInfo]]: + """Return a list of all weights associated with a given layer.""" + + def num_layers_config_key(self) -> str: + """Key in config that represents number of layers""" + return "num_hidden_layers" + + def num_layers(self, config: PretrainedConfig) -> int: + """Return the number of layers in a model.""" + return get_config_value(config, self.num_layers_config_key()) + + def all_weights(self, config: PretrainedConfig) -> List[WeightInfo]: + """Return all weights associated with a model.""" + num_layers = self.num_layers(config) + res = list(self.pre_weights(config)) + for layer_idx in range(num_layers): + res.extend(self.layer_weights(layer_idx, config)) + res.extend(self.post_weights(config)) + return res + + +class ConfiguredModuleArchitecture( + BaseModel, frozen=True, arbitrary_types_allowed=True +): + info: ModuleArchitecture + config: PretrainedConfig + weight_prefix: Optional[str] = None + + def num_layers(self) -> int: + return self.info.num_layers(self.config) + + def pre_weights(self) -> List[WeightInfo]: + return [ + _prefix_weight(w, self.weight_prefix) + for w in self.info.pre_weights(self.config) + ] + + def post_weights(self) -> List[WeightInfo]: + return [ + _prefix_weight(w, self.weight_prefix) + for w in self.info.post_weights(self.config) + ] + + def layer_weights(self, index: int) -> List[WeightInfo]: + return [ + _prefix_weight(w, self.weight_prefix) + for w in self.info.layer_weights(index, self.config) + ] + + def all_weights(self) -> List[WeightInfo]: + return [ + _prefix_weight(w, self.weight_prefix) + for w in self.info.all_weights(self.config) + ] + + +class ModuleDefinition(BaseModel, frozen=True, arbitrary_types_allowed=True): + architecture: ModuleArchitecture + weight_prefix: Optional[str] = None + subfolder: Optional[str] = None + + +class ModelArchitecture(BaseModel, frozen=True): + modules: Dict[str, ModuleDefinition] + architectures: List[str] + expected_model_type: str = Field(alias="model_type") + tagalong_files: Optional[List[str]] = None + vocab_size_config_key: Optional[str] = None + + def all_weights(self, config: PretrainedConfig) -> List[WeightInfo]: + res = [] + for module in self.modules.values(): + for weight_info in module.architecture.all_weights(config=config): + res.append(_prefix_weight(weight_info, module.weight_prefix)) + return res + + +class ConfiguredModelArchitecture(BaseModel, frozen=True, arbitrary_types_allowed=True): + info: ModelArchitecture + config: PretrainedConfig + + def all_weights(self) -> List[WeightInfo]: + return self.info.all_weights(self.config) + + def get_module(self, module_name: str) -> ConfiguredModuleArchitecture: + return ConfiguredModuleArchitecture( + info=self.info.modules[module_name].architecture, + config=self.config, + weight_prefix=self.info.modules[module_name].weight_prefix, + ) diff --git a/src/mindnlp/wizard/merge/architecture/json_definitions.py b/src/mindnlp/wizard/merge/architecture/json_definitions.py new file mode 100644 index 000000000..ff37fb871 --- /dev/null +++ b/src/mindnlp/wizard/merge/architecture/json_definitions.py @@ -0,0 +1,192 @@ +# Copyright (C) 2025 Arcee AI +# SPDX-License-Identifier: LGPL-3.0-only +# Modified for MindSpore/Ascend NPU by MindNLP Wizard contributors. +# +# instead of importlib.resources on the mergekit package. + +import json +import string +from pathlib import Path +from typing import Dict, List, Optional, Tuple + +from pydantic import BaseModel, Field +from transformers import PretrainedConfig +from typing_extensions import Literal + +from .base import ( + ModelArchitecture, + ModuleArchitecture, + ModuleDefinition, + WeightInfo, +) + +_DATA_DIR = (Path(__file__).resolve().parent.parent / "_data" / "architectures") + + +class JsonLayerTemplates(BaseModel, frozen=True): + weights: List[WeightInfo] + + +class JsonModuleArchDef(BaseModel, frozen=True): + expected_model_type: str = Field(alias="model_type") + architectures: List[str] + pre_weights: List[WeightInfo] + layer_templates: JsonLayerTemplates + post_weights: List[WeightInfo] + num_layers_config_key: Optional[str] = None + override_num_layers: Optional[int] = None + + +class JsonModuleArchitecture(ModuleArchitecture, BaseModel, frozen=True): + kind: Literal["module"] = "module" + definition: JsonModuleArchDef + + def _substitute( + self, + item: WeightInfo, + config: PretrainedConfig, + layer_idx: Optional[int] = None, + ) -> WeightInfo: + num_layers = self.num_layers(config) + + obj_dict = item.model_dump(mode="json", exclude_unset=True) + for key in obj_dict: + if isinstance(obj_dict[key], str): + obj_dict[key] = _template_substitution( + obj_dict[key], num_layers, layer_idx + ) + elif isinstance(obj_dict[key], list): + obj_dict[key] = [ + ( + _template_substitution(s, num_layers, layer_idx) + if isinstance(s, str) + else s + ) + for s in obj_dict[key] + ] + return type(item).model_validate(obj_dict) + + def name(self) -> str: + return self.definition.expected_model_type + + def pre_weights(self, config: PretrainedConfig) -> List[WeightInfo]: + return [ + self._substitute(wi, config=config) for wi in self.definition.pre_weights + ] + + def layer_weights( + self, index: int, config: PretrainedConfig + ) -> Optional[List[WeightInfo]]: + return [ + self._substitute(wi, config=config, layer_idx=index) + for wi in self.definition.layer_templates.weights + ] + + def post_weights(self, config: PretrainedConfig) -> List[WeightInfo]: + return [ + self._substitute(wi, config=config) for wi in self.definition.post_weights + ] + + def num_layers_config_key(self) -> str: + return self.definition.num_layers_config_key + + def num_layers(self, config): + if self.definition.override_num_layers is not None: + return self.definition.override_num_layers + return super().num_layers(config) + + +class JsonModuleDefinition(BaseModel, frozen=True): + architecture: JsonModuleArchDef + weight_prefix: Optional[str] = None + subfolder: Optional[str] = None + + +class JsonModularArchitectureDefinition(BaseModel, frozen=True): + kind: Literal["modular"] + modules: Dict[str, JsonModuleDefinition] + architectures: List[str] + expected_model_type: str = Field(alias="model_type") + tagalong_files: Optional[List[str]] = None + vocab_size_config_key: Optional[str] = None + + +class TemplateWithArithmetic(string.Template): + idpattern = r"(?a:[_a-z][_a-z0-9]*([+-]1)?)" + + +def _template_substitution( + template: str, num_layers: int, layer_idx: Optional[int] = None +) -> str: + if "{" not in template: + return template + + substitutions = { + "num_layers": num_layers, + "num_layers+1": num_layers + 1, + "num_layers-1": num_layers - 1, + } + + if layer_idx is not None: + substitutions.update( + { + "layer_index": layer_idx, + "layer_index+1": layer_idx + 1, + "layer_index-1": layer_idx - 1, + } + ) + + return TemplateWithArithmetic(template).substitute(substitutions) + + +def _load_architecture_json(text: str) -> ModelArchitecture: + data = json.loads(text) + kind = data.get("kind", "module") + if kind == "modular": + parsed = JsonModularArchitectureDefinition.model_validate_json(text) + return ModelArchitecture( + modules={ + k: ModuleDefinition( + architecture=JsonModuleArchitecture(definition=v.architecture), + weight_prefix=v.weight_prefix, + subfolder=v.subfolder, + ) + for k, v in parsed.modules.items() + }, + architectures=parsed.architectures, + model_type=parsed.expected_model_type, + tagalong_files=parsed.tagalong_files, + vocab_size_config_key=parsed.vocab_size_config_key, + ) + elif data.get("kind", "module") == "module": + module = JsonModuleArchitecture( + definition=JsonModuleArchDef.model_validate(data) + ) + return ModelArchitecture( + modules={"default": ModuleDefinition(architecture=module)}, + architectures=module.definition.architectures, + model_type=module.definition.expected_model_type, + ) + else: + raise RuntimeError(f"Unexpected architecture kind: {data['kind']}") + + +def _load_all_architectures() -> ( + Tuple[List[ModelArchitecture], Dict[str, List[ModelArchitecture]]] +): + architectures: List[ModelArchitecture] = [] + if _DATA_DIR.is_dir(): + for f in _DATA_DIR.iterdir(): + if f.is_file() and f.name.lower().endswith(".json"): + text = f.read_text() + architectures.append(_load_architecture_json(text)) + + name_to_arch: Dict[str, List[JsonModuleArchitecture]] = {} + for arch_info in architectures: + for arch_name in arch_info.architectures: + name_to_arch[arch_name] = name_to_arch.get(arch_name, []) + name_to_arch[arch_name].append(arch_info) + return architectures, name_to_arch + + +JSON_ARCHITECTURES, NAME_TO_ARCH = _load_all_architectures() diff --git a/src/mindnlp/wizard/merge/architecture/moe_defs.py b/src/mindnlp/wizard/merge/architecture/moe_defs.py new file mode 100644 index 000000000..daceeb096 --- /dev/null +++ b/src/mindnlp/wizard/merge/architecture/moe_defs.py @@ -0,0 +1,207 @@ +# Copyright (C) 2025 Arcee AI +# SPDX-License-Identifier: LGPL-3.0-only +# Modified for MindSpore/Ascend NPU by MindNLP Wizard contributors. +# + +from typing import ClassVar, List, Optional + +from pydantic import BaseModel +from transformers import PretrainedConfig + +from .base import ( + ModuleArchitecture, + WeightInfo, +) +from .json_definitions import NAME_TO_ARCH + +MISTRAL_INFO = NAME_TO_ARCH["MistralForCausalLM"][0] +MISTRAL_MODULE_ARCH = MISTRAL_INFO.modules["default"].architecture + + +class MixtralModuleArchitecture(ModuleArchitecture, BaseModel): + ARCHITECTURE_NAME: ClassVar[str] = "MixtralForCausalLM" + num_local_experts: int + + def name(self) -> str: + return "mixtral" + + @classmethod + def from_config(cls, config: PretrainedConfig): + return MixtralModuleArchitecture(num_local_experts=config.num_local_experts) + + def pre_weights(self, config: PretrainedConfig) -> List[WeightInfo]: + return MISTRAL_MODULE_ARCH.pre_weights(config) + + def post_weights(self, config: PretrainedConfig) -> List[WeightInfo]: + return MISTRAL_MODULE_ARCH.post_weights(config) + + def num_layers_config_key(self) -> str: + return MISTRAL_MODULE_ARCH.num_layers_config_key() + + def layer_weights( + self, index: int, config: PretrainedConfig + ) -> Optional[List[WeightInfo]]: + num_experts = self.num_local_experts + prefix = f"model.layers.{index}" + tensor_names = [] + for expert_idx in range(num_experts): + for param in ("w1", "w2", "w3"): + tensor_names.append( + prefix + f".block_sparse_moe.experts.{expert_idx}.{param}.weight" + ) + tensor_names.append(prefix + ".block_sparse_moe.gate.weight") + res = [] + for name in tensor_names: + res.append(WeightInfo(name=name)) + for weight_info in MISTRAL_MODULE_ARCH.layer_weights(index, config): + if ".mlp." in weight_info.name: + continue + res.append(weight_info) + return res + + +QWEN3_INFO = NAME_TO_ARCH["Qwen3ForCausalLM"][0] +QWEN3_MODULE_ARCH = QWEN3_INFO.modules["default"].architecture + + +class Qwen3MoeModuleArchitecture(ModuleArchitecture, BaseModel): + ARCHITECTURE_NAME: ClassVar[str] = "Qwen3MoeForCausalLM" + num_experts: int + + def name(self) -> str: + return "qwen3_moe" + + @classmethod + def from_config(cls, config: PretrainedConfig): + return Qwen3MoeModuleArchitecture(num_experts=config.num_experts) + + def pre_weights(self, config: PretrainedConfig) -> List[WeightInfo]: + return QWEN3_MODULE_ARCH.pre_weights(config) + + def post_weights(self, config: PretrainedConfig) -> List[WeightInfo]: + return QWEN3_MODULE_ARCH.post_weights(config) + + def num_layers_config_key(self) -> str: + return QWEN3_MODULE_ARCH.num_layers_config_key() + + def layer_weights( + self, index: int, config: PretrainedConfig + ) -> Optional[List[WeightInfo]]: + prefix = f"model.layers.{index}" + tensor_names = [] + for expert_idx in range(self.num_experts): + for param in ("up_proj", "gate_proj", "down_proj"): + tensor_names.append( + prefix + f".mlp.experts.{expert_idx}.{param}.weight" + ) + tensor_names.append(prefix + ".mlp.gate.weight") + res = [] + for name in tensor_names: + res.append(WeightInfo(name=name)) + for weight_info in QWEN3_MODULE_ARCH.layer_weights(index, config): + if ".mlp." in weight_info.name: + continue + res.append(weight_info) + return res + + +AFMOE_PARTIAL_INFO = NAME_TO_ARCH["_AfmoePartialForCausalLM"][0] +AFMOE_PARTIAL_MODULE_ARCH = AFMOE_PARTIAL_INFO.modules["default"].architecture + + +class AfmoeModuleArchitecture(ModuleArchitecture, BaseModel): + ARCHITECTURE_NAME: ClassVar[str] = "AfmoeForCausalLM" + num_experts: int + + def name(self) -> str: + return "afmoe" + + @classmethod + def from_config(cls, config: PretrainedConfig): + return AfmoeModuleArchitecture(num_experts=config.num_experts) + + def pre_weights(self, config: PretrainedConfig) -> List[WeightInfo]: + return AFMOE_PARTIAL_MODULE_ARCH.pre_weights(config) + + def post_weights(self, config: PretrainedConfig) -> List[WeightInfo]: + return AFMOE_PARTIAL_MODULE_ARCH.post_weights(config) + + def num_layers_config_key(self) -> str: + return AFMOE_PARTIAL_MODULE_ARCH.num_layers_config_key() + + def layer_weights( + self, index: int, config: PretrainedConfig + ) -> Optional[List[WeightInfo]]: + res = AFMOE_PARTIAL_MODULE_ARCH.layer_weights(index, config) or [] + prefix = f"model.layers.{index}" + for expert_idx in range(self.num_experts): + for param in ("up_proj", "gate_proj", "down_proj"): + res.append( + WeightInfo( + name=prefix + f".mlp.experts.{expert_idx}.{param}.weight", + optional=True, + ) + ) + return res + + +GLM4_INFO = NAME_TO_ARCH["Glm4MoeForCausalLM"][0] +GLM4_MODULE_ARCH = GLM4_INFO.modules["default"].architecture + + +class Glm4MoeModuleArchitecture(ModuleArchitecture, BaseModel): + ARCHITECTURE_NAME: ClassVar[str] = "Glm4MoeForCausalLM" + num_experts: int + + def name(self) -> str: + return "glm4_moe" + + @classmethod + def from_config(cls, config: PretrainedConfig): + return Glm4MoeModuleArchitecture(num_experts=config.n_routed_experts) + + def pre_weights(self, config: PretrainedConfig) -> List[WeightInfo]: + return GLM4_MODULE_ARCH.pre_weights(config) + + def post_weights(self, config: PretrainedConfig) -> List[WeightInfo]: + return GLM4_MODULE_ARCH.post_weights(config) + + def num_layers_config_key(self) -> str: + return GLM4_MODULE_ARCH.num_layers_config_key() + + def layer_weights( + self, index: int, config: PretrainedConfig + ) -> Optional[List[WeightInfo]]: + prefix = f"model.layers.{index}" + + if index < config.first_k_dense_replace: + return GLM4_MODULE_ARCH.layer_weights(index, config) + else: + tensor_names = [] + for expert_idx in range(self.num_experts): + tensor_names.append( + prefix + f".mlp.experts.{expert_idx}.gate_proj.weight" + ) + tensor_names.append( + prefix + f".mlp.experts.{expert_idx}.up_proj.weight" + ) + tensor_names.append( + prefix + f".mlp.experts.{expert_idx}.down_proj.weight" + ) + tensor_names.append(prefix + ".mlp.gate.weight") + tensor_names.append(prefix + ".mlp.gate.e_score_correction_bias") + shared_expert_names = [ + (prefix + ".mlp.shared_experts.gate_proj.weight", False), + (prefix + ".mlp.shared_experts.up_proj.weight", False), + (prefix + ".mlp.shared_experts.down_proj.weight", False), + ] + res = [] + for name in tensor_names: + res.append(WeightInfo(name=name)) + for name, optional in shared_expert_names: + res.append(WeightInfo(name=name, optional=optional)) + for weight_info in GLM4_MODULE_ARCH.layer_weights(index, config): + if ".mlp." in weight_info.name: + continue + res.append(weight_info) + return res diff --git a/src/mindnlp/wizard/merge/card.py b/src/mindnlp/wizard/merge/card.py new file mode 100644 index 000000000..42f1bd512 --- /dev/null +++ b/src/mindnlp/wizard/merge/card.py @@ -0,0 +1,226 @@ +# Copyright (C) 2025 Arcee AI +# SPDX-License-Identifier: LGPL-3.0-only +# Modified for MindSpore/Ascend NPU by MindNLP Wizard contributors. +# + +""" +Model card generation for merged models. + +and related helpers. +""" + +import os +from typing import Generator, List, Optional + +import huggingface_hub +import yaml +from huggingface_hub.utils import HFValidationError +from . import merge_methods +from .config import MergeConfiguration, ModelReference + +CARD_TEMPLATE = """--- +{metadata} +--- +# {name} + +This is a merge of pre-trained language models created using [Wizard](https://github.com/mindspore-lab/mindnlp). + +## Merge Details +### Merge Method + +This model was merged using the {merge_method} merge method{base_text}. + +### Models Merged + +The following models were included in the merge: +{model_list} + +### Configuration + +The following YAML configuration was used to produce this model: + +```yaml +{config_yaml} +``` +""" + +CARD_TEMPLATE_LORA = """--- +{metadata} +--- +# {name} + +This is a LoRA extracted from a language model. It was extracted using [Wizard](https://github.com/mindspore-lab/mindnlp). + +## LoRA Details + +{details} + +### Parameters + +The following command was used to extract this LoRA adapter: + +```sh +{invocation} +``` +""" + + +def is_hf(path: str) -> bool: + """ + Determines if the given path is a Hugging Face model repository. + + Args: + path: A string path to check. + """ + if path[0] in "/~" or path.count("/") > 1: + return False + if not os.path.exists(path): + return True + try: + return huggingface_hub.repo_exists(path, repo_type="model", token=False) + except HFValidationError: + return False + + +def extract_hf_paths(models: List[ModelReference]) -> Generator[str, None, None]: + """ + Yields all valid Hugging Face paths from a list of ModelReference objects. + + Args: + models: A list of ModelReference objects. + """ + for model in models: + if is_hf(model.model.path): + yield model.model.path + + if model.lora and is_hf(model.lora.path): + yield model.lora.path + + +def method_md(merge_method: str) -> str: + """ + Returns a markdown string for the given merge method. + + Args: + merge_method: A string indicating the merge method used. + """ + try: + method = merge_methods.get(merge_method) + except RuntimeError: + return merge_method + ref_url = method.reference_url() + name = method.pretty_name() or method.name() + if ref_url and ref_url.strip(): + return f"[{name}]({ref_url})" + return name + + +def maybe_link_hf(path: str) -> str: + """ + Convert a path to a clickable link if it's a Hugging Face model path. + + Args: + path: A string path to possibly convert to a link. + """ + if is_hf(path): + return f"[{path}](https://huggingface.co/{path})" + return path + + +def modelref_md(model: ModelReference) -> str: + """ + Generates markdown description for a ModelReference object. + + Args: + model: A ModelReference object. + + Returns: + A markdown formatted string describing the model reference. + """ + text = maybe_link_hf(model.model.path) + if model.lora: + text += " + " + maybe_link_hf(model.lora.path) + return text + + +def generate_card( + config: MergeConfiguration, + config_yaml: str, + name: Optional[str] = None, +) -> str: + """ + Generates a markdown card for a merged model configuration. + + Args: + config: A MergeConfiguration object. + config_yaml: YAML source text of the config. + name: An optional name for the model. + """ + if not name: + name = "Untitled Model (1)" + + hf_bases = list(extract_hf_paths(config.referenced_models())) + tags = ["wizard", "merge"] + + actual_base = config.base_model + if config.merge_method == "slerp": + actual_base = None + + base_text = "" + if actual_base: + base_text = f" using {modelref_md(actual_base)} as a base" + + model_bullets = [] + for model in config.referenced_models(): + if model == actual_base: + continue + + model_bullets.append("* " + modelref_md(model)) + + return CARD_TEMPLATE.format( + metadata=yaml.dump( + {"base_model": hf_bases, "tags": tags, "library_name": "transformers"} + ), + model_list="\n".join(model_bullets), + base_text=base_text, + merge_method=method_md(config.merge_method), + name=name, + config_yaml=config_yaml, + ) + + +def generate_card_lora( # pylint: disable=too-many-positional-arguments + base_ref: ModelReference, + finetuned_ref: ModelReference, + invocation: str, + name: str, + base_vocab_size: Optional[int] = None, + final_vocab_size: Optional[int] = None, +) -> str: + if not name: + name = "Untitled LoRA Model (1)" + + hf_bases = list(extract_hf_paths([base_ref, finetuned_ref])) + tags = ["wizard", "peft"] + + details = ( + f"This LoRA adapter was extracted from {modelref_md(finetuned_ref)} " + f"and uses {modelref_md(base_ref)} as a base." + ) + + if base_vocab_size and final_vocab_size and base_vocab_size != final_vocab_size: + verb = "extended" if final_vocab_size > base_vocab_size else "reduced" + details += ( + f"\n\n [!WARNING]\n> The vocabulary size has been {verb} from the base " + f"model's {base_vocab_size} to {final_vocab_size}. To load this adapter, " + f"you must first call `model.resize_token_embeddings({final_vocab_size})`." + ) + + return CARD_TEMPLATE_LORA.format( + metadata=yaml.dump( + {"base_model": hf_bases, "tags": tags, "library_name": "peft"} + ), + name=name, + details=details, + invocation=invocation, + ) diff --git a/src/mindnlp/wizard/merge/common.py b/src/mindnlp/wizard/merge/common.py new file mode 100644 index 000000000..51cca80b4 --- /dev/null +++ b/src/mindnlp/wizard/merge/common.py @@ -0,0 +1,592 @@ +# Originally from MergeKit (https://github.com/arcee-ai/mergekit) +# Copyright (C) 2025 Arcee AI +# SPDX-License-Identifier: LGPL-3.0-only +# Modified for MindSpore/Ascend NPU by MindNLP Wizard contributors. + +"""Common types shared across the merge package.""" + +import binascii +import logging +import os +import os.path +from typing import ( + Any, + Callable, + Dict, + Generic, + Mapping, + Optional, + Protocol, + Tuple, + Union, + get_args, +) + +import huggingface_hub +import immutables +import mindspore +from pydantic import BaseModel, model_serializer, model_validator +from pydantic_core import core_schema +from transformers import AutoConfig, PretrainedConfig +from typing_extensions import TypeVar + +LOG = logging.getLogger(__name__) + + +# --------------------------------------------------------------------------- +# Config helpers +# --------------------------------------------------------------------------- + +def set_config_value(config: PretrainedConfig, key: str, value: Any): + parts = key.split(".") + obj = config + for idx, part in enumerate(parts[:-1]): + if not hasattr(obj, part): + raise RuntimeError( + f"Config {config} has no attribute {'.'.join(parts[: idx + 1])}" + ) + obj = getattr(obj, part) + setattr(obj, parts[-1], value) + + +def get_config_value(config: PretrainedConfig, key: str) -> Any: + parts = key.split(".") + obj = config + for idx, part in enumerate(parts): + if not hasattr(obj, part): + raise RuntimeError( + f"Config {config} has no attribute {'.'.join(parts[: idx + 1])}" + ) + obj = getattr(obj, part) + return obj + + +# --------------------------------------------------------------------------- +# ModelPath / ModelReference +# --------------------------------------------------------------------------- + +class ModelPath(BaseModel, frozen=True): + path: str + revision: Optional[str] = None + + @model_validator(mode="before") + def validate_string(cls, value): + if isinstance(value, str): + at_ct = value.count("@") + if at_ct > 1: + raise RuntimeError(f"Invalid model path - multiple @: {value}") + elif at_ct == 1: + path, rev = value.split("@") + return {"path": path, "revision": rev} + else: + return {"path": value} + return value + + def __str__(self): + if self.revision: + return f"{self.path}@{self.revision}" + return self.path + + def _unique_id(self): + return ( + os.path.basename(self.path) + + "_" + + str(binascii.crc32(self.__str__().encode())) + ) + + +class ModelReference(BaseModel, frozen=True): + """A reference to a language model (hub path or local). + + Optionally includes a LoRA adapter path. + """ + + model: ModelPath + lora: Optional[ModelPath] = None + override_architecture: Optional[str] = None + + def merged( + self, + cache_dir: Optional[str] = None, + trust_remote_code: bool = False, + lora_merge_dtype: Optional[str] = None, + ) -> "ModelReference": + """Merge the LoRA adapter (if any) and return a new reference.""" + if not self.lora: + return self + + if not cache_dir: + raise RuntimeError("Need to specify cache dir to merge adapters") + + out_path = os.path.join( + cache_dir, + self.model._unique_id() + "_" + self.lora._unique_id(), + ) + + if not os.path.exists(out_path): + os.makedirs(out_path, exist_ok=True) + + config = self.config(trust_remote_code) + auto_cls = get_auto_cls(config.architectures[0]) + + logging.info("Loading %s for LoRA merge...", self.model) + model = auto_cls.from_pretrained( + self.model.path, + revision=self.model.revision, + ms_dtype=dtype_from_name(lora_merge_dtype), + trust_remote_code=trust_remote_code, + ) + + # MindSpore PEFT equivalent (mindnlp / mindpet) + try: + from mindnlp.peft import PeftModel + except ImportError: + raise ImportError( + "mindnlp.peft is required for LoRA merging. " + "Install mindnlp with PEFT support." + ) + model = PeftModel.from_pretrained( + model, + self.lora.path, + is_trainable=False, + ) + logging.info("Merging %s into %s", self.lora, self.model) + model = model.merge_and_unload() + model.save_pretrained(out_path, safe_serialization=True) + del model + + return ModelReference(model=ModelPath(path=out_path)) + + def config(self, trust_remote_code: bool = False) -> PretrainedConfig: + res = AutoConfig.from_pretrained( + self.model.path, + revision=self.model.revision, + trust_remote_code=trust_remote_code, + ) + if self.override_architecture: + res.architectures = [self.override_architecture] + return res + + def local_path( + self, cache_dir: Optional[str] = None, ignore_lora: bool = False + ) -> str: + if not ignore_lora: + assert ( + self.lora is None + ), "LoRA not merged - use .merged() to get a local path" + + path = self.model.path + if not os.path.exists(path): + has_safetensors = any( + fn.lower().endswith(".safetensors") + for fn in huggingface_hub.list_repo_files( + path, repo_type="model", revision=self.model.revision + ) + ) + patterns = ["tokenizer.model", "*.json"] + if has_safetensors: + patterns.append("*.safetensors") + else: + patterns.append("*.bin") + + path = huggingface_hub.snapshot_download( + path, + revision=self.model.revision, + cache_dir=cache_dir, + allow_patterns=patterns, + ) + return path + + def tensor_index(self, cache_dir: Optional[str] = None): + from .io import ShardedTensorIndex + + return ShardedTensorIndex.from_disk(self.local_path(cache_dir)) + + def lazy_loader( + self, cache_dir: Optional[str] = None, lazy_loader: bool = True + ): + from .io import LazyTensorLoader + + return LazyTensorLoader( + self.tensor_index(cache_dir), + lazy_loader=lazy_loader, + ) + + @model_validator(mode="before") + def validate_string(cls, value): + if isinstance(value, str): + chunks = value.split("+") + if len(chunks) == 1: + return {"model": value} + elif len(chunks) == 2: + return {"model": chunks[0], "lora": chunks[1]} + raise RuntimeError(f"Can't parse {value}") + return value + + @model_serializer() + def serialize(self): + if self.override_architecture is not None: + return { + "model": self.model, + "lora": self.lora, + "override_architecture": self.override_architecture, + } + res = str(self) + if '"' in res or " " in res: + return self + return res + + @classmethod + def parse(cls, value: str) -> "ModelReference": + return ModelReference.model_validate(value) + + def __str__(self) -> str: + if self.lora: + return f"{str(self.model)}+{str(self.lora)}" + return str(self.model) + + +# --------------------------------------------------------------------------- +# dtype helpers (MindSpore) +# --------------------------------------------------------------------------- + +def dtype_from_name(name: Optional[str]) -> Optional[mindspore.dtype]: + if not name: + return None + + for prefix in ("torch.", "mindspore.", "ms."): + if name.startswith(prefix): + name = name[len(prefix):] + break + + _MAP = { + "bfloat16": mindspore.bfloat16, + "float16": mindspore.float16, + "float32": mindspore.float32, + "float64": mindspore.float64, + "int32": mindspore.int32, + "int64": mindspore.int64, + } + if name in _MAP: + return _MAP[name] + raise RuntimeError(f'Unimplemented dtype "{name}"') + + +# --------------------------------------------------------------------------- +# parse_kmb +# --------------------------------------------------------------------------- + +def parse_kmb(value: Union[str, int]) -> int: + if isinstance(value, int): + return value + elif value.isnumeric(): + return int(value) + elif value[-1].lower() == "k": + return int(value[:-1]) * 1000 + elif value[-1].lower() == "m": + return int(value[:-1]) * 1000 * 1000 + elif value[-1].lower() == "b": + return int(value[:-1]) * 1000 * 1000 * 1000 + else: + raise ValueError(value) + + +# --------------------------------------------------------------------------- +# ImmutableMap +# --------------------------------------------------------------------------- + +T_K = TypeVar("T_K") +T_V = TypeVar("T_V") + + +class ImmutableMap(Generic[T_K, T_V]): + data: immutables.Map + + def __init__(self, data: Mapping): + self.data = data + + @classmethod + def __get_pydantic_core_schema__( + cls, source: Any, handler: Callable[[Any], core_schema.CoreSchema] + ) -> core_schema.CoreSchema: + instance_schema = core_schema.is_instance_schema(cls) + args = get_args(source) + if args: + dict_schema = handler(Dict[args[0], args[1]]) + else: + dict_schema = handler(Dict) + non_instance_schema = core_schema.with_info_after_validator_function( + lambda value, _info: immutables.Map(value), dict_schema + ) + return core_schema.union_schema([instance_schema, non_instance_schema]) + + def __iter__(self): + return self.data.__iter__() + + def __getitem__(self, key): + return self.data[key] + + def __len__(self) -> int: + return len(self.data) + + def keys(self): + return self.data.keys() + + def items(self): + return self.data.items() + + def values(self): + return self.data.values() + + +# --------------------------------------------------------------------------- +# Auto model class detection (framework-agnostic via transformers) +# --------------------------------------------------------------------------- + +ARCH_NAME_TO_AUTO_CLS: Dict[str, Any] = {} + +try: + import transformers + import transformers.models.auto.modeling_auto as tf_auto +except ImportError: + tf_auto = None + +if tf_auto is not None: + for map_name, cls_name in [ + ("MODEL_MAPPING_NAMES", "AutoModel"), + ( + "MODEL_FOR_AUDIO_CLASSIFICATION_MAPPING_NAMES", + "AutoModelForAudioClassification", + ), + ( + "MODEL_FOR_IMAGE_CLASSIFICATION_MAPPING_NAMES", + "AutoModelForImageClassification", + ), + ("MODEL_FOR_SPEECH_SEQ_2_SEQ_MAPPING_NAMES", "AutoModelForSpeechSeq2Seq"), + ( + "MODEL_FOR_SEQUENCE_CLASSIFICATION_MAPPING_NAMES", + "AutoModelForSequenceClassification", + ), + ("MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING_NAMES", "AutoModelForSeq2SeqLM"), + ( + "MODEL_FOR_TOKEN_CLASSIFICATION_MAPPING_NAMES", + "AutoModelForTokenClassification", + ), + ( + "MODEL_FOR_IMAGE_TEXT_TO_TEXT_MAPPING_NAMES", + "AutoModelForImageTextToText", + ), + ("MODEL_FOR_TEXT_TO_WAVEFORM_MAPPING_NAMES", "AutoModelForTextToWaveform"), + ("MODEL_FOR_MASKED_LM_MAPPING_NAMES", "AutoModelForMaskedLM"), + ("MODEL_FOR_CAUSAL_LM_MAPPING_NAMES", "AutoModelForCausalLM"), + ]: + cls = getattr(transformers, cls_name, None) + if cls is None: + logging.info("Could not find %s in transformers", cls_name) + continue + if hasattr(tf_auto, map_name): + name_to_arch_name = getattr(tf_auto, map_name) + for arch_name in name_to_arch_name.values(): + ARCH_NAME_TO_AUTO_CLS[arch_name] = cls + + +class AutoClassProtocol(Protocol): + def from_pretrained( + self, pretrained_model_name_or_path: str, *model_args, **kwargs + ): ... + + def from_config(self, config, *model_args, **kwargs): ... + + +def get_auto_cls(arch_name: str) -> AutoClassProtocol: + if arch_name in ARCH_NAME_TO_AUTO_CLS: + return ARCH_NAME_TO_AUTO_CLS[arch_name] + + if arch_name.endswith("ForMaskedLM"): + auto_cls = transformers.AutoModelForMaskedLM + elif arch_name.endswith("ForSequenceClassification"): + auto_cls = transformers.AutoModelForSequenceClassification + elif arch_name.endswith("ForTokenClassification"): + auto_cls = transformers.AutoModelForTokenClassification + else: + if not arch_name.endswith("ForCausalLM") or arch_name.endswith( + "LMHeadModel" + ): + logging.warning( + "Unknown model type %s — assuming AutoModelForCausalLM", + arch_name, + ) + auto_cls = transformers.AutoModelForCausalLM + return auto_cls + + +# --------------------------------------------------------------------------- +# Ascend NPU / accelerator helpers +# --------------------------------------------------------------------------- + +def get_ascend_device_count() -> int: + """Return the number of available Ascend NPU devices.""" + try: + import acl # pylint: disable=import-error # Ascend Computing Language + raw = acl.rt.get_device_count() + count = _normalize_device_count(raw) + if count > 0: + return count + except Exception as exc: + LOG.debug( + "ACL-based Ascend device probing failed (%s: %s)", + type(exc).__name__, + exc, + ) + try: + count = int(os.environ.get("ASCEND_DEVICE_NUM", "0")) + if count > 0: + return count + except ValueError as exc: + LOG.debug("Invalid ASCEND_DEVICE_NUM value (%s)", exc) + try: + device_list = mindspore.get_context("device_target") + if device_list == "Ascend": + return max(1, int(os.environ.get("RANK_SIZE", "1"))) + except Exception as exc: + LOG.debug( + "MindSpore context probing for Ascend failed (%s: %s)", + type(exc).__name__, + exc, + ) + return 0 + + +def get_accelerator_count(accelerator_name: Optional[str] = None) -> int: + """Return device count for the specified (or default) accelerator.""" + if accelerator_name is not None: + target, dev_id = _parse_accelerator(accelerator_name) + if dev_id is not None: + return 1 + else: + target = _default_accelerator() + + if target == "CPU": + return 1 + count = _probe_device_count(target) + if count > 0: + return count + + LOG.warning( + "Could not determine device count for accelerator '%s'; defaulting to 1", + target, + ) + return 1 + + +def get_accelerator_type(accelerator_name: Optional[str] = None) -> str: + if accelerator_name is not None: + return _parse_accelerator(accelerator_name)[0] + return _default_accelerator() + + +def _parse_accelerator(spec: str) -> Tuple[str, Optional[int]]: + parts = spec.split(":") + target = parts[0] + dev_id = int(parts[1]) if len(parts) > 1 else None + return target, dev_id + + +def _default_accelerator() -> str: + """Detect the default accelerator available on this system.""" + try: + target = mindspore.get_context("device_target") + if target and target != "CPU": + return target + except Exception as exc: + LOG.debug( + "MindSpore context accelerator probing failed (%s: %s)", + type(exc).__name__, + exc, + ) + + for candidate in ("Ascend", "GPU"): + if _accelerator_available(candidate): + return candidate + + if get_ascend_device_count() > 0: + return "Ascend" + return "CPU" + + +def _probe_device_count(target: str) -> int: + normalized_target = target.upper() + if normalized_target == "ASCEND": + count = get_ascend_device_count() + if count > 0: + return count + + try: + if hasattr(mindspore, "hal") and hasattr(mindspore.hal, "device_count"): + count = _normalize_device_count( + mindspore.hal.device_count(device_target=target) + ) + if count > 0: + return count + except Exception as exc: + LOG.debug( + "MindSpore HAL device_count probing failed for %s (%s: %s)", + target, + type(exc).__name__, + exc, + ) + + for env_key in ( + "RANK_SIZE", + "WORLD_SIZE", + "OMPI_COMM_WORLD_SIZE", + "DEVICE_NUM", + "ASCEND_DEVICE_NUM", + ): + raw = os.environ.get(env_key) + if not raw: + continue + try: + count = int(raw) + if count > 0: + return count + except ValueError: + LOG.debug("Invalid %s value=%r during device count probing", env_key, raw) + + return 0 + + +def _accelerator_available(target: str) -> bool: + try: + if hasattr(mindspore, "hal") and hasattr(mindspore.hal, "is_available"): + return bool(mindspore.hal.is_available(target)) + except Exception as exc: + LOG.debug( + "MindSpore HAL is_available probing failed for %s (%s: %s)", + target, + type(exc).__name__, + exc, + ) + return _probe_device_count(target) > 0 + + +def _normalize_device_count(raw: Any) -> int: + """Normalize vendor/runtime-specific device count return values.""" + if isinstance(raw, int): + return raw + if isinstance(raw, (tuple, list)): + for item in raw: + if isinstance(item, int): + return item + return 0 + if isinstance(raw, dict): + for key in ("count", "device_count", "num_devices"): + value = raw.get(key) + if isinstance(value, int): + return value + return 0 + try: + return int(raw) + except Exception: + LOG.debug("Unsupported device count return value: %r", raw) + return 0 diff --git a/src/mindnlp/wizard/merge/config.py b/src/mindnlp/wizard/merge/config.py new file mode 100644 index 000000000..9a0b9eb7e --- /dev/null +++ b/src/mindnlp/wizard/merge/config.py @@ -0,0 +1,315 @@ +# Originally from MergeKit (https://github.com/arcee-ai/mergekit) +# Copyright (C) 2025 Arcee AI +# SPDX-License-Identifier: LGPL-3.0-only +# Modified for MindSpore/Ascend NPU by MindNLP Wizard contributors. + +"""Configuration models and parameter readers for merge.""" + +from __future__ import annotations + +import re +from typing import Any, Dict, List, Literal, Optional, Tuple, Union + +import yaml +from pydantic import BaseModel, ConfigDict, model_validator + +from .common import ModelReference +from .tokenizer.config import TokenizerConfig + + +class ConditionalParameter(BaseModel, frozen=True): + """Conditionally-applied parameter value by tensor name filter.""" + + value: Any + filter: Optional[str] = None + + +ParameterSetting = Union[ + int, + float, + bool, + str, + ConditionalParameter, + List[int], + List[float], + List[bool], + List[str], + List[ConditionalParameter], +] + + +def _filter_match(tensor_name: str, pattern: str) -> bool: + if not pattern: + return True + if pattern == "*": + return True + try: + return re.search(pattern, tensor_name) is not None + except re.error: + # Keep a safe fallback when users pass plain substrings. + return pattern in tensor_name + + +def evaluate_setting( + tensor_name: str, + setting: Optional[ParameterSetting], + *, + t: float = 0.0, +) -> Any: + """Resolve a parameter setting for a specific tensor and interpolation point.""" + if setting is None: + return None + + if isinstance(setting, ConditionalParameter): + if setting.filter and not _filter_match(tensor_name, setting.filter): + return None + return evaluate_setting(tensor_name, setting.value, t=t) + + if isinstance(setting, list): + if not setting: + return None + + if all(isinstance(x, ConditionalParameter) for x in setting): + for item in setting: + value = evaluate_setting(tensor_name, item, t=t) + if value is not None: + return value + return None + + if all(isinstance(x, (int, float)) for x in setting): + if len(setting) == 1: + return setting[0] + pos = max(0.0, min(1.0, float(t))) * (len(setting) - 1) + left = int(pos) + right = min(left + 1, len(setting) - 1) + if left == right: + return setting[left] + frac = pos - left + return float(setting[left]) * (1.0 - frac) + float(setting[right]) * frac + + return setting[0] + + return setting + + +class InputModelDefinition(BaseModel, frozen=True): + model: ModelReference + parameters: Optional[Dict[str, ParameterSetting]] = None + + +class InputSliceDefinition(BaseModel, frozen=True): + model: ModelReference + layer_range: Tuple[int, int] + parameters: Optional[Dict[str, ParameterSetting]] = None + + +class OutputSliceDefinition(BaseModel, frozen=True): + sources: List[InputSliceDefinition] + base_model: Optional[ModelReference] = None + parameters: Optional[Dict[str, ParameterSetting]] = None + + +class OutputModuleDefinition(BaseModel): + model_config = ConfigDict(extra="allow") + + name: Optional[str] = None + models: Optional[List[InputModelDefinition]] = None + slices: Optional[List[OutputSliceDefinition]] = None + parameters: Optional[Dict[str, ParameterSetting]] = None + + +class MergeConfiguration(BaseModel): + """Top-level merge recipe.""" + + model_config = ConfigDict(extra="allow") + + merge_method: str + models: Optional[List[InputModelDefinition]] = None + slices: Optional[List[Union[OutputSliceDefinition, InputSliceDefinition]]] = None + modules: Optional[Dict[str, OutputModuleDefinition]] = None + parameters: Optional[Dict[str, ParameterSetting]] = None + base_model: Optional[ModelReference] = None + tokenizer_source: Optional[Union[Literal["union"], Literal["base"], ModelReference]] = None + tokenizer: Optional[TokenizerConfig] = None + chat_template: Optional[str] = None + dtype: Optional[str] = None + out_dtype: Optional[str] = None + # Extra compatibility fields used in architecture disambiguation. + architectures: Optional[List[str]] = None + model_type: Optional[str] = None + + @model_validator(mode="after") + def _normalize(self): + sources = int(bool(self.models)) + int(bool(self.slices)) + int(bool(self.modules)) + if sources != 1: + raise RuntimeError( + "Exactly one of models, slices, or modules must be specified" + ) + + if self.tokenizer_source is not None and self.tokenizer is not None: + if self.tokenizer.source != self.tokenizer_source: + raise RuntimeError( + "Cannot specify both tokenizer_source and tokenizer" + ) + + if self.slices and isinstance(self.slices[0], InputSliceDefinition): + self.slices = [ + OutputSliceDefinition(sources=[s]) # type: ignore[arg-type] + for s in self.slices # type: ignore[assignment] + ] + + if self.modules: + normalized: Dict[str, OutputModuleDefinition] = {} + for name, module in self.modules.items(): + if module.name is None: + module.name = name + normalized[name] = module + self.modules = normalized + + if self.tokenizer_source is not None and self.tokenizer is None: + self.tokenizer = TokenizerConfig(source=self.tokenizer_source) + + model_count = len(self.models) if self.models else 0 + if model_count == 0 and not self.slices and not self.modules: + raise RuntimeError("At least one model source must be provided") + + two_model_base_methods = {"slerp", "arcee_fusion", "nearswap"} + if self.merge_method in two_model_base_methods: + if self.base_model is None: + raise RuntimeError( + f"merge_method '{self.merge_method}' requires base_model" + ) + return self + + def referenced_models(self) -> List[ModelReference]: + models: List[ModelReference] = [] + + def _add(m: Optional[ModelReference]): + if m is not None and m not in models: + models.append(m) + + if self.models: + for model_in in self.models: + _add(model_in.model) + if self.slices: + for sl in self.slices: + for src in sl.sources: + _add(src.model) + if self.modules: + for module in self.modules.values(): + if module.models: + for model_in in module.models: + _add(model_in.model) + if module.slices: + for sl in module.slices: + for src in sl.sources: + _add(src.model) + _add(self.base_model) + if isinstance(self.tokenizer_source, ModelReference): + _add(self.tokenizer_source) + if self.tokenizer and isinstance(self.tokenizer.source, ModelReference): + _add(self.tokenizer.source) + return models + + def to_yaml(self) -> str: + return yaml.safe_dump( + self.model_dump(mode="json", exclude_none=True), + allow_unicode=True, + sort_keys=False, + ) + + +class ConfigReader(BaseModel): + """Contextual parameter resolver with merge precedence rules.""" + + model_config = ConfigDict(arbitrary_types_allowed=True, frozen=True) + + config: MergeConfiguration + t: float = 0.0 + tensor_name: str = "" + module: Optional[OutputModuleDefinition] = None + slice_out: Optional[OutputSliceDefinition] = None + + @property + def base_model(self) -> Optional[ModelReference]: + if self.slice_out and self.slice_out.base_model is not None: + return self.slice_out.base_model + return self.config.base_model + + def with_t(self, t: float) -> "ConfigReader": + return ConfigReader( + config=self.config, + t=t, + tensor_name=self.tensor_name, + module=self.module, + slice_out=self.slice_out, + ) + + def for_tensor(self, tensor_name: str) -> "ConfigReader": + return ConfigReader( + config=self.config, + t=self.t, + tensor_name=tensor_name, + module=self.module, + slice_out=self.slice_out, + ) + + def for_out_slice(self, out_slice: OutputSliceDefinition) -> "ConfigReader": + return ConfigReader( + config=self.config, + t=self.t, + tensor_name=self.tensor_name, + module=self.module, + slice_out=out_slice, + ) + + def parameter( + self, + name: str, + *, + model: Optional[ModelReference] = None, + required: bool = False, + default: Any = None, + ) -> Any: + if model is not None and self.slice_out is not None: + for source in self.slice_out.sources: + if source.model == model and source.parameters and name in source.parameters: + value = evaluate_setting( + self.tensor_name, + source.parameters[name], + t=self.t, + ) + if value is not None: + return value + + if self.slice_out and self.slice_out.parameters and name in self.slice_out.parameters: + value = evaluate_setting(self.tensor_name, self.slice_out.parameters[name], t=self.t) + if value is not None: + return value + + if self.module and self.module.parameters and name in self.module.parameters: + value = evaluate_setting(self.tensor_name, self.module.parameters[name], t=self.t) + if value is not None: + return value + + if self.config.parameters and name in self.config.parameters: + value = evaluate_setting(self.tensor_name, self.config.parameters[name], t=self.t) + if value is not None: + return value + + if required and default is None: + raise RuntimeError(f"Missing required parameter: {name}") + return default + + +__all__ = [ + "ConditionalParameter", + "ConfigReader", + "InputModelDefinition", + "InputSliceDefinition", + "MergeConfiguration", + "OutputModuleDefinition", + "OutputSliceDefinition", + "ParameterSetting", + "evaluate_setting", +] diff --git a/src/mindnlp/wizard/merge/dtype_policy.py b/src/mindnlp/wizard/merge/dtype_policy.py new file mode 100644 index 000000000..ed3e1f6e9 --- /dev/null +++ b/src/mindnlp/wizard/merge/dtype_policy.py @@ -0,0 +1,148 @@ +# Copyright (c) MindNLP Wizard contributors. +# Licensed under the Apache License, Version 2.0. + +"""Dtype policy helpers for merge execution.""" + +from __future__ import annotations + +import logging +from typing import Dict, Iterable, Optional + +import mindspore +import numpy + +LOG = logging.getLogger(__name__) + +HALF_DTYPES = (mindspore.bfloat16, mindspore.float16) +_WARNED_SAFE_PATHS: Dict[str, bool] = {} + + +def get_device_target(device_target: Optional[str] = None) -> str: + if device_target: + return str(device_target).strip().upper() + try: + target = mindspore.get_context("device_target") + if target: + return str(target).strip().upper() + except Exception: + pass + return "CPU" + + +def needs_safe_path( + ref_dtype: mindspore.dtype, + *, + device_target: Optional[str] = None, + promote_half_to_fp32: bool = True, +) -> bool: + if not promote_half_to_fp32: + return False + return ( + get_device_target(device_target) == "CPU" + and ref_dtype in HALF_DTYPES + ) + + +def choose_work_dtype( + ref_dtype: mindspore.dtype, + *, + device_target: Optional[str] = None, + promote_half_to_fp32: bool = True, +) -> mindspore.dtype: + if needs_safe_path( + ref_dtype, + device_target=device_target, + promote_half_to_fp32=promote_half_to_fp32, + ): + return mindspore.float32 + return ref_dtype + + +def cast_to_work(tensor: mindspore.Tensor, work_dtype: mindspore.dtype) -> mindspore.Tensor: + if tensor.dtype == work_dtype: + return tensor + return tensor.astype(work_dtype) + + +def cast_back(tensor: mindspore.Tensor, out_dtype: mindspore.dtype) -> mindspore.Tensor: + if tensor.dtype == out_dtype: + return tensor + return tensor.astype(out_dtype) + + +def cast_many_to_work( + tensors: Iterable[mindspore.Tensor], + work_dtype: mindspore.dtype, +) -> list[mindspore.Tensor]: + return [cast_to_work(t, work_dtype) for t in tensors] + + +def warn_safe_path_once( + op_name: str, + *, + in_dtype: mindspore.dtype, + work_dtype: mindspore.dtype, + device_target: Optional[str] = None, +) -> None: + if in_dtype == work_dtype: + return + key = f"{get_device_target(device_target)}::{op_name}::{in_dtype}->{work_dtype}" + if _WARNED_SAFE_PATHS.get(key): + return + _WARNED_SAFE_PATHS[key] = True + LOG.warning( + "Safe dtype path enabled for %s: %s -> %s on %s", + op_name, + in_dtype, + work_dtype, + get_device_target(device_target), + ) + + +# --------------------------------------------------------------------------- +# Centralised numpy <-> MindSpore conversion with bfloat16 safety +# --------------------------------------------------------------------------- + +def _is_ml_dtypes_bfloat16(arr: numpy.ndarray) -> bool: + """Check whether *arr* uses ``ml_dtypes.bfloat16``.""" + try: + import ml_dtypes + return arr.dtype == ml_dtypes.bfloat16 + except (ImportError, AttributeError): + return False + + +def numpy_to_mindspore(arr: numpy.ndarray) -> mindspore.Tensor: + """Convert a numpy array to a MindSpore tensor, handling ``bfloat16``. + + ``ml_dtypes.bfloat16`` is not a built-in numpy dtype, so + ``mindspore.Tensor(arr)`` may silently misinterpret the buffer. + This helper detects the dtype and passes an explicit + ``dtype=mindspore.bfloat16`` when needed. + """ + if _is_ml_dtypes_bfloat16(arr): + return mindspore.Tensor(arr, dtype=mindspore.bfloat16) + return mindspore.Tensor(arr) + + +def mindspore_to_numpy(tensor: mindspore.Tensor) -> numpy.ndarray: + """Convert a MindSpore tensor to a numpy array, handling ``bfloat16``. + + MindSpore's ``asnumpy()`` for bfloat16 tensors may return a raw + ``uint16`` view. This helper reinterprets such arrays as + ``ml_dtypes.bfloat16`` so that downstream consumers (safetensors, + etc.) see the correct dtype and metadata. + """ + arr = tensor.asnumpy() + if tensor.dtype == mindspore.bfloat16: + try: + import ml_dtypes + if arr.dtype != ml_dtypes.bfloat16: + arr = arr.view(ml_dtypes.bfloat16) + except ImportError: + LOG.warning( + "ml_dtypes is not installed — bfloat16 tensor will be " + "exported as raw %s; downstream tools may misinterpret it.", + arr.dtype, + ) + return arr diff --git a/src/mindnlp/wizard/merge/eval/__init__.py b/src/mindnlp/wizard/merge/eval/__init__.py new file mode 100644 index 000000000..7e22101aa --- /dev/null +++ b/src/mindnlp/wizard/merge/eval/__init__.py @@ -0,0 +1,6 @@ +# Copyright (c) MindNLP Wizard contributors. +# Licensed under the Apache License, Version 2.0. + +def get_mindspore_lm(): + from mindnlp.wizard.merge.eval.mindspore_lm import MindSporeLM + return MindSporeLM diff --git a/src/mindnlp/wizard/merge/eval/mindspore_lm.py b/src/mindnlp/wizard/merge/eval/mindspore_lm.py new file mode 100644 index 000000000..1fca8762c --- /dev/null +++ b/src/mindnlp/wizard/merge/eval/mindspore_lm.py @@ -0,0 +1,432 @@ +# Copyright (c) MindNLP Wizard contributors. +# Licensed under the Apache License, Version 2.0. + +"""MindSpore model backend for lm-evaluation-harness. + +Uses mindnlp.transformers (MindSpore) for model loading and inference on Ascend NPU. +Registered as model="mindspore" or model="mindnlp" in lm-eval. + +IMPORTANT: mindnlp must be imported BEFORE lm_eval so mindtorch properly replaces torch. +After importing mindnlp, patch torch.utils.collect_env.get_pretty_env_info for lm_eval compat. +""" +import copy +import logging +from typing import Dict, List, Optional, Tuple, Union + +import numpy as np + +from lm_eval import utils +from lm_eval.api.instance import Instance +from lm_eval.api.model import TemplateLM +from lm_eval.api.registry import register_model + +eval_logger = logging.getLogger(__name__) + + +def _pad_sequences(sequences: List[List[int]], pad_value: int = 0, + padding_side: str = "left") -> Tuple[np.ndarray, np.ndarray]: + max_len = max(len(s) for s in sequences) + padded = np.full((len(sequences), max_len), pad_value, dtype=np.int64) + mask = np.zeros((len(sequences), max_len), dtype=np.int64) + for i, s in enumerate(sequences): + if padding_side == "left": + offset = max_len - len(s) + padded[i, offset:] = s + mask[i, offset:] = 1 + else: + padded[i, :len(s)] = s + mask[i, :len(s)] = 1 + return padded, mask + + +@register_model("mindspore", "mindnlp") +class MindSporeLM(TemplateLM): + """lm-eval model backend using MindSpore / MindNLP on Ascend NPU.""" + + _DEFAULT_MAX_LENGTH = 2048 + + def __init__( # pylint: disable=too-many-positional-arguments + self, + pretrained: str, + dtype: str = "float16", + batch_size: Union[int, str] = 1, + max_length: Optional[int] = None, + device: Optional[str] = None, + trust_remote_code: bool = True, + add_bos_token: bool = False, + prefix_token_id: Optional[int] = None, + **kwargs, + ): + super().__init__() + + import mindspore + from mindnlp.transformers import AutoModelForCausalLM, AutoTokenizer # pylint: disable=no-name-in-module + from transformers import AutoConfig + + self._batch_size = int(batch_size) if str(batch_size).isdigit() else 1 + self._max_length = max_length + self.add_bos_token = add_bos_token + self.custom_prefix_token_id = prefix_token_id + + # Device must already be set via DEVICE_TARGET env var before import mindnlp. + # We only verify here, not set — setting here is too late for mindtorch init. + current_device = mindspore.get_context('device_target') + if current_device != 'Ascend': + eval_logger.warning( + f"[MindSporeLM] device_target is '{current_device}', not 'Ascend'. " + f"NPU will NOT be used. Set os.environ['DEVICE_TARGET']='Ascend' " + f"BEFORE importing mindnlp to enable NPU.") + eval_logger.info(f"[MindSporeLM] Loading model: {pretrained}, dtype={dtype}, device={current_device}") + + self._config = AutoConfig.from_pretrained( + pretrained, trust_remote_code=trust_remote_code) + + self.tokenizer = AutoTokenizer.from_pretrained( + pretrained, trust_remote_code=trust_remote_code) + + if self.tokenizer.pad_token_id is None: + if self.tokenizer.eos_token_id is not None: + self.tokenizer.pad_token_id = self.tokenizer.eos_token_id + else: + self.tokenizer.add_special_tokens({"pad_token": ""}) + + import torch as _torch + torch_dtype_map = { + "float16": _torch.float16, "float32": _torch.float32, + "bfloat16": _torch.bfloat16, "auto": "auto", + } + _torch_dtype = torch_dtype_map.get(dtype, "auto") + + self._model = AutoModelForCausalLM.from_pretrained( + pretrained, torch_dtype=_torch_dtype, trust_remote_code=trust_remote_code) + + if current_device == 'Ascend': + eval_logger.info("[MindSporeLM] Moving model to NPU...") + self._model = self._model.npu() + + if "gemma" in getattr(self._config, "model_type", ""): + self.add_bos_token = True + + self.vocab_size = self.tokenizer.vocab_size + eval_logger.info( + f"[MindSporeLM] Model loaded. type={getattr(self._config, 'model_type', '?')}, " + f"layers={getattr(self._config, 'num_hidden_layers', '?')}, " + f"vocab={self.vocab_size}") + + @property + def config(self): + return self._config + + @property + def model(self): + return self._model + + @property + def eot_token_id(self): + return self.tokenizer.eos_token_id + + @property + def prefix_token_id(self): + if self.custom_prefix_token_id is not None: + return self.custom_prefix_token_id + if self.tokenizer.bos_token_id is not None: + return self.tokenizer.bos_token_id + return self.tokenizer.eos_token_id + + @property + def max_length(self): + if self._max_length: + return self._max_length + for attr in ("n_positions", "max_position_embeddings", "n_ctx"): + if hasattr(self._config, attr): + return getattr(self._config, attr) + if hasattr(self.tokenizer, "model_max_length"): + if self.tokenizer.model_max_length == 1000000000000000019884624838656: + return self._DEFAULT_MAX_LENGTH + return self.tokenizer.model_max_length + return self._DEFAULT_MAX_LENGTH + + @property + def max_gen_toks(self) -> int: + return 256 + + @property + def batch_size(self): + return self._batch_size + + @property + def device(self): + return "Ascend" + + @property + def tokenizer_name(self) -> str: + return self.tokenizer.name_or_path.replace("/", "__") + + def tok_encode(self, string: str, left_truncate_len=None, + add_special_tokens=None) -> List[int]: + special_tokens_kwargs = {} + if add_special_tokens is None: + special_tokens_kwargs = {"add_special_tokens": False or self.add_bos_token} + else: + special_tokens_kwargs = {"add_special_tokens": add_special_tokens} + + encoding = self.tokenizer.encode(string, **special_tokens_kwargs) + if left_truncate_len: + encoding = encoding[-left_truncate_len:] + return encoding + + def tok_decode(self, tokens, skip_special_tokens=True): + if isinstance(tokens, int): + tokens = [tokens] + return self.tokenizer.decode(tokens, skip_special_tokens=skip_special_tokens) + + def _loglikelihood_tokens( + self, + requests: List[Tuple[Tuple[str, str], List[int], List[int]]], + disable_tqdm: bool = False, + override_bs: int = None, + ) -> List[Tuple[float, bool]]: + import torch + import torch.nn.functional as F + from tqdm import tqdm + + res = [] + + def _collate(req): + toks = req[1] + req[2] + return -len(toks), tuple(toks) + + re_ord = utils.Reorderer(requests, _collate) + batch_size = override_bs if override_bs is not None else self._batch_size + + chunks = list( + _chunks_list(re_ord.get_reordered(), batch_size)) + + pbar = tqdm(total=len(requests), disable=disable_tqdm, + desc="Running loglikelihood requests") + + for chunk in chunks: + inps_list = [] + cont_toks_list = [] + inplens = [] + + for _, context_enc, continuation_enc in chunk: + inp = (context_enc + continuation_enc)[-(self.max_length + 1):][:-1] + inps_list.append(inp) + cont_toks_list.append(continuation_enc) + inplens.append(len(inp)) + + padded_inps, attn_mask = _pad_sequences( + inps_list, + pad_value=self.tokenizer.pad_token_id, + padding_side="right", + ) + + input_ids = torch.tensor(padded_inps, dtype=torch.long).npu() + attention_mask = torch.tensor(attn_mask, dtype=torch.long).npu() + with torch.no_grad(): + outputs = self._model(input_ids=input_ids, attention_mask=attention_mask) + logits = outputs.logits if hasattr(outputs, "logits") else outputs[0] + multi_logits = F.log_softmax(logits.float(), dim=-1) + + for idx, ((request_str, ctx_tokens, _), inplen, cont_toks) in enumerate( + zip(chunk, inplens, cont_toks_list) + ): + contlen = len(cont_toks) + logits_row = multi_logits[idx, inplen - contlen: inplen, :] + + greedy_tokens = logits_row.argmax(dim=-1) + cont_toks_t = torch.tensor(cont_toks, dtype=torch.long).npu() + cont_logprobs = logits_row[torch.arange(contlen).npu(), cont_toks_t] + + logprob_sum = float(cont_logprobs.sum().cpu()) + is_greedy = bool((greedy_tokens == cont_toks_t).all().cpu()) + + answer = (logprob_sum, is_greedy) + res.append(answer) + + if request_str is not None: + self.cache_hook.add_partial("loglikelihood", request_str, answer) + pbar.update(1) + + pbar.close() + return re_ord.get_original(res) + + def loglikelihood_rolling( + self, requests: List[Instance], disable_tqdm: bool = False + ) -> List[float]: + from tqdm import tqdm + + loglikelihoods = [] + for (string,) in tqdm( + [req.args for req in requests], + disable=disable_tqdm, + desc="Running rolling loglikelihood", + ): + token_list = self.tok_encode(string) + rolling_windows = list( + map( + utils.make_disjoint_window, + utils.get_rolling_token_windows( + token_list=token_list, + prefix_token=self.prefix_token_id, + max_seq_len=self.max_length, + context_len=1, + ), + ) + ) + + windows_as_requests = [ + (None, ctx, cont) for ctx, cont in rolling_windows + ] + + nlls = self._loglikelihood_tokens( + windows_as_requests, disable_tqdm=True, + override_bs=self._batch_size, + ) + + total_nll = sum(nll for nll, _ in nlls) + loglikelihoods.append(total_nll) + self.cache_hook.add_partial("loglikelihood_rolling", (string,), total_nll) + + return loglikelihoods + + def generate_until( + self, requests: List[Instance], disable_tqdm: bool = False + ) -> List[str]: + from tqdm import tqdm + + res = [] + reqs = [req.args for req in requests] + + def _collate(x): + toks = self.tok_encode(x[0]) + return -len(toks), x[0] + + re_ord = utils.Reorderer(reqs, _collate) + + eos_str = self.tok_decode(self.eot_token_id, skip_special_tokens=False) + + pbar = tqdm(total=len(requests), disable=disable_tqdm, + desc="Running generate_until requests") + + for chunk in _chunks_list(re_ord.get_reordered(), self._batch_size): + contexts = [] + gen_kwargs_list = [] + for context, gen_kwargs in chunk: + contexts.append(context) + gen_kwargs_list.append(gen_kwargs) + + kwargs = copy.deepcopy(gen_kwargs_list[0]) if isinstance(gen_kwargs_list[0], dict) else {} + until = kwargs.pop("until", None) or [] + if eos_str and eos_str not in until: + until.append(eos_str) + max_gen_toks = kwargs.pop("max_gen_toks", self.max_gen_toks) + kwargs.pop("do_sample", None) + + max_ctx_len = self.max_length - max_gen_toks + assert max_ctx_len > 0 + + add_special = {"add_special_tokens": False or self.add_bos_token} + encodings = [ + self.tokenizer.encode(ctx, **add_special)[-max_ctx_len:] + for ctx in contexts + ] + + padded_inps, attn_mask = _pad_sequences( + encodings, pad_value=self.tokenizer.pad_token_id, padding_side="left") + + import torch + input_ids = torch.tensor(padded_inps, dtype=torch.long).npu() + attention_mask = torch.tensor(attn_mask, dtype=torch.long).npu() + + eos_ids = [] + if self.tokenizer.eos_token_id is not None: + eos_ids.append(self.tokenizer.eos_token_id) + for special_tok in ["<|im_end|>", "<|endoftext|>"]: + tok_id = self.tokenizer.convert_tokens_to_ids(special_tok) + if isinstance(tok_id, int) and tok_id != self.tokenizer.unk_token_id and tok_id not in eos_ids: + eos_ids.append(tok_id) + + generation_kwargs = { + "max_new_tokens": max_gen_toks, + "do_sample": False, + "pad_token_id": self.tokenizer.pad_token_id, + } + if eos_ids: + generation_kwargs["eos_token_id"] = eos_ids + + with torch.no_grad(): + output_ids = self._model.generate( + input_ids=input_ids, + attention_mask=attention_mask, + **generation_kwargs, + ) + + if isinstance(output_ids, tuple): + output_ids = output_ids[0] + + # MindTorch NPU tensors may expose `.numpy()` but require CPU init. + # Prefer robust conversion order to avoid assertion on device tensors. + if hasattr(output_ids, "asnumpy"): + output_ids_np = output_ids.asnumpy() + elif hasattr(output_ids, "cpu") and hasattr(output_ids, "numpy"): + output_ids_np = output_ids.cpu().numpy() + elif hasattr(output_ids, "numpy"): + try: + output_ids_np = output_ids.numpy() + except Exception: + if hasattr(output_ids, "cpu"): + output_ids_np = output_ids.cpu().numpy() + else: + raise + else: + output_ids_np = np.array(output_ids) + + for i, context in enumerate(contexts): + ctx_len = len(encodings[i]) + pad_offset = padded_inps.shape[1] - ctx_len + gen_start = pad_offset + ctx_len + + if output_ids_np.ndim == 1: + cont_toks = output_ids_np[gen_start:].tolist() + else: + cont_toks = output_ids_np[i, gen_start:].tolist() + + s = self.tokenizer.decode(cont_toks, skip_special_tokens=True) + + for term in until: + if len(term) > 0: + s = s.split(term)[0] + + res.append(s) + self.cache_hook.add_partial( + "generate_until", (context, gen_kwargs_list[i]), s) + pbar.update(1) + + pbar.close() + return re_ord.get_original(res) + + def apply_chat_template( + self, chat_history: List[Dict[str, str]], add_generation_prompt: bool = True + ) -> str: + try: + return self.tokenizer.apply_chat_template( + chat_history, + tokenize=False, + add_generation_prompt=add_generation_prompt, + continue_final_message=not add_generation_prompt, + ) + except Exception: + chat_history = [m for m in chat_history if m["role"] != "system"] + return self.tokenizer.apply_chat_template( + chat_history, + tokenize=False, + add_generation_prompt=add_generation_prompt, + continue_final_message=not add_generation_prompt, + ) + + +def _chunks_list(lst, n): + for i in range(0, len(lst), n): + yield lst[i: i + n] diff --git a/src/mindnlp/wizard/merge/evo/__init__.py b/src/mindnlp/wizard/merge/evo/__init__.py new file mode 100644 index 000000000..a397a9d30 --- /dev/null +++ b/src/mindnlp/wizard/merge/evo/__init__.py @@ -0,0 +1,4 @@ +# Copyright (C) 2025 Arcee AI +# SPDX-License-Identifier: LGPL-3.0-only +# Modified for MindSpore/Ascend NPU by MindNLP Wizard contributors. +# diff --git a/src/mindnlp/wizard/merge/evo/actors.py b/src/mindnlp/wizard/merge/evo/actors.py new file mode 100644 index 000000000..e660af383 --- /dev/null +++ b/src/mindnlp/wizard/merge/evo/actors.py @@ -0,0 +1,350 @@ +# Copyright (C) 2025 Arcee AI +# SPDX-License-Identifier: LGPL-3.0-only +# Modified for MindSpore/Ascend NPU by MindNLP Wizard contributors. +# + +import gc +import logging +import tempfile +from typing import Optional, Union + +import lm_eval # pylint: disable=import-error +import lm_eval.api.model # pylint: disable=import-error +import lm_eval.models.huggingface # pylint: disable=import-error +import lm_eval.tasks # pylint: disable=import-error +import mindspore # pylint: disable=import-error +import ray # pylint: disable=import-error +import ray.util.queue # pylint: disable=import-error +import ray.util.scheduling_strategies # pylint: disable=import-error +import transformers +from transformers.utils import is_flash_attn_2_available + +from ..architecture.base import ConfiguredModelArchitecture + +try: + import vllm +except ImportError: + vllm = None + +from ..architecture import arch_info_for_config +from ..common import get_accelerator_type +from ..config import MergeConfiguration +from .config import EvolMergeConfiguration +from .genome import InvalidGenotypeError, ModelGenome +from .helpers import _eval_model, evaluate_model, merge_model +from .monkeypatch import ( + NoInit, + monkeypatch_lmeval_shuffle, + monkeypatch_lmeval_vllm, +) +from ..graph import Executor +from ..io.tasks import LoaderCache, ReturnTensor +from ..merge import _model_out_config +from ..options import MergeOptions +from ..plan import MergePlanner + +LOG = logging.getLogger(__name__) + + +class MergeActorBase: + def __init__( # pylint: disable=too-many-positional-arguments + self, + config: EvolMergeConfiguration, + genome: ModelGenome, + merge_options: MergeOptions, + model_storage_path: Optional[str] = None, + vllm: bool = False, + batch_size: Optional[int] = None, + task_manager: Optional[lm_eval.tasks.TaskManager] = None, + quantization_config: Optional[transformers.BitsAndBytesConfig] = None, + ): + self.config = config + self.genome = genome + self.merge_options = merge_options + self.cache = LoaderCache() + self.cache.setup(merge_options) + self.model_storage_path = model_storage_path + self.vllm = vllm + self.batch_size = batch_size + self.task_manager = task_manager + self.quantization_config = quantization_config + + if config.shuffle: + monkeypatch_lmeval_shuffle() + + monkeypatch_lmeval_vllm() + + +@ray.remote(num_cpus=1, num_gpus=1.0) +class OnDiskMergeEvaluator(MergeActorBase): + """ + Merges models to disk then evaluates them in a separate process. + + Maximum compatibility and potential for parallelism, but higher overhead. + """ + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + + def evaluate_genotype( + self, + genotype: mindspore.Tensor, + ) -> dict: + gc.collect() + accelerator_type = get_accelerator_type(self.merge_options.device) + if accelerator_type == "Ascend": + try: + import acl # pylint: disable=import-error + acl.rt.reset_device(0) + except Exception as exc: + LOG.debug( + "Failed to reset Ascend device before evaluation (%s: %s)", + type(exc).__name__, + exc, + ) + LOG.info("Merging model") + merged_path = merge_model( + genotype, self.genome, self.model_storage_path, self.merge_options + ) + if not merged_path: + LOG.error("Model merge failed") + return {"score": None, "results": None} + + model_kwargs = {} + if self.quantization_config is not None: + model_kwargs["quantization_config"] = self.quantization_config + LOG.info(f"Model merged to {merged_path}") + return evaluate_model( + merged_path, + self.config.tasks, + num_fewshot=self.config.num_fewshot, + limit=self.config.limit, + vllm=self.vllm, + batch_size=self.batch_size, + task_manager=self.task_manager, + apply_chat_template=self.config.apply_chat_template, + fewshot_as_multiturn=self.config.fewshot_as_multiturn, + model_kwargs=model_kwargs, + ) + + +@ray.remote(num_cpus=1, num_gpus=1) +class InMemoryMergeEvaluator(MergeActorBase): + """ + Performs merges in memory, using a single model instance. + + This reduces overhead from disk I/O and model loading, but prevents + parallelism and may be slower for large models. + + Implementation is dark sorcery tampering with the internals of lm-eval, + transformers, and vLLM and may break at any time. + """ + + model: Union[ + lm_eval.models.huggingface.HFLM, lm_eval.models.vllm_causallms.VLLM, None + ] = None + arch_info: Optional[ConfiguredModelArchitecture] = None + + def __init__( + self, + *args, + vllm: bool = False, + **kwargs, + ): + super().__init__(*args, vllm=vllm, **kwargs) + + def _maybe_init_model(self, config: MergeConfiguration): + ai = arch_info_for_config(self.genome._input_config_example) + cfg_out = _model_out_config( + config, + ai, + trust_remote_code=self.merge_options.trust_remote_code, + ) + cfg_out.use_cache = True + cfg_out.ms_dtype = mindspore.bfloat16 + + if self.arch_info is not None: + different = False + for key in cfg_out.to_diff_dict(): + if key in ["architectures", "model_type"]: + continue + if key in ["use_cache", "ms_dtype", "torch_dtype"]: + continue + if key.endswith("_token_id"): + setattr(self.arch_info.config, key, getattr(cfg_out, key, None)) + continue + + if getattr(cfg_out, key) != getattr(self.arch_info.config, key, None): + LOG.warning(f"Config key {key} changed, reinitializing model") + different = True + break + + if not different: + return + + self.inner_model = None + + model_kwargs = { + "trust_remote_code": self.merge_options.trust_remote_code, + "torch_dtype": "bfloat16", + } + if is_flash_attn_2_available(): + model_kwargs["attn_implementation"] = "flash_attention_2" + + with NoInit(): + inner_model = transformers.AutoModelForCausalLM.from_config( + cfg_out, + **model_kwargs, + ) + inner_model.eval() + + if self.vllm: + with tempfile.TemporaryDirectory( + dir=self.model_storage_path, prefix="vllm" + ) as tempdir: + inner_model.save_pretrained( + tempdir, safe_serialization=True, out_shard_size=1_000_000_000_000 + ) + del inner_model + tokenizer_donor = self.genome.definition.base_model + if tokenizer_donor is None: + LOG.warning( + "Base model not set, using tokenizer from first model in genome" + ) + tokenizer_donor = self.genome.definition.models[0] + tok = transformers.AutoTokenizer.from_pretrained( + tokenizer_donor.model.path, use_fast=True + ) + tok.save_pretrained(tempdir) + + max_model_len = None + if ( + seq_len := getattr(cfg_out, "max_position_embeddings", None) + ) is not None: + max_model_len = seq_len + if (window_sz := getattr(cfg_out, "sliding_window", None)) is not None: + max_model_len = min(max_model_len or 1024, window_sz) + if max_model_len and max_model_len > 8192: + max_model_len = 8192 + LOG.warning(f"Clipping sequence length to {max_model_len}") + + accelerator_type = get_accelerator_type(self.merge_options.device) + mem_util = ( + 0.7 if accelerator_type in ["Ascend", "cuda", "xpu"] else 0.9 + ) + self.model = lm_eval.models.vllm_causallms.VLLM( + pretrained=tempdir, + batch_size=self.batch_size or "auto", + max_model_len=max_model_len, + gpu_memory_utilization=mem_util, + dtype="bfloat16", + device=self.merge_options.device, + trust_remote_code=self.merge_options.trust_remote_code, + ) + else: + self.model = lm_eval.models.huggingface.HFLM(pretrained=inner_model) + self.arch_info = ( + ConfiguredModelArchitecture( + info=ai, + config=cfg_out, + ) + if ai + else None + ) + LOG.info("Model initialized") + + def evaluate(self, genotype: mindspore.Tensor) -> dict: + try: + config = self.genome.genotype_merge_config(genotype) + except InvalidGenotypeError as e: + LOG.error("Invalid genotype", exc_info=e) + return {"score": None, "results": None} + + self._maybe_init_model(config) + + planner = MergePlanner( + config, + self.arch_info.info, + self.merge_options, + self.arch_info.config, + ) + + tasks = planner.plan_in_memory() + + model = self.model.model + if vllm is not None and isinstance(model, vllm.LLM): + assert ( + model.llm_engine.parallel_config.world_size == 1 + ), "Must be single GPU" + engine = model.llm_engine + if hasattr(engine, "model_executor"): + worker = engine.model_executor.worker + elif hasattr(engine, "driver_worker"): + worker = engine.driver_worker + else: + raise ValueError("Unknown LLM engine type") + model = worker.model_runner.model + param_dict = dict(model.named_parameters()) + + stacked_mapping = { + ".q_proj.": (".qkv_proj.", "q"), + ".k_proj.": (".qkv_proj.", "k"), + ".v_proj.": (".qkv_proj.", "v"), + ".gate_proj.": (".gate_up_proj.", 0), + ".up_proj.": (".gate_up_proj.", 1), + } + + accelerator_type = get_accelerator_type(self.merge_options.device) + executor = Executor( + tasks, + math_device=( + self.merge_options.device + if accelerator_type in ["Ascend", "cuda", "xpu"] + else "CPU" + ), + storage_device=( + self.merge_options.device + if accelerator_type in ["Ascend", "cuda", "xpu"] + else "CPU" + ), + ) + for tensor_task, value in executor.run(quiet=True): + assert isinstance(tensor_task, ReturnTensor) + name = tensor_task.weight_info.name + + if name in param_dict: + param_dict[name].set_data(value) + elif self.vllm: + stacked = False + for needle, (replacement, shard_id) in stacked_mapping.items(): + if needle in name: + target = name.replace(needle, replacement) + param = param_dict[target] + weight_loader = param.weight_loader + weight_loader(param, value, shard_id) + stacked = True + break + + if not stacked: + raise ValueError(f"Unknown parameter {name}") + else: + raise ValueError(f"Unknown parameter {name}") + + del value + + return _eval_model( + self.model, + self.config.tasks, + num_fewshot=self.config.num_fewshot, + limit=self.config.limit, + task_manager=self.task_manager, + batch_size=self.batch_size, + apply_chat_template=self.config.apply_chat_template, + fewshot_as_multiturn=self.config.fewshot_as_multiturn, + ) + + def evaluate_genotype( + self, + genotype: mindspore.Tensor, + ) -> dict: + return self.evaluate(genotype) diff --git a/src/mindnlp/wizard/merge/evo/config.py b/src/mindnlp/wizard/merge/evo/config.py new file mode 100644 index 000000000..5b0caa63d --- /dev/null +++ b/src/mindnlp/wizard/merge/evo/config.py @@ -0,0 +1,84 @@ +# Copyright (C) 2025 Arcee AI +# SPDX-License-Identifier: LGPL-3.0-only +# Modified for MindSpore/Ascend NPU by MindNLP Wizard contributors. +# + +import logging +from typing import List, Optional + +from pydantic import BaseModel, model_validator + +from .genome import ModelGenomeDefinition + + +class TaskConfiguration(BaseModel, frozen=True): + name: str + weight: float = 1.0 + metric: str = "acc,none" + + @model_validator(mode="before") + def validate_string(cls, value): + if isinstance(value, str): + return {"name": value} + return value + + +class EvolMergeConfiguration(BaseModel, frozen=True): + genome: ModelGenomeDefinition + tasks: List[TaskConfiguration] + limit: Optional[int] = None + num_fewshot: Optional[int] = None + shuffle: bool = False + random_init: bool = False + apply_chat_template: bool = True + fewshot_as_multiturn: bool = True + + +NAUGHTY_PREFIXES = [ + "mmlu", + "hendrycks", + "agieval", + "gsm8k", + "hellaswag", + "winogrande", + "arc_", + "ai2_arc", + "truthfulqa", + "bigbench", + "piqa", + "openbookqa", + "leaderboard", +] + + +def check_for_naughty_config(config: EvolMergeConfiguration, allow: bool = False): + """ + Check if the given configuration is naughty and should be disallowed. + + mergekit-evolve is perfectly set up to directly optimize against the test set + of common benchmarks, which just makes the world a worse place. There are + cases where this is useful but it deserves a giant honking warning. + """ + suffix = "" + if not allow: + suffix = ( + " To proceed, set the " + "--i-understand-the-depths-of-the-evils-i-am-unleashing flag." + ) + for task in config.tasks: + for prefix in NAUGHTY_PREFIXES: + if task.name.startswith(prefix): + if task.name.endswith("_train"): + continue + + message = ( + f"Task {task.name} is a common benchmark task. " + "Optimizing against this task directly is unsporting at best " + "and outright malicious at worst. Using mergekit-evolve to " + "game benchmarks will be a black mark on your name for a " + f"thousand generations.{suffix}" + ) + if not allow: + raise ValueError(message) + else: + logging.warning(message) diff --git a/src/mindnlp/wizard/merge/evo/genome.py b/src/mindnlp/wizard/merge/evo/genome.py new file mode 100644 index 000000000..84523b975 --- /dev/null +++ b/src/mindnlp/wizard/merge/evo/genome.py @@ -0,0 +1,374 @@ +# Copyright (C) 2025 Arcee AI +# SPDX-License-Identifier: LGPL-3.0-only +# Modified for MindSpore/Ascend NPU by MindNLP Wizard contributors. +# + +import logging +import os +from typing import Any, Dict, List, Optional, Union + +import mindspore # pylint: disable=import-error +from mindspore import ops # pylint: disable=import-error +import numpy as np +import transformers +from pydantic import BaseModel, model_validator + +from ..common import ModelReference +from ..config import MergeConfiguration + +METHOD_PARAM_MAPS = { + "linear": ["weight"], + "task_arithmetic": ["weight"], + "ties": ["weight", "density"], + "dare_ties": ["weight", "density"], + "slerp": ["t"], +} + + +class InvalidGenotypeError(RuntimeError): + pass + + +class ModelGenomeDefinition(BaseModel, frozen=True): + models: List[ModelReference] + merge_method: str + base_model: Optional[ModelReference] = None + tokenizer_source: Optional[str] = None + layer_granularity: int = 0 + normalize: Optional[bool] = None + allow_negative_weights: bool = False + filters: Optional[List[str]] = None + smooth: bool = False + + @model_validator(mode="after") + def validate(self): + assert self.merge_method in METHOD_PARAM_MAPS, "Invalid merge method" + + if self.merge_method in ["ties", "dare_ties", "task_arithmetic"]: + assert self.base_model is not None, "base_model is required for this method" + + if self.merge_method == "slerp": + assert not self.smooth, "smooth is not supported for slerp merge method" + assert ( + not self.filters + ), "tensor name filtering is not supported for slerp merge method" + + return self + + +class ModelGenome: + definition: ModelGenomeDefinition + num_layers: int + _input_config_example: transformers.PretrainedConfig + + def __init__( + self, definition: ModelGenomeDefinition, trust_remote_code: bool = False + ): + self.definition = definition + + self._input_config_example = self.definition.models[0].config( + trust_remote_code=trust_remote_code + ) + self.num_layers = self._input_config_example.num_hidden_layers + + assert ( + self.definition.layer_granularity < 1 + or self.num_layers % self.definition.layer_granularity == 0 + ), "Number of layers must be a multiple of layer_granularity" + + def initial_genotype(self, random: bool = False) -> mindspore.Tensor: + """Generate an initial genotype for the given number of layers.""" + if self.definition.layer_granularity > 0: + n_layer_groups = self.num_layers // self.definition.layer_granularity + else: + n_layer_groups = 1 + n_param_sets = len(self.definition.filters or []) + 1 + n_models = len(self.definition.models) + n_params = len(METHOD_PARAM_MAPS[self.definition.merge_method]) + + shape = (n_layer_groups, n_models, n_param_sets, n_params) + if random: + return ops.rand(*shape) + else: + x0_t = ops.zeros(*shape) + x0_t[:, :, :, 0] = 1 / n_models + if n_params > 1: + x0_t[:, :, :, 1:] = 1 + return x0_t + + def genotype_merge_config( + self, genotype: Union[mindspore.Tensor, np.ndarray] + ) -> MergeConfiguration: + """Convert a genotype tensor to a mergekit configuration.""" + + genotype = self._to_mindspore(genotype) + + (n_layer_groups, n_models, n_param_sets, n_params) = genotype.shape + if self.definition.layer_granularity > 0: + assert n_layer_groups * self.definition.layer_granularity == self.num_layers + assert n_models == len(self.definition.models) + assert n_params == len(METHOD_PARAM_MAPS[self.definition.merge_method]) + + if self.definition.merge_method == "slerp": + slices = self._slerp_slices(genotype) + models = None + else: + param_arrays = {} + for param_idx, param in enumerate( + METHOD_PARAM_MAPS[self.definition.merge_method] + ): + values = genotype[:, :, :, param_idx] + if param == "density": + values = ops.clamp(ops.abs(values), 0, 1) + if not self.definition.allow_negative_weights and param in [ + "weight", + "t", + ]: + values = ops.abs(values) + param_arrays[param] = values + + if self.definition.smooth: + slices = None + models = self._smooth_config_models(n_param_sets, param_arrays) + else: + models = None + slices = self._discrete_config_slices( + n_layer_groups, n_param_sets, param_arrays + ) + + normalize = self.definition.normalize + if normalize is None: + normalize = self.definition.merge_method in ["ties", "dare_ties", "linear"] + return MergeConfiguration.model_validate( + { + "merge_method": self.definition.merge_method, + "slices": slices, + "models": models, + "parameters": { + "normalize": normalize, + "int8_mask": True, + }, + "dtype": "bfloat16", + "base_model": self.definition.base_model, + "tokenizer_source": self.definition.tokenizer_source, + } + ) + + def _discrete_config_slices( + self, + n_layer_groups: int, + n_param_sets: int, + param_arrays: Dict[str, mindspore.Tensor], + ) -> List[Dict]: + """Generate merge config output slices for non-interpolated parameters.""" + slices = [] + layer_step = ( + self.definition.layer_granularity + if self.definition.layer_granularity > 0 + else self.num_layers + ) + for slice_idx in range(n_layer_groups): + sources = [] + for model_idx, model in enumerate(self.definition.models): + params = {} + if n_param_sets > 1: + for param, values in param_arrays.items(): + params[param] = [] + for set_idx in range(n_param_sets): + value = values[ + slice_idx, + model_idx, + set_idx, + ] + filter_ = (self.definition.filters + [None])[set_idx] + params[param].append( + {"filter": filter_, "value": value.item()} + ) + else: + for param, values in param_arrays.items(): + params[param] = values[ + slice_idx, + model_idx, + 0, + ].item() + + sources.append( + { + "model": model, + "layer_range": [ + slice_idx * layer_step, + (slice_idx + 1) * layer_step, + ], + "parameters": params, + } + ) + + if self.definition.base_model and ( + self.definition.base_model not in self.definition.models + ): + sources.append( + { + "model": self.definition.base_model, + "layer_range": [ + slice_idx * layer_step, + (slice_idx + 1) * layer_step, + ], + } + ) + slices.append({"sources": sources}) + return slices + + def _smooth_config_models( + self, n_param_sets: int, param_arrays: Dict[str, mindspore.Tensor] + ) -> List[Dict]: + """Generate merge config model section with parameter interpolation.""" + models = [] + for model_idx, model in enumerate(self.definition.models): + params = {} + if n_param_sets > 1: + for param, values in param_arrays.items(): + params[param] = [] + for set_idx in range(n_param_sets): + value = values[:, model_idx, set_idx] + filter_ = (self.definition.filters + [None])[set_idx] + params[param].append( + { + "filter": filter_, + "value": _unpack_single_element(value.asnumpy().tolist()), + } + ) + else: + for param, values in param_arrays.items(): + params[param] = _unpack_single_element( + values[:, model_idx, 0].asnumpy().tolist() + ) + + models.append( + { + "model": model, + "layer_range": [0, self.num_layers], + "parameters": params, + } + ) + + if self.definition.base_model and ( + self.definition.base_model not in self.definition.models + ): + models.append({"model": self.definition.base_model}) + return models + + def _slerp_slices(self, genotype: mindspore.Tensor) -> List[Dict]: + """Generate merge config output slices for SLERP. + + This method is a bit more complex because it requires choosing the + two models with the highest weight for each layer group and calculating + the interpolation parameter t. Parameter interpolation and component + splitting are not supported because it's too hard and I don't want to. + """ + n_layer_groups, n_models, _, _ = genotype.shape + layer_step = ( + self.definition.layer_granularity + if self.definition.layer_granularity > 0 + else self.num_layers + ) + slices = [] + for slice_idx in range(n_layer_groups): + s = { + "sources": [ + { + "model": self.definition.models[i], + "layer_range": [ + slice_idx * layer_step, + (slice_idx + 1) * layer_step, + ], + } + for i in range(n_models) + ] + } + + chosen_values, chosen_indices = ops.topk( + genotype[slice_idx, :, 0, 0], 2 + ) + t = ops.softmax(chosen_values, axis=-1)[1].item() + s["parameters"] = {"t": t} + s["base_model"] = self.definition.models[chosen_indices[0].item()] + s["sources"] = [ + s["sources"][chosen_indices[0].item()], + s["sources"][chosen_indices[1].item()], + ] + if self.definition.tokenizer_source: + s["sources"][0]["parameters"] = {"weight": 1 - t} + s["sources"][1]["parameters"] = {"weight": t} + + if self.definition.base_model and ( + self.definition.base_model not in self.definition.models + ): + s["sources"].append( + { + "model": self.definition.base_model, + "layer_range": [ + slice_idx * layer_step, + (slice_idx + 1) * layer_step, + ], + } + ) + + slices.append(s) + return slices + + def _to_mindspore( + self, genotype: Union[mindspore.Tensor, np.ndarray] + ) -> mindspore.Tensor: + """Convert a genotype to a MindSpore tensor of the correct shape.""" + if not isinstance(genotype, mindspore.Tensor): + genotype = mindspore.Tensor(genotype, dtype=mindspore.float32) + if len(genotype.shape) == 1: + num_layer_groups = ( + self.num_layers // self.definition.layer_granularity + if self.definition.layer_granularity > 0 + else 1 + ) + genotype = genotype.view( + num_layer_groups, + len(self.definition.models), + len(self.definition.filters or []) + 1, + -1, + ) + + if len(genotype.shape) != 4: + logging.error(f"Invalid genotype shape: {genotype.shape}") + raise InvalidGenotypeError( + "Invalid genotype shape - must be 4D tensor or 1D array" + ) + + return genotype + + def genotype_to_param_arrays( + self, genotype: Union[mindspore.Tensor, np.ndarray] + ) -> Dict[str, mindspore.Tensor]: + """Convert a genotype tensor to a dictionary of numpy arrays.""" + genotype = self._to_mindspore(genotype) + + res = {} + for idx, param_name in enumerate( + METHOD_PARAM_MAPS[self.definition.merge_method] + ): + for model_idx, model in enumerate(self.definition.models): + model_name = os.path.basename(model.model.path) + for set_idx, filter_ in enumerate( + (self.definition.filters or []) + [None] + ): + suffix = "" + if filter_ is not None: + suffix = f"_{filter_}" + res[f"{model_name}_{param_name}{suffix}"] = genotype[ + :, model_idx, set_idx, idx + ] + + return res + + +def _unpack_single_element(x: List) -> Any: + if len(x) == 1: + return x[0] + return x diff --git a/src/mindnlp/wizard/merge/evo/helpers.py b/src/mindnlp/wizard/merge/evo/helpers.py new file mode 100644 index 000000000..e39eec78a --- /dev/null +++ b/src/mindnlp/wizard/merge/evo/helpers.py @@ -0,0 +1,118 @@ +# Copyright (C) 2025 Arcee AI +# SPDX-License-Identifier: LGPL-3.0-only +# Modified for MindSpore/Ascend NPU by MindNLP Wizard contributors. +# + +import logging +import os +import shutil +import tempfile +from typing import Any, Dict, List, Optional, Union + +import lm_eval # pylint: disable=import-error +import lm_eval.api.model # pylint: disable=import-error +import lm_eval.models.huggingface # pylint: disable=import-error +import lm_eval.tasks # pylint: disable=import-error +import mindspore # pylint: disable=import-error +import ray # pylint: disable=import-error +import ray.util.queue # pylint: disable=import-error +import ray.util.scheduling_strategies # pylint: disable=import-error + +from .config import TaskConfiguration +from .genome import InvalidGenotypeError, ModelGenome +from .monkeypatch import monkeypatch_lmeval_vllm +from ..merge import run_merge +from ..options import MergeOptions + + +def _eval_model( + model: Union[str, lm_eval.api.model.LM], + tasks: List[TaskConfiguration], + model_args: Optional[Dict[str, Any]] = None, + task_manager: Optional[lm_eval.tasks.TaskManager] = None, + **kwargs, +) -> Dict[str, Any]: + results = lm_eval.simple_evaluate( + model=model, + model_args=model_args, + tasks=list({task.name for task in tasks}), + log_samples=False, + verbosity="WARNING", + task_manager=task_manager, + **kwargs, + ) + + logging.info(results["results"]) + res = 0 + for task in tasks: + res += results["results"][task.name][task.metric] * task.weight + return {"score": res, "results": results["results"]} + + +def evaluate_model( # pylint: disable=too-many-positional-arguments + merged_path: str, + tasks: List[TaskConfiguration], + num_fewshot: Optional[int], + limit: Optional[int], + vllm: bool, + batch_size: Optional[int] = None, + task_manager: Optional[lm_eval.tasks.TaskManager] = None, + model_kwargs: Optional[Dict[str, Any]] = None, + **kwargs, +) -> dict: + monkeypatch_lmeval_vllm() + try: + model_args = { + "pretrained": merged_path, + "dtype": "bfloat16", + **(model_kwargs or {}), + } + if vllm: + model_args["gpu_memory_utilization"] = 0.8 + model_args["tensor_parallel_size"] = 1 + model_args["batch_size"] = "auto" + model_args["max_model_len"] = 4096 + else: + model_args["use_cache"] = True + + res = _eval_model( + "vllm" if vllm else "huggingface", + tasks, + model_args, + num_fewshot=num_fewshot, + limit=limit, + batch_size=batch_size, + task_manager=task_manager, + **kwargs, + ) + return res + finally: + shutil.rmtree(merged_path) + + +evaluate_model_ray = ray.remote(num_cpus=1, num_gpus=1.0)(evaluate_model) + + +def merge_model( + genotype: mindspore.Tensor, + genome: ModelGenome, + model_storage_path: str, + merge_options: MergeOptions, +) -> str: + try: + cfg = genome.genotype_merge_config(genotype) + except InvalidGenotypeError as e: + logging.error("Invalid genotype", exc_info=e) + return None + os.makedirs(model_storage_path, exist_ok=True) + res = tempfile.mkdtemp(prefix="merged", dir=model_storage_path) + run_merge(cfg, out_path=res, options=merge_options) + return res + + +merge_model_ray = ray.remote( + num_cpus=1, + num_gpus=1, + max_retries=3, + retry_exceptions=[ConnectionError], +)(merge_model) diff --git a/src/mindnlp/wizard/merge/evo/monkeypatch.py b/src/mindnlp/wizard/merge/evo/monkeypatch.py new file mode 100644 index 000000000..45ce021ba --- /dev/null +++ b/src/mindnlp/wizard/merge/evo/monkeypatch.py @@ -0,0 +1,152 @@ +# Copyright (C) 2025 Arcee AI +# SPDX-License-Identifier: LGPL-3.0-only +# Modified for MindSpore/Ascend NPU by MindNLP Wizard contributors. +# + +import logging + +from mindspore.common import initializer # pylint: disable=import-error +import transformers + +LOG = logging.getLogger(__name__) + + +def monkeypatch_lmeval_shuffle(): + """Monkeypatch lm_eval to shuffle the dataset after downloading.""" + import lm_eval.api.task # pylint: disable=import-error + + if hasattr(lm_eval.api.task.Task, "_monkey_patched"): + return + + _old_task_dl = lm_eval.api.task.Task.download + + def _dl_shuffled(self: lm_eval.api.task.Task, *args, **kwargs): + _old_task_dl(self, *args, **kwargs) + self.dataset = self.dataset.shuffle() + + lm_eval.api.task.Task.download = _dl_shuffled + + _old_ct_dl = lm_eval.api.task.ConfigurableTask.download + + def _ct_dl_shuffled(self, *args, **kwargs): + _old_ct_dl(self, *args, **kwargs) + self.dataset = self.dataset.shuffle() + + lm_eval.api.task.ConfigurableTask.download = _ct_dl_shuffled + + lm_eval.api.task.Task._monkey_patched = True + print("monkey has been patched") + + +def monkeypatch_tqdm(lm_eval: bool = True, mergekit: bool = True): + """Patch lm_eval & wizard to use Ray's tqdm for progress bars.""" + + from ray.experimental.tqdm_ray import tqdm as tqdm_ray # pylint: disable=import-error + + def _tqdm_wrap(iterable=None, disable: bool = False, **kwargs): + if disable: + if iterable is not None: + return iterable + return lambda x: x + res = tqdm_ray(iterable=iterable, **kwargs, flush_interval_s=1.0) + res.refresh() + return res + + def _patch_lm_eval(): + import lm_eval # pylint: disable=import-error + + if hasattr(lm_eval, "_mk_tqdm_patched"): + return + + import lm_eval.api.metrics # pylint: disable=import-error + import lm_eval.api.model # pylint: disable=import-error + import lm_eval.api.task # pylint: disable=import-error + import lm_eval.models.huggingface # pylint: disable=import-error + import lm_eval.models.vllm_causallms # pylint: disable=import-error + + for module in ( + lm_eval.models.huggingface, + lm_eval.models.vllm_causallms, + lm_eval.api.model, + lm_eval.api.task, + lm_eval.api.metrics, + ): + setattr(module, "tqdm", _tqdm_wrap) + + lm_eval._mk_tqdm_patched = True + + if lm_eval: + _patch_lm_eval() + + if mergekit: + del mergekit + + from .. import graph as wizard_graph + from .. import merge as wizard_merge + from .. import tokenizer as wizard_tokenizer + + fake_module = type("fake_module", (), {"tqdm": staticmethod(_tqdm_wrap)})() + + wizard_graph.tqdm = fake_module + wizard_merge.tqdm = fake_module + wizard_tokenizer.tqdm = fake_module + + +def monkeypatch_lmeval_vllm(): + import lm_eval.models.vllm_causallms # pylint: disable=import-error + + lm_eval.models.vllm_causallms.VLLM.AUTO_MODEL_CLASS = ( + transformers.AutoModelForCausalLM + ) + + +class NoInit: + """Context manager that disables weight initialization for faster model + instantiation. Patches ``mindspore.common.initializer`` entry-points + used by transformers model construction so that allocated parameters are + left uninitialised (they will be overwritten by the merge anyway). + """ + + def __enter__(self): + def noop(*args, **kwargs): + pass + + self._originals = { + "kaiming_uniform_": getattr(initializer, "HeUniform", None), + "uniform_": getattr(initializer, "Uniform", None), + "normal_": getattr(initializer, "Normal", None), + } + + try: + import torch.nn.init as _init + + (k, u, n) = ( + _init.kaiming_uniform_, + _init.uniform_, + _init.normal_, + ) + _init.kaiming_uniform_ = noop + _init.uniform_ = noop + _init.normal_ = noop + self._torch_funcs = (k, u, n) + except ImportError: + self._torch_funcs = None + + transformers.modeling_utils._init_weights = False + + def __exit__(self, *args): + if self._torch_funcs is not None: + try: + import torch.nn.init as _init + + (k, u, n) = self._torch_funcs + _init.kaiming_uniform_ = k + _init.uniform_ = u + _init.normal_ = n + except ImportError: + LOG.debug( + "Torch is unavailable while leaving NoInit context; " + "skipping torch init restoration" + ) + + transformers.modeling_utils._init_weights = True diff --git a/src/mindnlp/wizard/merge/evo/strategy.py b/src/mindnlp/wizard/merge/evo/strategy.py new file mode 100644 index 000000000..e6d1c4baf --- /dev/null +++ b/src/mindnlp/wizard/merge/evo/strategy.py @@ -0,0 +1,317 @@ +# Copyright (C) 2025 Arcee AI +# SPDX-License-Identifier: LGPL-3.0-only +# Modified for MindSpore/Ascend NPU by MindNLP Wizard contributors. +# + +import asyncio +import logging +import os +from abc import ABC, abstractmethod +from typing import Dict, List, Optional, Tuple, Union + +import lm_eval.tasks # pylint: disable=import-error +import numpy as np +import ray # pylint: disable=import-error +import ray.util.queue # pylint: disable=import-error +import ray.util.scheduling_strategies # pylint: disable=import-error +import transformers + +from ..common import get_accelerator_count +from .actors import InMemoryMergeEvaluator, OnDiskMergeEvaluator +from .config import EvolMergeConfiguration +from .genome import ModelGenome +from .helpers import evaluate_model_ray, merge_model_ray +from ..options import MergeOptions + + +class EvaluationStrategyBase(ABC): + def __init__( # pylint: disable=too-many-positional-arguments + self, + config: EvolMergeConfiguration, + genome: ModelGenome, + merge_options: MergeOptions, + num_gpus: Optional[int] = None, + batch_size: Optional[int] = None, + task_search_path: Union[str, List[str], None] = None, + model_storage_path: Optional[str] = None, + quantization_config: Optional[transformers.BitsAndBytesConfig] = None, + ): + self.config = config + self.genome = genome + self.merge_options = merge_options + self.num_gpus = num_gpus or get_accelerator_count( + self.merge_options.device + ) + self.batch_size = batch_size + self.task_manager = lm_eval.tasks.TaskManager(include_path=task_search_path) + self.model_storage_path = model_storage_path + self.quantization_config = quantization_config + if self.model_storage_path: + os.makedirs(self.model_storage_path, exist_ok=True) + + @abstractmethod + def evaluate_genotypes(self, genotypes: List[np.ndarray]) -> List[dict]: + pass + + @abstractmethod + def evaluate_genotype(self, genotype: np.ndarray) -> dict: + pass + + +class ActorPoolEvaluationStrategy(EvaluationStrategyBase): + """ + Uses a fixed-size pool of actors to evaluate genotypes in parallel. + """ + + def __init__( + self, + *args, + in_memory: bool = False, + vllm: bool = False, + **kwargs, + ): + super().__init__(*args, **kwargs) + if in_memory: + self.actor_cls = InMemoryMergeEvaluator + else: + self.actor_cls = OnDiskMergeEvaluator + + self.actor_pool = ray.util.ActorPool( + [ + self.actor_cls.remote( + self.config, + self.genome, + self.merge_options, + model_storage_path=self.model_storage_path, + vllm=vllm, + batch_size=self.batch_size, + task_manager=self.task_manager, + quantization_config=self.quantization_config, + ) + for _ in range(self.num_gpus) + ] + ) + + def evaluate_genotypes(self, genotypes: List[np.ndarray]) -> List[dict]: + return list( + self.actor_pool.map( + lambda a, x: a.evaluate_genotype.remote(x), + genotypes, + ) + ) + + def evaluate_genotype(self, genotype: np.ndarray) -> dict: + return self.evaluate_genotypes([genotype])[0] + + +@ray.remote +class BufferedRayEvaluationStrategyActor: + def __init__( # pylint: disable=too-many-positional-arguments + self, + config: EvolMergeConfiguration, + genome: ModelGenome, + merge_options: MergeOptions, + vllm: bool = False, + num_gpus: Optional[int] = None, + batch_size: Optional[int] = None, + task_manager: Optional[lm_eval.tasks.TaskManager] = None, + model_storage_path: Optional[str] = None, + quantization_config: Optional[transformers.BitsAndBytesConfig] = None, + ): + self.config = config + self.genome = genome + self.merge_options = merge_options + self.vllm = vllm + self.num_gpus = num_gpus or get_accelerator_count( + self.merge_options.device + ) + self.input_queue = [] + self.batch_size = batch_size + self.task_manager = task_manager + self.model_storage_path = model_storage_path + self.quantization_config = quantization_config + self._shutdown = False + + async def evaluate_genotype(self, genotype: np.ndarray): + future_result = asyncio.Future() + self.input_queue.append((genotype, future_result)) + return await future_result + + async def process_queue(self): + merging: Dict[ray.ObjectRef, asyncio.Future] = {} + merged: List[Tuple[asyncio.Future, ray.ObjectRef]] = [] + evaluating: Dict[ray.ObjectRef, asyncio.Future] = {} + + logging.info("Starting processing loop") + + try: + while not self._shutdown: + while self.input_queue and (len(merging) + len(merged) < self.num_gpus): + genotype, future_result = self.input_queue.pop(0) + merging[ + merge_model_ray.remote( + genotype, + self.genome, + self.model_storage_path, + self.merge_options, + ) + ] = future_result + + while merged and len(evaluating) < self.num_gpus: + future_result, merged_path = merged.pop() + kwargs = {} + if self.quantization_config is not None: + kwargs["quantization_config"] = self.quantization_config + evaluating[ + evaluate_model_ray.remote( + merged_path, + self.config.tasks, + num_fewshot=self.config.num_fewshot, + limit=self.config.limit, + vllm=self.vllm, + batch_size=self.batch_size, + task_manager=self.task_manager, + apply_chat_template=self.config.apply_chat_template, + fewshot_as_multiturn=self.config.fewshot_as_multiturn, + **kwargs, + ) + ] = future_result + + ready, _ = ray.wait( + list(merging.keys()) + list(evaluating.keys()), + num_returns=1, + fetch_local=False, + timeout=1, + ) + for r in ready: + if r in merging: + future_result = merging.pop(r) + merged.append((future_result, r)) + elif r in evaluating: + future_result = evaluating.pop(r) + future_result.set_result(await r) + + if ( + not self.input_queue + and not merging + and not merged + and not evaluating + ): + await asyncio.sleep(1) + except Exception as e: + logging.error("Error in processing loop", exc_info=e) + raise + + async def shutdown(self): + self._shutdown = True + + +class BufferedRayEvaluationStrategy(EvaluationStrategyBase): + def __init__( + self, + *args, + vllm: bool = False, + in_memory: bool = False, + **kwargs, + ): + if in_memory: + raise ValueError("In-memory evaluation is not supported for buffered mode") + + super().__init__(*args, **kwargs) + self.actor = BufferedRayEvaluationStrategyActor.options( + max_concurrency=1000 + ).remote( + self.config, + self.genome, + self.merge_options, + model_storage_path=self.model_storage_path, + vllm=vllm, + num_gpus=self.num_gpus, + task_manager=self.task_manager, + batch_size=self.batch_size, + quantization_config=self.quantization_config, + ) + self.actor.process_queue.remote() + + def evaluate_genotypes(self, genotypes: List[np.ndarray]) -> List[dict]: + return ray.get([self.actor.evaluate_genotype.remote(x) for x in genotypes]) + + def evaluate_genotype(self, genotype: np.ndarray) -> dict: + return ray.get(self.actor.evaluate_genotype.remote(genotype)) + + +@ray.remote +def evaluate_genotype_serial( # pylint: disable=too-many-positional-arguments + genotype: np.ndarray, + config: EvolMergeConfiguration, + genome: ModelGenome, + merge_options: MergeOptions, + model_storage_path: Optional[str] = None, + vllm: bool = False, + batch_size: Optional[int] = None, + task_manager: Optional[lm_eval.tasks.TaskManager] = None, + quantization_config: Optional[transformers.BitsAndBytesConfig] = None, +): + pg = ray.util.placement_group([{"CPU": 1, "GPU": 1}], strategy="STRICT_PACK") + strat = ray.util.scheduling_strategies.PlacementGroupSchedulingStrategy( + placement_group=pg + ) + merged_path = merge_model_ray.options(scheduling_strategy=strat).remote( + genotype, genome, model_storage_path, merge_options + ) + if not merged_path: + return {"score": None, "results": None} + kwargs = {} + if quantization_config is not None: + kwargs["quantization_config"] = quantization_config + res = ray.get( + evaluate_model_ray.options(scheduling_strategy=strat).remote( + merged_path, + config.tasks, + num_fewshot=config.num_fewshot, + limit=config.limit, + vllm=vllm, + batch_size=batch_size, + task_manager=task_manager, + apply_chat_template=config.apply_chat_template, + fewshot_as_multiturn=config.fewshot_as_multiturn, + **kwargs, + ) + ) + ray.util.remove_placement_group(pg) + return res + + +class SerialEvaluationStrategy(EvaluationStrategyBase): + def __init__( + self, + *args, + vllm: bool = False, + in_memory: bool = False, + **kwargs, + ): + self.vllm = vllm + if in_memory: + raise ValueError("In-memory evaluation is not supported for serial mode") + super().__init__(*args, **kwargs) + + def evaluate_genotypes(self, genotypes: List[np.ndarray]) -> List[dict]: + return ray.get( + [ + evaluate_genotype_serial.remote( + x, + self.config, + self.genome, + self.merge_options, + model_storage_path=self.model_storage_path, + vllm=self.vllm, + batch_size=self.batch_size, + task_manager=self.task_manager, + quantization_config=self.quantization_config, + ) + for x in genotypes + ] + ) + + def evaluate_genotype(self, genotype: np.ndarray) -> dict: + return self.evaluate_genotypes([genotype])[0] diff --git a/src/mindnlp/wizard/merge/graph.py b/src/mindnlp/wizard/merge/graph.py new file mode 100644 index 000000000..d17b5f789 --- /dev/null +++ b/src/mindnlp/wizard/merge/graph.py @@ -0,0 +1,461 @@ +# Originally from MergeKit (https://github.com/arcee-ai/mergekit) +# Copyright (C) 2025 Arcee AI +# SPDX-License-Identifier: LGPL-3.0-only +# Modified for MindSpore/Ascend NPU by MindNLP Wizard contributors. + +"""Computational graph execution engine.""" + +import logging +import os +import resource +import time +from abc import ABC, abstractmethod +from typing import Any, Dict, Iterable, Iterator, List, Optional, Tuple, Union + +import mindspore +import networkx +import tqdm +from pydantic import BaseModel, ConfigDict +from typing_extensions import Generic, TypeVar + +ValueT = TypeVar("ValueT") +LOG = logging.getLogger(__name__) + + +class Task(ABC, BaseModel, Generic[ValueT], frozen=True): + """Abstract base class representing a task in a computational graph. + + Subclasses must implement ``arguments`` and ``execute``. + """ + model_config = ConfigDict(arbitrary_types_allowed=True) + + @abstractmethod + def arguments(self) -> Dict[str, "Task"]: + ... + + @abstractmethod + def execute(self, **kwargs) -> ValueT: + ... + + def priority(self) -> int: + return 0 + + def group_label(self) -> Optional[str]: + return None + + def uses_accelerator(self) -> bool: + return False + + def main_thread_only(self) -> bool: + return False + + def duplicate_per_gpu(self) -> bool: + return False + + def cost_hint(self) -> Optional[Dict[str, float]]: + return None + + +class TaskUniverse: + """Container for tasks and their dependency relationships.""" + + tasks: List[Task] + task_to_index: Dict[Task, int] + task_arguments: Dict[int, Dict[str, int]] + _type_id_to_index: Dict[Tuple[type, int], int] + + def __init__(self, tasks: Optional[Iterable[Task]] = None): + self.tasks = [] + self.task_to_index = {} + self.task_arguments = {} + self._type_id_to_index = {} + if tasks is not None: + for task in tasks: + self.add_task(task) + + def add_task(self, task: Task, recursive: bool = True) -> "TaskHandle": + _ti_key = (type(task), id(task)) + if _ti_key in self._type_id_to_index: + index = self._type_id_to_index[_ti_key] + assert ( + self.tasks[index] == task + ), "Task modified after being added to universe" + return TaskHandle(self, index) + + index = self.task_to_index.setdefault(task, len(self.tasks)) + if index < len(self.tasks): + return TaskHandle(self, index) + self.tasks.append(task) + self._type_id_to_index[_ti_key] = index + + if recursive: + self.task_arguments[index] = {} + for k, v in task.arguments().items(): + self.task_arguments[index][k] = self.add_task( + v, recursive=True + )._index + return TaskHandle(self, index) + + def get_handle(self, task: Task) -> Optional["TaskHandle"]: + if task not in self.task_to_index: + return None + return TaskHandle(self, self.task_to_index[task]) + + +class TaskHandle: + """Lightweight reference to a task within a :class:`TaskUniverse`.""" + + __slots__ = ["_universe", "_index"] + _universe: TaskUniverse + _index: int + + def __init__(self, universe: TaskUniverse, index: int): + self._universe = universe + self._index = index + + def task(self) -> Task: + return self._universe.tasks[self._index] + + def arguments(self) -> Dict[str, "TaskHandle"]: + return { + k: TaskHandle(self._universe, v) + for k, v in self._universe.task_arguments[self._index].items() + } + + def __eq__(self, other): + if not isinstance(other, TaskHandle): + return False + if self._index != other._index: + return False + if self._universe is not other._universe: + return False + return True + + def __hash__(self): + return self._index + + def __str__(self): + return f"TaskHandle({type(self.task()).__name__}, {self._index})" + + __repr__ = __str__ + + +class ExecutionSchedule: + """Ordered schedule of tasks with lifecycle annotations.""" + + tasks: List[TaskHandle] + last_use_index: Dict[TaskHandle, int] + + def __init__( + self, tasks: List[TaskHandle], last_use_index: Dict[TaskHandle, int] + ): + self.tasks = tasks + self.last_use_index = last_use_index + + +def build_schedule( + targets: List[TaskHandle], + cached_values: Dict[TaskHandle, Any], +) -> ExecutionSchedule: + """Build a topologically-sorted execution schedule.""" + if not targets: + return ExecutionSchedule(tasks=[], last_use_index={}) + + universe = targets[0]._universe + assert all( + t._universe is universe for t in targets + ), "All tasks must be from the same universe" + + dummy_handle = TaskHandle(universe, -1) + edge_tups: List[Tuple[TaskHandle, TaskHandle]] = [] + + explored: set = set() + to_explore: set = set(targets) + while to_explore: + task = to_explore.pop() + if task in explored: + continue + explored.add(task) + if task in (cached_values or {}): + continue + for dep in task.arguments().values(): + to_explore.add(dep) + edge_tups.append((dep, task)) + + for target in targets: + edge_tups.append((dummy_handle, target)) + + def _compare_key(node: TaskHandle) -> Tuple[str, int]: + if node._index < 0: + return ("", 0) + task = node.task() + return ( + task.group_label() or "", + -task.priority(), + ) + + graph = networkx.DiGraph(edge_tups) + schedule: List[TaskHandle] = [ + node + for node in networkx.lexicographical_topological_sort( + graph, key=_compare_key + ) + if (node != dummy_handle) and node not in (cached_values or {}) + ] + + last_use_index: Dict[TaskHandle, int] = {} + for idx, task in reversed(list(enumerate(schedule))): + for dep in task.arguments().values(): + if dep not in last_use_index: + last_use_index[dep] = idx + if task not in last_use_index: + last_use_index[task] = idx + for task in cached_values or {}: + if task not in last_use_index: + last_use_index[task] = len(schedule) + 1 + + return ExecutionSchedule(tasks=schedule, last_use_index=last_use_index) + + +# --------------------------------------------------------------------------- +# Device helpers for MindSpore +# --------------------------------------------------------------------------- + +def _parse_device(spec: str) -> Tuple[str, Optional[int]]: + """Parse a device specifier like ``"Ascend:0"`` or ``"CPU"``. + + Returns ``(target, device_id)`` where *target* is one of + ``"CPU"``, ``"Ascend"``, ``"GPU"`` and *device_id* may be ``None``. + """ + parts = str(spec).split(":") + target = parts[0] + device_id = int(parts[1]) if len(parts) > 1 else None + return target, device_id + + +def _runtime_device_target(spec: str) -> str: + """Map user device specs to runtime strings used by mindtorch Tensor.to().""" + target, device_id = _parse_device(spec) + key = str(target).strip().lower() + if key in ("ascend", "npu", "gpu", "cuda"): + base = "cuda" + elif key == "cpu": + base = "cpu" + else: + base = key + if device_id is not None and base != "cpu": + return f"{base}:{device_id}" + return base + + +def _mindspore_device_target(spec: str) -> str: + """Map user device specs to MindSpore move_to targets.""" + target, _ = _parse_device(spec) + key = str(target).strip().lower() + mapping = { + "cpu": "CPU", + "ascend": "Ascend", + "npu": "Ascend", + "gpu": "GPU", + "cuda": "GPU", + } + return mapping.get(key, str(target)) + + +class Executor: + """Schedule and execute a DAG of tasks. + + ``math_device`` / ``storage_device`` are device specifier strings such as + ``"CPU"`` or ``"Ascend:0"``. + """ + + math_device: str + storage_device: str + universe: TaskUniverse + targets: List[TaskHandle] + schedule: ExecutionSchedule + cached_values: Optional[Dict[TaskHandle, Any]] + + def __init__( + self, + targets: Union[List[Task], List[TaskHandle]], + math_device: str = "CPU", + storage_device: str = "CPU", + cached_values: Optional[Dict[TaskHandle, Any]] = None, + ): + self.cached_values = cached_values + self.math_device = math_device + self.storage_device = storage_device + self._task_metrics: List[Dict[str, Any]] = [] + + if targets and isinstance(targets[0], Task): + universe = TaskUniverse(targets) + targets = [universe.add_task(t) for t in targets] + elif targets and isinstance(targets[0], TaskHandle): + universe = targets[0]._universe + elif not targets: + universe = TaskUniverse() + else: + raise ValueError( + "Targets must be a list of Task or TaskHandle instances" + ) + + self.universe = universe + self.targets = targets + self.schedule = build_schedule( + targets, + cached_values=cached_values, + ) + + # ------------------------------------------------------------------ + + def _run( + self, + quiet: bool = False, + desc: Optional[str] = None, + ) -> Iterator[Tuple[TaskHandle, Any]]: + last_use_index = self.schedule.last_use_index + + values: Dict[TaskHandle, Any] = {} + if self.cached_values: + for task, value in self.cached_values.items(): + values[task] = value + + for idx, task_handle in ( + pbar := tqdm.tqdm( + list(enumerate(self.schedule.tasks)), + disable=quiet, + desc=desc or "Executing graph", + ) + ): + task = task_handle.task() + use_math_device = task.uses_accelerator() + + arguments: Dict[str, Any] = {} + for name, dep_handle in task_handle.arguments().items(): + value = values[dep_handle] + if use_math_device: + value = self._move_tensors(value, self.math_device) + arguments[name] = value + del value + + start_t = time.perf_counter() + res = task.execute(**arguments) + elapsed = time.perf_counter() - start_t + self._task_metrics.append( + { + "task": type(task).__name__, + "wait_ms": 0.0, + "run_ms": elapsed * 1000.0, + } + ) + del arguments + res = self._move_tensors(res, self.storage_device) + + values[task_handle] = res + del res + + if task_handle in self.targets: + yield (task_handle, values[task_handle]) + + expired = [ + key for key in values if idx >= last_use_index[key] + ] + for key in expired: + del values[key] + + del values + del pbar + + def run( + self, + quiet: bool = False, + desc: Optional[str] = None, + ) -> Iterator[Tuple[Task, Any]]: + for handle, value in self._run(quiet=quiet, desc=desc): + yield (handle.task(), value) + + def execute(self, desc: Optional[str] = None) -> None: + for _ in self.run(desc=desc): + pass + + def metrics_snapshot(self) -> Dict[str, Any]: + rss_mb = 0.0 + try: + rss_kb = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss + # Linux reports ru_maxrss in KB. + rss_mb = float(rss_kb) / 1024.0 + except Exception: + pass + return { + "executor": "single_device", + "pid": os.getpid(), + "task_count": len(self._task_metrics), + "tasks": list(self._task_metrics), + "queue_depth_samples": [], + "backpressure_trigger_count": 0, + "rss_peak_mb": rss_mb, + "npu_used_peak_mb": None, + } + + # ------------------------------------------------------------------ + # MindSpore tensor device movement + # ------------------------------------------------------------------ + + @staticmethod + def _move_tensors( + value: Any, + device: str, + non_blocking: Optional[bool] = None, + ) -> Any: + """Recursively move MindSpore tensors to *device*. + + MindSpore's device model is context-based. When both math and + storage targets are the same (the common single-NPU case) this is + effectively a no-op. Cross-device copies go through NumPy when the + source and target differ. + """ + if isinstance(value, mindspore.Tensor): + runtime_target = _runtime_device_target(device) + try: + if hasattr(value, "to"): + try: + return value.to( + device=runtime_target, + non_blocking=(runtime_target != "cpu"), + ) + except TypeError: + return value.to(runtime_target) + except Exception as exc: + LOG.debug( + "Tensor.to(%s) failed (%s: %s), falling back to move_to", + runtime_target, + type(exc).__name__, + exc, + ) + + ms_target = _mindspore_device_target(device) + try: + return value.move_to(ms_target) + except Exception as exc: + LOG.warning( + "Failed to move tensor to %s (%s: %s); keeping original tensor", + ms_target, + type(exc).__name__, + exc, + ) + return value + elif isinstance(value, dict): + return { + k: Executor._move_tensors(v, device, non_blocking) + for k, v in value.items() + } + elif isinstance(value, list): + return [ + Executor._move_tensors(v, device, non_blocking) for v in value + ] + elif isinstance(value, tuple): + return tuple( + Executor._move_tensors(v, device, non_blocking) for v in value + ) + return value diff --git a/src/mindnlp/wizard/merge/io/__init__.py b/src/mindnlp/wizard/merge/io/__init__.py new file mode 100644 index 000000000..69607768d --- /dev/null +++ b/src/mindnlp/wizard/merge/io/__init__.py @@ -0,0 +1,13 @@ +# Copyright (C) 2025 Arcee AI +# SPDX-License-Identifier: LGPL-3.0-only +# Modified for MindSpore/Ascend NPU by MindNLP Wizard contributors. + +from .lazy_tensor_loader import LazyTensorLoader, ShardedTensorIndex, ShardInfo +from .tensor_writer import TensorWriter + +__all__ = [ + "LazyTensorLoader", + "ShardedTensorIndex", + "ShardInfo", + "TensorWriter", +] diff --git a/src/mindnlp/wizard/merge/io/_device.py b/src/mindnlp/wizard/merge/io/_device.py new file mode 100644 index 000000000..8ae08a70d --- /dev/null +++ b/src/mindnlp/wizard/merge/io/_device.py @@ -0,0 +1,99 @@ +# Copyright (c) MindNLP Wizard contributors. +# Licensed under the Apache License, Version 2.0. + +"""Centralised device-movement helpers for tensor I/O. + +Every loader / writer that needs to move a tensor to a specific device +should import from here instead of duplicating the logic. +""" + +import logging +from typing import Optional + +import mindspore + +LOG = logging.getLogger(__name__) +_MOVE_WARNED_TARGETS: set = set() + + +def normalize_device_target(device: Optional[str]) -> Optional[str]: + """Normalize device spec strings to MindSpore target names (CPU/Ascend/GPU).""" + if not device: + return None + target = str(device).split(":", maxsplit=1)[0].strip().lower() + mapping = { + "cpu": "CPU", + "ascend": "Ascend", + "gpu": "GPU", + } + return mapping.get(target) + + +def runtime_device_target(device: Optional[str]) -> Optional[str]: + """Normalize to torch-style device strings used by mindtorch ``Tensor.to()``.""" + if not device: + return None + raw = str(device).strip().lower() + if ":" in raw: + target, index = raw.split(":", 1) + else: + target, index = raw, None + if target in ("ascend", "npu", "gpu", "cuda"): + base = "cuda" + elif target == "cpu": + base = "cpu" + else: + base = target + if index and base != "cpu": + return f"{base}:{index}" + return base + + +def move_tensor_to_device( + tensor: mindspore.Tensor, + device: Optional[str], + *, + caller: str = "", +) -> mindspore.Tensor: + """Best-effort device movement for a MindSpore tensor. + + Parameters + ---------- + tensor : mindspore.Tensor + device : str or None + caller : str + Free-form label included in warning messages for diagnostics. + """ + rt_target = runtime_device_target(device) + ms_target = normalize_device_target(device) + if not rt_target and not ms_target: + return tensor + + if rt_target and hasattr(tensor, "to"): + try: + return tensor.to(device=rt_target, non_blocking=(rt_target != "cpu")) + except TypeError: + try: + return tensor.to(rt_target) + except Exception: + pass + except Exception: + pass + + target = ms_target + if not target: + return tensor + try: + return tensor.move_to(target) + except Exception as exc: + if target not in _MOVE_WARNED_TARGETS: + label = f" ({caller})" if caller else "" + LOG.warning( + "Failed to move tensor to %s%s (%s: %s); keeping original tensor", + target, + label, + type(exc).__name__, + exc, + ) + _MOVE_WARNED_TARGETS.add(target) + return tensor diff --git a/src/mindnlp/wizard/merge/io/lazy_ckpt.py b/src/mindnlp/wizard/merge/io/lazy_ckpt.py new file mode 100644 index 000000000..273c67047 --- /dev/null +++ b/src/mindnlp/wizard/merge/io/lazy_ckpt.py @@ -0,0 +1,328 @@ +# Copyright (c) MindNLP Wizard contributors. +# Licensed under the Apache License, Version 2.0. + +""" +B-level lazy loader for MindSpore ``.ckpt`` checkpoint files. + +Scans the protobuf wire format to build an index of tensor locations +(file offsets and byte lengths), then reads only the requested tensor +bytes from disk on demand — no full file materialisation. + +Protobuf schema (from ``mindspore/ccsrc/utils/checkpoint.proto``):: + + message Checkpoint { + message Value { + required string tag = 1; + oneof value { + TensorProto tensor = 2; + MapTensorProto maptensor = 3; + } + } + repeated Value value = 1; + } + message TensorProto { + repeated int64 dims = 1; + required string tensor_type = 2; + required bytes tensor_content = 3; + } + +Large tensors may be sliced into multiple ``Value`` entries sharing +the same ``tag``; their ``tensor_content`` chunks are concatenated on read. +""" + +from __future__ import annotations + +import logging +import os +from dataclasses import dataclass +from typing import Dict, List, Optional, Tuple + +import numpy + +LOG = logging.getLogger(__name__) + + +class CkptFormatNotSupported(Exception): + """The .ckpt file uses features not handled by the lazy loader + (encryption, MapTensor, unknown wire types, …). + + Callers should fall back to ``DumbCkptLoader`` which delegates + to ``mindspore.load_checkpoint``. + """ + + +# ── MindSpore dtype string → numpy dtype ────────────────────────────────── + +_TENSOR_TYPE_TO_NUMPY = { + # str(mindspore.dtype) format — produced by native MindSpore Tensor + "Float32": numpy.float32, + "Float16": numpy.float16, + "Float64": numpy.float64, + "Int8": numpy.int8, + "Int16": numpy.int16, + "Int32": numpy.int32, + "Int64": numpy.int64, + "UInt8": numpy.uint8, + "UInt16": numpy.uint16, + "UInt32": numpy.uint32, + "UInt64": numpy.uint64, + "Bool": numpy.bool_, + "BFloat16": "bfloat16", + # repr(mindspore.dtype) format — produced when mindtorch wraps tensors + "mindspore.float32": numpy.float32, + "mindspore.float16": numpy.float16, + "mindspore.float64": numpy.float64, + "mindspore.int8": numpy.int8, + "mindspore.int16": numpy.int16, + "mindspore.int32": numpy.int32, + "mindspore.int64": numpy.int64, + "mindspore.uint8": numpy.uint8, + "mindspore.uint16": numpy.uint16, + "mindspore.uint32": numpy.uint32, + "mindspore.uint64": numpy.uint64, + "mindspore.bool": numpy.bool_, + "mindspore.bool_": numpy.bool_, + "mindspore.bfloat16": "bfloat16", +} + + +# ── Protobuf wire-format primitives ─────────────────────────────────────── + +def _read_varint(f) -> Optional[int]: + """Read a base-128 varint. Return *None* at EOF.""" + result = 0 + shift = 0 + while True: + b = f.read(1) + if not b: + return None + byte = b[0] + result |= (byte & 0x7F) << shift + if (byte & 0x80) == 0: + return result + shift += 7 + if shift > 63: + raise CkptFormatNotSupported( + "Varint exceeds 64 bits — file may be corrupted or encrypted" + ) + + +def _skip_wire_field(f, wire_type: int) -> None: + """Advance past one protobuf field value.""" + if wire_type == 0: + _read_varint(f) + elif wire_type == 1: + f.seek(8, os.SEEK_CUR) + elif wire_type == 2: + length = _read_varint(f) + f.seek(length, os.SEEK_CUR) + elif wire_type == 5: + f.seek(4, os.SEEK_CUR) + else: + raise CkptFormatNotSupported( + f"Unknown protobuf wire type {wire_type}" + ) + + +# ── Index data structures ───────────────────────────────────────────────── + +@dataclass +class CkptTensorSlice: + """One contiguous byte range inside a ``.ckpt`` file.""" + file_offset: int + byte_length: int + + +@dataclass +class CkptTensorEntry: + """All slices of a single parameter, plus its metadata.""" + slices: List[CkptTensorSlice] + dims: Tuple[int, ...] + tensor_type_str: str + + +# ── CkptIndex: scanner + on-demand reader ───────────────────────────────── + +class CkptIndex: + """Index of every tensor in a ``.ckpt`` file. + + Build with :meth:`from_file`, then call :meth:`read_tensor` to + materialise individual parameters as numpy arrays. + """ + + def __init__(self, file_path: str, entries: Dict[str, CkptTensorEntry]): + self.file_path = file_path + self.entries = entries + + @classmethod + def from_file(cls, file_path: str) -> "CkptIndex": + entries = _scan_ckpt_file(file_path) + return cls(file_path, entries) + + def read_tensor(self, key: str) -> numpy.ndarray: + if key not in self.entries: + raise KeyError(f"Tensor '{key}' not found in {self.file_path}") + + entry = self.entries[key] + np_dtype = _TENSOR_TYPE_TO_NUMPY.get(entry.tensor_type_str) + + if np_dtype is None: + raise CkptFormatNotSupported( + f"Unsupported tensor_type '{entry.tensor_type_str}' " + f"for key '{key}' in {self.file_path}" + ) + + chunks: List[bytes] = [] + with open(self.file_path, "rb") as f: + for slc in entry.slices: + f.seek(slc.file_offset) + chunks.append(f.read(slc.byte_length)) + + raw = b"".join(chunks) + + if np_dtype == "bfloat16": + import ml_dtypes + arr = numpy.frombuffer(raw, dtype=ml_dtypes.bfloat16).copy() + else: + arr = numpy.frombuffer(raw, dtype=np_dtype).copy() + + if entry.dims: + arr = arr.reshape(entry.dims) + return arr + + +# ── File scanner ────────────────────────────────────────────────────────── + +def _scan_ckpt_file(path: str) -> Dict[str, CkptTensorEntry]: + """Walk the protobuf wire format to index every tensor without + loading ``tensor_content`` into memory.""" + file_size = os.path.getsize(path) + entries: Dict[str, CkptTensorEntry] = {} + + try: + with open(path, "rb") as f: + while f.tell() < file_size: + key = _read_varint(f) + if key is None: + break + field_num = key >> 3 + wire_type = key & 0x7 + + if field_num == 1 and wire_type == 2: + value_len = _read_varint(f) + if value_len is None: + break + value_end = f.tell() + value_len + _parse_value(f, value_end, entries) + f.seek(value_end) + else: + _skip_wire_field(f, wire_type) + except (OSError, ValueError, UnicodeDecodeError) as exc: + raise CkptFormatNotSupported( + f"Failed to parse {path}: {type(exc).__name__}: {exc}" + ) from exc + + return entries + + +def _parse_value( + f, + value_end: int, + entries: Dict[str, CkptTensorEntry], +) -> None: + """Parse one ``Checkpoint.Value`` submessage.""" + tag: Optional[str] = None + + dims: List[int] = [] + tensor_type_str: Optional[str] = None + content_offset: Optional[int] = None + content_length: int = 0 + + while f.tell() < value_end: + sub_key = _read_varint(f) + if sub_key is None: + break + sub_field = sub_key >> 3 + sub_wire = sub_key & 0x7 + + if sub_field == 1 and sub_wire == 2: + # Value.tag (string) + str_len = _read_varint(f) + tag = f.read(str_len).decode("utf-8") + + elif sub_field == 2 and sub_wire == 2: + # Value.tensor (TensorProto embedded message) + tp_len = _read_varint(f) + tp_end = f.tell() + tp_len + dims, tensor_type_str, content_offset, content_length = ( + _parse_tensor_proto(f, tp_end) + ) + f.seek(tp_end) + + elif sub_field == 3 and sub_wire == 2: + raise CkptFormatNotSupported( + f"MapTensor entries are not supported by the lazy loader " + f"(tag='{tag}')" + ) + + else: + _skip_wire_field(f, sub_wire) + + if tag is None: + return + if content_offset is None or tensor_type_str is None: + LOG.debug("Skipping non-tensor entry '%s'", tag) + return + + slc = CkptTensorSlice(file_offset=content_offset, byte_length=content_length) + if tag in entries: + entries[tag].slices.append(slc) + else: + entries[tag] = CkptTensorEntry( + slices=[slc], + dims=tuple(dims), + tensor_type_str=tensor_type_str, + ) + + +def _parse_tensor_proto( + f, + tp_end: int, +) -> Tuple[List[int], Optional[str], Optional[int], int]: + """Parse a ``TensorProto`` submessage, returning + ``(dims, tensor_type_str, content_offset, content_length)``. + """ + dims: List[int] = [] + tensor_type_str: Optional[str] = None + content_offset: Optional[int] = None + content_length: int = 0 + + while f.tell() < tp_end: + tp_key = _read_varint(f) + if tp_key is None: + break + tp_field = tp_key >> 3 + tp_wire = tp_key & 0x7 + + if tp_field == 1 and tp_wire == 2: + # dims (packed repeated int64) + packed_len = _read_varint(f) + packed_end = f.tell() + packed_len + while f.tell() < packed_end: + dims.append(_read_varint(f)) + elif tp_field == 1 and tp_wire == 0: + # dims (single int64 varint) + dims.append(_read_varint(f)) + elif tp_field == 2 and tp_wire == 2: + # tensor_type (string, e.g. "Float32") + str_len = _read_varint(f) + tensor_type_str = f.read(str_len).decode("utf-8") + elif tp_field == 3 and tp_wire == 2: + # tensor_content (bytes — record offset, skip data) + content_length = _read_varint(f) + content_offset = f.tell() + f.seek(content_length, os.SEEK_CUR) + else: + _skip_wire_field(f, tp_wire) + + return dims, tensor_type_str, content_offset, content_length diff --git a/src/mindnlp/wizard/merge/io/lazy_tensor_loader.py b/src/mindnlp/wizard/merge/io/lazy_tensor_loader.py new file mode 100644 index 000000000..9937d840f --- /dev/null +++ b/src/mindnlp/wizard/merge/io/lazy_tensor_loader.py @@ -0,0 +1,232 @@ +# Copyright (C) 2025 Arcee AI +# SPDX-License-Identifier: LGPL-3.0-only +# Modified for MindSpore/Ascend NPU by MindNLP Wizard contributors. +# + +""" +Sharded tensor index and lazy loader. + +""" + +import json +import logging +import os +import os.path +import threading +from dataclasses import dataclass +from typing import Any, Dict, List, Optional + +import mindspore +import safetensors + +from ._device import move_tensor_to_device +from .loader import TensorLoader + +LOG = logging.getLogger(__name__) + +_FORMAT_FOR_EXTENSION = { + ".safetensors": "safetensors", + ".bin": "bin", + ".ckpt": "ckpt", +} + + +def _detect_format(path: str) -> str: + ext = os.path.splitext(path)[1].lower() + return _FORMAT_FOR_EXTENSION.get(ext, "bin") + + +@dataclass +class ShardInfo: + filename: str + contained_keys: List[str] + + +@dataclass +class ShardedTensorIndex: + base_path: str + format: str + tensor_paths: Dict[str, str] + shards: List[ShardInfo] + + @property + def is_safetensors(self) -> bool: + return self.format == "safetensors" + + @classmethod + def from_disk(cls, base_path: str) -> "ShardedTensorIndex": + model_path = None + for model_file_name in [ + "model.safetensors", + "mindspore_model.ckpt", + "pytorch_model.bin", + ]: + candidate_path = os.path.join(base_path, model_file_name) + if os.path.exists(candidate_path) or os.path.exists( + candidate_path + ".index.json" + ): + model_path = candidate_path + break + + if not model_path: + ckpt_files = [ + f for f in os.listdir(base_path) + if f.lower().endswith(".ckpt") + ] if os.path.isdir(base_path) else [] + if ckpt_files: + raise RuntimeError( + f"Found {len(ckpt_files)} .ckpt file(s) in {base_path} " + f"but no recognized entry point (mindspore_model.ckpt or " + f"*.index.json). If these are segmented checkpoints from " + f"distributed training, please consolidate them first " + f"using mindspore.parallel.load_segmented_checkpoints() " + f"and re-save as a single checkpoint." + ) + raise RuntimeError( + f"Unable to find model files at {base_path}" + ) + + fmt = _detect_format(model_path) + tensor_paths = None + shards: List[ShardInfo] = [] + + if os.path.exists(model_path + ".index.json"): + with open(model_path + ".index.json", "r", encoding="utf-8") as fd: + weight_map = json.load(fd)["weight_map"] + tensor_paths = weight_map + + shard_names = list( + sorted(set(tensor_paths[e] for e in tensor_paths)) + ) + for shard_name in shard_names: + info = ShardInfo( + shard_name, + [ + key + for key in tensor_paths + if tensor_paths[key] == shard_name + ], + ) + shards.append(info) + + return ShardedTensorIndex( + base_path=base_path, + format=fmt, + tensor_paths=tensor_paths, + shards=shards, + ) + elif os.path.exists(model_path): + return ShardedTensorIndex.from_file(model_path) + else: + raise RuntimeError( + f"Unable to find model files at {base_path}" + ) + + @classmethod + def from_file(cls, file_path: str) -> "ShardedTensorIndex": + if not os.path.exists(file_path): + raise FileNotFoundError(file_path) + + lower = file_path.lower() + shard_name = os.path.basename(file_path) + fmt = _detect_format(file_path) + + if lower.endswith(".safetensors"): + with safetensors.safe_open( + file_path, framework="numpy" + ) as st: + tensor_paths = {key: shard_name for key in st.keys()} + elif lower.endswith(".ckpt"): + from .lazy_ckpt import CkptIndex + idx = CkptIndex.from_file(file_path) + tensor_paths = {key: shard_name for key in idx.entries} + else: + from .lazy_unpickle import load_bin_lazy + index = load_bin_lazy(file_path) + tensor_paths = {key: shard_name for key in index} + + return ShardedTensorIndex( + base_path=os.path.dirname(file_path), + format=fmt, + tensor_paths=tensor_paths, + shards=[ShardInfo(shard_name, list(tensor_paths.keys()))], + ) + + +class LazyTensorLoader: + """Thread-safe loader that opens one shard at a time.""" + + index: ShardedTensorIndex + current_shard: Optional[TensorLoader] + lazy_loader: bool + ckpt_load_kwargs: Optional[Dict[str, Any]] + lock: threading.Lock + + def __init__( + self, + index: ShardedTensorIndex, + lazy_loader: bool = True, + ckpt_load_kwargs: Optional[Dict[str, Any]] = None, + ): + self.index = index + self.current_shard = None + self.lazy_loader = lazy_loader + self.ckpt_load_kwargs = ckpt_load_kwargs + self.lock = threading.Lock() + + def get_tensor( + self, + key: str, + device: str = "CPU", + aliases: Optional[List[str]] = None, + raise_on_missing: bool = True, + ) -> Optional[mindspore.Tensor]: + if aliases and key not in self.index.tensor_paths: + for alias in aliases: + if alias in self.index.tensor_paths: + key = alias + break + + with self.lock: + if ( + self.current_shard is None + or key not in self.current_shard.keys() + ): + if key not in self.index.tensor_paths: + if raise_on_missing: + raise KeyError(key) + return None + + self.current_shard = None + + shard_file = self.index.tensor_paths[key] + shard_full_path = os.path.join( + self.index.base_path, shard_file + ) + logging.debug("Opening shard %s", shard_full_path) + self.current_shard = TensorLoader.get( + shard_full_path, + use_lazy_loader=self.lazy_loader, + device=device, + ckpt_load_kwargs=self.ckpt_load_kwargs, + ) + + tensor = self.current_shard.get_tensor(key) + return move_tensor_to_device(tensor, device, caller="LazyTensorLoader") + + def flush(self): + with self.lock: + self.current_shard = None + + @classmethod + def from_disk( + cls, + base_path: str, + lazy_loader: bool = True, + ckpt_load_kwargs: Optional[Dict[str, Any]] = None, + ) -> "LazyTensorLoader": + return LazyTensorLoader( + ShardedTensorIndex.from_disk(base_path), + lazy_loader, + ckpt_load_kwargs=ckpt_load_kwargs, + ) diff --git a/src/mindnlp/wizard/merge/io/lazy_unpickle.py b/src/mindnlp/wizard/merge/io/lazy_unpickle.py new file mode 100644 index 000000000..c05d8e683 --- /dev/null +++ b/src/mindnlp/wizard/merge/io/lazy_unpickle.py @@ -0,0 +1,301 @@ +# Copyright (C) 2025 Arcee AI +# SPDX-License-Identifier: LGPL-3.0-only +# Modified for MindSpore/Ascend NPU by MindNLP Wizard contributors. +# + +""" +Lazy unpickler for PyTorch ``.bin`` checkpoint files. + +Reads the pickle metadata without materialising tensors. When a tensor +is actually needed, raw bytes are read from the ZIP archive and +reconstructed as a ``mindspore.Tensor`` via NumPy. +""" + +import codecs +import collections +import contextlib +import logging +import operator +import os +import pickle +import zipfile +from functools import reduce +from typing import Any, Dict, Optional, Tuple, Union + +import mindspore +import numpy + +from ..dtype_policy import numpy_to_mindspore +from ._device import move_tensor_to_device + +LOG = logging.getLogger(__name__) + +# --------------------------------------------------------------------------- +# dtype helpers — map PyTorch storage class names to numpy dtypes +# --------------------------------------------------------------------------- + +_STORAGE_DTYPE_MAP = { + "DoubleStorage": numpy.float64, + "FloatStorage": numpy.float32, + "HalfStorage": numpy.float16, + "LongStorage": numpy.int64, + "IntStorage": numpy.int32, + "ShortStorage": numpy.int16, + "CharStorage": numpy.int8, + "ByteStorage": numpy.uint8, + "BoolStorage": numpy.bool_, + "BFloat16Storage": "bfloat16", +} + +_NUMPY_DTYPE_BYTES = { + numpy.float64: 8, + numpy.float32: 4, + numpy.float16: 2, + numpy.int64: 8, + numpy.int32: 4, + numpy.int16: 2, + numpy.int8: 1, + numpy.uint8: 1, + numpy.bool_: 1, + "bfloat16": 2, +} + + +def _resolve_numpy_dtype(name: str): + """Return the numpy dtype (or ``'bfloat16'`` sentinel) for a storage name.""" + return _STORAGE_DTYPE_MAP.get(name, numpy.float32) + + +def _element_size(np_dtype) -> int: + return _NUMPY_DTYPE_BYTES.get(np_dtype, 4) + + +# --------------------------------------------------------------------------- +# Placeholder storage types used during unpickling +# --------------------------------------------------------------------------- + +class _StoragePlaceholder: + """Stand-in for ``torch.*Storage`` during pickle loading.""" + def __init__(self, np_dtype): + self.np_dtype = np_dtype + self.dtype = np_dtype + + +def _make_storage_placeholder(name: str): + np_dt = _resolve_numpy_dtype(name) + return _StoragePlaceholder(np_dt) + + +_STORAGE_PLACEHOLDERS = { + name: _make_storage_placeholder(name) + for name in _STORAGE_DTYPE_MAP +} + + +# --------------------------------------------------------------------------- +# DeferredLoad +# --------------------------------------------------------------------------- + +class DeferredLoad: + """Describes a not-yet-materialised tensor inside a ``.bin`` archive.""" + + def __init__(self, name: str, location: str, np_dtype): + self.name = name + self.location = location + self.np_dtype = np_dtype + self.file_offset: Optional[int] = None + self.shape: Optional[Tuple[int, ...]] = None + self.stride: Optional[Tuple[int, ...]] = None + self.requires_grad = False + + @staticmethod + def rebuild( + load: "DeferredLoad", + offset: int, + shape: Union[Tuple[int, ...], Any], + stride: Tuple[int, ...], + ) -> "DeferredLoad": + load.shape = tuple(shape) + load.stride = tuple(stride) + load.file_offset = offset * _element_size(load.np_dtype) + return load + + def execute( + self, + reader: "TorchArchiveReader", + map_location: Any = None, + ) -> mindspore.Tensor: + if self.shape is None or self.stride is None or self.file_offset is None: + raise RuntimeError( + f"DeferredLoad for '{self.name}' was not fully initialised" + ) + + total_params = reduce(operator.mul, self.shape, 1) + elem_sz = _element_size(self.np_dtype) + total_bytes = total_params * elem_sz + + f = reader.open_file(file_name=self.name, offset=self.file_offset) + raw = f.read(total_bytes) + + if self.np_dtype == "bfloat16": + import ml_dtypes # noqa: F811 + arr = numpy.frombuffer(raw, dtype=ml_dtypes.bfloat16).copy() + else: + arr = numpy.frombuffer(raw, dtype=self.np_dtype).copy() + + arr = arr.reshape(self.shape) + tensor = numpy_to_mindspore(arr) + return move_tensor_to_device(tensor, map_location, caller="LazyPickleLoader") + + +# --------------------------------------------------------------------------- +# Custom unpickler +# --------------------------------------------------------------------------- + +def _rebuild_tensor_v2_placeholder(storage, offset, shape, stride, *extra): + # PyTorch may pass additional trailing args such as requires_grad, + # backward_hooks, metadata, etc. We only need the storage/offset/shape/stride + # information to reconstruct a DeferredLoad placeholder. + load = DeferredLoad.rebuild(storage, offset, shape, stride) + if extra: + try: + load.requires_grad = bool(extra[0]) + except Exception as exc: + LOG.debug( + "Failed to parse requires_grad flag from pickle metadata (%s: %s)", + type(exc).__name__, + exc, + ) + return load + + +ACCEPTABLE_TYPES = { + ("torch._utils", "_rebuild_tensor_v2"): _rebuild_tensor_v2_placeholder, + ("collections", "OrderedDict"): collections.OrderedDict, + ("numpy.core.multiarray", "scalar"): numpy.core.multiarray.scalar, + ("numpy", "dtype"): numpy.core.multiarray.scalar, + ("_codecs", "encode"): codecs.encode, +} +for _sname, _placeholder in _STORAGE_PLACEHOLDERS.items(): + ACCEPTABLE_TYPES[("torch", _sname)] = _placeholder + + +class LazyTorchUnpickler(pickle.Unpickler): + def find_class(self, module: str, name: str) -> Any: + if (module, name) in ACCEPTABLE_TYPES: + return ACCEPTABLE_TYPES[(module, name)] + raise pickle.UnpicklingError(f"Unsupported type {module}.{name}") + + def persistent_load(self, pid: Any) -> Any: + if not isinstance(pid, tuple) or pid[0] != "storage": + raise RuntimeError( + f"Unpickling object with unexpected PID: {repr(pid)}" + ) + storage_type, key, location, _ = pid[1:] + if isinstance(storage_type, _StoragePlaceholder): + np_dtype = storage_type.np_dtype + else: + np_dtype = numpy.float32 + return DeferredLoad(name=key, location=location, np_dtype=np_dtype) + + +class LazyUnpickleModule: + """Drop-in ``pickle_module`` for ``torch.load`` replacement.""" + Unpickler = LazyTorchUnpickler + + @staticmethod + def load(*args, **kwargs): + return LazyTorchUnpickler(*args, **kwargs).load() + + +# --------------------------------------------------------------------------- +# Archive reader +# --------------------------------------------------------------------------- + +class TorchArchiveReader: + """Lazily reads files from a PyTorch ZIP archive.""" + + def __init__(self, path: str): + self.archive = zipfile.ZipFile(path, mode="r") + self.archive_name = os.path.basename( + os.path.normpath(path) + ).split(".")[0] + self.file_name: Optional[str] = None + self.file = None + + def open_file( + self, file_name: str, offset: int = 0 + ) -> zipfile.ZipExtFile: + if self.file_name != file_name or ( + self.file is not None and self.file.tell() > offset + ): + if self.file is not None: + self.file.close() + try: + fd = self.archive.open( + f"archive/data/{file_name}", mode="r" + ) + except Exception: + fd = self.archive.open( + f"{self.archive_name}/data/{file_name}", mode="r" + ) + self.file = fd + self.file_name = file_name + + skip_bytes = offset - self.file.tell() + if skip_bytes < 0: + raise RuntimeError( + f"Cannot seek backwards in zip stream (need {skip_bytes})" + ) + self.file.seek(skip_bytes, os.SEEK_CUR) + return self.file + + def close(self): + if self.file is not None: + self.file.close() + self.file = None + self.archive.close() + + +# --------------------------------------------------------------------------- +# Context manager for lazy loading (replaces torch_lazy_load) +# --------------------------------------------------------------------------- + +@contextlib.contextmanager +def lazy_load_context(): + """Context manager that monkey-patches ``pickle`` so that + ``pickle.load`` on a ``.bin`` file returns ``DeferredLoad`` placeholders + instead of real tensors. + """ + old_unpickler = pickle.Unpickler + old_load = pickle.load + try: + def load_mp(*args, **kwargs): + return LazyTorchUnpickler(*args, **kwargs).load() + + pickle.Unpickler = LazyTorchUnpickler + pickle.load = load_mp + yield + finally: + pickle.Unpickler = old_unpickler + pickle.load = old_load + + +def load_bin_lazy(path: str) -> Dict[str, DeferredLoad]: + """Load a ``.bin`` file lazily, returning ``{name: DeferredLoad}``.""" + with lazy_load_context(): + if zipfile.is_zipfile(path): + archive = zipfile.ZipFile(path, mode="r") + archive_name = os.path.basename(os.path.normpath(path)).split(".")[0] + try: + try: + with archive.open("archive/data.pkl", mode="r") as f: + return LazyTorchUnpickler(f).load() + except KeyError: + with archive.open(f"{archive_name}/data.pkl", mode="r") as f: + return LazyTorchUnpickler(f).load() + finally: + archive.close() + + with open(path, "rb") as f: + return LazyTorchUnpickler(f).load() diff --git a/src/mindnlp/wizard/merge/io/loader.py b/src/mindnlp/wizard/merge/io/loader.py new file mode 100644 index 000000000..6cedf62be --- /dev/null +++ b/src/mindnlp/wizard/merge/io/loader.py @@ -0,0 +1,189 @@ +# Copyright (C) 2025 Arcee AI +# SPDX-License-Identifier: LGPL-3.0-only +# Modified for MindSpore/Ascend NPU by MindNLP Wizard contributors. +# + +""" +Tensor loader abstraction and concrete implementations. + +``TensorLoader.get()`` is the factory method that picks the right backend +(safetensors / lazy-ckpt / lazy-pickle / eager-pickle / eager-ckpt) +for a given shard file. +""" + +from abc import ABC, abstractmethod +import logging +from typing import Any, Dict, Optional, Sequence + +import mindspore +import safetensors + +from ..dtype_policy import numpy_to_mindspore +from ._device import move_tensor_to_device +from .lazy_unpickle import ( + DeferredLoad, + TorchArchiveReader, + load_bin_lazy, +) + +LOG = logging.getLogger(__name__) + + +class TensorLoader(ABC): + """Base class for (potentially lazy) tensor loaders.""" + + @abstractmethod + def get_tensor(self, key: str) -> mindspore.Tensor: + ... + + @abstractmethod + def keys(self) -> Sequence[str]: + ... + + @classmethod + def get( + cls, + shard_path: str, + use_lazy_loader: bool = False, + device: Optional[str] = None, + ckpt_load_kwargs: Optional[Dict[str, Any]] = None, + *, + use_lazy_unpickle: Optional[bool] = None, + ) -> "TensorLoader": + if use_lazy_unpickle is not None and use_lazy_loader is False: + use_lazy_loader = use_lazy_unpickle + + lower = shard_path.lower() + + if lower.endswith(".safetensors"): + return SafetensorsLoader(shard_path, device=device) + + if lower.endswith(".ckpt"): + if use_lazy_loader: + try: + return LazyCkptLoader(shard_path, device=device) + except Exception as exc: + from .lazy_ckpt import CkptFormatNotSupported + if isinstance(exc, CkptFormatNotSupported): + LOG.warning( + "Lazy ckpt loading failed for %s (%s); " + "falling back to eager loading.", + shard_path, + exc, + ) + else: + raise + return DumbCkptLoader( + shard_path, + device=device, + ckpt_load_kwargs=ckpt_load_kwargs, + ) + + if use_lazy_loader: + return LazyPickleLoader(shard_path, device=device) + return DumbPytorchLoader(shard_path, device=device) + + +class SafetensorsLoader(TensorLoader): + """Load tensors from a safetensors file via numpy → MindSpore.""" + + def __init__(self, path: str, device: Optional[str] = None): + self._handle = safetensors.safe_open(path, framework="numpy") + self._keys = list(self._handle.keys()) + self._device = device + + def get_tensor(self, key: str) -> mindspore.Tensor: + arr = self._handle.get_tensor(key) + tensor = numpy_to_mindspore(arr) + return move_tensor_to_device(tensor, self._device, caller="SafetensorsLoader") + + def keys(self) -> Sequence[str]: + return self._keys + + +class LazyPickleLoader(TensorLoader): + """Lazy loader for ``.bin`` files — reads metadata up-front, data on demand.""" + + def __init__(self, path: str, device: Optional[str] = None): + self.zip_reader = TorchArchiveReader(path) + self.index: Dict[str, DeferredLoad] = load_bin_lazy(path) + self.device = device + + def get_tensor(self, key: str) -> mindspore.Tensor: + if key not in self.index: + raise KeyError(key) + return self.index[key].execute( + self.zip_reader, map_location=self.device + ) + + def keys(self) -> Sequence[str]: + return list(self.index.keys()) + + +class DumbPytorchLoader(TensorLoader): + """Eager full-load of a ``.bin`` file — highest memory, best compatibility.""" + + def __init__(self, path: str, device: Optional[str] = None): + self.zip_reader = TorchArchiveReader(path) + index = load_bin_lazy(path) + self.tensors: Dict[str, mindspore.Tensor] = { + key: dl.execute(self.zip_reader, map_location=device) + for key, dl in index.items() + } + + def get_tensor(self, key: str) -> mindspore.Tensor: + return self.tensors[key] + + def keys(self) -> Sequence[str]: + return list(self.tensors.keys()) + + +class LazyCkptLoader(TensorLoader): + """B-level lazy loader for ``.ckpt`` files. + + Scans protobuf wire format once to build an index, then reads + only the requested tensor bytes on demand. + """ + + def __init__(self, path: str, device: Optional[str] = None): + from .lazy_ckpt import CkptIndex + self._index = CkptIndex.from_file(path) + self._device = device + + def get_tensor(self, key: str) -> mindspore.Tensor: + arr = self._index.read_tensor(key) + tensor = numpy_to_mindspore(arr) + return move_tensor_to_device(tensor, self._device, caller="LazyCkptLoader") + + def keys(self) -> Sequence[str]: + return list(self._index.entries.keys()) + + +class DumbCkptLoader(TensorLoader): + """Eager loader for ``.ckpt`` files via ``mindspore.load_checkpoint``. + + Handles encrypted checkpoints, CRC checks, and other advanced + features that the protobuf wire-format scanner cannot support. + """ + + def __init__( + self, + path: str, + device: Optional[str] = None, + ckpt_load_kwargs: Optional[Dict[str, Any]] = None, + ): + kwargs = dict(ckpt_load_kwargs or {}) + param_dict = mindspore.load_checkpoint(path, **kwargs) + self.tensors: Dict[str, mindspore.Tensor] = {} + for key, value in param_dict.items(): + if isinstance(value, (mindspore.Tensor, mindspore.Parameter)): + t = move_tensor_to_device(value, device, caller="DumbCkptLoader") + self.tensors[key] = t + else: + LOG.debug("Skipping non-tensor entry '%s' (type=%s)", key, type(value).__name__) + + def get_tensor(self, key: str) -> mindspore.Tensor: + return self.tensors[key] + + def keys(self) -> Sequence[str]: + return list(self.tensors.keys()) diff --git a/src/mindnlp/wizard/merge/io/tasks.py b/src/mindnlp/wizard/merge/io/tasks.py new file mode 100644 index 000000000..a06249a79 --- /dev/null +++ b/src/mindnlp/wizard/merge/io/tasks.py @@ -0,0 +1,339 @@ +# Originally from MergeKit (https://github.com/arcee-ai/mergekit) +# Copyright (C) 2025 Arcee AI +# SPDX-License-Identifier: LGPL-3.0-only +# Modified for MindSpore/Ascend NPU by MindNLP Wizard contributors. + +"""IO-related tasks for the DAG execution engine.""" + +import os +import re +import threading +from typing import Dict, Optional, Tuple + +import mindspore + +from ..common import ImmutableMap, ModelReference, dtype_from_name +from ..graph import Task +from .lazy_tensor_loader import LazyTensorLoader +from .tensor_writer import TensorWriter + + +class LoaderCache: + """Thread-local singleton cache for :class:`LazyTensorLoader` instances.""" + + loaders: Dict[ModelReference, LazyTensorLoader] + lora_cache_dir: Optional[str] + hf_cache_dir: Optional[str] + lazy_loader: bool + trust_remote_code: bool + lora_merge_dtype: Optional[str] + + _instance = threading.local() + + def __new__(cls) -> "LoaderCache": + if not hasattr(cls._instance, "value"): + cls._instance.value = super(LoaderCache, cls).__new__(cls) + cls._instance.value.loaders = {} + cls._instance.value.lora_cache_dir = None + cls._instance.value.hf_cache_dir = None + cls._instance.value.lazy_loader = False + cls._instance.value.trust_remote_code = False + cls._instance.value.lora_merge_dtype = None + return cls._instance.value + + def get(self, model: ModelReference) -> LazyTensorLoader: + if model not in self.loaders: + merged = model.merged( + cache_dir=self.lora_cache_dir, + trust_remote_code=self.trust_remote_code, + lora_merge_dtype=self.lora_merge_dtype, + ) + self.loaders[model] = merged.lazy_loader( + cache_dir=self.hf_cache_dir, + lazy_loader=self.lazy_loader, + ) + return self.loaders[model] + + def flush_all(self): + for loader in self.loaders.values(): + loader.flush() + + def setup(self, options): + self.lora_cache_dir = options.lora_merge_cache + self.hf_cache_dir = options.transformers_cache + self.lazy_loader = options.lazy_loader + self.trust_remote_code = options.trust_remote_code + self.lora_merge_dtype = options.lora_merge_dtype + + +shard_name_re = re.compile(r"model\-([0-9]+)-of-([0-9]+)") + + +def _normalized_shard_name(path: str) -> str: + name, _ext = os.path.splitext(os.path.basename(path)) + name = name.lower().replace("pytorch_model", "model").replace("mindspore_model", "model") + m = shard_name_re.search(name) + if m: + frac = int(m.group(1)) / int(m.group(2)) + name = f"model-{int(frac * 100):03d}pct" + return name + + +def _dtype_factor(dtype_name: Optional[str]) -> float: + if not dtype_name: + return 1.0 + name = str(dtype_name).lower() + if "float32" in name or name == "fp32": + return 2.0 + if "bfloat16" in name or "float16" in name or name in ("bf16", "fp16"): + return 1.0 + if "int8" in name: + return 0.5 + return 1.0 + + +class LoadTensor(Task[Optional[mindspore.Tensor]]): + model: ModelReference + tensor: str + dtype: Optional[str] = None + device: Optional[str] = None + optional: bool = False + aliases: Optional[Tuple[str, ...]] = None + tied_names: Optional[Tuple[str, ...]] = None + per_gpu: bool = False + + def arguments(self) -> Dict[str, Task]: + return {} + + def _resolve_name(self, loader: LazyTensorLoader) -> Optional[str]: + all_names = ( + [self.tensor] + + list(self.aliases or []) + + list(self.tied_names or []) + ) + for name in all_names: + if name in loader.index.tensor_paths: + return name + return None + + def execute(self) -> Optional[mindspore.Tensor]: + loader = LoaderCache().get(self.model) + name = self._resolve_name(loader) + if not name: + if not self.optional: + raise RuntimeError( + f"Tensor {self.tensor} required but not present in model {self.model}" + ) + return None + + x = loader.get_tensor(name, device=self.device or "CPU") + if self.dtype: + target_dtype = dtype_from_name(self.dtype) + if target_dtype is not None and target_dtype != x.dtype: + x = x.astype(target_dtype) + return x + + def priority(self) -> int: + return -1000 + + def group_label(self) -> Optional[str]: + loader = LoaderCache().get(self.model) + tensor_name = self._resolve_name(loader) + if tensor_name is None: + return None + tensor_path = loader.index.tensor_paths.get(tensor_name) + if tensor_path: + # Expose shard/file locality to the scheduler. + return f"{str(self.model)}::{_normalized_shard_name(tensor_path)}" + return tensor_name + + def duplicate_per_gpu(self): + return self.per_gpu + + def cost_hint(self): + return { + "read": 1.0, + "bytes_in": float(64 * 1024 * 1024), + "dtype_factor": _dtype_factor(self.dtype), + "fanout": 1.0, + } + + +class GatherTensors(Task[Dict[ModelReference, mindspore.Tensor]]): + weight_info: ImmutableMap + dtype: Optional[str] = None + device: Optional[str] = None + + def arguments(self) -> Dict[str, Task]: + return { + f"{str(model)}:{wi.name}": LoadTensor( + model=model, + tensor=wi.name, + dtype=wi.force_dtype or self.dtype, + device=self.device, + optional=wi.optional, + aliases=wi.aliases, + tied_names=wi.tied_names, + ) + for (model, wi) in self.weight_info.items() + } + + def group_label(self) -> Optional[str]: + return max( + t.group_label() or "" for t in self.arguments().values() + ) + + def priority(self) -> int: + return -10 + + def cost_hint(self): + fan_in = max(1, len(self.weight_info)) + dtype_name = self.dtype + return { + "read": float(fan_in), + "compute": 0.25, + "bytes_in": float(fan_in * 64 * 1024 * 1024), + "dtype_factor": _dtype_factor(dtype_name), + "fanout": float(fan_in), + } + + def execute( + self, **kwargs + ) -> Dict[ModelReference, mindspore.Tensor]: + key2model = { + f"{str(model)}:{wi.name}": model + for (model, wi) in self.weight_info.items() + } + return { + key2model[key]: kwargs[key] + for key in key2model + if kwargs[key] is not None + } + + +class TensorWriterTask(Task[TensorWriter]): + out_path: str + max_shard_size: int + output_format: str = "safetensors" + override_basename: Optional[str] = None + use_async: bool = False + write_threads: int = 1 + + def arguments(self) -> Dict[str, Task]: + return {} + + def execute(self, **_kwargs) -> TensorWriter: + return TensorWriter( + self.out_path, + max_shard_size=self.max_shard_size, + output_format=self.output_format, + override_basename=self.override_basename, + use_async=self.use_async, + max_write_threads=self.write_threads, + ) + + def priority(self): + return 10000 + + def main_thread_only(self): + return True + + def cost_hint(self): + return {"write": 0.5, "bytes_out": float(16 * 1024 * 1024)} + + +class SaveTensor(Task[None]): + tensor_name: str + tensor_task: Task + writer_task: TensorWriterTask + clone: bool + optional: bool = False + dtype: Optional[str] = None + force_main_thread: bool = False + + def arguments(self) -> Dict[str, Task]: + return {"writer": self.writer_task, "tensor": self.tensor_task} + + def priority(self) -> int: + return 1000 + + def group_label(self) -> Optional[str]: + return self.tensor_task.group_label() + + def main_thread_only(self): + return self.force_main_thread + + def cost_hint(self): + return { + "write": 1.0, + "bytes_out": float(64 * 1024 * 1024), + "dtype_factor": _dtype_factor(self.dtype), + "fanout": 1.0, + } + + def execute( + self, + writer: TensorWriter, + tensor: Optional[mindspore.Tensor], + ) -> None: + if tensor is None: + if not self.optional: + raise RuntimeError( + f"No value for required tensor {self.tensor_name}" + ) + return + if self.dtype: + target = dtype_from_name(self.dtype) + if target is not None and target != tensor.dtype: + tensor = tensor.astype(target) + writer.save_tensor( + name=self.tensor_name, tensor=tensor, clone=self.clone + ) + + +class FinalizeModel(Task[None]): + tensor_save_tasks: Tuple[Task, ...] + writer_task: TensorWriterTask + + def arguments(self) -> Dict[str, Task]: + return { + "writer": self.writer_task, + **{ + f"_unused_{idx}": t + for idx, t in enumerate(self.tensor_save_tasks) + }, + } + + def execute(self, writer: TensorWriter, **kwargs) -> None: + writer.finalize() + + def main_thread_only(self): + return True + + def cost_hint(self): + return {"write": 0.3, "bytes_out": float(8 * 1024 * 1024)} + + +class ReturnTensor(Task[mindspore.Tensor]): + weight_info: object # WeightInfo — resolved at import time + tensor_task: Task[mindspore.Tensor] + dtype: Optional[str] = None + + def arguments(self) -> Dict[str, Task]: + return {"tensor": self.tensor_task} + + def priority(self) -> int: + return 10000 + + def group_label(self) -> Optional[str]: + return self.tensor_task.group_label() + + def cost_hint(self): + return {"compute": 0.1, "bytes_in": float(8 * 1024 * 1024)} + + def execute(self, tensor: mindspore.Tensor) -> mindspore.Tensor: + if self.dtype: + target = dtype_from_name(self.dtype) + if target is not None and target != tensor.dtype: + tensor = tensor.astype(target) + return tensor diff --git a/src/mindnlp/wizard/merge/io/tensor_writer.py b/src/mindnlp/wizard/merge/io/tensor_writer.py new file mode 100644 index 000000000..2616d83df --- /dev/null +++ b/src/mindnlp/wizard/merge/io/tensor_writer.py @@ -0,0 +1,265 @@ +# Copyright (C) 2025 Arcee AI +# SPDX-License-Identifier: LGPL-3.0-only +# Modified for MindSpore/Ascend NPU by MindNLP Wizard contributors. +# + +""" +Sharded tensor writer. + +Streams tensors to safetensors or MindSpore ckpt shards on disk. +""" + +import json +import logging +import os +import threading +from concurrent.futures import Future, ThreadPoolExecutor +from typing import Dict, List, Optional + +import mindspore +import numpy +import safetensors.numpy + +from ..dtype_policy import mindspore_to_numpy, numpy_to_mindspore + +LOG = logging.getLogger(__name__) + +_BASENAME_MAP = { + "safetensors": "model", + "ckpt": "mindspore_model", +} + + +class TensorWriter: + out_path: str + override_basename: Optional[str] + max_shard_size: int + output_format: str + use_async: bool + + shards_written: int + weight_map: Dict[str, str] + current_shard: Dict[str, mindspore.Tensor] + current_shard_size: int + + _lock: threading.RLock + _executor: Optional[ThreadPoolExecutor] + _write_futures: List[Future] + + def __init__( # pylint: disable=too-many-positional-arguments + self, + out_path: str, + max_shard_size: int = 1000 * 1000 * 1000 * 5, + output_format: str = "safetensors", + override_basename: Optional[str] = None, + use_async: bool = False, + max_write_threads: int = 1, + *, + safe_serialization: Optional[bool] = None, + ) -> None: + os.makedirs(out_path, exist_ok=True) + + if safe_serialization is not None and output_format == "safetensors": + output_format = "safetensors" if safe_serialization else "bin" + if output_format not in ("safetensors", "ckpt"): + raise ValueError( + f"Unsupported output_format '{output_format}'. " + f"Use 'safetensors' or 'ckpt'." + ) + + self.out_path = out_path + self.override_basename = override_basename + self.max_shard_size = max_shard_size + self.output_format = output_format + self.use_async = use_async + + self.shards_written = 0 + self.weight_map: Dict[str, str] = {} + self.current_shard: Dict[str, mindspore.Tensor] = {} + self.current_shard_size = 0 + self.total_size = 0 + + self._lock = threading.RLock() + self._write_futures = [] + if self.use_async: + self._executor = ThreadPoolExecutor( + max_workers=max_write_threads + ) + + def __enter__(self): + return self + + def __exit__(self, exc_type, exc_val, exc_tb): + self.finalize() + + def save_tensor( + self, name: str, tensor: mindspore.Tensor, clone: bool = False + ): + if clone: + tensor = numpy_to_mindspore(mindspore_to_numpy(tensor).copy()) + + tensor_size = int(tensor.nbytes) + + with self._lock: + if ( + self.current_shard + and self.max_shard_size > 0 + and self.current_shard_size + tensor_size + > self.max_shard_size + ): + self._flush_current_shard() + + self.current_shard[name] = tensor + self.current_shard_size += tensor_size + + def _flush_current_shard(self): + if not self.current_shard: + return + + shard_to_write = self.current_shard + shard_index = self.shards_written + + self.total_size += self.current_shard_size + self.current_shard = {} + self.current_shard_size = 0 + self.shards_written += 1 + + prefix, extension = self._get_name_components() + shard_name = f"{prefix}-{shard_index + 1}.{extension}" + shard_path = os.path.join(self.out_path, shard_name) + for key in shard_to_write: + self.weight_map[key] = shard_name + + if self.use_async: + LOG.info( + "Dispatching shard #%d to be written.", shard_index + 1, + ) + future = self._executor.submit( + self._write_shard_task, + shard_to_write, + shard_index, + shard_path, + ) + self._write_futures.append(future) + else: + self._write_shard_task( + shard_data=shard_to_write, + shard_index=shard_index, + shard_path=shard_path, + ) + + def _write_shard_task( + self, + shard_data: Dict[str, mindspore.Tensor], + shard_index: int, + shard_path: str, + ): + LOG.info("Writing shard #%d...", shard_index + 1) + if self.output_format == "safetensors": + self._save_st(shard_data, shard_path) + elif self.output_format == "ckpt": + self._save_ckpt(shard_data, shard_path) + else: + raise RuntimeError( + f"Cannot write output format '{self.output_format}'. " + f"Use 'safetensors' or 'ckpt'." + ) + LOG.info("Finished writing shard #%d.", shard_index + 1) + + def finalize(self): + with self._lock: + self._flush_current_shard() + + if self.use_async: + if self._write_futures: + LOG.info( + "Waiting for %d shard(s) to finish writing...", + len(self._write_futures), + ) + for future in self._write_futures: + future.result() + LOG.info("All shards have been written to disk.") + self._write_futures.clear() + self._executor.shutdown() + + with self._lock: + LOG.info("Finalizing shard names and creating index file.") + prefix, extension = self._get_name_components() + total_shards = self.shards_written + + name_remap: Dict[str, str] = {} + if total_shards == 1: + name_remap[f"{prefix}-1.{extension}"] = ( + f"{prefix}.{extension}" + ) + else: + for idx in range(total_shards): + old_name = f"{prefix}-{idx + 1}.{extension}" + new_name = f"{prefix}-{idx + 1:05d}-of-{total_shards:05d}.{extension}" + name_remap[old_name] = new_name + + for old_name, new_name in name_remap.items(): + old_path = os.path.join(self.out_path, old_name) + new_path = os.path.join(self.out_path, new_name) + os.rename(old_path, new_path) + + if total_shards > 1: + for key in self.weight_map: + self.weight_map[key] = name_remap.get( + self.weight_map[key], self.weight_map[key] + ) + + index_filename = f"{prefix}.{extension}.index.json" + index_path = os.path.join(self.out_path, index_filename) + with open(index_path, "w", encoding="utf-8") as f: + content = { + "metadata": { + "total_size": self.total_size, + "mergekit_version": "0.1.4", + }, + "weight_map": self.weight_map, + } + json.dump(content, f, indent=2) + + def _get_name_components(self): + if self.override_basename: + basename = self.override_basename + else: + basename = _BASENAME_MAP.get(self.output_format, "model") + return basename, self.output_format + + def _save_st(self, shard_data: dict, shard_path: str): + np_data = {} + for key, tensor in shard_data.items(): + arr = mindspore_to_numpy(tensor) + if not arr.flags["C_CONTIGUOUS"]: + arr = numpy.ascontiguousarray(arr) + np_data[key] = arr + + def _do_save(sd): + safetensors.numpy.save_file( + sd, shard_path, metadata={"format": "np"} + ) + + try: + _do_save(np_data) + except RuntimeError as e: + if ( + len(e.args) > 0 + and isinstance(e.args[0], str) + and "share memory" in e.args[0] + ): + LOG.warning( + "Duplicated tensors detected — cloning before save." + ) + np_data = {k: v.copy() for k, v in np_data.items()} + _do_save(np_data) + else: + raise + + def _save_ckpt(self, shard_data: dict, shard_path: str): + param_list = [ + {"name": key, "data": tensor} + for key, tensor in shard_data.items() + ] + mindspore.save_checkpoint(param_list, shard_path, format="ckpt") diff --git a/src/mindnlp/wizard/merge/merge.py b/src/mindnlp/wizard/merge/merge.py new file mode 100644 index 000000000..939d32616 --- /dev/null +++ b/src/mindnlp/wizard/merge/merge.py @@ -0,0 +1,489 @@ +# Originally from MergeKit (https://github.com/arcee-ai/mergekit) +# Copyright (C) 2025 Arcee AI +# SPDX-License-Identifier: LGPL-3.0-only +# Modified for MindSpore/Ascend NPU by MindNLP Wizard contributors. + +"""Main merge entry point.""" + +import importlib +import importlib.resources +import json +import logging +import os +import shutil +import statistics +from collections import Counter +from typing import Any, Dict, List, Optional, Tuple + +import tqdm +import transformers + +from ._data import chat_templates +from .architecture import ModelArchitecture, get_architecture_info +from .card import generate_card +from .common import ModelReference, set_config_value +from .config import MergeConfiguration +from .graph import Executor +from .io.tasks import LoaderCache +from .multigpu_executor import MultiDeviceExecutor +from .options import MergeOptions +from .plan import MergePlanner +from .preflight import run_merge_preflight +from .tokenizer import TokenizerInfo + +LOG = logging.getLogger(__name__) + + +def run_merge( + merge_config: MergeConfiguration, + out_path: str, + options: MergeOptions, + config_source: Optional[str] = None, +): + if options.random_seed is not None: + transformers.trainer_utils.set_seed(options.random_seed) + + if not merge_config.models and not merge_config.slices and not merge_config.modules: + raise RuntimeError("No output requested") + + run_merge_preflight(merge_config=merge_config, options=options) + + arch_info = get_architecture_info(merge_config, options) + loader_cache = LoaderCache() + loader_cache.setup(options=options) + + cfg_out = _model_out_config( + merge_config, arch_info, trust_remote_code=options.trust_remote_code + ) + + for model in ( + pbar := tqdm.tqdm( + merge_config.referenced_models(), + desc="Warmup loader cache", + disable=options.quiet, + ) + ): + loader_cache.get(model) + del pbar + + LOG.info("Planning operations") + targets = MergePlanner( + merge_config, + arch_info, + options=options, + out_model_config=cfg_out, + ).plan_to_disk(out_path=out_path) + + if options.multi_npu: + exec = MultiDeviceExecutor( + targets=targets, + storage_device=None if options.low_cpu_memory else "CPU", + ) + else: + exec = Executor( + targets=targets, + math_device=options.device, + storage_device=options.device if options.low_cpu_memory else "CPU", + ) + + tokenizer = None + for _task, value in exec.run(quiet=options.quiet): + if isinstance(value, TokenizerInfo): + tokenizer = value.tokenizer + + if hasattr(exec, "metrics_snapshot"): + _write_execution_report( + out_path=out_path, + metrics=exec.metrics_snapshot(), + metadata={}, + ) + + if tokenizer: + pad_to_multiple_of = None + if merge_config.tokenizer and merge_config.tokenizer.pad_to_multiple_of: + pad_to_multiple_of = merge_config.tokenizer.pad_to_multiple_of + _update_config_vocab( + cfg_out, arch_info, tokenizer, pad_to_multiple_of=pad_to_multiple_of + ) + + LOG.info("Saving config") + cfg_out.save_pretrained(out_path) + + if options.write_model_card: + if not config_source: + config_source = merge_config.to_yaml() + + card_md = generate_card( + config=merge_config, + config_yaml=config_source, + name=os.path.basename(out_path), + ) + with open(os.path.join(out_path, "README.md"), "w", encoding="utf-8") as fp: + fp.write(card_md) + + with open( + os.path.join(out_path, "wizard_config.yml"), "w", encoding="utf-8" + ) as fp: + fp.write(config_source) + + if tokenizer is not None: + LOG.info("Saving tokenizer") + _set_chat_template(tokenizer, merge_config) + tokenizer.save_pretrained(out_path, safe_serialization=True) + else: + if options.copy_tokenizer: + try: + _copy_tokenizer(merge_config, out_path, options=options) + except Exception as e: + LOG.error( + "Failed to copy tokenizer. The merge was still successful, " + "just copy it from somewhere else.", + exc_info=e, + ) + elif merge_config.chat_template: + LOG.warning( + "Chat template specified but no tokenizer found. " + "Chat template will not be saved." + ) + + _copy_tagalong_files( + merge_config, + out_path, + files=arch_info.tagalong_files or [], + options=options, + ) + + if getattr(arch_info, "post_fill_parameters", False): + logging.info( + "Filling missing parameters from base model %s into new directory", + arch_info.post_fill_parameters, + ) + try: + from .scripts.fill_missing_params import copy_and_fill_missing_params + + copy_and_fill_missing_params( + base_model_repo_id=arch_info.post_fill_parameters, + sub_model_dir=out_path, + ) + logging.info("Deleting initial merge directory: %s", out_path) + shutil.rmtree(out_path) + except ImportError: + LOG.warning( + "fill_missing_params script not available — skipping post-fill step" + ) + + +def _set_chat_template( + tokenizer: transformers.PreTrainedTokenizerBase, + merge_config: MergeConfiguration, + trust_remote_code: bool = False, +): + chat_template = merge_config.chat_template + if not chat_template: + return + + if chat_template == "auto": + model_templates = [] + for model in merge_config.referenced_models(): + try: + tok = transformers.AutoTokenizer.from_pretrained( + model.model.path, + revision=model.model.revision, + trust_remote_code=trust_remote_code, + ) + template = tok.chat_template + if isinstance(template, dict): + template = template.get("default", None) + if template: + model_templates.append(template.strip()) + except Exception as e: + LOG.warning("Unable to load tokenizer for %s", model, exc_info=e) + + if not model_templates: + return + + chat_template = Counter(model_templates).most_common(1)[0][0] + LOG.info("Auto-selected chat template: %s", chat_template) + + elif ( + t := importlib.resources.files(chat_templates).joinpath( + chat_template + ".jinja" + ) + ).is_file(): + chat_template = t.read_text() + + elif len(chat_template) < 20 or "{" not in chat_template: + raise RuntimeError(f"Invalid chat template: {chat_template}") + + tokenizer.chat_template = chat_template + + +def _get_donor_model( + merge_config: MergeConfiguration, + options: MergeOptions, +) -> Tuple[ModelReference, str]: + donor_model = merge_config.base_model or (merge_config.referenced_models()[0]) + donor_local_path = donor_model.merged( + cache_dir=options.lora_merge_cache, + trust_remote_code=options.trust_remote_code, + lora_merge_dtype=options.lora_merge_dtype, + ).local_path(cache_dir=options.transformers_cache) + if not donor_local_path: + raise RuntimeError(f"Unable to find local path for {donor_model}") + return donor_model, donor_local_path + + +def _copy_tagalong_files( + merge_config: MergeConfiguration, + out_path: str, + files: List[str], + options: MergeOptions, +): + donor_model, donor_local_path = _get_donor_model(merge_config, options=options) + + for file_name in files: + fp = os.path.join(donor_local_path, file_name) + if os.path.exists(fp): + LOG.info("Copying %s from %s", file_name, donor_model) + shutil.copy( + fp, + os.path.join(out_path, file_name), + ) + + +def _copy_tokenizer( + merge_config: MergeConfiguration, out_path: str, options: MergeOptions +): + donor_model, donor_local_path = _get_donor_model(merge_config, options=options) + + # MergeKit-compatible warning: when base_model is used as tokenizer donor, + # chat templates from instruct/chat models will not be inherited unless + # users explicitly set `chat_template: auto` or `tokenizer_source`. + if not merge_config.chat_template and merge_config.base_model and not merge_config.tokenizer_source: + try: + donor_tok = transformers.AutoTokenizer.from_pretrained( + donor_model.model.path, + revision=donor_model.model.revision, + trust_remote_code=options.trust_remote_code, + ) + donor_has_template = bool(getattr(donor_tok, "chat_template", None)) + other_has_template = False + for model in merge_config.referenced_models(): + if model == donor_model: + continue + try: + tok = transformers.AutoTokenizer.from_pretrained( + model.model.path, + revision=model.model.revision, + trust_remote_code=options.trust_remote_code, + ) + if getattr(tok, "chat_template", None): + other_has_template = True + break + except Exception: + continue + if (not donor_has_template) and other_has_template: + LOG.warning( + "Base model tokenizer is being used as donor and does not define a " + "chat template, while another input model does. This matches " + "MergeKit default behavior, but the merged output may lose the " + "instruct/chat template. Consider setting `chat_template: auto` " + "or `tokenizer_source` to the instruct/chat model." + ) + except Exception: + pass + + if ( + (not merge_config.chat_template) + and os.path.exists(os.path.join(donor_local_path, "tokenizer_config.json")) + and ( + os.path.exists(os.path.join(donor_local_path, "tokenizer.json")) + or os.path.exists(os.path.join(donor_local_path, "tokenizer.model")) + ) + ): + LOG.info("Copying tokenizer from %s", donor_model) + + for file_name in [ + "tokenizer_config.json", + "special_tokens_map.json", + "tokenizer.json", + "tokenizer.model", + "added_tokens.json", + "merges.txt", + "chat_template.jinja", + "generation_config.json", + ]: + if os.path.exists(os.path.join(donor_local_path, file_name)): + shutil.copy( + os.path.join(donor_local_path, file_name), + os.path.join(out_path, file_name), + ) + + return + + LOG.info("Reserializing tokenizer from %s", donor_model) + tokenizer = transformers.AutoTokenizer.from_pretrained( + donor_model.model.path, + revision=donor_model.model.revision, + trust_remote_code=options.trust_remote_code, + ) + _set_chat_template(tokenizer, merge_config) + tokenizer.save_pretrained(out_path, safe_serialization=True) + + +def _model_out_config( + config: MergeConfiguration, + arch_info: ModelArchitecture, + trust_remote_code: bool = False, +) -> transformers.PretrainedConfig: + """Return a configuration for the resulting model.""" + if config.base_model: + res = config.base_model.config(trust_remote_code=trust_remote_code) + else: + res = config.referenced_models()[0].config(trust_remote_code=trust_remote_code) + if config.out_dtype: + res.torch_dtype = config.out_dtype + elif config.dtype: + res.torch_dtype = config.dtype + + module_layers = {} + for module_name in arch_info.modules: + if config.modules and module_name in config.modules: + module_def = config.modules.get(module_name) + if module_def and module_def.slices: + module_layers[module_name] = sum( + [ + s.sources[0].layer_range[1] - s.sources[0].layer_range[0] + for s in module_def.slices + ] + ) + elif config.slices: + module_layers[module_name] = sum( + [ + s.sources[0].layer_range[1] - s.sources[0].layer_range[0] + for s in config.slices + ] + ) + + if module_layers: + for module_name in module_layers: + if module_name not in arch_info.modules: + LOG.warning( + "Module %s in config but not in architecture info", + module_name, + ) + continue + module_info = arch_info.modules[module_name] + cfg_key = module_info.architecture.num_layers_config_key() + if not cfg_key: + if module_layers[module_name] > 0: + LOG.warning( + "Module %s has no configuration key for number of layers, " + "but the number of layers is not zero.", + module_name, + ) + continue + try: + set_config_value(res, cfg_key, module_layers[module_name]) + except Exception as e: + LOG.warning( + "Unable to set number of layers for module %s in output config " + "- you may need to manually correct it.", + module_name, + exc_info=e, + ) + + return res + + +def _update_config_vocab( + config: transformers.PretrainedConfig, + arch_info: ModelArchitecture, + tokenizer: transformers.PreTrainedTokenizerBase, + pad_to_multiple_of: Optional[int] = None, +): + vocab_size = len(tokenizer.get_vocab()) + if pad_to_multiple_of and vocab_size % pad_to_multiple_of: + vocab_size = vocab_size + pad_to_multiple_of - (vocab_size % pad_to_multiple_of) + try: + set_config_value( + config, arch_info.vocab_size_config_key or "vocab_size", vocab_size + ) + except Exception as e: + LOG.warning( + "Unable to set vocabulary size in output config " + "- you may need to manually correct it.", + exc_info=e, + ) + + +__all__ = ["MergeOptions", "run_merge"] + + +def _write_execution_report( + out_path: str, + metrics: Dict[str, Any], + metadata: Dict[str, Any], +) -> None: + os.makedirs(out_path, exist_ok=True) + task_runs = [float(t.get("run_ms", 0.0)) for t in metrics.get("tasks", [])] + queue_depth = [int(x) for x in metrics.get("queue_depth_samples", [])] + summary = { + "meta": metadata, + "executor": metrics.get("executor"), + "task_count": metrics.get("task_count", 0), + "run_ms_avg": statistics.fmean(task_runs) if task_runs else 0.0, + "run_ms_p95": _p95(task_runs) if task_runs else 0.0, + "queue_depth_avg": statistics.fmean(queue_depth) if queue_depth else 0.0, + "queue_depth_peak": max(queue_depth) if queue_depth else 0, + "backpressure_trigger_count": int( + metrics.get("backpressure_trigger_count", 0) + ), + "rss_peak_mb": float(metrics.get("rss_peak_mb", 0.0)), + "npu_used_peak_mb": metrics.get("npu_used_peak_mb"), + "island_assignment": metrics.get("island_assignment", []), + "tasks": metrics.get("tasks", []), + } + + json_path = os.path.join(out_path, "wizard_execution_report.json") + with open(json_path, "w", encoding="utf-8") as fp: + json.dump(summary, fp, ensure_ascii=False, indent=2) + + md_path = os.path.join(out_path, "wizard_execution_report.md") + with open(md_path, "w", encoding="utf-8") as fp: + fp.write("# Wizard Execution Report\n\n") + fp.write(f"- executor: `{summary['executor']}`\n") + fp.write(f"- task_count: `{summary['task_count']}`\n") + fp.write(f"- run_ms_avg: `{summary['run_ms_avg']:.3f}`\n") + fp.write(f"- run_ms_p95: `{summary['run_ms_p95']:.3f}`\n") + fp.write(f"- queue_depth_avg: `{summary['queue_depth_avg']:.3f}`\n") + fp.write(f"- queue_depth_peak: `{summary['queue_depth_peak']}`\n") + fp.write( + f"- backpressure_trigger_count: `{summary['backpressure_trigger_count']}`\n" + ) + fp.write(f"- rss_peak_mb: `{summary['rss_peak_mb']:.3f}`\n") + fp.write(f"- npu_used_peak_mb: `{summary['npu_used_peak_mb']}`\n") + fp.write( + f"- island_assignment_count: `{len(summary['island_assignment'])}`\n" + ) + fp.write("\n## Schedule Metadata\n\n") + for key, value in metadata.items(): + fp.write(f"- {key}: `{value}`\n") + if summary["island_assignment"]: + fp.write("\n## Island Assignment\n\n") + for item in summary["island_assignment"]: + fp.write( + "- device: `{device}`, task_count: `{task_count}`, dominant_locality_key: `{dominant}`\n".format( + device=item.get("device"), + task_count=item.get("task_count"), + dominant=item.get("dominant_locality_key"), + ) + ) + + +def _p95(values: List[float]) -> float: + if not values: + return 0.0 + ordered = sorted(float(x) for x in values) + idx = min(len(ordered) - 1, max(0, int(round((len(ordered) - 1) * 0.95)))) + return ordered[idx] diff --git a/src/mindnlp/wizard/merge/merge_methods/__init__.py b/src/mindnlp/wizard/merge/merge_methods/__init__.py new file mode 100644 index 000000000..6d27e3746 --- /dev/null +++ b/src/mindnlp/wizard/merge/merge_methods/__init__.py @@ -0,0 +1,53 @@ +# Copyright (C) 2025 Arcee AI +# SPDX-License-Identifier: LGPL-3.0-only +# Modified for MindSpore/Ascend NPU by MindNLP Wizard contributors. +# +# +import logging + +LOG = logging.getLogger(__name__) + +try: + from . import multislerp # noqa: F401 +except Exception as e: + LOG.warning("Failed to import merge method module multislerp", exc_info=e) + +try: + from . import nearswap # noqa: F401 +except Exception as e: + LOG.warning("Failed to import merge method module nearswap", exc_info=e) + +try: + from . import ram # noqa: F401 +except Exception as e: + LOG.warning("Failed to import merge method module ram", exc_info=e) + +try: + from . import sce # noqa: F401 +except Exception as e: + LOG.warning("Failed to import merge method module sce", exc_info=e) + +from .base import MergeMethod # pylint: disable=wrong-import-position +from .registry import REGISTERED_MERGE_METHODS # pylint: disable=wrong-import-position + +try: + from .generalized_task_arithmetic import GeneralizedTaskArithmeticMerge +except Exception as e: + LOG.warning( + "Failed to import generalized_task_arithmetic merge methods", exc_info=e + ) + GeneralizedTaskArithmeticMerge = None + + +def get(method: str) -> MergeMethod: + if method in REGISTERED_MERGE_METHODS: + return REGISTERED_MERGE_METHODS[method] + raise RuntimeError(f"Unimplemented merge method {method}") + + +__all__ = [ + "MergeMethod", + "get", + "GeneralizedTaskArithmeticMerge", + "REGISTERED_MERGE_METHODS", +] diff --git a/src/mindnlp/wizard/merge/merge_methods/arcee_fusion.py b/src/mindnlp/wizard/merge/merge_methods/arcee_fusion.py new file mode 100644 index 000000000..036ab3309 --- /dev/null +++ b/src/mindnlp/wizard/merge/merge_methods/arcee_fusion.py @@ -0,0 +1,160 @@ +# Copyright (C) 2025 Arcee AI +# SPDX-License-Identifier: LGPL-3.0-only +# Modified for MindSpore/Ascend NPU by MindNLP Wizard contributors. +# + +from typing import Dict, List, Optional + +import mindspore # pylint: disable=import-error +from mindspore import ops # pylint: disable=import-error +from typing_extensions import override + +from ..architecture.base import WeightInfo +from ..common import ModelReference +from ..dtype_policy import choose_work_dtype +from ..graph import Task +from ..safe_ops import safe_abs +from .base import ( + ConfigParameterDef, + MergeMethod, + MergeTensorInput, +) +from .rectify_embed import rectify_embed_sizes + + +class DynamicThresholdFusion: + def approximate_quantiles(self, tensor, q): + flat_tensor = tensor.reshape(-1) + + if flat_tensor.numel() > 1e6: + perm = ops.randperm(flat_tensor.numel())[:1000000] + flat_tensor = flat_tensor[perm] + + sorted_tensor, _ = ops.sort(flat_tensor) + + quantile_indices = (q * (sorted_tensor.numel() - 1)).astype(mindspore.int64) + + return sorted_tensor[quantile_indices] + + def calculate_dynamic_threshold(self, importance_scores): + median = self.approximate_quantiles( + importance_scores, mindspore.Tensor([0.5]) + )[0] + q1, q3 = self.approximate_quantiles( + importance_scores, mindspore.Tensor([0.25, 0.75]) + ) + + iqr = q3 - q1 + + dynamic_threshold = median + 1.5 * iqr + + return dynamic_threshold + + def compute_fusion_mask(self, importance_scores): + threshold = self.calculate_dynamic_threshold(importance_scores) + fusion_mask = (importance_scores >= threshold).astype(mindspore.float32) + return fusion_mask, threshold + + +class ArceeFusionMergeTask(Task[mindspore.Tensor]): + gather_tensors: MergeTensorInput + base_model: ModelReference + weight_info: WeightInfo + split_pieces: int = 1 + max_tensor_mem_gb: Optional[float] = None + + def uses_accelerator(self) -> bool: + return True + + def arguments(self) -> Dict[str, Task]: + return {"tensors": self.gather_tensors} + + def execute(self, tensors: Dict[ModelReference, mindspore.Tensor]) -> mindspore.Tensor: + if len(tensors) == 1: + return list(tensors.values())[0] + elif len(tensors) != 2: + raise RuntimeError("ArceeFusion merge expects exactly two models") + elif self.base_model not in tensors: + raise RuntimeError("Base model not in input tensors") + + [a, b] = list(tensors.items()) + if a[0] != self.base_model: + [a, b] = [b, a] + lhs, rhs = a[1], b[1] + threshold = self.max_tensor_mem_gb + if ( + threshold is not None + and self.split_pieces > 1 + and lhs.ndim >= 1 + and int(lhs.nbytes) > int(float(threshold) * (1024**3)) + and int(lhs.shape[0]) >= self.split_pieces + ): + total = int(lhs.shape[0]) + outputs = [] + for piece_idx in range(self.split_pieces): + start = (total * piece_idx) // self.split_pieces + end = (total * (piece_idx + 1)) // self.split_pieces + if end <= start: + continue + outputs.append(self._merge_core(lhs[start:end], rhs[start:end])) + if outputs: + return ops.concat(outputs, axis=0) + return self._merge_core(lhs, rhs) + + def _merge_core( + self, lhs: mindspore.Tensor, rhs: mindspore.Tensor + ) -> mindspore.Tensor: + prepped_tensors = [lhs, rhs] + out_dtype = prepped_tensors[0].dtype + work_dtype = choose_work_dtype(out_dtype) + rectify_embed_sizes(self.weight_info, prepped_tensors) + prepped_tensors = [t.astype(work_dtype) for t in prepped_tensors] + importance_scores = self._compute_importance(prepped_tensors[1], prepped_tensors[0]) + dynamic_threshold_fusion = DynamicThresholdFusion() + fusion_mask, _threshold = dynamic_threshold_fusion.compute_fusion_mask(importance_scores) + delta = prepped_tensors[1] - prepped_tensors[0] + masked_delta = delta * fusion_mask + fused = prepped_tensors[0] + masked_delta + return fused.astype(out_dtype) + + def _compute_importance( + self, params: mindspore.Tensor, base_params: mindspore.Tensor, eps: float = 1e-8 + ) -> mindspore.Tensor: + diff = safe_abs(params - base_params, out_dtype=params.dtype, op_name="arcee.diff_abs") + p = ops.softmax(params.astype(mindspore.float32), axis=-1) + eps + q = ops.softmax(base_params.astype(mindspore.float32), axis=-1) + eps + kl_div = (p * ops.log(p / q)).sum(axis=-1) + return diff * kl_div.unsqueeze(-1) + + +class ArceeFusionMerge(MergeMethod): + def name(self) -> str: + return "arcee_fusion" + + @override + def pretty_name(self) -> Optional[str]: + return "Arcee Fusion" + + @override + def reference_url(self) -> Optional[str]: + return "https://arcee.ai" + + def parameters(self) -> List[ConfigParameterDef]: + return [] + + def make_task( # pylint: disable=too-many-positional-arguments + self, + output_weight: WeightInfo, + tensors: MergeTensorInput, + base_model: Optional[ModelReference], + split_pieces: int = 1, + max_tensor_mem_gb: Optional[float] = None, + **kwargs, + ) -> Task[mindspore.Tensor]: + return ArceeFusionMergeTask( + gather_tensors=tensors, + weight_info=output_weight, + base_model=base_model, + split_pieces=split_pieces, + max_tensor_mem_gb=max_tensor_mem_gb, + ) diff --git a/src/mindnlp/wizard/merge/merge_methods/base.py b/src/mindnlp/wizard/merge/merge_methods/base.py new file mode 100644 index 000000000..0ef8b8a44 --- /dev/null +++ b/src/mindnlp/wizard/merge/merge_methods/base.py @@ -0,0 +1,71 @@ +# Copyright (C) 2025 Arcee AI +# SPDX-License-Identifier: LGPL-3.0-only +# Modified for MindSpore/Ascend NPU by MindNLP Wizard contributors. +# + +from abc import ABC, abstractmethod +from typing import Any, Dict, List, Optional, Union + +import mindspore # pylint: disable=import-error +from pydantic import BaseModel +from typing_extensions import TypeAlias + +from ..architecture.base import WeightInfo +from ..common import ImmutableMap, ModelReference +from ..graph import Task +from ..io.tasks import GatherTensors +from ..tokenizer import PermutedEmbeddings + + +class TensorDictWrapper(Task[Dict[ModelReference, mindspore.Tensor]]): + tensors: ImmutableMap[ModelReference, Task[mindspore.Tensor]] + + def arguments(self) -> Dict[str, Task]: + return { + k.model_dump_json( + exclude_none=True, exclude_defaults=True, round_trip=True + ): v + for k, v in self.tensors.items() + } + + def execute(self, **kwargs) -> Dict[ModelReference, mindspore.Tensor]: + return {ModelReference.model_validate_json(k): v for k, v in kwargs.items()} + + +MergeTensorInput: TypeAlias = Union[ + GatherTensors, PermutedEmbeddings, TensorDictWrapper +] + + +class ConfigParameterDef(BaseModel): + name: str + required: bool = False + default_value: Any = None + + +class MergeMethod(ABC): + def tensor_parameters(self) -> List[ConfigParameterDef]: + return [] + + def parameters(self) -> List[ConfigParameterDef]: + return [] + + @abstractmethod + def name(self) -> str: ... + + def pretty_name(self) -> Optional[str]: + return None + + def reference_url(self) -> Optional[str]: + return None + + @abstractmethod + def make_task( + self, + *, + output_weight: WeightInfo, + tensors: MergeTensorInput, + parameters: ImmutableMap[str, Any], + tensor_parameters: ImmutableMap[ModelReference, ImmutableMap[str, Any]], + base_model: Optional[ModelReference], + ) -> Task: ... diff --git a/src/mindnlp/wizard/merge/merge_methods/easy_define.py b/src/mindnlp/wizard/merge/merge_methods/easy_define.py new file mode 100644 index 000000000..6ce9b3878 --- /dev/null +++ b/src/mindnlp/wizard/merge/merge_methods/easy_define.py @@ -0,0 +1,347 @@ +# Copyright (C) 2025 Arcee AI +# SPDX-License-Identifier: LGPL-3.0-only +# Modified for MindSpore/Ascend NPU by MindNLP Wizard contributors. +# + +import inspect +import typing +from typing import Any, Dict, List, Optional + +import mindspore # pylint: disable=import-error +import pydantic +from pydantic import Field +from typing_extensions import Callable # pylint: disable=no-name-in-module + +from ..architecture.base import WeightInfo +from ..common import ImmutableMap, ModelReference +from ..graph import Task +from .base import ( + ConfigParameterDef, + MergeMethod, + MergeTensorInput, +) +from .registry import REGISTERED_MERGE_METHODS + +STANDARD_KWARGS = {"output_weight", "base_model"} + + +def __merge_method( + func: Callable, + name: str, + reference_url: Optional[str] = None, + pretty_name: Optional[str] = None, +) -> Callable: + use_base_tensor_arg = False + require_base_tensor = False + used_kwargs = set() + parameters: List[ConfigParameterDef] = [] + tensor_parameters: List[ConfigParameterDef] = [] + sig = inspect.signature(func) + if "tensors" not in sig.parameters: + raise ValueError("Merge methods must have a 'tensors' parameter") + tensor_param = sig.parameters["tensors"] + if ( + (tensor_param.annotation is None) + or (not hasattr(tensor_param.annotation, "__origin__")) + or not ( + tensor_param.annotation.__origin__ == list + and tensor_param.annotation.__args__ == (mindspore.Tensor,) + ) + ): + raise ValueError("'tensors' must be annotated with List[mindspore.Tensor]") + if "base_tensor" in sig.parameters: + bt_param = sig.parameters["base_tensor"] + if bt_param.annotation == mindspore.Tensor: + require_base_tensor = True + elif ( + hasattr(bt_param.annotation, "__origin__") + and bt_param.annotation.__origin__ == typing.Union + and bt_param.annotation.__args__ in ( # pylint: disable=consider-using-in + (mindspore.Tensor, type(None)), + (type(None), mindspore.Tensor), + ) + ): + require_base_tensor = False + else: + raise ValueError( + "'base_tensor' must be annotated either mindspore.Tensor or Optional[mindspore.Tensor]" + ) + use_base_tensor_arg = True + for arg, arg_info in sig.parameters.items(): + if arg in ("base_tensor", "tensors"): + continue + if arg in STANDARD_KWARGS: + used_kwargs.add(arg) + else: + if arg_info.annotation is None: + raise ValueError( + "All merge method arguments must have type annotations" + ) + elif arg_info.annotation in (int, float, bool): + default_value = arg_info.default + if default_value == inspect.Parameter.empty: + default_value = None + required = True + else: + required = False + parameters.append( + ConfigParameterDef( + name=arg, required=required, default_value=default_value + ) + ) + elif ( + hasattr(arg_info.annotation, "__origin__") + and arg_info.annotation.__origin__ == list + and arg_info.annotation.__args__[0] in (float, int) + ): + default_value = arg_info.default + if default_value == inspect.Parameter.empty: + default_value = None + required = True + else: + required = False + if (not required) and (not isinstance(default_value, (int, float))): + raise ValueError( + f"Unexpected default for presumed tensor parameter '{arg}' - should be single number, got {repr(default_value)}" + ) + tensor_parameters.append( + ConfigParameterDef( + name=arg, required=required, default_value=default_value + ) + ) + + tt_fields = {} + tt_fields["gather_tensors"] = (MergeTensorInput, Field(...)) + tt_fields["split_pieces"] = (int, Field(default=1)) + tt_fields["max_tensor_mem_gb"] = (Optional[float], Field(default=None)) + if ("base_model" in used_kwargs) or use_base_tensor_arg: + bm_ty = ModelReference if require_base_tensor else Optional[ModelReference] + field_kwargs = {"default": None} if not require_base_tensor else {} + tt_fields["base_model"] = (bm_ty, Field(**field_kwargs)) + if "output_weight" in used_kwargs: + tt_fields["output_weight"] = (WeightInfo, Field(...)) + if parameters: + tt_fields["parameters"] = (ImmutableMap[str, Any], Field(...)) + if tensor_parameters: + tt_fields["tensor_parameters"] = ( + ImmutableMap[ModelReference, ImmutableMap[str, Any]], + Field(...), + ) + + def _arguments(self) -> Dict[str, Task]: + return {"tensors": self.gather_tensors} + + def _group_label(self) -> Optional[str]: + return self.gather_tensors.group_label() + + def _uses_accelerator(self) -> bool: + return True + + def _execute(self, tensors: Dict[ModelReference, mindspore.Tensor], **_kwargs): + model_refs = set(tensors.keys()) + base_model = getattr(self, "base_model", None) + if base_model and base_model in model_refs: + model_refs.remove(base_model) + if not use_base_tensor_arg: + model_refs = [base_model] + list(model_refs) + else: + model_refs = list(model_refs) + base_tensor = tensors.get(base_model, None) + tensors = [tensors[key] for key in model_refs] + inner_kwargs = {} + for key in used_kwargs: + inner_kwargs[key] = getattr(self, key) + if use_base_tensor_arg: + inner_kwargs["base_tensor"] = base_tensor + if require_base_tensor and (inner_kwargs["base_tensor"] is None): + raise ValueError("Base model tensor required but not present") + for key in parameters: + inner_kwargs[key.name] = self.parameters[key.name] + for key in tensor_parameters: + inner_kwargs[key.name] = [ + self.tensor_parameters[ref][key.name] for ref in model_refs + ] + + if ( + self.max_tensor_mem_gb is not None + and self.split_pieces > 1 + and tensors + and tensors[0].ndim >= 1 + ): + threshold_bytes = int(float(self.max_tensor_mem_gb) * (1024**3)) + first_tensor_bytes = int(tensors[0].nbytes) + if first_tensor_bytes > threshold_bytes and tensors[0].shape[0] >= self.split_pieces: + piece_outputs: List[mindspore.Tensor] = [] + total = int(tensors[0].shape[0]) + base_tensor_piece = inner_kwargs.get("base_tensor", None) + for piece_idx in range(self.split_pieces): + start = (total * piece_idx) // self.split_pieces + end = (total * (piece_idx + 1)) // self.split_pieces + if end <= start: + continue + piece_tensors = [t[start:end] for t in tensors] + piece_kwargs = dict(inner_kwargs) + if use_base_tensor_arg and base_tensor_piece is not None: + piece_kwargs["base_tensor"] = base_tensor_piece[start:end] + piece_outputs.append(func(tensors=piece_tensors, **piece_kwargs)) + + if piece_outputs: + return mindspore.ops.concat(piece_outputs, axis=0) + + return func(tensors=tensors, **inner_kwargs) + + tt_name = f"{name.title().replace(' ', '')}MergeTask" + tt_cls = pydantic.create_model(tt_name, __base__=Task[mindspore.Tensor], **tt_fields) + setattr(tt_cls, "arguments", _arguments) + setattr(tt_cls, "group_label", _group_label) + setattr(tt_cls, "uses_accelerator", _uses_accelerator) + setattr(tt_cls, "execute", _execute) + tt_cls.__abstractmethods__ = frozenset() + + mm_fields = {} + + def _make_task( + self, + *, + output_weight: WeightInfo, + tensors: MergeTensorInput, + parameters: ImmutableMap[str, Any], + tensor_parameters: ImmutableMap[ModelReference, ImmutableMap[str, Any]], + base_model: Optional[ModelReference], + split_pieces: int = 1, + max_tensor_mem_gb: Optional[float] = None, + **_kwargs, + ) -> Task: + tt_kwargs = {"gather_tensors": tensors} + tt_kwargs["split_pieces"] = split_pieces + tt_kwargs["max_tensor_mem_gb"] = max_tensor_mem_gb + if "base_model" in tt_fields: + tt_kwargs["base_model"] = base_model + if "output_weight" in tt_fields: + tt_kwargs["output_weight"] = output_weight + if "parameters" in tt_fields: + tt_kwargs["parameters"] = parameters + if "tensor_parameters" in tt_fields: + tt_kwargs["tensor_parameters"] = tensor_parameters + return tt_cls(**tt_kwargs) + + mm_fields["make_task"] = _make_task + + def _name(self) -> str: + return name + + mm_fields["name"] = _name + + def _pretty_name(self) -> Optional[str]: + return pretty_name + + mm_fields["pretty_name"] = _pretty_name + + def _reference_url(self) -> Optional[str]: + return reference_url + + mm_fields["reference_url"] = _reference_url + + def _tensor_parameters(self) -> List[ConfigParameterDef]: + return tensor_parameters + + mm_fields["tensor_parameters"] = _tensor_parameters + + def _parameters(self) -> List[ConfigParameterDef]: + return parameters + + mm_fields["parameters"] = _parameters + + mm_name = f"{name.title().replace(' ', '')}MergeMethod" + mm_cls = type(mm_name, (MergeMethod,), mm_fields) + REGISTERED_MERGE_METHODS[name] = mm_cls() # pylint: disable=abstract-class-instantiated + return func + + +def merge_method( + name: str, + reference_url: Optional[str] = None, + pretty_name: Optional[str] = None, +) -> Callable: + """Decorator for registering custom model merging algorithms. + + Enables creation of new merge algorithms that can be specified in merge configurations + and executed through mergekit's processing pipeline. Handles parameter validation, task + creation, and registration in the mergekit system. + + Args: + name: Unique identifier for the merge method (lowercase, snake_case recommended) + reference_url: Optional URL to paper/documentation explaining the method (used in generated READMEs) + pretty_name: Human-readable display name (used in generated READMEs) + + Returns: + A decorator that registers the function as a merge method implementation + + Notes: + The decorated function must meet these requirements: + - First parameter must be `tensors: List[mindspore.Tensor]` + - Must return a single `mindspore.Tensor` + - All parameters must have type annotations + + Key behavioral considerations: + + *Base Model Handling:* + - If the method includes a `base_tensor` parameter: + * `mindspore.Tensor` annotation: Requires `base_model` in config, receives its tensor + * `Optional[mindspore.Tensor]` annotation: `base_model` optional, `None` if not provided + * Non-base model tensors passed in `tensors` list + - Without `base_tensor` parameter: + * Base model tensor (if specified) will be first in `tensors` list + + *Parameter Types:* + - Standard parameters (auto-populated): + * `base_tensor`: Tensor from base model (type determines requirement) + * `output_weight`: WeightInfo with output configuration + * `base_model`: ModelReference if using base model logic + - Scalar parameters (global config): + * `float`, `int`, or `bool` types specified in top-level `parameters` + - Tensor parameters (per-model weights): + * Annotated as `List[float]` or `List[int]` + * Configured per-model in their `parameters` section + * Collected into lists ordered by input models + + Example: + ```python + @merge_method( + name="average", + pretty_name="Simple Average", + reference_url="https://example.com/mean-merge" + ) + def average_merge( + tensors: List[mindspore.Tensor], # Input tensors to merge + weights: List[float], # Per-model weights (tensor parameter) + normalize: bool = True # Scalar parameter + ) -> mindspore.Tensor: + if normalize: + weights = [w / sum(weights) for w in weights] + return sum(t * w for t, w in zip(tensors, weights)) + ``` + + This would enable merge configurations like: + ```yaml + merge_method: average + models: + - model: model_a + parameters: + weights: 0.3 + - model: model_b + parameters: + weights: [0.6, 0.8] + parameters: + normalize: true + ``` + + Raises: + ValueError: If function signature doesn't meet requirements + TypeError: For invalid parameter annotations + """ + + def _wrap(func: Callable) -> Callable: + return __merge_method(func, name, reference_url, pretty_name) + + return _wrap diff --git a/src/mindnlp/wizard/merge/merge_methods/generalized_task_arithmetic.py b/src/mindnlp/wizard/merge/merge_methods/generalized_task_arithmetic.py new file mode 100644 index 000000000..b313f5fc4 --- /dev/null +++ b/src/mindnlp/wizard/merge/merge_methods/generalized_task_arithmetic.py @@ -0,0 +1,345 @@ +# Copyright (C) 2025 Arcee AI +# SPDX-License-Identifier: LGPL-3.0-only +# Modified for MindSpore/Ascend NPU by MindNLP Wizard contributors. +# + +import logging +from enum import Enum +from typing import Any, Dict, List, Optional, Tuple + +import numpy as np +import mindspore # pylint: disable=import-error +from mindspore import ops # pylint: disable=import-error +from pydantic import BaseModel +from typing_extensions import Literal, override + +from ..architecture.base import WeightInfo +from ..common import ImmutableMap, ModelReference +from ..graph import Task +from ..safe_ops import safe_abs, safe_mul, safe_stack, safe_sum, safe_where +from .base import ( + ConfigParameterDef, + MergeMethod, + MergeTensorInput, +) +from ..sparsify import RescaleNorm, SparsificationMethod, sparsify + + +class ConsensusMethod(str, Enum): + count = "count" + sum = "sum" + + +class GeneralizedTaskArithmeticMerge(MergeMethod, BaseModel, frozen=True): + consensus_method: Optional[ConsensusMethod] + sparsification_method: Optional[SparsificationMethod] + default_normalize: bool + default_rescale: bool + method_name: str + method_pretty_name: Optional[str] + method_reference_url: Optional[str] + + def name(self) -> str: + return self.method_name + + @override + def pretty_name(self) -> Optional[str]: + return self.method_pretty_name + + @override + def reference_url(self) -> Optional[str]: + return self.method_reference_url + + def parameters(self) -> List[ConfigParameterDef]: + return [ + ConfigParameterDef(name="int8_mask", required=False, default_value=False), + ConfigParameterDef( + name="normalize", required=False, default_value=self.default_normalize + ), + ConfigParameterDef( + name="rescale", required=False, default_value=self.default_rescale + ), + ConfigParameterDef(name="lambda", required=False, default_value=1.0), + ] + + def tensor_parameters(self) -> List[ConfigParameterDef]: + res = [ + ConfigParameterDef(name="weight", required=True), + ConfigParameterDef(name="density", required=False, default_value=1.0), + ] + if self.sparsification_method == SparsificationMethod.magnitude_outliers: + res.append( + ConfigParameterDef( + name="gamma", + default_value=0.01, + ) + ) + if self.sparsification_method == SparsificationMethod.della_magprune: + res.append( + ConfigParameterDef( + name="epsilon", + default_value=0.15, + ) + ) + return res + + def make_task( # pylint: disable=too-many-positional-arguments + self, + output_weight: WeightInfo, + tensors: MergeTensorInput, + base_model: Optional[ModelReference], + parameters: ImmutableMap[str, Any], + tensor_parameters: ImmutableMap[ModelReference, ImmutableMap[str, Any]], + split_pieces: int = 1, + max_tensor_mem_gb: Optional[float] = None, + ) -> Task: + return GTATask( + method=self, + tensors=tensors, + base_model=base_model, + tensor_parameters=tensor_parameters, + int8_mask=parameters["int8_mask"], + normalize=parameters["normalize"], + lambda_=parameters["lambda"], + rescale_norm=RescaleNorm.l1 if parameters["rescale"] else None, + weight_info=output_weight, + split_pieces=split_pieces, + max_tensor_mem_gb=max_tensor_mem_gb, + ) + + +class GTATask(Task[mindspore.Tensor]): + method: GeneralizedTaskArithmeticMerge + tensors: MergeTensorInput + base_model: ModelReference + weight_info: WeightInfo + tensor_parameters: ImmutableMap[ModelReference, Any] + int8_mask: bool + normalize: bool + lambda_: float + rescale_norm: Optional[RescaleNorm] + split_pieces: int = 1 + max_tensor_mem_gb: Optional[float] = None + + def uses_accelerator(self) -> bool: + return True + + def arguments(self) -> Dict[str, Task]: + return {"tensors": self.tensors} + + def execute( + self, + tensors: Dict[ModelReference, mindspore.Tensor], + **_kwargs, + ) -> mindspore.Tensor: + if not tensors: + raise RuntimeError("No tensors provided to GTATask") + + first = next(iter(tensors.values())) + threshold = self.max_tensor_mem_gb + if ( + threshold is not None + and self.split_pieces > 1 + and first.ndim >= 1 + and int(first.nbytes) > int(float(threshold) * (1024**3)) + and int(first.shape[0]) >= self.split_pieces + ): + total = int(first.shape[0]) + outputs = [] + for piece_idx in range(self.split_pieces): + start = (total * piece_idx) // self.split_pieces + end = (total * (piece_idx + 1)) // self.split_pieces + if end <= start: + continue + piece_tensors = {k: v[start:end] for k, v in tensors.items()} + outputs.append(self._execute_core(piece_tensors)) + if outputs: + return ops.concat(outputs, axis=0) + + return self._execute_core(tensors) + + def _execute_core(self, tensors: Dict[ModelReference, mindspore.Tensor]) -> mindspore.Tensor: + tvs, base = get_task_vectors( + self.weight_info, + self.base_model, + dict(tensors), + tensor_parameters=self.tensor_parameters.data, + ) + if not tvs: + return base + + out_dtype = base.dtype + work_dtype = mindspore.float32 + + if self.method.sparsification_method: + for tv_info in tvs: + kwargs = {} + if "gamma" in tv_info: + kwargs["gamma"] = tv_info["gamma"] + + if "epsilon" in tv_info: + kwargs["epsilon"] = tv_info["epsilon"] + + tv_info["delta"] = sparsify( + tv_info["delta"], + density=tv_info["density"], + method=self.method.sparsification_method, + rescale_norm=self.rescale_norm, + **kwargs, + ) + + deltas = safe_stack( + [tv["delta"] for tv in tvs], + axis=0, + out_dtype=work_dtype, + op_name="gta.stack", + ) + + weight_list = [tv["weight"] for tv in tvs] + w_np = np.array(weight_list, dtype=np.float32) + for _ in range(len(deltas.shape) - 1): + w_np = np.expand_dims(w_np, axis=-1) + weights = mindspore.Tensor(w_np, dtype=work_dtype) + + weighted_deltas = safe_mul( + deltas, weights, out_dtype=work_dtype, op_name="gta.mul" + ) + + if self.method.consensus_method: + mask_dtype = mindspore.int8 if self.int8_mask else work_dtype + mask = get_mask( + weighted_deltas, + method=self.method.consensus_method, + mask_dtype=mask_dtype, + ) + + mixed_delta = safe_sum( + safe_mul( + weighted_deltas, mask, out_dtype=work_dtype, op_name="gta.mask_mul" + ), + axis=0, + out_dtype=work_dtype, + op_name="gta.mask_sum", + ) + divisor = safe_sum( + safe_mul(weights, mask, out_dtype=work_dtype, op_name="gta.div_mul"), + axis=0, + out_dtype=work_dtype, + op_name="gta.div_sum", + ) + zero_fill = mindspore.Tensor(1.0, dtype=work_dtype) + divisor = safe_where( + ops.equal(divisor, 0), + zero_fill, + divisor, + out_dtype=work_dtype, + op_name="gta.divisor_fill_zero", + ) + else: + mixed_delta = safe_sum( + weighted_deltas, axis=0, out_dtype=work_dtype, op_name="gta.sum" + ) + divisor = safe_sum(weights, axis=0, out_dtype=work_dtype, op_name="gta.wsum") + tiny = mindspore.Tensor(1e-8, dtype=work_dtype) + one = mindspore.Tensor(1.0, dtype=work_dtype) + divisor = safe_where( + ops.less(safe_abs(divisor, out_dtype=work_dtype, op_name="gta.divisor_abs"), tiny), + one, + divisor, + out_dtype=work_dtype, + op_name="gta.divisor_fill_tiny", + ) + + if self.normalize: + mixed_delta = ops.div(mixed_delta, divisor) + + if self.lambda_ != 1: + lambda_tensor = mindspore.Tensor(float(self.lambda_), dtype=work_dtype) + mixed_delta = ops.mul(mixed_delta, lambda_tensor) + + result = ops.add(base.astype(work_dtype), mixed_delta) + + return result.astype(out_dtype) + + def group_label(self) -> Optional[str]: + return self.tensors.group_label() + + +def get_task_vectors( + weight_info: WeightInfo, + base_model: ModelReference, + tensors: ImmutableMap[ModelReference, mindspore.Tensor], + tensor_parameters: ImmutableMap[ModelReference, ImmutableMap[str, Any]], +) -> Tuple[List[Dict[str, Any]], mindspore.Tensor]: + keys = list(tensors.keys()) + base = tensors[base_model] + + parameter_name = weight_info.name + + # Compute deltas in float32 unconditionally to avoid unsupported + # half-precision arithmetic and intermediate overflow. + work_dtype = mindspore.float32 + base_work = base.astype(work_dtype) + + res = [] + for model in keys: + if model == base_model: + continue + + x = tensors[model].astype(base.dtype) + if x.shape != base.shape: + if weight_info.is_embed: + x = x[: base.shape[0], : base.shape[1]] + logging.warning(f"Using submatrix of {model}:{parameter_name}") + else: + logging.warning( + f"skipping {model}:{parameter_name} due to size mismatch" + ) + continue + + delta = x.astype(work_dtype) - base_work + del x + del tensors[model] + + d = {} + d["model"] = model + d["delta"] = delta + for p in tensor_parameters[model]: + d[p] = tensor_parameters[model][p] + res.append(d) + return res, base + + +def get_mask( + delta: mindspore.Tensor, + method: Literal["sum", "count"] = "sum", + mask_dtype: Optional[mindspore.dtype] = None, +): + """Returns a mask determining which delta vectors should be merged + into the final model. + + For the methodology described in the TIES paper use 'sum'. For a + simpler naive count of signs, use 'count'.""" + if mask_dtype is None: + mask_dtype = delta.dtype + + sign = ops.sign(delta).astype(mask_dtype) + + # Use explicit typed constants to avoid MindSpore CPU rejecting + # mixed-type Mul/Sub (e.g. int8 × int64). + _two = mindspore.Tensor(2, dtype=mask_dtype) + _one = mindspore.Tensor(1, dtype=mask_dtype) + + if method == "sum": + sign_weight = safe_sum(delta, axis=0, out_dtype=delta.dtype, op_name="gta.mask_sign_sum") + ge_zero = ops.greater_equal(sign_weight, mindspore.Tensor(0, dtype=sign_weight.dtype)) + majority_sign = ops.sub(ops.mul(ge_zero.astype(mask_dtype), _two), _one) + del sign_weight + elif method == "count": + sign_sum = safe_sum(sign, axis=0, out_dtype=sign.dtype, op_name="gta.mask_count_sum") + ge_zero = ops.greater_equal(sign_sum, mindspore.Tensor(0, dtype=sign_sum.dtype)) + majority_sign = ops.sub(ops.mul(ge_zero.astype(mask_dtype), _two), _one) + else: + raise RuntimeError(f'Unimplemented mask method "{method}"') + + return ops.equal(sign, majority_sign) diff --git a/src/mindnlp/wizard/merge/merge_methods/karcher.py b/src/mindnlp/wizard/merge/merge_methods/karcher.py new file mode 100644 index 000000000..e8c289efc --- /dev/null +++ b/src/mindnlp/wizard/merge/merge_methods/karcher.py @@ -0,0 +1,219 @@ +# Copyright (C) 2025 Arcee AI +# SPDX-License-Identifier: LGPL-3.0-only +# Modified for MindSpore/Ascend NPU by MindNLP Wizard contributors. +# + +from typing import Any, Dict, List, Optional + +import mindspore # pylint: disable=import-error +from mindspore import ops # pylint: disable=import-error +from typing_extensions import override + +from ..architecture.base import WeightInfo +from ..common import ImmutableMap, ModelReference +from ..dtype_policy import choose_work_dtype +from ..graph import Task +from .base import ( + ConfigParameterDef, + MergeMethod, + MergeTensorInput, +) +from .rectify_embed import rectify_embed_sizes + + +class KarcherTask(Task[mindspore.Tensor]): + """ + Task for merging model weights using the Riemannian (Karcher) mean algorithm. + + The Karcher mean provides a geometrically meaningful way to average points on a manifold, + which is particularly useful for merging model weights that can be interpreted as points + on a hypersphere. + """ + + gather_tensors: MergeTensorInput + weight_info: WeightInfo + max_iter: int + tol: float + split_pieces: int = 1 + max_tensor_mem_gb: Optional[float] = None + + def uses_accelerator(self) -> bool: + return True + + def arguments(self) -> Dict[str, Task]: + return {"tensors": self.gather_tensors} + + def execute(self, tensors: Dict[ModelReference, mindspore.Tensor]) -> mindspore.Tensor: + if len(tensors) == 1: + return list(tensors.values())[0] + + model_tensors = list(tensors.values()) + threshold = self.max_tensor_mem_gb + if ( + threshold is not None + and self.split_pieces > 1 + and model_tensors + and model_tensors[0].ndim >= 1 + and int(model_tensors[0].nbytes) > int(float(threshold) * (1024**3)) + and int(model_tensors[0].shape[0]) >= self.split_pieces + ): + total = int(model_tensors[0].shape[0]) + outputs = [] + for piece_idx in range(self.split_pieces): + start = (total * piece_idx) // self.split_pieces + end = (total * (piece_idx + 1)) // self.split_pieces + if end <= start: + continue + outputs.append(self._merge_core([t[start:end] for t in model_tensors])) + if outputs: + return ops.concat(outputs, axis=0) + return self._merge_core(model_tensors) + + def _merge_core(self, model_tensors: List[mindspore.Tensor]) -> mindspore.Tensor: + + for i in range(1, len(model_tensors)): + rectify_embed_sizes(self.weight_info, [model_tensors[0], model_tensors[i]]) + + alphas = [1.0 / len(model_tensors)] * len(model_tensors) + + return karcher_merge_tensors(model_tensors, alphas, max_iter=self.max_iter, tol=self.tol) + + def group_label(self) -> Optional[str]: + return self.gather_tensors.group_label() + + +class KarcherMerge(MergeMethod): + """ + Implementation of the Karcher mean merge method. + + This method merges model weights using the Riemannian (Karcher) mean concept, + which provides a geometrically meaningful way to average points on a manifold. + """ + + def name(self) -> str: + return "karcher" + + @override + def pretty_name(self) -> Optional[str]: + return "Karcher Mean" + + @override + def reference_url(self) -> Optional[str]: + return "https://en.wikipedia.org/wiki/Karcher_mean" + + def parameters(self) -> List[ConfigParameterDef]: + return [ + ConfigParameterDef(name="max_iter", required=False, default_value=10), + ConfigParameterDef(name="tol", required=False, default_value=1e-5), + ] + + def make_task( + self, + *, + output_weight: WeightInfo, + tensors: MergeTensorInput, + parameters: ImmutableMap[str, Any], + tensor_parameters: ImmutableMap[ModelReference, ImmutableMap[str, Any]], + base_model: Optional[ModelReference], + split_pieces: int = 1, + max_tensor_mem_gb: Optional[float] = None, + **_kwargs, + ) -> Task: + max_iter = parameters["max_iter"] if "max_iter" in parameters else 10 + tol = parameters["tol"] if "tol" in parameters else 1e-5 + + return KarcherTask( + gather_tensors=tensors, + weight_info=output_weight, + max_iter=max_iter, + tol=tol, + split_pieces=split_pieces, + max_tensor_mem_gb=max_tensor_mem_gb, + ) + + +def karcher_merge_tensors(tensors, alphas, max_iter=10, tol=1e-5): + """ + Implements weight fusion based on the Riemannian (Karcher) mean concept. + + Args: + tensors: List of tensors to merge + alphas: List of weights for each tensor + max_iter: Maximum number of iterations for the Karcher mean algorithm + tol: Convergence tolerance + + Returns: + Merged tensor using Karcher mean algorithm + """ + if len(tensors) == 1: + return tensors[0] + out_dtype = tensors[0].dtype + work_dtype = choose_work_dtype(out_dtype) + tensors = [t.astype(work_dtype) for t in tensors] + + norms = [] + units = [] + for t in tensors: + t_float = t.astype(mindspore.float32) + n = ops.norm(t_float) + n_val = n.asnumpy().item() + if n_val == 0.0: + norms.append(0.0) + units.append(ops.zeros_like(t)) + else: + norms.append(n_val) + units.append((t / n).astype(work_dtype)) + + valid_indices = [i for i, n in enumerate(norms) if n > tol] + if not valid_indices: + return ops.zeros_like(tensors[0]) + + valid_alphas = [alphas[i] for i in valid_indices] + alpha_sum = sum(valid_alphas) + normalized_alphas = [a / alpha_sum for a in valid_alphas] + valid_units = [units[i] for i in valid_indices] + + u = ops.zeros_like(valid_units[0]) + for a, ui in zip(normalized_alphas, valid_units): + u = u + a * ui + norm_u = ops.norm(u.astype(mindspore.float32)).asnumpy().item() + if norm_u < tol: + u = valid_units[0].copy() + else: + u = (u / norm_u).astype(u.dtype) + + for _ in range(max_iter): + T = ops.zeros_like(u) + for a, ui in zip(normalized_alphas, valid_units): + dot = ops.clamp( + ( + u.flatten().astype(mindspore.float32) + * ui.flatten().astype(mindspore.float32) + ).sum(), + -1.0, + 1.0, + ) + theta = ops.acos(dot) + theta_val = theta.asnumpy().item() + if theta_val < tol: + continue + sin_theta = ops.sin(theta) + T = T + a * (theta / sin_theta) * (ui - dot * u) + + norm_T = ops.norm(T.astype(mindspore.float32)) + if norm_T.asnumpy().item() < tol: + break + + cos_norm_T = ops.cos(norm_T) + sin_norm_T = ops.sin(norm_T) + u = (cos_norm_T * u + sin_norm_T * (T / norm_T)).astype(u.dtype) + + u_norm = ops.norm(u.astype(mindspore.float32)) + if u_norm.asnumpy().item() > tol: + u = (u / u_norm).astype(u.dtype) + + s = 0.0 + for a, n in zip(alphas, norms): + s += a * n + + return (s * u).astype(out_dtype) diff --git a/src/mindnlp/wizard/merge/merge_methods/linear.py b/src/mindnlp/wizard/merge/merge_methods/linear.py new file mode 100644 index 000000000..6dbe46ef8 --- /dev/null +++ b/src/mindnlp/wizard/merge/merge_methods/linear.py @@ -0,0 +1,131 @@ +# Copyright (C) 2025 Arcee AI +# SPDX-License-Identifier: LGPL-3.0-only +# Modified for MindSpore/Ascend NPU by MindNLP Wizard contributors. +# + +from typing import Any, Dict, List, Optional + +import mindspore # pylint: disable=import-error +from mindspore import ops # pylint: disable=import-error +from typing_extensions import override + +from ..architecture.base import WeightInfo +from ..common import ImmutableMap, ModelReference +from ..dtype_policy import choose_work_dtype +from ..graph import Task +from ..safe_ops import safe_stack +from .base import ( + ConfigParameterDef, + MergeMethod, + MergeTensorInput, +) +from .rectify_embed import rectify_embed_sizes + + +class LinearMergeTask(Task[mindspore.Tensor]): + gather_tensors: MergeTensorInput + tensor_parameters: ImmutableMap[ModelReference, ImmutableMap[str, Any]] + normalize: bool + weight_info: WeightInfo + split_pieces: int = 1 + max_tensor_mem_gb: Optional[float] = None + + def uses_accelerator(self) -> bool: + return True + + def arguments(self) -> Dict[str, Task]: + return {"tensors": self.gather_tensors} + + def execute( + self, tensors: Dict[ModelReference, mindspore.Tensor], **_kwargs + ) -> mindspore.Tensor: + keys = list(tensors.keys()) + tensor_list = [tensors[key] for key in keys] + weights = [self.tensor_parameters[key]["weight"] for key in keys] + + threshold = self.max_tensor_mem_gb + if ( + threshold is not None + and self.split_pieces > 1 + and tensor_list + and tensor_list[0].ndim >= 1 + and int(tensor_list[0].nbytes) > int(float(threshold) * (1024**3)) + and int(tensor_list[0].shape[0]) >= self.split_pieces + ): + total = int(tensor_list[0].shape[0]) + outputs = [] + for piece_idx in range(self.split_pieces): + start = (total * piece_idx) // self.split_pieces + end = (total * (piece_idx + 1)) // self.split_pieces + if end <= start: + continue + piece_tensors = [t[start:end] for t in tensor_list] + outputs.append(self._merge(piece_tensors, weights)) + if outputs: + return ops.concat(outputs, axis=0) + + return self._merge(tensor_list, weights) + + def _merge( + self, tensors: List[mindspore.Tensor], weights: List[float] + ) -> mindspore.Tensor: + rectify_embed_sizes(self.weight_info, tensors) + unique_shapes = set(t.shape for t in tensors) + if len(unique_shapes) != 1: + raise RuntimeError( + f"Tensor size mismatch for {self.weight_info.name}, sizes: {list(unique_shapes)}" + ) + out_dtype = tensors[0].dtype + work_dtype = choose_work_dtype(out_dtype) + stacked = safe_stack(tensors, axis=0, out_dtype=work_dtype, op_name="linear.stack") + weight_tensor = mindspore.Tensor(weights, dtype=work_dtype) + while len(weight_tensor.shape) < len(stacked.shape): + weight_tensor = weight_tensor.unsqueeze(-1) + res = (weight_tensor * stacked).sum(axis=0) + if self.normalize: + res = res / weight_tensor.sum(axis=0) + return res.astype(out_dtype) + + def group_label(self) -> Optional[str]: + return self.gather_tensors.group_label() + + +class LinearMerge(MergeMethod): + def name(self) -> str: + return "linear" + + @override + def pretty_name(self) -> Optional[str]: + return "Linear" + + @override + def reference_url(self) -> Optional[str]: + return "https://arxiv.org/abs/2203.05482" + + def parameters(self) -> List[ConfigParameterDef]: + return [ + ConfigParameterDef(name="normalize", required=False, default_value=True), + ] + + def tensor_parameters(self) -> List[ConfigParameterDef]: + return [ConfigParameterDef(name="weight", required=True)] + + def make_task( + self, + *, + output_weight: WeightInfo, + tensors: MergeTensorInput, + parameters: Dict[str, Any], + tensor_parameters: ImmutableMap[ModelReference, ImmutableMap[str, Any]], + split_pieces: int = 1, + max_tensor_mem_gb: Optional[float] = None, + **_kwargs, + ) -> Task: + return LinearMergeTask( + gather_tensors=tensors, + tensor_parameters=tensor_parameters, + normalize=parameters["normalize"], + weight_info=output_weight, + split_pieces=split_pieces, + max_tensor_mem_gb=max_tensor_mem_gb, + ) diff --git a/src/mindnlp/wizard/merge/merge_methods/model_stock.py b/src/mindnlp/wizard/merge/merge_methods/model_stock.py new file mode 100644 index 000000000..e4efd0029 --- /dev/null +++ b/src/mindnlp/wizard/merge/merge_methods/model_stock.py @@ -0,0 +1,181 @@ +# Copyright (C) 2025 Arcee AI +# SPDX-License-Identifier: LGPL-3.0-only +# Modified for MindSpore/Ascend NPU by MindNLP Wizard contributors. +# + +import logging +from typing import Any, Dict, List, Optional + +import mindspore # pylint: disable=import-error +from mindspore import ops # pylint: disable=import-error +from typing_extensions import override + +from ..architecture.base import WeightInfo +from ..common import ImmutableMap, ModelReference +from ..dtype_policy import choose_work_dtype +from ..graph import Task +from ..safe_ops import safe_norm, safe_stack +from .base import ( + ConfigParameterDef, + MergeMethod, + MergeTensorInput, +) +from .rectify_embed import rectify_embed_sizes + + +class ModelStockMergeTask(Task[mindspore.Tensor]): + gather_tensors: MergeTensorInput + base_model: ModelReference + weight_info: WeightInfo + filter_wise: bool = False + split_pieces: int = 1 + max_tensor_mem_gb: Optional[float] = None + + def uses_accelerator(self) -> bool: + return True + + def arguments(self) -> Dict[str, Task]: + return {"tensors": self.gather_tensors} + + def execute(self, tensors: Dict[ModelReference, mindspore.Tensor]) -> mindspore.Tensor: + if len(tensors) == 1 and self.base_model in tensors: + return tensors[self.base_model] + if len(tensors) < 3: + if self.weight_info.optional: + logging.warning( + f"Optional weight {self.weight_info.name} not present in enough models, discarding" + ) + return None + + raise ValueError( + "ModelStockMerge requires at least 3 models (base plus two+ others)" + ) + + w_0, ws = self.get_rectified_weights(tensors) + threshold = self.max_tensor_mem_gb + if ( + threshold is not None + and self.split_pieces > 1 + and self.filter_wise + and w_0.ndim >= 1 + and int(w_0.nbytes) > int(float(threshold) * (1024**3)) + and int(w_0.shape[0]) >= self.split_pieces + ): + total = int(w_0.shape[0]) + outputs = [] + for piece_idx in range(self.split_pieces): + start = (total * piece_idx) // self.split_pieces + end = (total * (piece_idx + 1)) // self.split_pieces + if end <= start: + continue + outputs.append( + self._merge_core( + w_0[start:end], + [w[start:end] for w in ws], + filter_wise=True, + ) + ) + if outputs: + return ops.concat(outputs, axis=0) + + return self._merge_core(w_0, ws, filter_wise=self.filter_wise) + + def _merge_core( + self, w_0: mindspore.Tensor, ws: List[mindspore.Tensor], *, filter_wise: bool + ) -> mindspore.Tensor: + out_dtype = w_0.dtype + work_dtype = choose_work_dtype(out_dtype) + w_0 = w_0.astype(work_dtype) + ws = [w.astype(work_dtype) for w in ws] + out_shape = w_0.shape + + if filter_wise: + if w_0.ndim == 1: + w_0 = w_0.unsqueeze(0) + ws = [w.unsqueeze(0) for w in ws] + else: + w_0 = w_0.reshape(-1) + ws = [w.reshape(-1) for w in ws] + + offsets = [w - w_0 for w in ws] + cos_thetas = [] + for i, w_0_offset in enumerate(offsets): + for j in range(i + 1, len(offsets)): + w_1_offset = offsets[j] + norm_product = safe_norm( + w_0_offset, + axis=-1, + out_dtype=mindspore.float32, + op_name="model_stock.norm0", + ) * safe_norm( + w_1_offset, + axis=-1, + out_dtype=mindspore.float32, + op_name="model_stock.norm1", + ) + cos_theta = ( + (w_0_offset * w_1_offset).sum(axis=-1) / norm_product.clamp(min=1e-6) + ).clamp(-1, 1) + cos_thetas.append(cos_theta) + + cos_theta = safe_stack( + cos_thetas, out_dtype=work_dtype, op_name="model_stock.stack" + ).mean(axis=0).unsqueeze(-1) + N = len(ws) + t = (N * cos_theta) / (1 + (N - 1) * cos_theta) + w_avg = sum(ws) / len(ws) + w_h = t * w_avg + (1 - t) * w_0 + return w_h.reshape(out_shape).astype(out_dtype) + + def get_rectified_weights(self, tensors: Dict[ModelReference, mindspore.Tensor]): + if self.base_model not in tensors: + raise ValueError("Base model tensor not found") + + all_weights = [tensors[self.base_model]] + [ + tensors[k] for k in tensors if k != self.base_model + ] + rectify_embed_sizes(self.weight_info, all_weights) + w_0 = all_weights[0] + ws = all_weights[1:] + return w_0, ws + + def group_label(self) -> Optional[str]: + return self.gather_tensors.group_label() + + +class ModelStockMerge(MergeMethod): + def name(self) -> str: + return "model_stock" + + @override + def pretty_name(self) -> Optional[str]: + return "Model Stock" + + @override + def reference_url(self): + return "https://arxiv.org/abs/2403.19522" + + def parameters(self) -> List[ConfigParameterDef]: + return [ + ConfigParameterDef(name="filter_wise", required=False, default_value=False) + ] + + def make_task( + self, + *, + output_weight: WeightInfo, + tensors: MergeTensorInput, + base_model: Optional[ModelReference], + parameters: ImmutableMap[str, Any], + split_pieces: int = 1, + max_tensor_mem_gb: Optional[float] = None, + **_kwargs, + ) -> Task: + return ModelStockMergeTask( + gather_tensors=tensors, + base_model=base_model, + weight_info=output_weight, + filter_wise=parameters["filter_wise"], + split_pieces=split_pieces, + max_tensor_mem_gb=max_tensor_mem_gb, + ) diff --git a/src/mindnlp/wizard/merge/merge_methods/multislerp.py b/src/mindnlp/wizard/merge/merge_methods/multislerp.py new file mode 100644 index 000000000..32a3fb42c --- /dev/null +++ b/src/mindnlp/wizard/merge/merge_methods/multislerp.py @@ -0,0 +1,108 @@ +# Copyright (C) 2025 Arcee AI +# SPDX-License-Identifier: LGPL-3.0-only +# Modified for MindSpore/Ascend NPU by MindNLP Wizard contributors. +# + +from typing import List, Optional + +import mindspore # pylint: disable=import-error +from mindspore import ops # pylint: disable=import-error + +from ..dtype_policy import choose_work_dtype +from ..safe_ops import safe_norm, safe_stack +from .easy_define import merge_method + + +@merge_method( + name="multislerp", + pretty_name="Multi-SLERP", + reference_url="https://goddard.blog/posts/multislerp-wow-what-a-cool-idea", +) +def multislerp( + tensors: List[mindspore.Tensor], + weight: List[float], + base_tensor: Optional[mindspore.Tensor] = None, + normalize_weights: bool = True, + eps: float = 1e-8, +): + """ + Implements barycentric interpolation on a hypersphere. + + The approach: + 1. Project points onto a tangent space at their weighted Euclidean mean. + 2. Perform the interpolation in the tangent space. + 3. Project the result back to the hypersphere. + + Limitations: + - The weighted sum of the input tensors must not be zero. + - The tensors must not be all parallel or antiparallel. + + Args: + tensors: List of tensors to interpolate + weight: List of weights for each tensor + base_tensor: Optional tensor defining the origin of the hypersphere + normalize_weights: If True, the weights will be normalized to sum to 1 + eps: Small constant for numerical stability + """ + if len(tensors) == 1: + return tensors[0] + + out_dtype = tensors[0].dtype + work_dtype = choose_work_dtype(out_dtype) + tensors = safe_stack(tensors, axis=0, out_dtype=work_dtype, op_name="multislerp.stack") + if base_tensor is not None: + tensors = tensors - base_tensor.astype(work_dtype) + + tensors_flat = tensors.reshape(tensors.shape[0], -1) + + weights = mindspore.Tensor(weight, dtype=work_dtype) + if normalize_weights: + weights = weights / weights.sum() + + norms = safe_norm( + tensors_flat, + axis=-1, + keepdims=True, + out_dtype=mindspore.float32, + op_name="multislerp.norm", + ) + unit_tensors = tensors_flat / (norms + eps) + + mean = (unit_tensors * weights.reshape(-1, 1)).sum(0) + mean_norm = safe_norm(mean, out_dtype=mindspore.float32, op_name="multislerp.mean_norm") + if mean_norm < eps: + if tensors.shape[0] == 2: + res = (tensors[0] * weights[0] + tensors[1] * weights[1]).reshape( + tensors.shape[1:] + ) + if base_tensor is not None: + res = res + base_tensor.astype(work_dtype) + return res.astype(out_dtype) + raise ValueError( + "The weighted sum of the input tensors is zero. This occurs when " + "antipodal vectors or sets of vectors have weights that exactly " + "balance out (e.g., vectors a,-a with equal weights). Try using " + "different weights if you have antipodal vectors." + ) + mean = mean / mean_norm + + dots = (unit_tensors * mean).sum(-1, keepdims=True) + tangent_vectors = unit_tensors - dots * mean + + tangent_result = (tangent_vectors * weights.reshape(-1, 1)).sum(0) + + tangent_norm = safe_norm( + tangent_result, out_dtype=mindspore.float32, op_name="multislerp.tangent_norm" + ) + eps + result = mean * ops.cos(tangent_norm) + tangent_result * ( + ops.sin(tangent_norm) / tangent_norm + ) + + avg_norm = (norms.squeeze(-1) * weights).sum() + result = result * avg_norm + result = result.reshape(tensors.shape[1:]) + + if base_tensor is not None: + result = result + base_tensor.astype(work_dtype) + + return result.astype(out_dtype) diff --git a/src/mindnlp/wizard/merge/merge_methods/nearswap.py b/src/mindnlp/wizard/merge/merge_methods/nearswap.py new file mode 100644 index 000000000..a8028f264 --- /dev/null +++ b/src/mindnlp/wizard/merge/merge_methods/nearswap.py @@ -0,0 +1,36 @@ +# Copyright (C) 2025 Arcee AI +# SPDX-License-Identifier: LGPL-3.0-only +# Modified for MindSpore/Ascend NPU by MindNLP Wizard contributors. +# + +from typing import List + +import mindspore # pylint: disable=import-error + +from ..dtype_policy import choose_work_dtype +from ..safe_ops import safe_abs +from .easy_define import merge_method + + +@merge_method( + name="nearswap", + pretty_name="NearSwap", + reference_url="https://huggingface.co/alchemonaut/QuartetAnemoi-70B-t0.0001", +) +def nearswap_merge( + tensors: List[mindspore.Tensor], base_tensor: mindspore.Tensor, t: float +) -> mindspore.Tensor: + if not tensors: + return base_tensor + if len(tensors) != 1: + raise RuntimeError( + "NearSwap merge expects exactly two models, one base and one other" + ) + out_dtype = base_tensor.dtype + work_dtype = choose_work_dtype(out_dtype) + a = base_tensor.astype(work_dtype) + b = tensors[0].astype(work_dtype) + + absdiff = safe_abs(a - b, out_dtype=work_dtype, op_name="nearswap.absdiff") + weight = (t / absdiff.clamp(min=1e-6)).clamp(min=0, max=1) + return (weight * b + (1 - weight) * a).astype(out_dtype) diff --git a/src/mindnlp/wizard/merge/merge_methods/nuslerp.py b/src/mindnlp/wizard/merge/merge_methods/nuslerp.py new file mode 100644 index 000000000..7abb673f8 --- /dev/null +++ b/src/mindnlp/wizard/merge/merge_methods/nuslerp.py @@ -0,0 +1,208 @@ +# Copyright (C) 2025 Arcee AI +# SPDX-License-Identifier: LGPL-3.0-only +# Modified for MindSpore/Ascend NPU by MindNLP Wizard contributors. +# + +from typing import Any, Dict, List, Optional + +import mindspore # pylint: disable=import-error +from mindspore import ops # pylint: disable=import-error +from typing_extensions import override + +from ..architecture.base import WeightInfo +from ..common import ImmutableMap, ModelReference +from ..dtype_policy import choose_work_dtype +from ..graph import Task +from ..safe_ops import safe_norm, safe_where +from .base import ( + ConfigParameterDef, + MergeMethod, + MergeTensorInput, +) +from .rectify_embed import rectify_embed_sizes + + +class NuSlerpTask(Task[mindspore.Tensor]): + gather_tensors: MergeTensorInput + tensor_parameters: ImmutableMap[ModelReference, ImmutableMap[str, Any]] + weight_info: WeightInfo + row_wise: bool + flatten: bool + base_model: Optional[ModelReference] + split_pieces: int = 1 + max_tensor_mem_gb: Optional[float] = None + + def uses_accelerator(self) -> bool: + return True + + def arguments(self) -> Dict[str, Task]: + return {"tensors": self.gather_tensors} + + def execute(self, tensors: Dict[ModelReference, mindspore.Tensor]) -> mindspore.Tensor: + if len(tensors) == 1: + return list(tensors.values())[0] + + if self.base_model is not None: + if len(tensors) != 3: + raise RuntimeError( + "NuSlerp base model can not be one of the two models to merge" + ) + base_tensor = tensors.pop(self.base_model) + else: + base_tensor = None + + keys = list(tensors.keys()) + tensors = [tensors[key] for key in keys] + weights = [self.tensor_parameters[key]["weight"] for key in keys] + + if len(tensors) != 2: + raise RuntimeError( + "NuSlerp merge expects exactly two models (plus optional base model)" + ) + + if abs(sum(weights)) < 1e-6: + t = 0.5 + else: + t = weights[1] / sum(weights) + + lhs, rhs = tensors[0], tensors[1] + threshold = self.max_tensor_mem_gb + if ( + threshold is not None + and self.split_pieces > 1 + and lhs.ndim >= 1 + and int(lhs.nbytes) > int(float(threshold) * (1024**3)) + and int(lhs.shape[0]) >= self.split_pieces + ): + total = int(lhs.shape[0]) + outputs = [] + for piece_idx in range(self.split_pieces): + start = (total * piece_idx) // self.split_pieces + end = (total * (piece_idx + 1)) // self.split_pieces + if end <= start: + continue + base_piece = base_tensor[start:end] if base_tensor is not None else None + outputs.append(self._merge_core(lhs[start:end], rhs[start:end], t, base_piece)) + if outputs: + return ops.concat(outputs, axis=0) + return self._merge_core(lhs, rhs, t, base_tensor) + + def _merge_core( + self, + lhs: mindspore.Tensor, + rhs: mindspore.Tensor, + t: float, + base_tensor: Optional[mindspore.Tensor], + ) -> mindspore.Tensor: + prepped = [lhs, rhs] + if base_tensor is not None: + prepped.append(base_tensor) + rectify_embed_sizes(self.weight_info, prepped) + if base_tensor is not None: + base = prepped[2] + return base + nuslerp( + t, + prepped[0] - base, + prepped[1] - base, + dim=0 if self.row_wise else -1, + flatten=self.flatten, + ) + return nuslerp( + t, + prepped[0], + prepped[1], + dim=0 if self.row_wise else -1, + flatten=self.flatten, + ) + + +class NuSlerpMerge(MergeMethod): + def name(self) -> str: + return "nuslerp" + + @override + def pretty_name(self): + return "NuSLERP" + + def parameters(self) -> List[ConfigParameterDef]: + return [ + ConfigParameterDef( + name="nuslerp_row_wise", + required=False, + default_value=False, + ), + ConfigParameterDef( + name="nuslerp_flatten", + required=False, + default_value=True, + ), + ] + + def tensor_parameters(self) -> List[ConfigParameterDef]: + return [ConfigParameterDef(name="weight", required=True)] + + def make_task( + self, + *, + output_weight: WeightInfo, + tensors: MergeTensorInput, + base_model: Optional[ModelReference], + parameters: ImmutableMap[str, Any], + tensor_parameters: ImmutableMap[ModelReference, ImmutableMap[str, Any]], + split_pieces: int = 1, + max_tensor_mem_gb: Optional[float] = None, + **_kwargs, + ) -> Task: + return NuSlerpTask( + gather_tensors=tensors, + tensor_parameters=tensor_parameters, + weight_info=output_weight, + row_wise=parameters["nuslerp_row_wise"], + flatten=parameters["nuslerp_flatten"], + base_model=base_model, + split_pieces=split_pieces, + max_tensor_mem_gb=max_tensor_mem_gb, + ) + + +def nuslerp( # pylint: disable=too-many-positional-arguments + t: float, + v0: mindspore.Tensor, + v1: mindspore.Tensor, + dim: int = -1, + eps: float = 1e-8, + flatten: bool = False, +): + out_dtype = v0.dtype + work_dtype = choose_work_dtype(out_dtype) + v0 = v0.astype(work_dtype) + v1 = v1.astype(work_dtype) + out_shape = v0.shape + + def _normalize(x: mindspore.Tensor, eps: float = 1e-7) -> mindspore.Tensor: + return x / safe_norm( + x, axis=-1, keepdims=True, out_dtype=mindspore.float32, op_name="nuslerp.norm" + ).clamp(min=eps) + + if flatten: + v0 = v0.reshape(-1) + v1 = v1.reshape(-1) + elif dim != -1: + v0 = v0.transpose(dim, -1) + v1 = v1.transpose(dim, -1) + + v0_u = _normalize(v0) + v1_u = _normalize(v1) + + cos_theta = (v0_u * v1_u).sum(axis=-1, keepdims=True) + theta = ops.acos(cos_theta.clamp(-1, 1)) + sin_theta = ops.sin(theta) + + colinear = sin_theta.abs() < eps + res = (ops.sin((1 - t) * theta) * v0 + ops.sin(t * theta) * v1) / sin_theta + fallback = (1 - t) * v0 + t * v1 + res = safe_where(colinear, fallback, res, out_dtype=work_dtype, op_name="nuslerp.where") + + if dim != -1 and not flatten: + res = res.transpose(dim, -1) + return res.reshape(out_shape).astype(out_dtype) diff --git a/src/mindnlp/wizard/merge/merge_methods/passthrough.py b/src/mindnlp/wizard/merge/merge_methods/passthrough.py new file mode 100644 index 000000000..16c86a117 --- /dev/null +++ b/src/mindnlp/wizard/merge/merge_methods/passthrough.py @@ -0,0 +1,63 @@ +# Copyright (C) 2025 Arcee AI +# SPDX-License-Identifier: LGPL-3.0-only +# Modified for MindSpore/Ascend NPU by MindNLP Wizard contributors. +# + +from typing import Any, Dict, List, Optional + +import mindspore # pylint: disable=import-error +from typing_extensions import override + +from ..common import ImmutableMap, ModelReference +from ..graph import Task +from ..safe_ops import safe_mul +from .base import ( + ConfigParameterDef, + MergeMethod, + MergeTensorInput, +) + + +class PassthroughMergeTask(Task[mindspore.Tensor]): + gather_tensors: MergeTensorInput + tensor_parameters: ImmutableMap[ModelReference, ImmutableMap[str, Any]] + + def arguments(self) -> Dict[str, Task]: + return {"tensors": self.gather_tensors} + + def execute(self, tensors: Dict[ModelReference, mindspore.Tensor]) -> mindspore.Tensor: + if len(tensors) != 1: + raise RuntimeError("Passthrough merge expects exactly one tensor") + + model, tensor = list(tensors.items())[0] + scale = self.tensor_parameters[model].data.get("scale", None) + if scale is not None: + tensor = safe_mul(tensor, scale, out_dtype=tensor.dtype, op_name="passthrough.scale") + + return tensor + + def group_label(self) -> Optional[str]: + return self.gather_tensors.group_label() + + +class PassthroughMerge(MergeMethod): + def name(self) -> str: + return "passthrough" + + @override + def pretty_name(self) -> Optional[str]: + return "Passthrough" + + def tensor_parameters(self) -> List[ConfigParameterDef]: + return [ConfigParameterDef(name="scale", required=False, default_value=None)] + + def make_task( + self, + *, + tensors: MergeTensorInput, + tensor_parameters: ImmutableMap[ModelReference, ImmutableMap[str, Any]], + **kwargs, + ) -> Task: + return PassthroughMergeTask( + gather_tensors=tensors, tensor_parameters=tensor_parameters + ) diff --git a/src/mindnlp/wizard/merge/merge_methods/ram.py b/src/mindnlp/wizard/merge/merge_methods/ram.py new file mode 100644 index 000000000..9b28110a4 --- /dev/null +++ b/src/mindnlp/wizard/merge/merge_methods/ram.py @@ -0,0 +1,117 @@ +# Copyright (C) 2025 Arcee AI +# SPDX-License-Identifier: LGPL-3.0-only +# Modified for MindSpore/Ascend NPU by MindNLP Wizard contributors. +# + +from typing import List, Tuple + +import mindspore # pylint: disable=import-error + +from ..dtype_policy import choose_work_dtype +from ..safe_ops import safe_stack +from .easy_define import merge_method + + +@merge_method( + name="ram", + pretty_name="Reinforced Agent Merging", + reference_url="https://arxiv.org/abs/2601.13572", +) +def ram_merge( + tensors: List[mindspore.Tensor], + base_tensor: mindspore.Tensor, + epsilon: float = 1e-5, +) -> mindspore.Tensor: + if not tensors: + return base_tensor + work_dtype = choose_work_dtype(base_tensor.dtype) + base_work = base_tensor.astype(work_dtype) + tensors_work = [t.astype(work_dtype) for t in tensors] + + ( + tv_flat, + nonzero_mask, + contrib_counts, + overlap_mask, + unique_mask, + ) = _prepare_ram_vectors(tensors_work, base_work, epsilon) + + tv_flat_z = tv_flat * nonzero_mask + merged_tv_flat = ( + (tv_flat_z * unique_mask) + + (tv_flat_z * overlap_mask / contrib_counts.clamp(min=1)) + ).sum(axis=0, keepdims=True) + + return (base_work + merged_tv_flat.reshape(base_tensor.shape)).astype( + base_tensor.dtype + ) + + +@merge_method( + name="ramplus_tl", + pretty_name="Reinforced Agent Merging Plus (Tensor-Local)", + reference_url="https://arxiv.org/abs/2601.13572", +) +def ramplus_tl_merge( + tensors: List[mindspore.Tensor], + base_tensor: mindspore.Tensor, + r: float = 0.1, + alpha: float = 0.2, + epsilon: float = 1e-5, +) -> mindspore.Tensor: + if not tensors: + return base_tensor + work_dtype = choose_work_dtype(base_tensor.dtype) + base_work = base_tensor.astype(work_dtype) + tensors_work = [t.astype(work_dtype) for t in tensors] + + ( + tv_flat, + nonzero_mask, + contrib_counts, + overlap_mask, + unique_mask, + ) = _prepare_ram_vectors(tensors_work, base_work, epsilon) + + shared_counts = (nonzero_mask & overlap_mask).sum(axis=1) + unique_counts = (nonzero_mask & unique_mask).sum(axis=1) + rho = shared_counts / unique_counts.clamp(min=epsilon) + lambda_ = 1 + r * rho.clamp(min=0, max=alpha) + + tv_flat_z = tv_flat * nonzero_mask + merged_tv_flat = ( + (tv_flat_z * unique_mask * lambda_.unsqueeze(-1)) + + (tv_flat_z * overlap_mask / contrib_counts.clamp(min=1)) + ).sum(axis=0, keepdims=True) + merged_tv = merged_tv_flat.reshape(base_tensor.shape) + return (base_work + merged_tv).astype(base_tensor.dtype) + + +def _prepare_ram_vectors( + tensors: List[mindspore.Tensor], base_tensor: mindspore.Tensor, epsilon: float +) -> Tuple[mindspore.Tensor, mindspore.Tensor, mindspore.Tensor, mindspore.Tensor, mindspore.Tensor]: + """ + Helper function to compute task vectors, masks, and counts shared by RAM methods. + Returns: + tv_flat: Flattened task vectors + nonzero_mask: Mask of values > epsilon + contrib_counts: Count of models contributing to each parameter + overlap_mask: Mask where counts > 1 + unique_mask: Mask where counts == 1 + """ + task_vectors = safe_stack( + [t - base_tensor for t in tensors], + axis=0, + out_dtype=base_tensor.dtype, + op_name="ram.stack", + ) + tv_flat = task_vectors.reshape(len(tensors), -1) + + nonzero_mask = tv_flat.abs() > epsilon + + contrib_counts = nonzero_mask.sum(axis=0, keepdims=True) + + overlap_mask = contrib_counts > 1 + unique_mask = contrib_counts == 1 + + return tv_flat, nonzero_mask, contrib_counts, overlap_mask, unique_mask diff --git a/src/mindnlp/wizard/merge/merge_methods/rectify_embed.py b/src/mindnlp/wizard/merge/merge_methods/rectify_embed.py new file mode 100644 index 000000000..44b9eb847 --- /dev/null +++ b/src/mindnlp/wizard/merge/merge_methods/rectify_embed.py @@ -0,0 +1,36 @@ +# Copyright (C) 2025 Arcee AI +# SPDX-License-Identifier: LGPL-3.0-only +# Modified for MindSpore/Ascend NPU by MindNLP Wizard contributors. +# + + +import logging +from typing import List + +import mindspore # pylint: disable=import-error + +from ..architecture.base import WeightInfo + + +def rectify_embed_sizes(weight_info: WeightInfo, tensors: List[mindspore.Tensor]): + if weight_info.is_embed and all(len(t.shape) == 2 for t in tensors): + # special case - if lm_head.weight or embed_tokens.weight have a size + # mismatch, take the largest common submatrix of all of them + if take_common_submatrix(tensors): + logging.warning( + f"Using common submatrix of size {tensors[0].shape} for {weight_info.name}" + ) + + +def take_common_submatrix(tensors: List[mindspore.Tensor]) -> bool: + min_size = [None, None] + for t in tensors: + for idx in range(2): + if min_size[idx] is None or t.shape[idx] < min_size[idx]: + min_size[idx] = t.shape[idx] + + if not all(t.shape == tuple(min_size) for t in tensors): + for idx in range(len(tensors)): + tensors[idx] = tensors[idx][: min_size[0], : min_size[1]] + return True + return False diff --git a/src/mindnlp/wizard/merge/merge_methods/registry.py b/src/mindnlp/wizard/merge/merge_methods/registry.py new file mode 100644 index 000000000..09e9c0232 --- /dev/null +++ b/src/mindnlp/wizard/merge/merge_methods/registry.py @@ -0,0 +1,146 @@ +# Copyright (C) 2025 Arcee AI +# SPDX-License-Identifier: LGPL-3.0-only +# Modified for MindSpore/Ascend NPU by MindNLP Wizard contributors. +# +# +import logging +from typing import Dict, List + +from .base import MergeMethod +from ..sparsify import SparsificationMethod + +LOG = logging.getLogger(__name__) + +STATIC_MERGE_METHODS: List[MergeMethod] = [] + +try: + from .linear import LinearMerge + STATIC_MERGE_METHODS.append(LinearMerge()) +except Exception as e: + LOG.warning("Failed to register merge method linear", exc_info=e) + +try: + from .slerp import SlerpMerge + STATIC_MERGE_METHODS.append(SlerpMerge()) +except Exception as e: + LOG.warning("Failed to register merge method slerp", exc_info=e) + +try: + from .nuslerp import NuSlerpMerge + STATIC_MERGE_METHODS.append(NuSlerpMerge()) +except Exception as e: + LOG.warning("Failed to register merge method nuslerp", exc_info=e) + +try: + from .passthrough import PassthroughMerge + STATIC_MERGE_METHODS.append(PassthroughMerge()) +except Exception as e: + LOG.warning("Failed to register merge method passthrough", exc_info=e) + +try: + from .model_stock import ModelStockMerge + STATIC_MERGE_METHODS.append(ModelStockMerge()) +except Exception as e: + LOG.warning("Failed to register merge method model_stock", exc_info=e) + +try: + from .arcee_fusion import ArceeFusionMerge + STATIC_MERGE_METHODS.append(ArceeFusionMerge()) +except Exception as e: + LOG.warning("Failed to register merge method arcee_fusion", exc_info=e) + +try: + from .karcher import KarcherMerge + STATIC_MERGE_METHODS.append(KarcherMerge()) +except Exception as e: + LOG.warning("Failed to register merge method karcher", exc_info=e) + +try: + from .generalized_task_arithmetic import ( + ConsensusMethod, + GeneralizedTaskArithmeticMerge, + ) + # generalized task arithmetic methods + STATIC_MERGE_METHODS.extend([ + GeneralizedTaskArithmeticMerge( + consensus_method=None, + sparsification_method=None, + default_normalize=False, + default_rescale=False, + method_name="task_arithmetic", + method_pretty_name="Task Arithmetic", + method_reference_url="https://arxiv.org/abs/2212.04089", + ), + GeneralizedTaskArithmeticMerge( + consensus_method=ConsensusMethod.sum, + sparsification_method=SparsificationMethod.magnitude, + default_normalize=True, + default_rescale=False, + method_name="ties", + method_pretty_name="TIES", + method_reference_url="https://arxiv.org/abs/2306.01708", + ), + GeneralizedTaskArithmeticMerge( + consensus_method=ConsensusMethod.sum, + sparsification_method=SparsificationMethod.random, + default_normalize=False, + default_rescale=True, + method_name="dare_ties", + method_pretty_name="DARE TIES", + method_reference_url="https://arxiv.org/abs/2311.03099", + ), + GeneralizedTaskArithmeticMerge( + consensus_method=None, + sparsification_method=SparsificationMethod.random, + default_normalize=False, + default_rescale=True, + method_name="dare_linear", + method_pretty_name="Linear DARE", + method_reference_url="https://arxiv.org/abs/2311.03099", + ), + GeneralizedTaskArithmeticMerge( + consensus_method=None, + sparsification_method=SparsificationMethod.magnitude_outliers, + default_normalize=False, + default_rescale=False, + method_name="breadcrumbs", + method_pretty_name="Model Breadcrumbs", + method_reference_url="https://arxiv.org/abs/2312.06795", + ), + GeneralizedTaskArithmeticMerge( + consensus_method=ConsensusMethod.sum, + sparsification_method=SparsificationMethod.magnitude_outliers, + default_normalize=False, + default_rescale=False, + method_name="breadcrumbs_ties", + method_pretty_name="Model Breadcrumbs with TIES", + method_reference_url="https://arxiv.org/abs/2312.06795", + ), + GeneralizedTaskArithmeticMerge( + consensus_method=ConsensusMethod.sum, + sparsification_method=SparsificationMethod.della_magprune, + default_normalize=True, + default_rescale=True, + method_name="della", + method_pretty_name="DELLA", + method_reference_url="https://arxiv.org/abs/2406.11617", + ), + GeneralizedTaskArithmeticMerge( + consensus_method=None, + sparsification_method=SparsificationMethod.della_magprune, + default_normalize=False, + default_rescale=True, + method_name="della_linear", + method_pretty_name="Linear DELLA", + method_reference_url="https://arxiv.org/abs/2406.11617", + ), + ]) +except Exception as e: + LOG.warning( + "Failed to register generalized_task_arithmetic-derived merge methods", + exc_info=e, + ) + +REGISTERED_MERGE_METHODS: Dict[str, MergeMethod] = { + method.name(): method for method in STATIC_MERGE_METHODS +} diff --git a/src/mindnlp/wizard/merge/merge_methods/sce.py b/src/mindnlp/wizard/merge/merge_methods/sce.py new file mode 100644 index 000000000..5ce34abb5 --- /dev/null +++ b/src/mindnlp/wizard/merge/merge_methods/sce.py @@ -0,0 +1,96 @@ +# Copyright (C) 2025 Arcee AI +# SPDX-License-Identifier: LGPL-3.0-only +# Modified for MindSpore/Ascend NPU by MindNLP Wizard contributors. +# + +from typing import List, Optional + +import mindspore # pylint: disable=import-error +from mindspore import ops # pylint: disable=import-error +import numpy as np + +from ..dtype_policy import choose_work_dtype +from ..safe_ops import safe_mul, safe_stack, safe_sum +from .easy_define import merge_method +from .generalized_task_arithmetic import ( + get_mask as sign_consensus_mask, +) + + +@merge_method( + name="sce", + pretty_name="SCE", + reference_url="https://arxiv.org/abs/2408.07990", +) +def sce_merge( + tensors: List[mindspore.Tensor], + base_tensor: mindspore.Tensor, + int8_mask: bool = False, + select_topk: float = 1.0, +) -> mindspore.Tensor: + if not tensors: + return base_tensor + work_dtype = choose_work_dtype(base_tensor.dtype) + base_work = base_tensor.astype(work_dtype) + mask_dtype = mindspore.int8 if int8_mask else work_dtype + task_vectors = safe_stack( + [t.astype(work_dtype) - base_work for t in tensors], + axis=0, + out_dtype=work_dtype, + op_name="sce.stack", + ) + + if select_topk < 1: + mask = sce_mask(task_vectors, select_topk, mask_dtype) + task_vectors = safe_mul( + task_vectors, mask.unsqueeze(0), out_dtype=work_dtype, op_name="sce.mask_mul" + ) + + erase_mask = sign_consensus_mask(task_vectors, method="sum", mask_dtype=mask_dtype) + + tv_weights = sce_weight(task_vectors) + while tv_weights.ndim < task_vectors.ndim: + tv_weights = tv_weights.unsqueeze(-1) + + erased_weights = safe_mul( + tv_weights, erase_mask, out_dtype=work_dtype, op_name="sce.erase_mul" + ) + merged_tv = safe_sum( + safe_mul(task_vectors, erased_weights, out_dtype=work_dtype, op_name="sce.tv_mul"), + axis=0, + out_dtype=work_dtype, + op_name="sce.tv_sum", + ) + final_tv = merged_tv / safe_sum( + erased_weights, axis=0, out_dtype=work_dtype, op_name="sce.weight_sum" + ).clamp(min=1e-6) + + return (base_work + final_tv).astype(base_tensor.dtype) + + +def sce_weight(tvs: mindspore.Tensor) -> mindspore.Tensor: + weights = ops.mean((tvs.astype(mindspore.float32)) ** 2, axis=list(range(1, tvs.ndim))) + weight_sum = float(ops.sum(weights).asnumpy().item()) + if abs(weight_sum) < 1e-6: + return ops.ones_like(weights) / weights.shape[0] + return weights / weight_sum + + +def sce_mask( + tvs: mindspore.Tensor, density: float, mask_dtype: Optional[mindspore.dtype] = None +): + if density <= 0: + return ops.zeros_like(tvs, dtype=mask_dtype) + if density >= 1: + return ops.ones_like(tvs, dtype=mask_dtype) + + var = ops.var(tvs.astype(mindspore.float32), axis=0, ddof=0) + nonzero = int((var != 0).sum().asnumpy().item()) + k = int(nonzero * density) + if k == 0: + return ops.zeros_like(tvs, dtype=mask_dtype) + + _, indices = ops.topk(var.abs().reshape(-1), k=k, largest=True) + mask_np = np.zeros(var.shape, dtype=np.float32) + mask_np.reshape(-1)[indices.asnumpy()] = 1.0 + return mindspore.Tensor(mask_np, dtype=mask_dtype) diff --git a/src/mindnlp/wizard/merge/merge_methods/slerp.py b/src/mindnlp/wizard/merge/merge_methods/slerp.py new file mode 100644 index 000000000..b035fe07a --- /dev/null +++ b/src/mindnlp/wizard/merge/merge_methods/slerp.py @@ -0,0 +1,175 @@ +# Copyright (C) 2025 Arcee AI +# SPDX-License-Identifier: LGPL-3.0-only +# Modified for MindSpore/Ascend NPU by MindNLP Wizard contributors. +# + +from typing import Any, Dict, List, Optional + +import mindspore # pylint: disable=import-error +from mindspore import ops # pylint: disable=import-error +from typing_extensions import override + +from ..architecture.base import WeightInfo +from ..common import ImmutableMap, ModelReference +from ..graph import Task +from .base import ( + ConfigParameterDef, + MergeMethod, + MergeTensorInput, +) +from .rectify_embed import rectify_embed_sizes + + +class SlerpTask(Task[mindspore.Tensor]): + gather_tensors: MergeTensorInput + base_model: ModelReference + t: float + weight_info: WeightInfo + split_pieces: int = 1 + max_tensor_mem_gb: Optional[float] = None + + def uses_accelerator(self) -> bool: + return True + + def arguments(self) -> Dict[str, Task]: + return {"tensors": self.gather_tensors} + + def execute(self, tensors: Dict[ModelReference, mindspore.Tensor]) -> mindspore.Tensor: + if len(tensors) == 1: + return list(tensors.values())[0] + elif len(tensors) != 2: + raise RuntimeError("Slerp merge expects exactly two models") + elif self.base_model not in tensors: + raise RuntimeError("Base model not in input tensors") + + [a, b] = list(tensors.items()) + if a[0] != self.base_model: + [a, b] = [b, a] + lhs, rhs = a[1], b[1] + threshold = self.max_tensor_mem_gb + if ( + threshold is not None + and self.split_pieces > 1 + and lhs.ndim >= 1 + and int(lhs.nbytes) > int(float(threshold) * (1024**3)) + and int(lhs.shape[0]) >= self.split_pieces + ): + total = int(lhs.shape[0]) + outputs = [] + for piece_idx in range(self.split_pieces): + start = (total * piece_idx) // self.split_pieces + end = (total * (piece_idx + 1)) // self.split_pieces + if end <= start: + continue + outputs.append(self._merge_core(lhs[start:end], rhs[start:end])) + if outputs: + return ops.concat(outputs, axis=0) + return self._merge_core(lhs, rhs) + + def _merge_core( + self, lhs: mindspore.Tensor, rhs: mindspore.Tensor + ) -> mindspore.Tensor: + prepped_tensors = [lhs, rhs] + rectify_embed_sizes(self.weight_info, prepped_tensors) + return slerp(self.t, prepped_tensors[0], prepped_tensors[1]).astype( + prepped_tensors[0].dtype + ) + + def group_label(self) -> Optional[str]: + return self.gather_tensors.group_label() + + +class SlerpMerge(MergeMethod): + def name(self) -> str: + return "slerp" + + @override + def pretty_name(self) -> Optional[str]: + return "SLERP" + + @override + def reference_url(self): + return "https://en.wikipedia.org/wiki/Slerp" + + def parameters(self) -> List[ConfigParameterDef]: + return [ConfigParameterDef(name="t", required=True)] + + def make_task( + self, + *, + output_weight: WeightInfo, + tensors: MergeTensorInput, + parameters: ImmutableMap[str, Any], + base_model: Optional[ModelReference], + split_pieces: int = 1, + max_tensor_mem_gb: Optional[float] = None, + **_kwargs, + ) -> Task: + return SlerpTask( + gather_tensors=tensors, + base_model=base_model, + weight_info=output_weight, + t=parameters["t"], + split_pieces=split_pieces, + max_tensor_mem_gb=max_tensor_mem_gb, + ) + + +def lerp( + t: float, v0: mindspore.Tensor, v1: mindspore.Tensor +) -> mindspore.Tensor: + return (1 - t) * v0 + t * v1 + + +def slerp( + t: float, + v0: mindspore.Tensor, + v1: mindspore.Tensor, + DOT_THRESHOLD: float = 0.9995, + eps: float = 1e-8, +) -> mindspore.Tensor: + """ + Spherical linear interpolation + + From: https://gist.github.com/dvschultz/3af50c40df002da3b751efab1daddf2c + Args: + t (float/np.ndarray): Float value between 0.0 and 1.0 + v0 (mindspore.Tensor): Starting vector + v1 (mindspore.Tensor): Final vector + DOT_THRESHOLD (float): Threshold for considering the two vectors as + colinear. Not recommended to alter this. + Returns: + v2 (mindspore.Tensor): Interpolation vector between v0 and v1 + """ + v0_f = v0.astype(mindspore.float32) + v1_f = v1.astype(mindspore.float32) + + v0_copy = v0_f.copy() + v1_copy = v1_f.copy() + + v0_norm = normalize(v0_f, eps) + v1_norm = normalize(v1_f, eps) + + dot = (v0_norm * v1_norm).sum() + + if ops.abs(dot) > DOT_THRESHOLD: + return lerp(t, v0_copy, v1_copy) + + theta_0 = ops.acos(dot) + sin_theta_0 = ops.sin(theta_0) + + theta_t = theta_0 * t + sin_theta_t = ops.sin(theta_t) + + s0 = ops.sin(theta_0 - theta_t) / sin_theta_0 + s1 = sin_theta_t / sin_theta_0 + res = s0 * v0_copy + s1 * v1_copy + + return res + + +def normalize(v: mindspore.Tensor, eps: float) -> mindspore.Tensor: + norm_v = ops.norm(v.astype(mindspore.float32)) + if norm_v > eps: + v = v / norm_v + return v diff --git a/src/mindnlp/wizard/merge/moe/__init__.py b/src/mindnlp/wizard/merge/moe/__init__.py new file mode 100644 index 000000000..0ea1a882f --- /dev/null +++ b/src/mindnlp/wizard/merge/moe/__init__.py @@ -0,0 +1,31 @@ +# Copyright (C) 2025 Arcee AI +# SPDX-License-Identifier: LGPL-3.0-only +# Modified for MindSpore/Ascend NPU by MindNLP Wizard contributors. +# + +from typing import List + +from .arch import MoEOutputArchitecture +from .deepseek import DeepseekMoE +from .mixtral import MixtralMoE + +ALL_OUTPUT_ARCHITECTURES: List[MoEOutputArchitecture] = [MixtralMoE(), DeepseekMoE()] + +try: + from .qwen import QwenMoE +except ImportError: + pass +else: + ALL_OUTPUT_ARCHITECTURES.append(QwenMoE()) + +try: + from .qwen3 import Qwen3MoE +except ImportError: + pass +else: + ALL_OUTPUT_ARCHITECTURES.append(Qwen3MoE()) + +__all__ = [ + "ALL_OUTPUT_ARCHITECTURES", + "MoEOutputArchitecture", +] diff --git a/src/mindnlp/wizard/merge/moe/arch.py b/src/mindnlp/wizard/merge/moe/arch.py new file mode 100644 index 000000000..24657b7f8 --- /dev/null +++ b/src/mindnlp/wizard/merge/moe/arch.py @@ -0,0 +1,40 @@ +# Copyright (C) 2025 Arcee AI +# SPDX-License-Identifier: LGPL-3.0-only +# Modified for MindSpore/Ascend NPU by MindNLP Wizard contributors. +# + +from abc import ABC, abstractmethod +from typing import List, Optional + +import mindspore # pylint: disable=import-error + +from .config import MoEMergeConfig +from ..options import MergeOptions + + +class MoEOutputArchitecture(ABC): + @abstractmethod + def name(self) -> str: + """Return a human-readable name for the architecture.""" + + @abstractmethod + def supports_config( + self, + config: MoEMergeConfig, + explain: bool = False, + trust_remote_code: bool = False, + ) -> bool: + """Return whether this architecture supports the given config. + + If `explain` is True, log an explanation of why the config is not supported.""" + + @abstractmethod + def write_model( # pylint: disable=too-many-positional-arguments + self, + out_path: str, + config: MoEMergeConfig, + merge_options: MergeOptions, + router_weights: List[mindspore.Tensor], + shared_router_weights: Optional[List[mindspore.Tensor]] = None, + ): + """Write the config and tensors for the output MoE to the given path.""" diff --git a/src/mindnlp/wizard/merge/moe/common.py b/src/mindnlp/wizard/merge/moe/common.py new file mode 100644 index 000000000..e19302ad8 --- /dev/null +++ b/src/mindnlp/wizard/merge/moe/common.py @@ -0,0 +1,103 @@ +# Copyright (C) 2025 Arcee AI +# SPDX-License-Identifier: LGPL-3.0-only +# Modified for MindSpore/Ascend NPU by MindNLP Wizard contributors. +# + +import logging +from typing import Dict, Optional, Tuple + +import mindspore # pylint: disable=import-error +from mindspore import ops # pylint: disable=import-error +import tqdm + +from ..architecture import WeightInfo +from ..common import ModelReference, dtype_from_name +from ..io import LazyTensorLoader, TensorWriter +from ..options import MergeOptions +from .config import Expert, MoEMergeConfig + + +def initialize_io( + config: MoEMergeConfig, + out_path: str, + merge_options: MergeOptions, +) -> Tuple[Dict[ModelReference, LazyTensorLoader], LazyTensorLoader, TensorWriter]: + base_model = config.base_model + loaders: Dict[ModelReference, LazyTensorLoader] = {} + for model in tqdm.tqdm( + [base_model] + [e.source_model for e in config.experts], desc="Warm up loaders" + ): + loaders[model] = model.lazy_loader( + cache_dir=merge_options.transformers_cache, + lazy_loader=merge_options.lazy_loader, + ) + + base_loader = loaders.get(base_model) + writer = TensorWriter( + out_path=out_path, + max_shard_size=merge_options.out_shard_size, + output_format=merge_options.output_format, + use_async=merge_options.async_write, + max_write_threads=merge_options.write_threads, + ) + + return loaders, base_loader, writer + + +def select_dtype( + config: MoEMergeConfig, base_cfg +) -> Optional[mindspore.dtype]: + out_dtype = None + if config.dtype: + out_dtype = dtype_from_name(config.dtype) + + if out_dtype is None and hasattr(base_cfg, "torch_dtype") and base_cfg.torch_dtype: + out_dtype = dtype_from_name(str(base_cfg.torch_dtype)) + return out_dtype + + +def noise_and_scale( + tensor: mindspore.Tensor, expert: Expert, is_residual: bool = False +) -> mindspore.Tensor: + if expert.noise_scale is not None: + noise = ops.randn_like(tensor) * expert.noise_scale + tensor = tensor + noise + if is_residual and expert.residual_scale is not None: + tensor = tensor * expert.residual_scale + return tensor + + +def copy_tensor_out( # pylint: disable=too-many-positional-arguments + weight_info: WeightInfo, + loader: LazyTensorLoader, + writer: TensorWriter, + expert: Optional[Expert] = None, + is_residual: bool = False, + output_name: Optional[str] = None, + out_dtype: Optional[mindspore.dtype] = None, + clone: bool = False, +): + out_tensor_name = output_name or weight_info.name + aliases = weight_info.aliases or [] + if not weight_info.optional: + aliases += weight_info.tied_names or [] + try: + tensor = loader.get_tensor( + weight_info.name, + aliases=aliases, + ) + except KeyError: + tensor = None + if tensor is None: + if weight_info.optional: + return + logging.error(f"Missing weight: {weight_info.name} / {out_tensor_name}") + raise KeyError(out_tensor_name) + + if expert: + tensor = noise_and_scale(tensor, expert, is_residual=is_residual) + writer.save_tensor( + out_tensor_name, + tensor.astype(out_dtype) if out_dtype is not None else tensor, + clone=clone, + ) diff --git a/src/mindnlp/wizard/merge/moe/config.py b/src/mindnlp/wizard/merge/moe/config.py new file mode 100644 index 000000000..a51d7dcdf --- /dev/null +++ b/src/mindnlp/wizard/merge/moe/config.py @@ -0,0 +1,87 @@ +# Copyright (C) 2025 Arcee AI +# SPDX-License-Identifier: LGPL-3.0-only +# Modified for MindSpore/Ascend NPU by MindNLP Wizard contributors. +# + +import logging +from typing import List, Optional + +from pydantic import BaseModel + +from ..common import ModelReference + + +class Expert(BaseModel): + """ + Defines a model to be used as a set of layerwise experts in a MoE model. + """ + + source_model: ModelReference + + positive_prompts: Optional[List[str]] = None + negative_prompts: Optional[List[str]] = None + noise_scale: Optional[float] = None + residual_scale: Optional[float] = None + + +class MoEMergeConfig(BaseModel): + """ + Configuration for merging a set of "expert" models into a MoE model. + """ + + base_model: ModelReference + experts: List[Expert] + gate_mode: str = ( + "hidden" # possible values: "hidden", "cheap_embed", "random", "uniform_random" + ) + # "hidden" uses hidden state vectors for the given prompts for each layer + # "cheap_embed" uses the average of token embeddings for the prompts, same for each layer + # "random" is random + # "uniform_random" matches default initialization for mindspore.nn.Dense + dtype: Optional[str] = None + experts_per_token: int = 2 + shared_experts: Optional[List[Expert]] = None + architecture: Optional[str] = None + + +def is_bad_config(config: MoEMergeConfig, allow_all_same: bool = False) -> bool: + if config.experts_per_token < 1: + logging.error("Experts per token must be >= 1") + return True + + if len(config.experts) < config.experts_per_token: + logging.error("Must include at least as many experts as experts_per_token.") + return True + + if config.gate_mode == "random": + return False + + for expert_idx, expert in enumerate(config.experts): + if not expert.positive_prompts: + logging.error(f"Expert {expert_idx} has no positive prompts.") + return True + + def prompt_tup(e: Expert): + return (tuple(e.positive_prompts), tuple(e.negative_prompts or [])) + + p_first = prompt_tup(config.experts[0]) + if all(prompt_tup(e) == p_first for e in config.experts[1:]): + logging.error( + "Your positive and negative prompts are identical for all experts. This will not produce a functioning MoE." + ) + logging.error( + "For each expert, `positive_prompts` must contain one or more example prompt reflecting what should be routed to that expert." + ) + return True + + if not allow_all_same: + if all( + e.source_model == config.experts[0].source_model for e in config.experts[1:] + ): + logging.error( + "All of your expert models are the same. This will produce " + "a model that uses more resources but gives the exact same output. " + "If you plan to train the model after merging, proceed with the " + "--i-understand-this-is-not-useful-without-training flag." + ) + return True diff --git a/src/mindnlp/wizard/merge/moe/deepseek.py b/src/mindnlp/wizard/merge/moe/deepseek.py new file mode 100644 index 000000000..84cd1a9c9 --- /dev/null +++ b/src/mindnlp/wizard/merge/moe/deepseek.py @@ -0,0 +1,188 @@ +# Copyright (C) 2025 Arcee AI +# SPDX-License-Identifier: LGPL-3.0-only +# Modified for MindSpore/Ascend NPU by MindNLP Wizard contributors. +# + +import json +import logging +import os +from typing import Dict, List, Optional + +import mindspore # pylint: disable=import-error +import tqdm + +try: + from mindnlp.transformers import PretrainedConfig +except ImportError: + from transformers import PretrainedConfig + +from ..architecture import arch_info_for_config +from .arch import MoEOutputArchitecture +from .common import copy_tensor_out, initialize_io, select_dtype +from .config import MoEMergeConfig +from ..options import MergeOptions + + +class DeepseekMoE(MoEOutputArchitecture): + def name(self) -> str: + return "DeepSeek MoE" + + def supports_config( + self, + config: MoEMergeConfig, + explain: bool = False, + trust_remote_code: bool = False, + ) -> bool: + if config.shared_experts: + if len(config.shared_experts) > 1: + if explain: + logging.warning( + "DeepSeek MoE merge does not support more than one shared expert" + ) + return False + + if ( + config.shared_experts[0].positive_prompts + or config.shared_experts[0].negative_prompts + ): + if explain: + logging.warning( + "DeepSeek MoE merge does not support gating shared experts" + ) + return False + + model_types = [] + for model_ref in ( + [config.base_model] + + [e.source_model for e in config.experts] + + [e.source_model for e in (config.shared_experts or [])] + ): + model_cfg = model_ref.config(trust_remote_code=trust_remote_code) + model_types.append(model_cfg.model_type) + + if len(set(model_types)) != 1: + if explain: + logging.warning( + "Deepseek MoE requires all input models to have the same architecture" + ) + return False + if model_types[0] not in ("llama", "mistral"): + if explain: + logging.warning( + "Deepseek MoE requires all input models to be Llama or Mistral models" + ) + return False + return True + + def _generate_config( + self, + base_config: PretrainedConfig, + num_experts: int, + shared_experts: Optional[int] = None, + experts_per_token: Optional[int] = None, + ) -> Dict: + if shared_experts and shared_experts > 1: + raise NotImplementedError( + "Shared experts must be 0 or 1 for DeepSeek output" + ) + + res = base_config.to_dict() + res["architectures"] = ["DeepseekForCausalLM"] + res["model_type"] = "deepseek" + res["n_routed_experts"] = num_experts + res["n_shared_experts"] = shared_experts or None + res["num_experts_per_tok"] = experts_per_token or (1 if shared_experts else 2) + res["first_k_dense_replace"] = 0 + res["moe_layer_freq"] = 1 + res["scoring_func"] = "softmax" + res["norm_topk_prob"] = True + res["moe_intermediate_size"] = res["intermediate_size"] + res["auto_map"] = { + "AutoConfig": "deepseek-ai/deepseek-moe-16b-base--configuration_deepseek.DeepseekConfig", + "AutoModel": "deepseek-ai/deepseek-moe-16b-base--modeling_deepseek.DeepseekModel", + "AutoModelForCausalLM": "deepseek-ai/deepseek-moe-16b-base--modeling_deepseek.DeepseekForCausalLM", + } + return res + + def write_model( # pylint: disable=too-many-positional-arguments + self, + out_path: str, + config: MoEMergeConfig, + merge_options: MergeOptions, + router_weights: List[mindspore.Tensor], + shared_router_weights: Optional[List[mindspore.Tensor]] = None, + ): + base_model = config.base_model + base_cfg = base_model.config(trust_remote_code=merge_options.trust_remote_code) + + out_dtype = select_dtype(config, base_cfg) + out_cfg = self._generate_config( + base_cfg, + len(config.experts), + len(config.shared_experts or []), + config.experts_per_token, + ) + if out_dtype is not None: + out_cfg["torch_dtype"] = str(out_dtype) + with open(os.path.join(out_path, "config.json"), "w", encoding="utf-8") as f: + json.dump(out_cfg, f, indent=4) + + shared_def = config.shared_experts[0] if config.shared_experts else None + + loaders, base_loader, writer = initialize_io(config, out_path, merge_options) + shared_loader = loaders.get(shared_def.source_model) if shared_def else None + for weight_info in tqdm.tqdm( + arch_info_for_config(base_cfg).all_weights(base_cfg), + desc="Weights", + ): + tensor_name = weight_info.name + if ".mlp." in tensor_name: + for expert_idx, expert in enumerate(config.experts): + expert_name = tensor_name.replace( + ".mlp.", f".mlp.experts.{expert_idx}." + ) + expert_loader = loaders.get(expert.source_model) + copy_tensor_out( + weight_info, + expert_loader, + writer, + expert=expert, + is_residual="down_proj" in tensor_name, + output_name=expert_name, + out_dtype=out_dtype, + clone=merge_options.clone_tensors, + ) + + if shared_def is not None: + copy_tensor_out( + weight_info, + shared_loader, + writer, + expert=shared_def, + is_residual="down_proj" in tensor_name, + output_name=tensor_name.replace( + ".mlp.", ".mlp.shared_experts." + ), + out_dtype=out_dtype, + clone=merge_options.clone_tensors, + ) + else: + copy_tensor_out( + weight_info, + base_loader, + writer, + out_dtype=out_dtype, + clone=merge_options.clone_tensors, + ) + + for layer_idx, weight in enumerate( + tqdm.tqdm(router_weights, desc="Router weights") + ): + tensor = weight.astype(out_dtype) if out_dtype is not None else weight + writer.save_tensor( + f"model.layers.{layer_idx}.mlp.gate.weight", + tensor, + clone=merge_options.clone_tensors, + ) + + writer.finalize() diff --git a/src/mindnlp/wizard/merge/moe/mixtral.py b/src/mindnlp/wizard/merge/moe/mixtral.py new file mode 100644 index 000000000..008a82484 --- /dev/null +++ b/src/mindnlp/wizard/merge/moe/mixtral.py @@ -0,0 +1,174 @@ +# Copyright (C) 2025 Arcee AI +# SPDX-License-Identifier: LGPL-3.0-only +# Modified for MindSpore/Ascend NPU by MindNLP Wizard contributors. +# + +import logging +from typing import List, Optional + +import mindspore # pylint: disable=import-error +import tqdm + +try: + from mindnlp.transformers import PretrainedConfig, MistralConfig, MixtralConfig +except ImportError: + from transformers import PretrainedConfig, MistralConfig, MixtralConfig + +from ..architecture import WeightInfo +from ..architecture import NAME_TO_ARCH +from ..options import MergeOptions +from .arch import MoEOutputArchitecture +from .common import copy_tensor_out, initialize_io, select_dtype +from .config import MoEMergeConfig + +MISTRAL_INFO = NAME_TO_ARCH["MistralForCausalLM"][0] + + +class MixtralMoE(MoEOutputArchitecture): + def name(self) -> str: + return "Mixtral" + + def supports_config( + self, + config: MoEMergeConfig, + explain: bool = False, + trust_remote_code: bool = False, + ) -> bool: + if config.shared_experts: + if explain: + logging.warning("Mixtral does not support shared experts") + return False + + model_types = [] + for model_ref in [config.base_model] + [e.source_model for e in config.experts]: + model_cfg = model_ref.config(trust_remote_code=trust_remote_code) + model_types.append(model_cfg.model_type) + + if len(set(model_types)) != 1: + if explain: + logging.warning( + "Mixtral requires all input models to have the same architecture" + ) + return False + if model_types[0] not in ("llama", "mistral"): + if explain: + logging.warning( + "Mixtral requires all input models to be Llama or Mistral models" + ) + return False + return True + + def _generate_config( + self, + base_config: PretrainedConfig, + num_experts: int, + shared_experts: Optional[int] = None, + experts_per_token: Optional[int] = None, + ) -> PretrainedConfig: + if shared_experts: + raise NotImplementedError("Shared experts not supported for Mixtral output") + + if not isinstance(base_config, MistralConfig): + base_cfg_mistral = MistralConfig(**base_config.to_dict()) + base_cfg_mistral.sliding_window = None + base_cfg_mistral.max_position_embeddings = ( + base_config.max_position_embeddings + ) + base_config = base_cfg_mistral + + out_cfg = MixtralConfig(**base_config.to_dict()) + out_cfg.architectures = ["MixtralForCausalLM"] + out_cfg.num_local_experts = num_experts + out_cfg.num_experts_per_tok = experts_per_token or 2 + out_cfg.sliding_window = None + + if (out_cfg.num_local_experts & (out_cfg.num_local_experts - 1)) != 0: + logging.warning( + f"Your model has {out_cfg.num_local_experts} experts, which is " + "not a power of two. The model will not be usable in llama.cpp." + ) + return out_cfg + + def _remap_weight_name(self, weight: WeightInfo) -> str: + if ".mlp." not in weight.name: + return weight.name + + res = weight.name + for needle, replacement in [ + (".mlp.gate_proj", ".block_sparse_moe.experts.{expert_idx}.w1"), + (".mlp.down_proj", ".block_sparse_moe.experts.{expert_idx}.w2"), + (".mlp.up_proj", ".block_sparse_moe.experts.{expert_idx}.w3"), + ]: + res = res.replace(needle, replacement) + return res + + def _router_weight_name(self, layer_idx: int) -> str: + return f"model.layers.{layer_idx}.block_sparse_moe.gate.weight" + + def write_model( # pylint: disable=too-many-positional-arguments + self, + out_path: str, + config: MoEMergeConfig, + merge_options: MergeOptions, + router_weights: List[mindspore.Tensor], + shared_router_weights: Optional[List[mindspore.Tensor]] = None, + ): + base_model = config.base_model + base_cfg = base_model.config(trust_remote_code=merge_options.trust_remote_code) + + assert len(router_weights) == base_cfg.num_hidden_layers, ( + f"Expected {base_cfg.num_hidden_layers} router weights, " + f"got {len(router_weights)}" + ) + + out_dtype = select_dtype(config, base_cfg) + out_cfg = self._generate_config( + base_cfg, + len(config.experts), + len(config.shared_experts or []), + config.experts_per_token, + ) + if out_dtype is not None: + out_cfg.torch_dtype = out_dtype + out_cfg.save_pretrained(out_path) + + loaders, base_loader, writer = initialize_io(config, out_path, merge_options) + for weight_info in tqdm.tqdm( + MISTRAL_INFO.all_weights(base_cfg), + desc="Weights", + ): + tensor_name = self._remap_weight_name(weight_info) + if "{expert_idx}" in tensor_name: + for expert_index, expert in enumerate(config.experts): + expert_name = tensor_name.replace("{expert_idx}", str(expert_index)) + expert_loader = loaders.get(expert.source_model) + copy_tensor_out( + weight_info, + expert_loader, + writer, + expert=expert, + out_dtype=out_dtype, + output_name=expert_name, + clone=merge_options.clone_tensors, + is_residual="down_proj" in tensor_name, + ) + else: + copy_tensor_out( + weight_info, + base_loader, + writer, + out_dtype=out_dtype, + clone=merge_options.clone_tensors, + ) + + for layer_idx, weight in enumerate( + tqdm.tqdm(router_weights, desc="Router weights") + ): + tensor = weight.astype(out_dtype) if out_dtype is not None else weight + writer.save_tensor( + self._router_weight_name(layer_idx), + tensor, + clone=merge_options.clone_tensors, + ) + + writer.finalize() diff --git a/src/mindnlp/wizard/merge/moe/qwen.py b/src/mindnlp/wizard/merge/moe/qwen.py new file mode 100644 index 000000000..d662d3554 --- /dev/null +++ b/src/mindnlp/wizard/merge/moe/qwen.py @@ -0,0 +1,204 @@ +# Copyright (C) 2025 Arcee AI +# SPDX-License-Identifier: LGPL-3.0-only +# Modified for MindSpore/Ascend NPU by MindNLP Wizard contributors. +# + +import logging +from typing import List, Optional + +import mindspore # pylint: disable=import-error +from mindspore import ops # pylint: disable=import-error +import tqdm + +try: + from mindnlp.transformers import PretrainedConfig + from mindnlp.transformers.models.qwen2_moe import Qwen2MoeConfig +except ImportError: + from transformers import PretrainedConfig + from transformers.models.qwen2_moe import Qwen2MoeConfig + +from ..architecture import NAME_TO_ARCH +from .arch import MoEOutputArchitecture +from .common import copy_tensor_out, initialize_io, select_dtype +from .config import MoEMergeConfig +from ..options import MergeOptions + +QWEN2_INFO = NAME_TO_ARCH["Qwen2ForCausalLM"][0] + + +class QwenMoE(MoEOutputArchitecture): + def name(self) -> str: + return "Qwen MoE" + + def supports_config( + self, + config: MoEMergeConfig, + explain: bool = False, + trust_remote_code: bool = False, + ) -> bool: + if len(config.shared_experts or []) != 1: + if explain: + logging.warning("Qwen MoE merge requires exactly one shared expert") + return False + + if ( + config.gate_mode != "random" + and not config.shared_experts[0].positive_prompts + ): + if explain: + logging.warning("Qwen MoE requires the shared expert to have prompts") + return False + + model_types = [] + for model_ref in ( + [config.base_model] + + [e.source_model for e in config.experts] + + [e.source_model for e in (config.shared_experts or [])] + ): + model_cfg = model_ref.config(trust_remote_code=trust_remote_code) + model_types.append(model_cfg.model_type) + + if len(set(model_types)) != 1: + if explain: + logging.warning( + "Qwen MoE requires all input models to have the same architecture" + ) + return False + if model_types[0] not in ("llama", "mistral", "qwen2"): + if explain: + logging.warning( + "Qwen MoE requires all input models to be Qwen2, Llama or Mistral models" + ) + return False + return True + + def _generate_config( + self, + base_config: PretrainedConfig, + num_experts: int, + experts_per_token: Optional[int] = None, + ) -> Qwen2MoeConfig: + out_cfg = Qwen2MoeConfig(**base_config.to_dict()) + out_cfg.architectures = ["Qwen2MoeForCausalLM"] + out_cfg.num_experts = num_experts + out_cfg.num_experts_per_tok = experts_per_token or 2 + out_cfg.decoder_sparse_step = 1 + out_cfg.norm_topk_prob = True + out_cfg.sliding_window = None + out_cfg.use_sliding_window = False + out_cfg.shared_expert_intermediate_size = out_cfg.intermediate_size + out_cfg.moe_intermediate_size = out_cfg.intermediate_size + + if (out_cfg.num_experts & (out_cfg.num_experts - 1)) != 0: + logging.warning( + f"Your model has {out_cfg.num_experts} experts, which is " + "not a power of two. The model will not be usable in llama.cpp." + ) + return out_cfg + + def write_model( # pylint: disable=too-many-positional-arguments + self, + out_path: str, + config: MoEMergeConfig, + merge_options: MergeOptions, + router_weights: List[mindspore.Tensor], + shared_router_weights: Optional[List[mindspore.Tensor]] = None, + ): + base_model = config.base_model + base_cfg = base_model.config(trust_remote_code=merge_options.trust_remote_code) + + out_dtype = select_dtype(config, base_cfg) + out_cfg = self._generate_config( + base_cfg, + len(config.experts), + config.experts_per_token, + ) + if out_dtype is not None: + out_cfg.torch_dtype = out_dtype + out_cfg.save_pretrained(out_path) + + shared_def = config.shared_experts[0] + + loaders, base_loader, writer = initialize_io(config, out_path, merge_options) + shared_loader = loaders.get(shared_def.source_model) if shared_def else None + for weight_info in tqdm.tqdm( + QWEN2_INFO.all_weights(base_cfg), + desc="Weights", + ): + tensor_name = weight_info.name + if ".mlp." in tensor_name: + for expert_idx, expert in enumerate(config.experts): + expert_name = tensor_name.replace( + ".mlp.", f".mlp.experts.{expert_idx}." + ) + expert_loader = loaders.get(expert.source_model) + copy_tensor_out( + weight_info, + expert_loader, + writer, + expert=expert, + is_residual="down_proj" in tensor_name, + output_name=expert_name, + out_dtype=out_dtype, + clone=merge_options.clone_tensors, + ) + + copy_tensor_out( + weight_info, + shared_loader, + writer, + expert=shared_def, + is_residual="down_proj" in tensor_name, + output_name=tensor_name.replace(".mlp.", ".mlp.shared_expert."), + out_dtype=out_dtype, + clone=merge_options.clone_tensors, + ) + else: + try: + tensor = base_loader.get_tensor( + tensor_name, aliases=weight_info.aliases + ) + except KeyError: + if tensor_name.endswith("_proj.bias"): + head_dim = out_cfg.hidden_size // out_cfg.num_attention_heads + num_heads = ( + out_cfg.num_key_value_heads + if ( + tensor_name.endswith("k_proj.bias") + or tensor_name.endswith("v_proj.bias") + ) + else out_cfg.num_attention_heads + ) + tensor = ops.zeros(num_heads * head_dim, dtype=out_dtype) + elif weight_info.optional: + continue + else: + raise + + writer.save_tensor( + tensor_name, + tensor.astype(out_dtype) if out_dtype is not None else tensor, + clone=merge_options.clone_tensors, + ) + + for layer_idx, weight in enumerate( + tqdm.tqdm(router_weights, desc="Router weights") + ): + tensor = weight.astype(out_dtype) if out_dtype is not None else weight + writer.save_tensor( + f"model.layers.{layer_idx}.mlp.gate.weight", + tensor, + clone=merge_options.clone_tensors, + ) + shared_tensor = ( + shared_router_weights[layer_idx].astype(out_dtype) + if out_dtype is not None + else shared_router_weights[layer_idx] + ) + writer.save_tensor( + f"model.layers.{layer_idx}.mlp.shared_expert_gate.weight", + shared_tensor, + clone=merge_options.clone_tensors, + ) + + writer.finalize() diff --git a/src/mindnlp/wizard/merge/moe/qwen3.py b/src/mindnlp/wizard/merge/moe/qwen3.py new file mode 100644 index 000000000..b6508a35b --- /dev/null +++ b/src/mindnlp/wizard/merge/moe/qwen3.py @@ -0,0 +1,144 @@ +# Copyright (C) 2025 Arcee AI +# SPDX-License-Identifier: LGPL-3.0-only +# Modified for MindSpore/Ascend NPU by MindNLP Wizard contributors. +# + +import logging +from typing import List, Optional + +import mindspore # pylint: disable=import-error +import tqdm + +try: + from mindnlp.transformers import PretrainedConfig + from mindnlp.transformers.models.qwen3_moe import Qwen3MoeConfig +except ImportError: + from transformers import PretrainedConfig + from transformers.models.qwen3_moe import Qwen3MoeConfig + +from ..architecture import NAME_TO_ARCH +from .arch import MoEOutputArchitecture +from .common import copy_tensor_out, initialize_io, select_dtype +from .config import MoEMergeConfig +from ..options import MergeOptions + +QWEN3_INFO = NAME_TO_ARCH["Qwen3ForCausalLM"][0] + + +class Qwen3MoE(MoEOutputArchitecture): + def name(self) -> str: + return "Qwen3 MoE" + + def supports_config( + self, + config: MoEMergeConfig, + explain: bool = False, + trust_remote_code: bool = False, + ) -> bool: + if len(config.shared_experts or []) != 0: + if explain: + logging.warning("Qwen3 MoE merge does not support shared experts") + return False + + for model_ref in ( + [config.base_model] + + [e.source_model for e in config.experts] + + [e.source_model for e in (config.shared_experts or [])] + ): + model_cfg = model_ref.config(trust_remote_code=trust_remote_code) + if model_cfg.model_type != "qwen3": + if explain: + logging.warning("Qwen3 MoE only supports Qwen3 input models") + return False + return True + + def _generate_config( + self, + base_config: PretrainedConfig, + num_experts: int, + experts_per_token: Optional[int] = None, + ) -> Qwen3MoeConfig: + out_cfg = Qwen3MoeConfig(**base_config.to_dict()) + out_cfg.architectures = ["Qwen3MoeForCausalLM"] + out_cfg.num_experts = num_experts + out_cfg.num_experts_per_tok = experts_per_token or 2 + out_cfg.decoder_sparse_step = 1 + out_cfg.norm_topk_prob = True + out_cfg.moe_intermediate_size = out_cfg.intermediate_size + + if (out_cfg.num_experts & (out_cfg.num_experts - 1)) != 0: + logging.warning( + f"Your model has {out_cfg.num_experts} experts, which is " + "not a power of two. The model will not be usable in llama.cpp." + ) + return out_cfg + + def write_model( # pylint: disable=too-many-positional-arguments + self, + out_path: str, + config: MoEMergeConfig, + merge_options: MergeOptions, + router_weights: List[mindspore.Tensor], + shared_router_weights: Optional[List[mindspore.Tensor]] = None, + ): + base_model = config.base_model + base_cfg = base_model.config(trust_remote_code=merge_options.trust_remote_code) + + out_dtype = select_dtype(config, base_cfg) + out_cfg = self._generate_config( + base_cfg, + len(config.experts), + config.experts_per_token, + ) + if out_dtype is not None: + out_cfg.torch_dtype = out_dtype + out_cfg.save_pretrained(out_path) + + loaders, base_loader, writer = initialize_io(config, out_path, merge_options) + for weight_info in tqdm.tqdm( + QWEN3_INFO.all_weights(base_cfg), + desc="Weights", + ): + tensor_name = weight_info.name + if ".mlp." in tensor_name: + for expert_idx, expert in enumerate(config.experts): + expert_name = tensor_name.replace( + ".mlp.", f".mlp.experts.{expert_idx}." + ) + expert_loader = loaders.get(expert.source_model) + copy_tensor_out( + weight_info, + expert_loader, + writer, + expert=expert, + is_residual="down_proj" in tensor_name, + output_name=expert_name, + out_dtype=out_dtype, + clone=merge_options.clone_tensors, + ) + else: + tensor = base_loader.get_tensor( + tensor_name, + aliases=weight_info.aliases, + raise_on_missing=not weight_info.optional, + ) + if tensor is None: + continue + + writer.save_tensor( + tensor_name, + tensor.astype(out_dtype) if out_dtype is not None else tensor, + clone=merge_options.clone_tensors, + ) + + for layer_idx, weight in enumerate( + tqdm.tqdm(router_weights, desc="Router weights") + ): + tensor = weight.astype(out_dtype) if out_dtype is not None else weight + writer.save_tensor( + f"model.layers.{layer_idx}.mlp.gate.weight", + tensor, + clone=merge_options.clone_tensors, + ) + + writer.finalize() diff --git a/src/mindnlp/wizard/merge/moe/router.py b/src/mindnlp/wizard/merge/moe/router.py new file mode 100644 index 000000000..2f6ff84a6 --- /dev/null +++ b/src/mindnlp/wizard/merge/moe/router.py @@ -0,0 +1,198 @@ +# Copyright (C) 2025 Arcee AI +# SPDX-License-Identifier: LGPL-3.0-only +# Modified for MindSpore/Ascend NPU by MindNLP Wizard contributors. +# +# Key changes: +# - torch.no_grad() → mindspore.ops.stop_gradient / removed (PYNATIVE_MODE has no grad by default) +# - torch.linalg.cond() → manual SVD-based implementation +# - torch.nn.functional.one_hot() → ops.one_hot() +# - torch.randn/rand → ops.randn/rand +# - AutoModelForCausalLM → try mindnlp.transformers first + +import logging +import math +from typing import Dict, List + +import mindspore # pylint: disable=import-error +from mindspore import ops # pylint: disable=import-error +import tqdm + +try: + from mindnlp.transformers import ( + AutoModelForCausalLM, + PreTrainedTokenizerBase, + BatchEncoding, + ) +except ImportError: + from transformers import AutoModelForCausalLM, PreTrainedTokenizerBase, BatchEncoding + +from ..common import ModelReference +from .config import Expert + + +def get_hidden_states( + model, + tokenized: BatchEncoding, + average: bool = True, +) -> List[mindspore.Tensor]: + output = model( + **dict(tokenized.items()), + output_hidden_states=True, + return_dict=True, + ) + hidden_states = ops.stack( + output.hidden_states[:-1] + ) # (num_layers, batch_size, seq_len, hidden_size) + if average: + hidden_states = hidden_states.sum(axis=2) / hidden_states.shape[2] + else: + hidden_states = hidden_states[:, :, -1, :] + return hidden_states.sum(axis=1) / hidden_states.shape[1] + + +def get_cheap_embedding( + embed: mindspore.Tensor, + tokenized: Dict[str, mindspore.Tensor], + num_layers: int, + vocab_size: int, +) -> mindspore.Tensor: + input_ids = tokenized["input_ids"] + onehot = ops.one_hot( + input_ids, vocab_size, mindspore.Tensor(1.0), mindspore.Tensor(0.0) + ) # (batch_size, seq_len, vocab_size) + h = onehot.float() @ embed.float() # (batch_size, seq_len, hidden_size) + embedded = ( + (h * tokenized["attention_mask"].unsqueeze(-1)) + .sum(axis=1) + .sum(axis=0, keepdims=True) + ) # (1, hidden_size) + norm_val = ops.norm(embedded, dim=-1, keepdim=True).clamp(min=1e-8) + res = embedded / norm_val # (1, hidden_size) + return res.repeat(num_layers, 1) + + +def tokenize_prompts( + prompts: List[str], tokenizer: PreTrainedTokenizerBase +): + return tokenizer( + [(tokenizer.bos_token or "") + p for p in prompts], + return_tensors="ms", + padding=True, + add_special_tokens=False, + ) + + +def get_gate_params( # pylint: disable=too-many-positional-arguments + model_ref: ModelReference, + tokenizer: PreTrainedTokenizerBase, + experts: List[Expert], + mode: str = "hidden", + load_in_4bit: bool = False, + load_in_8bit: bool = False, + lazy_loader: bool = False, + trust_remote_code: bool = False, + device: str = "auto", +): + gate_vecs = [] + _do_it = None + + model_cfg = model_ref.config(trust_remote_code=trust_remote_code) + + if mode == "random": + return ops.randn( + (model_cfg.num_hidden_layers, len(experts), model_cfg.hidden_size) + ) + elif mode == "uniform_random": + in_features = model_cfg.hidden_size + scale = math.sqrt(1.0 / in_features) + return ( + ops.rand( + (model_cfg.num_hidden_layers, len(experts), model_cfg.hidden_size) + ) + * 2 + * scale + - scale + ) + elif mode == "cheap_embed": + embed = model_ref.lazy_loader(lazy_loader=lazy_loader).get_tensor( + "model.embed_tokens.weight" + ) + + def _do_it(tokenized): # pylint: disable=function-redefined + return get_cheap_embedding( + embed, + tokenized, + num_layers=model_cfg.num_hidden_layers, + vocab_size=model_cfg.vocab_size, + ) + + elif mode in ("hidden", "hidden_avg", "hidden_last"): + model = AutoModelForCausalLM.from_pretrained( + model_ref.model.path, + revision=model_ref.model.revision, + ms_dtype=mindspore.bfloat16, + device_map=device, + low_cpu_mem_usage=True, + load_in_4bit=load_in_4bit, + load_in_8bit=load_in_8bit, + trust_remote_code=trust_remote_code, + ) + + def _do_it(tokenized): # pylint: disable=function-redefined + return get_hidden_states( + model, tokenized=tokenized, average=mode == "hidden_avg" + ) + + gate_vecs = [] + for expert in tqdm.tqdm(experts, desc="expert prompts"): + hidden_states = _do_it(tokenize_prompts(expert.positive_prompts, tokenizer)) + if expert.negative_prompts: + hidden_states -= _do_it( + tokenize_prompts(expert.negative_prompts, tokenizer) + ) + + norm_val = ops.norm(hidden_states, ord=2, dim=-1, keepdim=True).clamp(min=1e-8) + hidden_states = hidden_states / norm_val + gate_vecs.append(hidden_states) + gate_vecs = ops.stack(gate_vecs, axis=0) # (num_expert, num_layer, hidden_size) + return gate_vecs.permute(1, 0, 2) + + +def _cond_via_svd(matrix: mindspore.Tensor) -> mindspore.Tensor: + """Compute the condition number of a matrix via SVD (replaces torch.linalg.cond).""" + _, svd_values, _ = ops.svd(matrix) + return svd_values.max() / svd_values.min() + + +def warn_degenerate_gates(gate_vecs: mindspore.Tensor, threshold: float = 5.0): + degen_indices = [] + num_layers, _num_experts, _hidden_size = gate_vecs.shape + for idx in range(num_layers): + c = _cond_via_svd(gate_vecs[idx, :, :].float()) + if c > threshold: + degen_indices.append(idx) + + if degen_indices: + if len(degen_indices) == 1: + layer_str = f"layer {degen_indices[0]}" + verb = "has" + elif len(degen_indices) == 2: + layer_str = f"layers {' and '.join(map(str, degen_indices))}" + verb = "have" + elif len(degen_indices) >= num_layers: + layer_str = "ALL layers" + verb = "have" + else: + layer_str = ( + "layers " + + ", ".join(map(str, degen_indices[:-1])) + + ", and " + + str(degen_indices[-1]) + ) + verb = "have" + + logging.warning( + f"{layer_str} {verb} degenerate routing parameters " + "- your prompts may be too similar." + ) + logging.warning("One or more experts will be underutilized in your model.") diff --git a/src/mindnlp/wizard/merge/multigpu_executor.py b/src/mindnlp/wizard/merge/multigpu_executor.py new file mode 100644 index 000000000..98e2fccd7 --- /dev/null +++ b/src/mindnlp/wizard/merge/multigpu_executor.py @@ -0,0 +1,585 @@ +# Copyright (C) 2025 Arcee AI +# SPDX-License-Identifier: LGPL-3.0-only +# Modified for MindSpore/Ascend NPU by MindNLP Wizard contributors. +# +# Device strings ("Ascend:0", "CPU") replace torch.device objects. + +""" +Implementation of multi-device parallel task execution. + +Handles distribution of parallelizable tasks across multiple devices (NPUs) +while respecting: +- Main-thread-only task requirements +- Task dependency graphs +- Device assignment of connected task components +- Intermediate result storage locations +""" + +import concurrent.futures +import contextlib +import logging +import os +import queue +import resource +import threading +from collections import Counter, defaultdict +from typing import Any, Dict, Iterator, List, Optional, Set, Tuple + +import networkx as nx +import tqdm + +from .common import ( + get_accelerator_count, + get_accelerator_type, +) +from .graph import ( + Executor, + Task, + TaskHandle, + TaskUniverse, + build_schedule, +) + +LOG = logging.getLogger(__name__) + + +# --------------------------------------------------------------------------- +# Stream helpers — thin wrappers around mindspore.hal.Stream when available +# --------------------------------------------------------------------------- + +def _make_stream(device: str): + """Create a device stream if the runtime supports it, else return *None*.""" + try: + import mindspore + return mindspore.hal.Stream(device=device) + except Exception as exc: + LOG.debug( + "Could not create stream for %s (%s: %s); falling back to default stream", + device, + type(exc).__name__, + exc, + ) + return None + + +@contextlib.contextmanager +def _stream_context(stream): + """Context manager that activates *stream* if it is not ``None``.""" + if stream is not None: + try: + import mindspore + with mindspore.hal.StreamCtx(stream): + yield + return + except Exception as exc: + LOG.debug( + "Could not enter stream context (%s: %s); using default context", + type(exc).__name__, + exc, + ) + yield + + +def _synchronize_device(device: str): + """Block until all operations on *device* have finished.""" + try: + import mindspore + mindspore.hal.synchronize(device) + except Exception as exc: + LOG.debug( + "Device synchronization failed for %s (%s: %s)", + device, + type(exc).__name__, + exc, + ) + + +# --------------------------------------------------------------------------- +# Multi-device executor +# --------------------------------------------------------------------------- + +class MultiDeviceExecutor: + """ + Execute computational tasks in parallel across multiple devices (NPUs). + + This class analyzes the dependency structure of a task graph and distributes + the workload across available devices while respecting: + 1. Tasks requiring main thread execution + 2. Tasks that need to be duplicated on each device + 3. Task dependencies and data locality + 4. Memory management for intermediate results + + It automatically partitions the task graph into leading tasks (main thread, + pre-parallel), parallel tasks (distributed across devices), and trailing + tasks (main thread, post-parallel). + + Attributes: + num_devices: Number of devices to utilize (None = all available) + storage_device: Device string for storing tensors between stages + targets: Final output tasks to retain results for + """ + + def __init__( + self, + targets: List[Task], + num_devices: Optional[int] = None, + storage_device: Optional[str] = None, + ): + """ + Initialize the executor with a list of target tasks. + + This performs initial task graph analysis, including: + - Finding tasks that must run on the main thread before parallel execution + - Finding tasks that must run on the main thread after parallel execution + - Partitioning parallel tasks into islands that can run independently + - Assigning islands to devices using a load-balancing approach + + Args: + targets: List of final target tasks to execute + num_devices: Number of devices to utilize (None = all available) + storage_device: Device string for storing intermediate results + between execution stages (e.g. ``"CPU"``, + ``"Ascend:0"``) + """ + self.results: Dict[TaskHandle, Any] = {} + self.storage_device = storage_device + self._metric_lock = threading.Lock() + self._task_metrics: List[Dict[str, Any]] = [] + self._queue_depth_samples: List[int] = [] + self._assignment_metrics: List[Dict[str, Any]] = [] + + self.accelerator_type = get_accelerator_type() + if num_devices is None: + num_devices = get_accelerator_count() + LOG.info( + "Using %d %s device(s) for parallel execution", + num_devices, self.accelerator_type, + ) + + self.universe = TaskUniverse(targets) + self.targets = {self.universe.get_handle(t) for t in targets} + self.serial_schedule = build_schedule( + list(self.targets), + {}, + ) + ordered_handles = self.serial_schedule.tasks + + self.per_device_tasks = { + t for t in ordered_handles if t.task().duplicate_per_gpu() + } + leading_tasks = self._find_leading_tasks(ordered_handles) + trailing_tasks = self._find_trailing_tasks(ordered_handles) + self.trailing_main_handles = [ + t for t in ordered_handles if t in trailing_tasks + ] + self.leading_main_handles = [ + t for t in ordered_handles if t in leading_tasks + ] + + self.trailing_dependencies: Set[TaskHandle] = set() + for task_handle in self.trailing_main_handles: + self.trailing_dependencies.update(task_handle.arguments().values()) + + parallel_handles = [ + t + for t in ordered_handles + if ( + t not in trailing_tasks + and t not in leading_tasks + and t not in self.per_device_tasks + ) + ] + LOG.info( + "Task breakdown: %d leading, %d duplicated per-device, " + "%d parallel, %d trailing", + len(self.leading_main_handles), + len(self.per_device_tasks), + len(parallel_handles), + len(self.trailing_main_handles), + ) + if any(t.task().main_thread_only() for t in parallel_handles): + raise RuntimeError( + "Main-thread-only tasks must be either leading or trailing" + ) + if any(t.task().main_thread_only() for t in self.per_device_tasks): + raise RuntimeError( + "Tasks can not be both per-device and main-thread-only" + ) + self.device_assignments = self._assign_islands_to_devices( + parallel_handles, num_devices + ) + + self.task_completion_queue: queue.Queue = queue.Queue() + self.done_event = threading.Event() + + # ------------------------------------------------------------------ + # Public API + # ------------------------------------------------------------------ + + def run(self, quiet: bool = False) -> Iterator[Tuple[Task, Any]]: + """ + Execute all tasks and yield target results. + + Yields: + Iterator[Tuple[Task, Any]]: Task and result pairs + """ + with tqdm.tqdm( + total=len(self.serial_schedule.tasks), + disable=quiet, + desc="Executing graph", + ) as pbar: + # Phase 1: leading (main-thread, pre-parallel) tasks + if self.leading_main_handles: + exec = Executor( + self.leading_main_handles, + math_device=self.storage_device or "CPU", + storage_device=self.storage_device or "CPU", + ) + for task_handle, result in exec._run(quiet=True): + pbar.update() + self.results[task_handle] = result + with self._metric_lock: + self._task_metrics.extend(exec.metrics_snapshot()["tasks"]) + + results_snapshot = dict(self.results) + + def update_progress(): + while not self.done_event.is_set(): + try: + task_idx, result = self.task_completion_queue.get( + timeout=0.1 + ) + task_handle = TaskHandle(self.universe, task_idx) + self.results[task_handle] = result + pbar.update() + except queue.Empty: + continue + + progress_thread = threading.Thread(target=update_progress) + progress_thread.start() + + # Phase 2: parallel tasks distributed across devices + with concurrent.futures.ThreadPoolExecutor() as executor: + futures = [] + for device, island_task_handles in ( + self.device_assignments.items() + ): + futures.append( + executor.submit( + self._device_worker, + task_list=( + list(self.per_device_tasks) + + island_task_handles + ), + cached_values=results_snapshot, + device=device, + quiet=True, + ) + ) + + for future in concurrent.futures.as_completed(futures): + if ex := future.exception(): + self.done_event.set() + executor.shutdown(wait=False) + raise ex + + self.done_event.set() + progress_thread.join() + + # Phase 3: trailing (main-thread, post-parallel) tasks + if self.trailing_main_handles: + exec = Executor( + self.trailing_main_handles, + math_device=self.storage_device or "CPU", + storage_device=self.storage_device or "CPU", + cached_values=dict(self.results), + ) + for task_handle, result in exec._run(quiet=True): + pbar.update() + if task_handle in self.targets: + self.results[task_handle] = result + with self._metric_lock: + self._task_metrics.extend(exec.metrics_snapshot()["tasks"]) + + for task_handle, result in self.results.items(): + if task_handle in self.targets: + yield task_handle.task(), result + + def execute(self) -> None: + """Execute all tasks and discard results.""" + for _ in self.run(quiet=False): + pass + + # ------------------------------------------------------------------ + # Task graph analysis helpers + # ------------------------------------------------------------------ + + def _find_trailing_tasks( + self, tasks: List[TaskHandle] + ) -> Set[TaskHandle]: + """ + Identify tasks that must execute AFTER parallel device tasks complete. + + A task is considered "trailing" if: + - It requires main thread execution (task.main_thread_only() is True) + - All tasks dependent on it are also trailing tasks (recursive condition) + - OR it has no dependents (terminal task) + + Args: + tasks: List of task handles to analyze + + Returns: + Set[TaskHandle]: Set of tasks that should be executed after + parallel processing + """ + dependants: Dict[TaskHandle, Set[TaskHandle]] = defaultdict(set) + for task_idx, arg_indices in self.universe.task_arguments.items(): + for dep_idx in arg_indices.values(): + dependants[TaskHandle(self.universe, dep_idx)].add( + TaskHandle(self.universe, task_idx) + ) + + trailing_tasks: Set[TaskHandle] = set() + to_explore = {t for t in tasks if not dependants[t]} + while to_explore: + task_handle = to_explore.pop() + task = task_handle.task() + if not task.main_thread_only(): + continue + if all(d in trailing_tasks for d in dependants[task_handle]): + trailing_tasks.add(task_handle) + to_explore.update(task_handle.arguments().values()) + return trailing_tasks + + def _find_leading_tasks( + self, tasks: List[TaskHandle] + ) -> Set[TaskHandle]: + """ + Identify tasks that must execute BEFORE parallel device tasks. + + A task is considered "leading" if: + - It requires main thread execution (task.main_thread_only() is True) + - It has no dependencies, or all its dependencies are also leading tasks + + Args: + tasks: List of task handles to analyze + + Returns: + Set[TaskHandle]: Set of tasks that should be executed before + parallel processing + """ + leading_tasks: Set[TaskHandle] = set() + for task_handle in tasks: + task = task_handle.task() + if not task.main_thread_only(): + continue + args = task_handle.arguments() + if args and any( + dep not in leading_tasks for dep in args.values() + ): + continue + leading_tasks.add(task_handle) + return leading_tasks + + # ------------------------------------------------------------------ + # Island assignment & device worker + # ------------------------------------------------------------------ + + def _assign_islands_to_devices( + self, tasks: List[TaskHandle], num_devices: int + ) -> Dict[str, List[TaskHandle]]: + """ + Assign task islands to devices for parallel execution. + + This method partitions the parallel task graph into independent + subgraphs (islands) that can be executed independently on different + devices. It uses a load-balancing approach to distribute islands + across available devices. + + Task islands are identified as weakly connected components in the + task dependency graph, meaning groups of tasks that are connected + through dependencies but don't have dependencies outside their group. + + Args: + tasks: List of parallel tasks to assign to devices + num_devices: Number of available devices + + Returns: + Dict[str, List[TaskHandle]]: Mapping from device strings + (e.g. ``"Ascend:0"``) to assigned tasks + """ + task_set = set(tasks) + + edge_list = [] + for task_handle in tasks: + for dep_handle in task_handle.arguments().values(): + if dep_handle in task_set: + edge_list.append( + (dep_handle._index, task_handle._index) + ) + + island_graph = nx.DiGraph() + island_graph.add_nodes_from([t._index for t in tasks]) + island_graph.add_edges_from(edge_list) + islands: List[Set[int]] = list( + nx.weakly_connected_components(island_graph) + ) + LOG.info("Found %d islands in parallel task graph", len(islands)) + + assignments: Dict[str, List[TaskHandle]] = {} + assignment_metrics: List[Dict[str, Any]] = [] + island_items: List[Tuple[int, int, str, List[TaskHandle]]] = [] + for island in islands: + if not island: + continue + island_tasks = [ + TaskHandle(self.universe, idx) for idx in island + ] + key_hist = Counter( + self._task_locality_key(t) + for t in island_tasks + if self._task_locality_key(t) is not None + ) + dominant_key = "" + dominant_count = 0 + if key_hist: + dominant_key, dominant_count = key_hist.most_common(1)[0] + island_items.append( + ( + len(island_tasks), + dominant_count, + dominant_key, + island_tasks, + ) + ) + + # Large islands first, then strong-locality islands first. + island_items.sort( + key=lambda item: (item[0], item[1]), + reverse=True, + ) + + device_load: Dict[str, int] = defaultdict(int) + device_locality: Dict[str, Counter] = defaultdict(Counter) + for _size, _dom_ct, dominant_key, island_tasks in island_items: + def _assign_key(i: int, _dk=dominant_key): + device = f"{self.accelerator_type}:{i}" + locality_hit = ( + device_locality[device].get(_dk, 0) + if _dk + else 0 + ) + # Prefer locality first, then load balance. + return ( + 0 if locality_hit > 0 else 1, + device_load[device], + -locality_hit, + i, + ) + + device_idx = min(range(num_devices), key=_assign_key) + device = f"{self.accelerator_type}:{device_idx}" + assignments[device] = ( + assignments.get(device, []) + island_tasks + ) + device_load[device] += len(island_tasks) + for th in island_tasks: + key = self._task_locality_key(th) + if key: + device_locality[device][key] += 1 + assignment_metrics.append( + { + "device": device, + "task_count": len(island_tasks), + "dominant_locality_key": dominant_key or None, + } + ) + self._assignment_metrics = assignment_metrics + return assignments + + def _task_locality_key(self, task_handle: TaskHandle) -> Optional[str]: + label = task_handle.task().group_label() + if not label: + return None + label = str(label).lower() + if "::" in label: + # Prefer explicit file/shard locality emitted by loader tasks. + label = label.split("::", 1)[1] + # Collapse shard-style names to a stable prefix for locality grouping. + if "model-" in label and "-of-" in label: + return label.split("-of-")[0] + return label + + def _device_worker( + self, + task_list: List[TaskHandle], + cached_values: Dict[TaskHandle, Any], + device: str, + quiet: bool, + ): + """ + Execute a set of tasks on a single device. + + This method runs as a thread worker for a specific device. It + creates an execution stream on the assigned device, runs the tasks, + and queues results back to the main thread. Only results needed for + target tasks or trailing tasks are retained. + + Args: + task_list: List of tasks to execute on this device + cached_values: Values of previously-executed dependent tasks + device: Device string (e.g. ``"Ascend:0"``) + quiet: Whether to suppress progress bar output + """ + LOG.debug("Device %s starting", device) + stream = _make_stream(device) + with _stream_context(stream): + exec = Executor( + targets=task_list, + math_device=device, + storage_device=self.storage_device or device, + cached_values=cached_values, + ) + count = 0 + for task_handle, result in exec._run(quiet=quiet): + count += 1 + if not ( + task_handle in self.targets + or task_handle in self.trailing_dependencies + ): + result = None + queue_depth = self.task_completion_queue.qsize() + with self._metric_lock: + self._queue_depth_samples.append(int(queue_depth)) + self.task_completion_queue.put((task_handle._index, result)) + with self._metric_lock: + self._task_metrics.extend(exec.metrics_snapshot()["tasks"]) + _synchronize_device(device) + LOG.debug("Device %s done", device) + + def metrics_snapshot(self) -> Dict[str, Any]: + rss_mb = 0.0 + try: + rss_kb = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss + rss_mb = float(rss_kb) / 1024.0 + except Exception: + pass + with self._metric_lock: + tasks = list(self._task_metrics) + qdepth = list(self._queue_depth_samples) + assignment = list(self._assignment_metrics) + return { + "executor": "multi_device", + "pid": os.getpid(), + "task_count": len(tasks), + "tasks": tasks, + "queue_depth_samples": qdepth, + "island_assignment": assignment, + "backpressure_trigger_count": 0, + "rss_peak_mb": rss_mb, + "npu_used_peak_mb": None, + } + + +# Backward-compatible alias +MultiGPUExecutor = MultiDeviceExecutor diff --git a/src/mindnlp/wizard/merge/options.py b/src/mindnlp/wizard/merge/options.py new file mode 100644 index 000000000..45f8b04da --- /dev/null +++ b/src/mindnlp/wizard/merge/options.py @@ -0,0 +1,309 @@ +# Originally from MergeKit (https://github.com/arcee-ai/mergekit) +# Copyright (C) 2025 Arcee AI +# SPDX-License-Identifier: LGPL-3.0-only +# Modified for MindSpore/Ascend NPU by MindNLP Wizard contributors. + +"""CLI option handling for merge commands.""" + +from __future__ import annotations + +import functools +import logging +import warnings +import random +from dataclasses import dataclass +from typing import Any, Callable, Dict, Optional + +import click + +LOG = logging.getLogger(__name__) + +VALID_OUTPUT_FORMATS = ("safetensors", "ckpt") + + +def _parse_device_spec(device: str) -> tuple[str, Optional[int]]: + raw = str(device).strip() + if ":" in raw: + target, idx = raw.split(":", 1) + try: + return target.strip(), int(idx) + except ValueError: + return target.strip(), None + return raw, None + + +def _to_mindspore_target(device: str) -> str: + key = device.strip().lower() + mapping = { + "cpu": "CPU", + "ascend": "Ascend", + "npu": "Ascend", + "gpu": "GPU", + "cuda": "GPU", + } + return mapping.get(key, device) + + +def _parse_size(value: Any) -> int: + if isinstance(value, int): + return value + text = str(value).strip().upper() + if text.isdigit(): + return int(text) + units = { + "KB": 10**3, + "MB": 10**6, + "GB": 10**9, + "TB": 10**12, + "K": 10**3, + "M": 10**6, + "G": 10**9, + "T": 10**12, + "B": 1, + } + for suffix, mul in units.items(): + if text.endswith(suffix): + num = text[: -len(suffix)].strip() + return int(float(num) * mul) + return int(float(text)) + + +@dataclass(frozen=True) +class MergeOptions: + device: str = "CPU" + strict_device_detect: bool = False + allow_crimes: bool = False + transformers_cache: Optional[str] = None + lora_merge_cache: Optional[str] = None + lora_merge_dtype: Optional[str] = None + lazy_loader: bool = False + trust_remote_code: bool = False + random_seed: Optional[int] = None + quiet: bool = False + verbosity: int = 0 + + multi_npu: bool = False + low_cpu_memory: bool = False + read_to_npu: bool = False + + out_shard_size: int = 5 * 10**9 + output_format: str = "safetensors" + safe_serialization: bool = True + async_write: bool = False + write_threads: int = 1 + clone_tensors: bool = False + max_tensor_mem_gb: Optional[float] = None + split_pieces: int = 1 + ckpt_load_kwargs: Optional[Dict[str, Any]] = None + + copy_tokenizer: bool = True + write_model_card: bool = True + + def __post_init__(self) -> None: + if self.max_tensor_mem_gb is not None and self.max_tensor_mem_gb <= 0: + raise ValueError("max_tensor_mem_gb must be > 0") + if self.split_pieces < 1: + raise ValueError("split_pieces must be >= 1") + + resolved_format = self.output_format + if resolved_format == "safetensors" and not self.safe_serialization: + warnings.warn( + "--no-safe-serialization is deprecated. " + "Use --output-format ckpt instead.", + DeprecationWarning, + stacklevel=2, + ) + resolved_format = "bin" + + if resolved_format == "bin": + raise ValueError( + "Output format 'bin' (PyTorch pickle) is not supported by " + "MindSpore Wizard. Use --output-format safetensors (default) " + "or --output-format ckpt instead." + ) + + if resolved_format not in VALID_OUTPUT_FORMATS: + raise ValueError( + f"Invalid output_format '{resolved_format}'. " + f"Must be one of: {', '.join(VALID_OUTPUT_FORMATS)}" + ) + object.__setattr__(self, "output_format", resolved_format) + + if str(self.device).lower() != "auto": + return + + from . import common as common_mod + + try: + detected = common_mod.get_accelerator_type() + object.__setattr__(self, "device", detected) + except Exception as exc: + msg = f"Automatic device detection failed ({type(exc).__name__}: {exc})" + if self.strict_device_detect: + raise RuntimeError(msg) from exc + LOG.warning("%s; falling back to CPU", msg) + object.__setattr__(self, "device", "CPU") + + def apply_global_options(self) -> None: + level = logging.WARNING + if self.quiet: + level = logging.ERROR + elif self.verbosity >= 2: + level = logging.DEBUG + elif self.verbosity == 1: + level = logging.INFO + logging.getLogger().setLevel(level) + + if self.random_seed is not None: + random.seed(self.random_seed) + try: + import numpy as np + + np.random.seed(self.random_seed) + except Exception: + pass + try: + import mindspore + + mindspore.set_seed(self.random_seed) + except Exception: + pass + + try: + import mindspore + + target, device_id = _parse_device_spec(self.device) + ms_target = _to_mindspore_target(target) + mindspore.set_context(device_target=ms_target) + if device_id is not None and ms_target != "CPU": + mindspore.set_context(device_id=device_id) + except Exception as exc: + msg = ( + f"Failed to apply MindSpore device context for {self.device} " + f"({type(exc).__name__}: {exc})" + ) + LOG.warning(msg) + requested_target, _ = _parse_device_spec(self.device) + if _to_mindspore_target(requested_target) != "CPU": + raise RuntimeError( + msg + + ". Requested accelerator is unavailable; aborting instead " + + "of falling back to CPU." + ) from exc + + +class PrettyPrintHelp(click.Command): + """Help command wrapper.""" + + +_MERGE_OPTION_NAMES = [ + "device", + "strict_device_detect", + "allow_crimes", + "transformers_cache", + "lora_merge_cache", + "lora_merge_dtype", + "lazy_loader", + "trust_remote_code", + "random_seed", + "quiet", + "verbosity", + "multi_npu", + "low_cpu_memory", + "read_to_npu", + "out_shard_size", + "output_format", + "safe_serialization", + "async_write", + "write_threads", + "clone_tensors", + "max_tensor_mem_gb", + "split_pieces", + "copy_tokenizer", + "write_model_card", +] + + +def _lazy_loader_from_cli(ctx, param, value): + """Accept both --lazy-loader and deprecated --lazy-unpickle.""" + return value + + +def add_merge_options(func: Callable[..., Any]) -> Callable[..., Any]: + """Attach shared merge CLI options and inject `merge_options`.""" + + @functools.wraps(func) + def wrapper(*args, **kwargs): + option_kwargs: Dict[str, Any] = {} + + # Handle deprecated --lazy-unpickle → lazy_loader + if "lazy_unpickle" in kwargs: + val = kwargs.pop("lazy_unpickle") + if val: + warnings.warn( + "--lazy-unpickle is deprecated. Use --lazy-loader instead.", + DeprecationWarning, + stacklevel=2, + ) + option_kwargs["lazy_loader"] = kwargs.pop("lazy_loader", False) or val + for name in _MERGE_OPTION_NAMES: + if name in kwargs: + option_kwargs[name] = kwargs.pop(name) + option_kwargs["out_shard_size"] = _parse_size(option_kwargs["out_shard_size"]) + kwargs["merge_options"] = MergeOptions(**option_kwargs) + return func(*args, **kwargs) + + opts = [ + click.option("--device", default="CPU", show_default=True), + click.option( + "--strict-device-detect/--no-strict-device-detect", + "strict_device_detect", + default=False, + ), + click.option("--allow-crimes/--no-allow-crimes", "allow_crimes", default=False), + click.option("--transformers-cache", default=None), + click.option("--lora-merge-cache", default=None), + click.option("--lora-merge-dtype", default=None), + click.option("--lazy-loader/--no-lazy-loader", "lazy_loader", default=False), + click.option( + "--lazy-unpickle/--no-lazy-unpickle", + "lazy_unpickle", + default=False, + hidden=True, + ), + click.option("--trust-remote-code/--no-trust-remote-code", "trust_remote_code", default=False), + click.option("--random-seed", type=int, default=None), + click.option("-q", "--quiet/--no-quiet", default=False), + click.option("-v", "--verbose", "verbosity", count=True), + click.option("--multi-npu/--no-multi-npu", "multi_npu", default=False), + click.option("--low-cpu-memory/--no-low-cpu-memory", "low_cpu_memory", default=False), + click.option("--read-to-npu/--no-read-to-npu", "read_to_npu", default=False), + click.option("--out-shard-size", default="5G", show_default=True), + click.option( + "--output-format", + "output_format", + type=click.Choice(VALID_OUTPUT_FORMATS, case_sensitive=False), + default="safetensors", + show_default=True, + help="Tensor output format.", + ), + click.option( + "--safe-serialization/--no-safe-serialization", + "safe_serialization", + default=True, + hidden=True, + ), + click.option("--async-write/--no-async-write", "async_write", default=False), + click.option("--write-threads", type=int, default=1, show_default=True), + click.option("--clone-tensors/--no-clone-tensors", "clone_tensors", default=False), + click.option("--max-tensor-mem-gb", type=float, default=None), + click.option("--split-pieces", type=int, default=1, show_default=True), + click.option("--copy-tokenizer/--no-copy-tokenizer", "copy_tokenizer", default=True), + click.option("--write-model-card/--no-write-model-card", "write_model_card", default=True), + ] + for opt in reversed(opts): + wrapper = opt(wrapper) + return wrapper + + +__all__ = ["MergeOptions", "PrettyPrintHelp", "add_merge_options", "VALID_OUTPUT_FORMATS"] diff --git a/src/mindnlp/wizard/merge/plan.py b/src/mindnlp/wizard/merge/plan.py new file mode 100644 index 000000000..20a9011e7 --- /dev/null +++ b/src/mindnlp/wizard/merge/plan.py @@ -0,0 +1,389 @@ +# Originally from MergeKit (https://github.com/arcee-ai/mergekit) +# Copyright (C) 2025 Arcee AI +# SPDX-License-Identifier: LGPL-3.0-only +# Modified for MindSpore/Ascend NPU by MindNLP Wizard contributors. + +"""Merge planner: module/slice/layer/tensor planning.""" + +import logging +import inspect +from functools import lru_cache +from typing import Any, List, Optional, Tuple + +from . import merge_methods +from .architecture import ( + ConfiguredModuleArchitecture, + ModelArchitecture, + WeightInfo, +) +from .architecture.base import ConfiguredModelArchitecture +from .common import ImmutableMap, ModelReference +from .config import ( + ConfigReader, + InputSliceDefinition, + MergeConfiguration, + OutputModuleDefinition, + OutputSliceDefinition, +) +from .graph import Task +from .io.tasks import ( + FinalizeModel, + GatherTensors, + LoaderCache, + ReturnTensor, + SaveTensor, + TensorWriterTask, +) +from .merge_methods import MergeMethod +from .options import MergeOptions +from .tokenizer import BuildTokenizer, PermutedEmbeddings + + +class MergePlanner: + config: MergeConfiguration + arch_info: ModelArchitecture + options: MergeOptions + out_model_config: Any + _method: MergeMethod + _tensors: List[Tuple[WeightInfo, Task]] + _current_module_layers: int = 0 + _tokenizer_task: Optional[BuildTokenizer] = None + + def __init__( + self, + config: MergeConfiguration, + arch_info: ModelArchitecture, + options: MergeOptions, + out_model_config: Any, + ): + self.config = config + self.arch_info = arch_info + self.options = options + self.out_model_config = out_model_config + self._method = merge_methods.get(config.merge_method) + self._tensors = [] + + token_cfg = {} + tokenizer_source = config.tokenizer_source + if config.tokenizer is not None: + token_cfg = config.tokenizer.tokens or {} + tokenizer_source = config.tokenizer.source + if tokenizer_source is not None: + self._tokenizer_task = BuildTokenizer( + base_model=config.base_model, + referenced_models=tuple(config.referenced_models()), + tokenizer_source=tokenizer_source, + trust_remote_code=options.trust_remote_code, + add_tokens=tuple(token_cfg.keys()), + ) + + def _out_module_arch(self, module: str) -> ConfiguredModuleArchitecture: + module_def = self.arch_info.modules[module] + return ConfiguredModuleArchitecture( + info=module_def.architecture, + config=self.out_model_config, + weight_prefix=module_def.weight_prefix, + ) + + @lru_cache + def _model_arch(self, model: ModelReference): + return ConfiguredModelArchitecture( + info=self.arch_info, + config=model.config(trust_remote_code=self.options.trust_remote_code), + ) + + def normalize_config(self): + base_model = self.config.base_model + + if self.config.models: + self.config.modules = {} + for module_name in self.arch_info.modules: + self.config.modules[module_name] = OutputModuleDefinition( + name=module_name, models=self.config.models + ) + self.config.models = None + + if self.config.slices: + if len(self.arch_info.modules) != 1: + raise RuntimeError( + "Model has multiple modules, must use modules: config syntax " + "to work with slices" + ) + module_name = list(self.arch_info.modules.keys())[0] + self.config.modules = { + module_name: OutputModuleDefinition(slices=self.config.slices) + } + self.config.slices = None + + for module_name in self.config.modules: + module_out = self.config.modules[module_name] + module_arch = self.arch_info.modules[module_name].architecture + + if module_out.models: + slices_in = [] + base_included = False + + for model_in in module_out.models: + if base_model and model_in.model == base_model: + base_included = True + + model_cfg = model_in.model.config( + trust_remote_code=self.options.trust_remote_code + ) + num_layers = module_arch.num_layers(model_cfg) + slices_in.append( + InputSliceDefinition( + layer_range=[0, num_layers], + model=model_in.model, + parameters=model_in.parameters, + ) + ) + + if base_model and not base_included: + logging.info( + "Base model specified but not in input models - adding" + ) + base_cfg = base_model.config( + trust_remote_code=self.options.trust_remote_code + ) + num_layers = module_arch.num_layers(base_cfg) + slices_in.append( + InputSliceDefinition( + layer_range=[0, num_layers], + model=base_model, + ) + ) + + module_out.slices = [OutputSliceDefinition(sources=slices_in)] + module_out.models = None + + def plan_tensor( + self, + weight: WeightInfo, + weights_in: List[WeightInfo], + models: List[ModelReference], + cfg_reader: ConfigReader, + ): + if weight.optional: + any_weight = False + for model, w_in in zip(models, weights_in): + index = LoaderCache().get(model).index + if any( + name in index.tensor_paths + for name in [w_in.name] + (w_in.aliases or []) + ): + any_weight = True + break + + if not any_weight: + logging.info("Skipping optional weight %s", weight.name) + return + + tensor_merge_method = self._method + cfg_g = cfg_reader.for_tensor(weight.name) + global_params = {} + for p in tensor_merge_method.parameters(): + global_params[p.name] = cfg_g.parameter( + p.name, model=None, required=p.required, default=p.default_value + ) + + base_model = cfg_reader.base_model + + tensor_params = {} + for model, weight_in in zip(models, weights_in): + is_base = model == base_model + tensor_params[model] = {} + cfg_m = cfg_reader.for_tensor(weight_in.name) + for p in tensor_merge_method.tensor_parameters(): + tensor_params[model][p.name] = cfg_m.parameter( + p.name, + model=model, + required=p.required and not is_base, + default=p.default_value, + ) + + gather_tensors = GatherTensors( + weight_info=ImmutableMap(data=dict(zip(models, weights_in))), + dtype=self.config.dtype, + device=self.options.device if self.options.read_to_npu else None, + ) + + tensor_input_task = gather_tensors + if self._tokenizer_task and weight.is_embed: + token_cfg = {} + pad_to_multiple = None + if cfg_reader.config.tokenizer: + token_cfg = cfg_reader.config.tokenizer.tokens + pad_to_multiple = cfg_reader.config.tokenizer.pad_to_multiple_of + tensor_input_task = PermutedEmbeddings( + gather_tensors=gather_tensors, + tokenizer_task=self._tokenizer_task, + tokens=token_cfg, + pad_to_multiple_of=pad_to_multiple, + base_model=base_model, + ) + + make_task_kwargs = dict( + output_weight=weight, + tensors=tensor_input_task, + parameters=ImmutableMap(data=global_params), + tensor_parameters=ImmutableMap( + data={ + key: ImmutableMap(data=tensor_params[key]) for key in tensor_params + } + ), + base_model=base_model, + ) + make_task_sig = inspect.signature(tensor_merge_method.make_task) + if "split_pieces" in make_task_sig.parameters: + make_task_kwargs["split_pieces"] = self.options.split_pieces + if "max_tensor_mem_gb" in make_task_sig.parameters: + make_task_kwargs["max_tensor_mem_gb"] = self.options.max_tensor_mem_gb + + tensor_task = tensor_merge_method.make_task(**make_task_kwargs) + self._tensors.append((weight, tensor_task)) + + def plan_layer( # pylint: disable=too-many-positional-arguments + self, + sources: List[InputSliceDefinition], + layer_offset: int, + t: float, + cfg_reader: ConfigReader, + module_name: str, + ): + module_arch = self._out_module_arch(module_name) + weights_out: List[WeightInfo] = module_arch.layer_weights( + index=self._current_module_layers, + ) + weights_in: List[List[WeightInfo]] = [ + self._model_arch(s.model) + .get_module(module_name) + .layer_weights(index=s.layer_range[0] + layer_offset) + for s in sources + ] + + for idx, w_o in enumerate(weights_out): + self.plan_tensor( + weight=w_o, + weights_in=[weights_in[j][idx] for j in range(len(weights_in))], + models=[s.model for s in sources], + cfg_reader=cfg_reader.with_t(t), + ) + + self._current_module_layers += 1 + + def plan_slice( + self, + definition: OutputSliceDefinition, + module_def: OutputModuleDefinition, + module_name: str, + ): + slice_lengths = [ + s.layer_range[1] - s.layer_range[0] for s in definition.sources + ] + if not all(s == slice_lengths[0] for s in slice_lengths): + raise RuntimeError( + "All inputs to a slice must contain the same number of layers" + ) + num_layers = slice_lengths[0] + + cfg_reader = ConfigReader( + config=self.config, slice_out=definition, t=0, module=module_def + ) + for idx in range(num_layers): + if num_layers > 1: + t = idx / (num_layers - 1) + else: + t = 1 + + self.plan_layer( + definition.sources, + layer_offset=idx, + t=t, + cfg_reader=cfg_reader, + module_name=module_name, + ) + + def plan_module(self, module_name: str, definition: OutputModuleDefinition): + self._current_module_layers = 0 + + module_arch = self._out_module_arch(module_name) + config_reader = ConfigReader(config=self.config, t=0, module=definition) + + for weight_info in module_arch.pre_weights(): + self.plan_tensor( + weight_info, + [weight_info] * len(definition.slices[0].sources), + [s.model for s in definition.slices[0].sources], + config_reader.for_tensor(tensor_name=weight_info.name).for_out_slice( + definition.slices[0] + ), + ) + + for out_slice in definition.slices: + self.plan_slice( + out_slice, + module_def=definition, + module_name=module_name, + ) + + for weight_info in module_arch.post_weights(): + self.plan_tensor( + weight_info, + [weight_info] * len(definition.slices[0].sources), + [s.model for s in definition.slices[-1].sources], + config_reader.for_tensor(tensor_name=weight_info.name).for_out_slice( + definition.slices[-1] + ), + ) + + def plan_to_disk(self, out_path: str) -> List[Task]: + """Plan the merge to be streamed to disk, returning a list of tasks.""" + self._plan() + + writer_task = TensorWriterTask( + out_path=out_path, + max_shard_size=self.options.out_shard_size, + output_format=self.options.output_format, + use_async=self.options.async_write, + write_threads=self.options.write_threads, + ) + save_tasks = [] + for weight, tensor_task in self._tensors: + save_tasks.append( + SaveTensor( + tensor_name=weight.name, + tensor_task=tensor_task, + writer_task=writer_task, + clone=self.options.clone_tensors, + optional=weight.optional, + dtype=weight.force_dtype or self.config.out_dtype or self.config.dtype, + ) + ) + finalize = FinalizeModel( + tensor_save_tasks=tuple(save_tasks), writer_task=writer_task + ) + + res = save_tasks + [finalize] + if self._tokenizer_task: + res.append(self._tokenizer_task) + return res + + def plan_in_memory(self) -> List[ReturnTensor]: + """Plan the merge to be performed in memory.""" + self._plan() + return [ + ReturnTensor( + weight_info=w, + tensor_task=t, + dtype=w.force_dtype or self.config.out_dtype or self.config.dtype, + ) + for w, t in self._tensors + ] + + def _plan(self): + self.normalize_config() + self._tasks = [] + + for module_name in self.config.modules: + self.plan_module(module_name, self.config.modules[module_name]) diff --git a/src/mindnlp/wizard/merge/preflight.py b/src/mindnlp/wizard/merge/preflight.py new file mode 100644 index 000000000..f802bed27 --- /dev/null +++ b/src/mindnlp/wizard/merge/preflight.py @@ -0,0 +1,101 @@ +# Copyright (c) MindNLP Wizard contributors. +# Licensed under the Apache License, Version 2.0. + +"""Preflight checks for merge runtime safety.""" + +from __future__ import annotations + +import logging + + +import mindspore +from mindspore import ops + +from .common import dtype_from_name +from .config import MergeConfiguration +from .dtype_policy import choose_work_dtype +from .options import MergeOptions +from .safe_ops import safe_abs, safe_mul, safe_stack, safe_sum, safe_where + +LOG = logging.getLogger(__name__) + +_HALF_DTYPE_NAMES = {"bfloat16", "float16", "bf16", "fp16", "half"} +_METHODS_REQUIRING_HALF_PRECHECK = { + "task_arithmetic", + "ties", + "dare_ties", + "dare_linear", + "breadcrumbs", + "breadcrumbs_ties", + "della", + "della_linear", + "sce", +} + + +def run_merge_preflight(merge_config: MergeConfiguration, options: MergeOptions) -> None: + """Run quick probes to fail fast on unsafe runtime combinations.""" + merge_method = (merge_config.merge_method or "").strip().lower() + dtype_name = ((merge_config.out_dtype or merge_config.dtype or "")).strip().lower() + if merge_method not in _METHODS_REQUIRING_HALF_PRECHECK: + return + if dtype_name not in _HALF_DTYPE_NAMES: + return + _probe_half_precision_math(merge_method=merge_method, dtype_name=dtype_name) + + +def _probe_half_precision_math(*, merge_method: str, dtype_name: str) -> None: + target = "UNKNOWN" + try: + target = str(mindspore.get_context("device_target")) + except Exception: + pass + + test_dtype = dtype_from_name(dtype_name) or mindspore.float16 + work_dtype = choose_work_dtype(test_dtype) + try: + a = mindspore.Tensor([1.0, -2.0, 3.0], dtype=test_dtype) + b = mindspore.Tensor([2.0, 4.0, -1.0], dtype=test_dtype) + stacked = safe_stack([a, b], axis=0, out_dtype=work_dtype, op_name="preflight.stack") + weights = mindspore.Tensor([[0.5], [0.5]], dtype=work_dtype) + weighted = safe_mul(stacked, weights, out_dtype=work_dtype, op_name="preflight.mul") + merged = safe_sum(weighted, axis=0, out_dtype=work_dtype, op_name="preflight.sum") + + sign = ops.sign(weighted) + mask = ops.equal( + sign, + safe_where( + ops.greater_equal( + safe_sum(sign, axis=0, out_dtype=work_dtype, op_name="preflight.sign_sum"), + mindspore.Tensor(0, dtype=work_dtype), + ), + mindspore.Tensor(1, dtype=work_dtype), + mindspore.Tensor(-1, dtype=work_dtype), + out_dtype=work_dtype, + op_name="preflight.majority", + ), + ) + masked = safe_mul(weighted, mask, out_dtype=work_dtype, op_name="preflight.mask_mul") + divisor = safe_sum(mask.astype(work_dtype), axis=0, out_dtype=work_dtype, op_name="preflight.divisor") + divisor = safe_where( + ops.less(safe_abs(divisor, out_dtype=work_dtype, op_name="preflight.abs"), mindspore.Tensor(1e-8, dtype=work_dtype)), + mindspore.Tensor(1, dtype=work_dtype), + divisor, + out_dtype=work_dtype, + op_name="preflight.divisor_fix", + ) + _ = ops.add(merged, ops.div(safe_sum(masked, axis=0, out_dtype=work_dtype, op_name="preflight.mask_sum"), divisor)) + except Exception as exc: + raise RuntimeError( + "Wizard merge preflight failed for half precision runtime. " + f"merge_method={merge_method}, dtype={dtype_name}, device_target={target}. " + "This indicates an unsafe compute path before actual merge. " + "Please inspect wizard safe_ops / generalized_task_arithmetic path." + ) from exc + + LOG.info( + "Wizard preflight passed: merge_method=%s dtype=%s device_target=%s", + merge_method, + dtype_name, + target, + ) diff --git a/src/mindnlp/wizard/merge/safe_ops.py b/src/mindnlp/wizard/merge/safe_ops.py new file mode 100644 index 000000000..3fa8db8e1 --- /dev/null +++ b/src/mindnlp/wizard/merge/safe_ops.py @@ -0,0 +1,160 @@ +# Copyright (c) MindNLP Wizard contributors. +# Licensed under the Apache License, Version 2.0. + +"""Safe operator wrappers.""" + +from __future__ import annotations + +from typing import Iterable, Optional, Union + +import mindspore +from mindspore import ops + +from .dtype_policy import ( + cast_back, + cast_many_to_work, + cast_to_work, + choose_work_dtype, + warn_safe_path_once, +) + + +def _resolve_work_dtype( + ref_dtype: mindspore.dtype, + *, + op_name: str, + out_dtype: Optional[mindspore.dtype], + device_target: Optional[str], +) -> tuple[mindspore.dtype, mindspore.dtype]: + result_dtype = out_dtype or ref_dtype + work_dtype = choose_work_dtype(result_dtype, device_target=device_target) + warn_safe_path_once( + op_name, + in_dtype=result_dtype, + work_dtype=work_dtype, + device_target=device_target, + ) + return result_dtype, work_dtype + + +def safe_stack( + tensors: Iterable[mindspore.Tensor], + *, + axis: int = 0, + out_dtype: Optional[mindspore.dtype] = None, + device_target: Optional[str] = None, + op_name: str = "stack", +) -> mindspore.Tensor: + tensors = list(tensors) + if not tensors: + raise ValueError("safe_stack expects non-empty tensors") + result_dtype, work_dtype = _resolve_work_dtype( + tensors[0].dtype, + op_name=op_name, + out_dtype=out_dtype, + device_target=device_target, + ) + stacked = ops.stack(cast_many_to_work(tensors, work_dtype), axis=axis) + return cast_back(stacked, result_dtype) + + +def safe_sum( + tensor: mindspore.Tensor, + *, + axis=None, + keepdims: bool = False, + out_dtype: Optional[mindspore.dtype] = None, + device_target: Optional[str] = None, + op_name: str = "sum", +) -> mindspore.Tensor: + result_dtype, work_dtype = _resolve_work_dtype( + tensor.dtype, + op_name=op_name, + out_dtype=out_dtype, + device_target=device_target, + ) + reduced = ops.sum(cast_to_work(tensor, work_dtype), dim=axis, keepdim=keepdims) + return cast_back(reduced, result_dtype) + + +def safe_mul( + lhs: mindspore.Tensor, + rhs: Union[mindspore.Tensor, float, int], + *, + out_dtype: Optional[mindspore.dtype] = None, + device_target: Optional[str] = None, + op_name: str = "mul", +) -> mindspore.Tensor: + if not isinstance(rhs, mindspore.Tensor): + rhs = mindspore.Tensor(rhs, dtype=lhs.dtype) + result_dtype, work_dtype = _resolve_work_dtype( + out_dtype or lhs.dtype, + op_name=op_name, + out_dtype=out_dtype or lhs.dtype, + device_target=device_target, + ) + res = ops.mul(cast_to_work(lhs, work_dtype), cast_to_work(rhs, work_dtype)) + return cast_back(res, result_dtype) + + +def safe_abs( + tensor: mindspore.Tensor, + *, + out_dtype: Optional[mindspore.dtype] = None, + device_target: Optional[str] = None, + op_name: str = "abs", +) -> mindspore.Tensor: + result_dtype, work_dtype = _resolve_work_dtype( + out_dtype or tensor.dtype, + op_name=op_name, + out_dtype=out_dtype or tensor.dtype, + device_target=device_target, + ) + res = ops.abs(cast_to_work(tensor, work_dtype)) + return cast_back(res, result_dtype) + + +def safe_norm( + tensor: mindspore.Tensor, + *, + axis=None, + keepdims: bool = False, + out_dtype: Optional[mindspore.dtype] = None, + device_target: Optional[str] = None, + op_name: str = "norm", +) -> mindspore.Tensor: + result_dtype, work_dtype = _resolve_work_dtype( + out_dtype or tensor.dtype, + op_name=op_name, + out_dtype=out_dtype or tensor.dtype, + device_target=device_target, + ) + res = ops.norm( + cast_to_work(tensor, work_dtype), + dim=axis, + keepdim=keepdims, + ) + return cast_back(res, result_dtype) + + +def safe_where( + condition: mindspore.Tensor, + x: mindspore.Tensor, + y: mindspore.Tensor, + *, + out_dtype: Optional[mindspore.dtype] = None, + device_target: Optional[str] = None, + op_name: str = "where", +) -> mindspore.Tensor: + result_dtype, work_dtype = _resolve_work_dtype( + out_dtype or x.dtype, + op_name=op_name, + out_dtype=out_dtype or x.dtype, + device_target=device_target, + ) + res = ops.where( + condition, + cast_to_work(x, work_dtype), + cast_to_work(y, work_dtype), + ) + return cast_back(res, result_dtype) diff --git a/src/mindnlp/wizard/merge/scripts/__init__.py b/src/mindnlp/wizard/merge/scripts/__init__.py new file mode 100644 index 000000000..6babb0e71 --- /dev/null +++ b/src/mindnlp/wizard/merge/scripts/__init__.py @@ -0,0 +1,4 @@ +# Originally from MergeKit (https://github.com/arcee-ai/mergekit) +# Copyright (C) 2025 Arcee AI +# SPDX-License-Identifier: LGPL-3.0-only +# Modified for MindSpore/Ascend NPU by MindNLP Wizard contributors. diff --git a/src/mindnlp/wizard/merge/scripts/bakllama.py b/src/mindnlp/wizard/merge/scripts/bakllama.py new file mode 100644 index 000000000..5a0ae1901 --- /dev/null +++ b/src/mindnlp/wizard/merge/scripts/bakllama.py @@ -0,0 +1,75 @@ +# Copyright (C) 2025 Arcee AI +# SPDX-License-Identifier: LGPL-3.0-only +# Modified for MindSpore/Ascend NPU by MindNLP Wizard contributors. +# pylint: disable=no-value-for-parameter + +from typing import List, Optional + +import click +import yaml +from pydantic import BaseModel + +from ..config import ( + ConditionalParameter, + InputSliceDefinition, + MergeConfiguration, +) +from ..merge import run_merge +from ..options import MergeOptions + + +class LayerSlice(BaseModel): + model: str + start: int + end: int + scale: Optional[float] = None + + +class BakllamaConfig(BaseModel): + layer_slices: List[LayerSlice] + embedding_source: Optional[str] = None + lm_head_source: Optional[str] = None + + +@click.command("wizard-bakllama") +@click.argument("config_path", type=click.Path(exists=True, dir_okay=False)) +@click.argument("out_path", type=str) +@click.option( + "--clone-tensors/--no-clone-tensors", + type=bool, + is_flag=True, + help="Clone tensors before saving, to allow multiple occurrences of the same layer", + default=False, +) +@click.option("--fp16/--no-fp16", type=bool, default=False) +def main( + config_path: str, + out_path: str, + clone_tensors: bool, + fp16: bool, +): + """Wrapper for using legacy bakllama configuration files.""" + with open(config_path, "r", encoding="utf-8") as file: + config = BakllamaConfig.model_validate(yaml.safe_load(file)) + + slices = [] + for s in config.layer_slices: + parameters = {} + if s.scale is not None: + parameters["scale"] = ConditionalParameter( + value=s.scale, filter="down_proj" + ) + slices.append( + InputSliceDefinition( + model=s.model, layer_range=(s.start, s.end), parameters=parameters + ) + ) + + merge_config = MergeConfiguration( + merge_method="passthrough", slices=slices, dtype="float16" if fp16 else None + ) + run_merge(merge_config, out_path, MergeOptions(clone_tensors=clone_tensors)) + + +if __name__ == "__main__": + main() diff --git a/src/mindnlp/wizard/merge/scripts/evolve.py b/src/mindnlp/wizard/merge/scripts/evolve.py new file mode 100644 index 000000000..7ec078080 --- /dev/null +++ b/src/mindnlp/wizard/merge/scripts/evolve.py @@ -0,0 +1,422 @@ +# Copyright (C) 2025 Arcee AI +# SPDX-License-Identifier: LGPL-3.0-only +# Modified for MindSpore/Ascend NPU by MindNLP Wizard contributors. +# pylint: disable=no-value-for-parameter + +import importlib.util +import logging +import os +import time +from typing import List, Optional + +import click +import cma # pylint: disable=import-error +import mindspore # pylint: disable=import-error +import numpy as np +import pandas +import ray # pylint: disable=import-error +import tqdm +import transformers +import yaml + +try: + import wandb +except ImportError: + wandb = None + + +from ..common import ModelReference +from ..evo.config import ( + EvolMergeConfiguration, + ModelGenomeDefinition, + check_for_naughty_config, +) +from ..evo.genome import ModelGenome +from ..evo.strategy import ( + ActorPoolEvaluationStrategy, + BufferedRayEvaluationStrategy, + SerialEvaluationStrategy, +) +from ..merge import run_merge +from ..options import MergeOptions + + +@click.command("wizard-evolve") +@click.argument("genome-config-path", type=str) +@click.option("--max-fevals", type=int, default=100) +@click.option("--vllm/--no-vllm", is_flag=True, default=False, help="Use vLLM") +@click.option( + "--strategy", + "-s", + type=click.Choice(["pool", "buffered", "serial"]), + default="pool", + help="Evaluation scheduling strategy", +) +@click.option( + "--in-memory/--no-in-memory", + is_flag=True, + default=False, + help="Use in-memory merge & evaluation", +) +@click.option( + "--storage-path", + type=str, + help="Path to storage accessible to all nodes for model storage", + required=True, +) +@click.option("--num-npus", type=int, help="Number of NPUs to use across all nodes") +@click.option("--merge-npu/--no-merge-npu", is_flag=True, default=True) +@click.option("--trust-remote-code/--no-trust-remote-code", is_flag=True, default=False) +@click.option("--allow-crimes/--no-allow-crimes", is_flag=True, default=False) +@click.option("--random-seed", type=int, default=0) +@click.option("--batch-size", type=int, default=None, help="Batch size for evaluation") +@click.option("--sigma0", type=float, default=1 / 6, help="Initial sigma for CMA-ES") +@click.option("use_wandb", "--wandb/--no-wandb", is_flag=True, default=False) +@click.option("--wandb-project", type=str, help="Wandb project name") +@click.option("--wandb-entity", type=str, help="Wandb entity name") +@click.option( + "--task-search-path", + type=str, + multiple=True, + help="Path to search for lmeval tasks", +) +@click.option( + "--i-understand-the-depths-of-the-evils-i-am-unleashing", + "allow_benchmark_tasks", + is_flag=True, + default=False, + help="Allow benchmark tasks as objectives", +) +@click.option( + "--save-final-model/--no-save-final-model", + is_flag=True, + default=True, + help="Save the final merged model", +) +@click.option( + "--reshard/--no-reshard", + is_flag=True, + default=True, + help="Convert models to single-shard safetensors for faster merge", +) +@click.option( + "--timeout", + type=float, + default=None, + help="Maximum time to run the optimization in seconds", +) +@click.option( + "--load-in-8bit", + is_flag=True, + default=False, + help="Evaluate models at 8-bit precision", +) +@click.option( + "--load-in-4bit", + is_flag=True, + default=False, + help="Evaluate models at 4-bit precision", +) +@click.option( + "--force-population-size", + type=int, + default=None, + help="Force a specific initial population size for CMA-ES", +) +def main( # pylint: disable=too-many-positional-arguments + genome_config_path: str, + max_fevals: int, + vllm: bool, + strategy: str, + in_memory: bool, + storage_path: str, + num_npus: Optional[int], + merge_npu: bool, + trust_remote_code: bool, + allow_crimes: bool, + random_seed: int, + batch_size: Optional[int], + sigma0: float, + use_wandb: bool, + wandb_project: Optional[str], + wandb_entity: Optional[str], + task_search_path: List[str], + allow_benchmark_tasks: bool, + save_final_model: bool, + reshard: bool, + timeout: Optional[float], + load_in_8bit: bool, + load_in_4bit: bool, + force_population_size: Optional[int], +): + config = EvolMergeConfiguration.model_validate( + yaml.safe_load(open(genome_config_path, "r", encoding="utf-8")) + ) + + check_for_naughty_config(config, allow=allow_benchmark_tasks) + + if load_in_4bit and load_in_8bit: + raise ValueError("Cannot load models in both 4-bit and 8-bit") + + if load_in_4bit or load_in_8bit: + if vllm: + raise ValueError("Cannot use vLLM with 4-bit or 8-bit models") + if in_memory: + raise ValueError("Cannot use in-memory mode with 4-bit or 8-bit models") + if not importlib.util.find_spec("bitsandbytes"): + raise RuntimeError("bitsandbytes is not installed") + + bnb_config = transformers.BitsAndBytesConfig( + load_in_8bit=load_in_8bit, + load_in_4bit=load_in_4bit, + bnb_4bit_compute_dtype="bfloat16", + bnb_4bit_quant_type="nf4", + bnb_4bit_use_double_quant=True, + ) + else: + bnb_config = None + + if use_wandb: + if not wandb: + raise RuntimeError("wandb is not installed") + run = wandb.init( + project=wandb_project or "wizard-evolve", + entity=wandb_entity, + config=config.model_dump(mode="json"), + ) + else: + run = None + + merge_options = MergeOptions( + transformers_cache=os.path.join(storage_path, "transformers_cache"), + lora_merge_cache=os.path.join(storage_path, "lora_merge_cache"), + multi_npu=merge_npu, + low_cpu_memory=merge_npu and not in_memory, + out_shard_size=1_000_000_000_000, + trust_remote_code=trust_remote_code, + allow_crimes=allow_crimes, + random_seed=random_seed, + quiet=True, + read_to_npu=merge_npu and not in_memory, + copy_tokenizer=True, + safe_serialization=True, + ) + + if reshard: + resharded_models = [] + resharded_base = None + for model in tqdm.tqdm(config.genome.models, desc="Resharding models"): + resharded_models.append( + _reshard_model( + model, + storage_path, + merge_options.lora_merge_cache, + trust_remote_code, + ) + ) + if config.genome.base_model is not None: + resharded_base = _reshard_model( + config.genome.base_model, + storage_path, + merge_options.lora_merge_cache, + trust_remote_code, + ) + else: + resharded_models = config.genome.models + resharded_base = config.genome.base_model + + genome = ModelGenome( + ModelGenomeDefinition.model_validate( + { + **config.genome.model_dump( + exclude=[ + "models", + "base_model", + ] + ), + "models": resharded_models, + "base_model": resharded_base, + } + ), + trust_remote_code=trust_remote_code, + ) + + if strategy == "pool": + strat_cls = ActorPoolEvaluationStrategy + elif strategy == "buffered": + strat_cls = BufferedRayEvaluationStrategy + elif strategy == "serial": + strat_cls = SerialEvaluationStrategy + else: + raise ValueError(f"Unknown strategy {strategy}") + + strat = strat_cls( + config, + genome, + merge_options, + num_gpus=num_npus, + vllm=vllm, + in_memory=in_memory, + model_storage_path=os.path.join(storage_path, "merged"), + batch_size=batch_size, + task_search_path=task_search_path, + quantization_config=bnb_config, + ) + + x0 = genome.initial_genotype(random=config.random_init).view(-1).asnumpy() + xbest = x0 + xbest_cost = np.inf + + def progress_callback(es: cma.CMAEvolutionStrategy): + nonlocal xbest, xbest_cost + + res = es.result + if use_wandb and run is not None: + best_params = genome.genotype_to_param_arrays(res.xbest) + mean_params = genome.genotype_to_param_arrays(res.xfavorite) + run.log( + { + "best_score": -res.fbest, + "best_genome": wandb.Table(data=pandas.DataFrame(best_params)), + "mean_genome": wandb.Table(data=pandas.DataFrame(mean_params)), + "mean_std": genome.genotype_to_param_arrays(res.stds), + "evaluations": res.evaluations, + }, + commit=True, + step=res.evaluations, + ) + + if res.fbest < xbest_cost: + xbest = res.xbest + xbest_cost = res.fbest + print(f"New best score: {-xbest_cost:.4f}") + best_yaml = genome.genotype_merge_config(xbest).to_yaml() + with open(os.path.join(storage_path, "best_config.yaml"), "w") as f: + f.write(best_yaml) + print(f"Merge configuration:\n{best_yaml}") + + if use_wandb: + art = wandb.Artifact("best_config", type="merge_config") + art.add_file(os.path.join(storage_path, "best_config.yaml")) + run.log_artifact(art) + + def parallel_evaluate(x: List[np.ndarray]) -> List[float]: + print(f"Received {len(x)} genotypes") + res = strat.evaluate_genotypes(x) + + if use_wandb: + res = list(res) + score_mean = np.mean([r["score"] for r in res]) + score_std = np.std([r["score"] for r in res]) + run.log( + { + "population/score_mean": score_mean, + "population/score_std": score_std, + }, + commit=False, + ) + for task in res[0]["results"]: + for metric in res[0]["results"][task]: + values = [r["results"][task][metric] for r in res] + values = [v for v in values if v is not None] + if not values or all(isinstance(v, str) for v in values): + continue + + mean = np.mean(values) + max_val = max(values) + min_val = min(values) + + metric_pretty = metric.replace(",none", "") + if metric_pretty.endswith("_stderr"): + continue + + run.log( + { + f"population/{task}_{metric_pretty}_mean": mean, + f"population/{task}_{metric_pretty}_max": max_val, + f"population/{task}_{metric_pretty}_min": min_val, + }, + commit=False, + ) + + return [-x["score"] for x in res] + + try: + cma_opts = {"maxfevals": max_fevals, "timeout": timeout} + if force_population_size is not None: + cma_opts["popsize"] = force_population_size + xbest, es = cma.fmin2( + None, + parallel_objective=parallel_evaluate, + x0=x0, + sigma0=sigma0, + options=cma_opts, + callback=progress_callback, + ) + xbest_cost = es.result.fbest + except KeyboardInterrupt: + ray.shutdown() + + print("!!! OPTIMIZATION COMPLETE !!!") + print(f"Best cost: {xbest_cost:.4f}") + print() + + time.sleep(1.0) + + genome_pretty = ModelGenome(config.genome, trust_remote_code=trust_remote_code) + best_config = genome_pretty.genotype_merge_config(xbest) + print("Best merge configuration:") + print(best_config.to_yaml()) + + if save_final_model: + print("Saving final model...") + run_merge(best_config, os.path.join(storage_path, "final_model"), merge_options) + + +def _reshard_model( + model: ModelReference, + storage_path: str, + merge_cache: Optional[str], + trust_remote_code: bool, +) -> ModelReference: + merged = model.merged( + cache_dir=merge_cache, + trust_remote_code=trust_remote_code, + lora_merge_dtype="bfloat16", + ) + out_path = os.path.join( + storage_path, + "input_models", + merged.model._unique_id(), + ) + + if os.path.exists(out_path): + logging.info(f"Using existing resharded model at {out_path}") + return ModelReference(model=out_path) + + model_hf = transformers.AutoModelForCausalLM.from_pretrained( + merged.model.path, + revision=merged.model.revision, + trust_remote_code=trust_remote_code, + ms_dtype=mindspore.bfloat16, + cache_dir=os.path.join(storage_path, "transformers_cache"), + ) + model_hf.save_pretrained( + out_path, safe_serialization=True, out_shard_size=1_000_000_000_000 + ) + try: + tokenizer = transformers.AutoTokenizer.from_pretrained( + model.model.path, + revision=model.model.revision, + trust_remote_code=trust_remote_code, + use_fast=True, + ) + tokenizer.save_pretrained(out_path) + except Exception as e: + logging.warning(f"Could not save tokenizer for {model.model}", exc_info=e) + + return ModelReference(model=out_path) + + +if __name__ == "__main__": + main() diff --git a/src/mindnlp/wizard/merge/scripts/extract_lora.py b/src/mindnlp/wizard/merge/scripts/extract_lora.py new file mode 100644 index 000000000..272be590c --- /dev/null +++ b/src/mindnlp/wizard/merge/scripts/extract_lora.py @@ -0,0 +1,594 @@ +# Copyright (C) 2025 Arcee AI +# SPDX-License-Identifier: LGPL-3.0-only +# Modified for MindSpore/Ascend NPU by MindNLP Wizard contributors. +# pylint: disable=no-value-for-parameter + +import json +import logging +import os +import re +import sys +from typing import Any, Dict, List, Optional, Tuple + +import click +import mindspore # pylint: disable=import-error +from mindspore import nn as ms_nn # pylint: disable=import-error +from mindspore import ops # pylint: disable=import-error +import tqdm +import transformers +from pydantic import BaseModel + +from ..architecture import WeightInfo, arch_info_for_config +from ..card import generate_card_lora +from ..common import ModelReference, get_auto_cls +from ..graph import Executor, Task +from ..io.tasks import FinalizeModel, LoadTensor, SaveTensor, TensorWriterTask +from ..io.tensor_writer import TensorWriter +from ..multigpu_executor import MultiDeviceExecutor +from ..options import MergeOptions, PrettyPrintHelp, add_merge_options + +LOG = logging.getLogger("extract_lora") + + +def _get_npu_device(): + """Get available Ascend/NPU device string.""" + try: + import mindspore # pylint: disable=reimported,import-error + mindspore.set_context(device_target="Ascend") + return "Ascend" + except Exception as exc: + LOG.debug( + "Ascend context is unavailable for extract_lora (%s: %s); using CPU fallback", + type(exc).__name__, + exc, + ) + return "CPU" + + +def _build_executor(tasks: List[Task], merge_options: MergeOptions, fallback_device: str): + if merge_options.multi_npu: + return MultiDeviceExecutor( + tasks, + storage_device="CPU" if not merge_options.low_cpu_memory else None, + ) + return Executor( + tasks, + math_device=merge_options.device or fallback_device, + storage_device=( + merge_options.device if merge_options.low_cpu_memory else "CPU" + ), + ) + + +@click.command("wizard-extract-lora", cls=PrettyPrintHelp) +@click.option( + "--model", + required=True, + help="Fine-tuned model path", +) +@click.option( + "--base-model", + required=True, + help="Base model path", +) +@click.option( + "--out-path", + required=True, + help="Output path for extracted LoRA adapter", +) +@click.option( + "--max-rank", + type=int, + default=128, + help="Maximum rank for LoRA decomposition", +) +@click.option( + "--distribute-scale/--no-distribute-scale", + is_flag=True, + default=True, + help="Distribute scale between A and B matrices", +) +@click.option( + "--embed-lora/--no-embed-lora", + is_flag=True, + default=False, + help="Extract LoRA weights for embeddings (vs. in modules_to_save)", +) +@click.option( + "--save-module", + "modules_to_save", + type=str, + multiple=True, + default=[], + help="Save the specified module(s) at full rank", +) +@click.option( + "--exclude-regex", + "-e", + "exclude_regexes", + type=str, + multiple=True, + help="Exclude modules matching the specified regex", +) +@click.option( + "--include-regex", + "-i", + "include_regexes", + type=str, + multiple=True, + help="Include modules matching the specified regex", +) +@click.option( + "--sv-epsilon", + type=float, + default=0, + help="Threshold for singular values to discard", + show_default=True, +) +@click.option( + "--skip-undecomposable", + is_flag=True, + help="Skip saving undecomposable modules", + default=False, +) +@add_merge_options +def main( # pylint: disable=too-many-positional-arguments + base_model: str, + model: str, + out_path: str, + max_rank: int, + distribute_scale: bool, + embed_lora: bool, + modules_to_save: List[str], + exclude_regexes: List[str], + include_regexes: List[str], + sv_epsilon: float, + skip_undecomposable: bool, + merge_options: MergeOptions, +): + merge_options.apply_global_options() + + if not modules_to_save: + modules_to_save = [] + + base_model_ref = ModelReference.model_validate(base_model) + model_ref = ModelReference.model_validate(model) + plan_result = plan_extraction( + base_model_ref=base_model_ref.merged( + cache_dir=merge_options.lora_merge_cache, + trust_remote_code=merge_options.trust_remote_code, + lora_merge_dtype=merge_options.lora_merge_dtype, + ), + model_ref=model_ref.merged( + cache_dir=merge_options.lora_merge_cache, + trust_remote_code=merge_options.trust_remote_code, + lora_merge_dtype=merge_options.lora_merge_dtype, + ), + modules_to_save=modules_to_save, + out_path=out_path, + options=merge_options, + max_rank=max_rank, + distribute_scale=distribute_scale, + embed_lora=embed_lora, + exclude_regexes=exclude_regexes, + include_regexes=include_regexes, + sv_epsilon=sv_epsilon, + skip_undecomposable=skip_undecomposable, + ) + + tasks = plan_result.tasks + device = _get_npu_device() + executor = _build_executor(tasks, merge_options, device) + + module_real_ranks = {} + for task, result in executor.run(): + if isinstance(task, TaskVectorDecompositionTask): + module_real_ranks[task.weight_info.name.removesuffix(".weight")] = result[ + 0 + ].shape[0] + + real_max_rank = max(module_real_ranks.values()) + config_dict = make_config_dict( + base_ref=base_model_ref, + max_rank=real_max_rank, + modules_to_save=modules_to_save, + target_modules=list( + set(key.split(".")[-1] for key in module_real_ranks.keys()) + ), + module_ranks=module_real_ranks, + ) + with open(os.path.join(out_path, "adapter_config.json"), "w") as f: + json.dump(config_dict, f, indent=4) + + invocation = " ".join(sys.argv) + with open(os.path.join(out_path, "README.md"), "w", encoding="utf-8") as f: + f.write( + generate_card_lora( + base_model_ref, + model_ref, + invocation, + os.path.basename(out_path), + base_vocab_size=plan_result.base_vocab_size, + final_vocab_size=plan_result.final_vocab_size, + ) + ) + + LOG.info(f"LoRA adapter extracted to {out_path}") + + +def make_config_dict( + base_ref: ModelReference, + max_rank: int, + modules_to_save: List[str], + target_modules: List[str], + module_ranks: Dict[str, int], +): + different_ranked = {k: v for k, v in module_ranks.items() if v != max_rank} + return { + "base_model_name_or_path": base_ref.model.path, + "peft_type": "LORA", + "use_rslora": False, + "target_modules": target_modules, + "modules_to_save": modules_to_save, + "task_type": "CAUSAL_LM", + "r": max_rank, + "lora_alpha": max_rank, + "rank_pattern": different_ranked, + "alpha_pattern": different_ranked, + "lora_dropout": 0.0, + "fan_in_fan_out": False, + "inference_mode": True, + } + + +class TaskVectorDecompositionTask(Task[Tuple[mindspore.Tensor, mindspore.Tensor]]): + weight_info: WeightInfo + input_task: Task + max_rank: int + distribute_scale: bool = True + transpose: bool = False + sv_epsilon: float = 0 + + def arguments(self) -> Dict[str, Any]: + return {"task_vector": self.input_task} + + def execute(self, task_vector: mindspore.Tensor) -> Tuple[mindspore.Tensor, mindspore.Tensor]: + if self.transpose: + task_vector = task_vector.T + out_dtype = task_vector.dtype + u, s, vh = ops.svd( + task_vector.astype(mindspore.float32), full_matrices=False + ) + rank = min(self.max_rank, s.shape[0]) + if self.sv_epsilon > 0: + rank = min(int((s > self.sv_epsilon).sum()), rank) + if self.distribute_scale: + sqrt_s = ops.diag(ops.sqrt(s[:rank])) + scale_a = sqrt_s + scale_b = sqrt_s + else: + scale_a = ops.diag(s[:rank]) + scale_b = ops.eye(rank, rank, dtype=mindspore.float32) + sqrt_s = ops.diag(ops.sqrt(s[:rank])) + weight_a = ops.matmul(scale_a, vh[:rank]) + weight_b = ops.matmul(u[:, :rank], scale_b) + + return weight_a.astype(out_dtype), weight_b.astype(out_dtype) + + def group_label(self) -> Optional[str]: + return self.input_task.group_label() + + def uses_accelerator(self): + return True + + +class TaskVectorTask(Task[mindspore.Tensor]): + base_tensor: Task + model_tensor: Task + + def arguments(self) -> Dict[str, Any]: + return {"base": self.base_tensor, "model": self.model_tensor} + + def execute(self, base: mindspore.Tensor, model: mindspore.Tensor) -> mindspore.Tensor: + return model - base + + def group_label(self): + return max( + self.base_tensor.group_label() or "", self.model_tensor.group_label() or "" + ) + + def uses_accelerator(self): + return True + + +class LoRAModuleSaveTask(Task): + weight_info: WeightInfo + writer_task: TensorWriterTask + model_ref: ModelReference + decomposition_task: TaskVectorDecompositionTask + + def arguments(self) -> Dict[str, Any]: + return {"writer": self.writer_task, "decomp": self.decomposition_task} + + def execute( + self, writer: TensorWriter, decomp: Tuple[mindspore.Tensor, mindspore.Tensor] + ) -> None: + weight_a, weight_b = decomp + if weight_a is None or weight_b is None: + if not self.weight_info.optional: + raise RuntimeError( + f"No SVD decomposition for required weight {self.weight_info.name}" + ) + return + lora_type = "lora_embedding" if self.decomposition_task.transpose else "lora" + lora_suffix = ".weight" if not self.decomposition_task.transpose else "" + base_name = self.weight_info.name.removesuffix(".weight") + writer.save_tensor( + f"base_model.model.{base_name}.{lora_type}_A{lora_suffix}", weight_a + ) + writer.save_tensor( + f"base_model.model.{base_name}.{lora_type}_B{lora_suffix}", weight_b + ) + + def priority(self) -> int: + return 1000 + + def group_label(self) -> Optional[str]: + return self.decomposition_task.group_label() + + +def _wi_load(model_ref: ModelReference, weight_info: WeightInfo) -> LoadTensor: + return LoadTensor( + model=model_ref, + tensor=weight_info.name, + dtype=weight_info.force_dtype, + optional=weight_info.optional, + aliases=weight_info.aliases, + tied_names=weight_info.tied_names, + ) + + +def _make_dummy_model( + model_ref: ModelReference, trust_remote_code: bool = False +) -> transformers.PreTrainedModel: + model_cfg = transformers.AutoConfig.from_pretrained( + model_ref.model.path, + revision=model_ref.model.revision, + trust_remote_code=trust_remote_code, + ) + auto_cls = get_auto_cls(model_cfg.architectures[0]) + res = auto_cls.from_config(model_cfg, trust_remote_code=trust_remote_code) + return res + + +class PlanResults(BaseModel): + tasks: List[Task] + base_vocab_size: int + final_vocab_size: int + + +def plan_extraction( # pylint: disable=too-many-positional-arguments + base_model_ref: ModelReference, + model_ref: ModelReference, + modules_to_save: List[str], + out_path: str, + options: MergeOptions, + max_rank: int, + distribute_scale: bool = True, + embed_lora: bool = False, + exclude_regexes: Optional[List[str]] = None, + include_regexes: Optional[List[str]] = None, + sv_epsilon: float = 0, + skip_undecomposable: bool = False, +) -> PlanResults: + targets = [] + writer_task = TensorWriterTask( + out_path=out_path, + override_basename="adapter_model", + max_shard_size=-1, + output_format=options.output_format, + use_async=options.async_write, + max_write_threads=options.write_threads, + ) + + name_to_wi = all_weights_map(model_ref, options) + dummy_base = _make_dummy_model(base_model_ref, options.trust_remote_code) + dummy_model = _make_dummy_model(model_ref, options.trust_remote_code) + + embed_in = dummy_model.get_input_embeddings() + embed_out = dummy_model.get_output_embeddings() + + ft_vocab = embed_in.weight.shape[0] + base_vocab = dummy_base.get_input_embeddings().weight.shape[0] + if ft_vocab != base_vocab and embed_lora: + LOG.warning( + f"Vocabulary size mismatch: fine-tuned model has {ft_vocab} tokens, base model has {base_vocab} tokens" + ) + LOG.warning("Enforcing embeddings in modules_to_save, embed_lora=False") + embed_lora = False + del dummy_base + + warned_modules = set() + + def _should_extract(name: str) -> bool: + if include_regexes and not any(re.search(r, name) for r in include_regexes): + return False + if any(re.search(r, name) for r in exclude_regexes): + return False + return True + + for name, module in tqdm.tqdm( + list(dummy_model.named_modules()), desc="Planning operations" + ): + module_type = type(module).__name__ + is_embed_module = isinstance(module, ms_nn.Embedding) or module_type == "Embedding" + is_supported_module = isinstance( + module, (ms_nn.Dense, ms_nn.Conv1d, ms_nn.Conv2d, ms_nn.Embedding) + ) or module_type in {"Linear", "Conv1d", "Conv2d", "Conv1D", "Embedding"} + wi = name_to_wi.get(name + ".weight") + bias_wi = name_to_wi.get(name + ".bias") + if wi is None: + if hasattr(module, "weight"): + LOG.warning( + f"Weight {name} present in model but not in architecture info" + ) + wi = WeightInfo( + name=name + ".weight", + optional=True, + is_embed=is_embed_module, + ) + else: + continue + + if ( + (not embed_lora) + and ( + module == embed_in + or module == embed_out + or is_embed_module + ) + and not any(re.search(r, name) for r in exclude_regexes or []) + ): + key = name.split(".")[-1] + if key not in modules_to_save: + LOG.warning(f"Adding {key} to modules_to_save") + modules_to_save.append(key) + + if name in modules_to_save or (name.split(".")[-1] in modules_to_save): + LOG.info(f"Planning to save {name} at full rank") + targets.extend(plan_module_to_save(model_ref, writer_task, wi, bias_wi)) + elif _should_extract(name): + if is_supported_module: + LOG.info(f"Planning LoRA extraction for {name}") + targets.extend( + plan_lora_module( + base_model_ref, + model_ref, + wi, + bias_wi, + writer_task, + max_rank, + distribute_scale, + transpose=is_embed_module, + sv_epsilon=sv_epsilon, + ) + ) + else: + key = name.split(".")[-1] + message = ( + f"{key} has unsupported module type {type(module).__name__} - " + + ("skipping" if skip_undecomposable else "saving at full rank") + ) + if not skip_undecomposable: + if key not in modules_to_save: + modules_to_save.append(key) + targets.extend( + plan_module_to_save(model_ref, writer_task, wi, bias_wi) + ) + if key not in warned_modules: + LOG.warning(message) + warned_modules.add(key) + + save_tasks = [t for t in targets if isinstance(t, (SaveTensor, LoRAModuleSaveTask))] + finalize = FinalizeModel(tensor_save_tasks=save_tasks, writer_task=writer_task) + return PlanResults( + tasks=targets + [finalize], + base_vocab_size=base_vocab, + final_vocab_size=ft_vocab, + ) + + +def plan_lora_module( # pylint: disable=too-many-positional-arguments + base_model_ref: ModelReference, + model_ref: ModelReference, + wi: WeightInfo, + bias_wi: Optional[WeightInfo], + writer_task: TensorWriterTask, + max_rank: int, + distribute_scale: bool = True, + transpose: bool = False, + sv_epsilon: float = 0, +) -> List[Task]: + targets = [] + base_load_task = _wi_load(base_model_ref, wi) + model_load_task = _wi_load(model_ref, wi) + tv_task = TaskVectorTask(base_tensor=base_load_task, model_tensor=model_load_task) + decomp_task = TaskVectorDecompositionTask( + weight_info=wi, + input_task=tv_task, + max_rank=max_rank, + distribute_scale=distribute_scale, + transpose=transpose, + sv_epsilon=sv_epsilon, + ) + targets.append(decomp_task) + targets.append( + LoRAModuleSaveTask( + weight_info=wi, + writer_task=writer_task, + model_ref=model_ref, + decomposition_task=decomp_task, + ) + ) + if bias_wi is not None: + base_bias_load_task = _wi_load(base_model_ref, bias_wi) + model_bias_load_task = _wi_load(model_ref, bias_wi) + tv_bias_task = TaskVectorTask( + base_tensor=base_bias_load_task, model_tensor=model_bias_load_task + ) + base_bias_name = bias_wi.name.removesuffix(".bias") + name_out = f"base_model.model.{base_bias_name}.lora_B.bias" + targets.append( + SaveTensor( + tensor_name=name_out, + tensor_task=tv_bias_task, + writer_task=writer_task, + optional=bias_wi.optional, + clone=False, + ) + ) + return targets + + +def plan_module_to_save( + model_ref: ModelReference, + writer_task: TensorWriterTask, + wi: WeightInfo, + bias_wi: Optional[WeightInfo], +): + save_tasks = [] + load_task = _wi_load(model_ref, wi) + save_task = SaveTensor( + tensor_name=f"base_model.model.{wi.name}", + tensor_task=load_task, + writer_task=writer_task, + optional=wi.optional, + clone=False, + ) + save_tasks.append(save_task) + if bias_wi is not None: + bias_load_task = _wi_load(model_ref, bias_wi) + bias_save_task = SaveTensor( + tensor_name=f"base_model.model.{bias_wi.name}", + tensor_task=bias_load_task, + writer_task=writer_task, + optional=bias_wi.optional, + clone=False, + ) + save_tasks.append(bias_save_task) + return save_tasks + + +def all_weights_map( + model_ref: ModelReference, options: MergeOptions +) -> Dict[str, WeightInfo]: + name_to_wi = {} + model_cfg = model_ref.config(trust_remote_code=options.trust_remote_code) + arch_info = arch_info_for_config(model_cfg) + for wi in arch_info.all_weights(model_cfg): + name_to_wi[wi.name] = wi + return name_to_wi + + +if __name__ == "__main__": + main() diff --git a/src/mindnlp/wizard/merge/scripts/fill_missing_params.py b/src/mindnlp/wizard/merge/scripts/fill_missing_params.py new file mode 100644 index 000000000..6581567d4 --- /dev/null +++ b/src/mindnlp/wizard/merge/scripts/fill_missing_params.py @@ -0,0 +1,367 @@ +# Copyright (C) 2025 Arcee AI +# SPDX-License-Identifier: LGPL-3.0-only +# Modified for MindSpore/Ascend NPU by MindNLP Wizard contributors. +# pylint: disable=no-value-for-parameter + +import logging +import shutil +from pathlib import Path +from typing import List, Optional, Tuple + +import click +import mindspore # pylint: disable=import-error +from huggingface_hub import snapshot_download +from safetensors import safe_open +from tqdm import tqdm + +from ..io.lazy_tensor_loader import ShardedTensorIndex +from ..io.tensor_writer import TensorWriter + +DEFAULT_SHARD_SIZE = 5 * 1024**3 + + +def load_tensor_from_file(tensor_name: str, tensor_file: str = None) -> mindspore.Tensor: + """ + Load a specific tensor from a .safetensors file. + + :param tensor_name: The name of the tensor to load. + :param tensor_file: The .safetensors file that contains the tensor. + :return: The loaded tensor as a MindSpore tensor. + """ + with safe_open(tensor_file, framework="numpy", device="cpu") as f: + if tensor_name in f.keys(): + np_tensor = f.get_tensor(tensor_name) + return mindspore.Tensor(np_tensor) + else: + raise ValueError( + f"Tensor '{tensor_name}' not found in file '{tensor_file}'" + ) + + +def load_tensor_from_index(tensor_name: str, index: ShardedTensorIndex) -> mindspore.Tensor: + """ + Load a specific tensor from a ShardedTensorIndex. + + :param tensor_name: The name of the tensor to load. + :param index: The ShardedTensorIndex containing the tensor. + :return: The loaded tensor as a MindSpore tensor. + """ + return load_tensor_from_file( + tensor_name, Path(index.base_path) / index.tensor_paths[tensor_name] + ) + + +def copy_and_fill_missing_params( + base_model_repo_id: str, + sub_model_dir: str, + max_shard_size: int = DEFAULT_SHARD_SIZE, + output_dir: str = None, +): + """ + Merge submodel weights into a base model and fill in missing parameters. + + Use Case: + Given a submodel (e.g., a language model) that is structurally identical to a subset of a + larger base model (e.g., a vision-language model). + The submodel contains only a subset of the weights (e.g., for the language model part), + while the base model contains all weights required for the complete architecture. + + This function replaces the shared parameters in the base model with those from the submodel, + facilitating testing after generating submodel parameters through merging. + + Parameters: + base_model_repo_id (str): + The path to the base model's directory or its Hugging Face repository ID. + sub_model_dir (str): + The path to the submodel's directory containing the merged weights. + max_shard_size (int, optional): + The maximum shard size for saving model weights, in bytes. Defaults to 5 GiB. + output_dir (str, optional): + The directory to save the final merged model. + + Returns: + pathlib.Path: The path to the directory where the final merged model is saved. + """ + output_dir = ( + Path(sub_model_dir).parent + / f"{Path(base_model_repo_id).stem}--{Path(sub_model_dir).stem}" + if output_dir is None + else Path(output_dir) + ) + output_dir.mkdir(parents=True, exist_ok=True) + + base_dir = ParameterNamesUtils.resolve_model_directory(base_model_repo_id) + files_to_copy = [ + item + for item in base_dir.rglob("*") + if item.is_file() and item.suffix not in {".safetensors", ".bin"} + ] + + with tqdm( + total=len(files_to_copy), desc="Copying non-parameter files", unit="file" + ) as pbar: + for item in files_to_copy: + target_path = output_dir / item.relative_to(base_dir) + target_path.parent.mkdir(parents=True, exist_ok=True) + shutil.copy2(item, target_path) + pbar.update(1) + + base_param_names = ParameterNamesUtils.get_model_parameter_names(base_model_repo_id) + submodel_param_names = ParameterNamesUtils.get_model_parameter_names(sub_model_dir) + + assert len(base_param_names) > len(submodel_param_names), ( + f"Base model must have more parameters than the submodel. " + f"Base: {len(base_param_names)}, Submodel: {len(submodel_param_names)}" + ) + + prefix = ParameterNamesUtils.find_prefix(base_param_names, submodel_param_names) + common_param_names = ParameterNamesUtils.find_common_ordered_names( + [base_param_names, submodel_param_names], ["", prefix] + ) + + base_index = ShardedTensorIndex.from_disk(str(base_dir)) + submodel_index = ShardedTensorIndex.from_disk( + str(ParameterNamesUtils.resolve_model_directory(sub_model_dir)) + ) + + writer = TensorWriter( + out_path=str(output_dir), max_shard_size=max_shard_size, output_format="safetensors" + ) + + for name, tensor_path in tqdm( + base_index.tensor_paths.items(), + total=len(base_index.tensor_paths), + desc="Merging tensors", + unit="tensor", + ): + tensor = load_tensor_from_index(name, base_index) + + if name in common_param_names: + submodel_name = ParameterNamesUtils.strip_prefix(name, prefix) + submodel_tensor = load_tensor_from_index(submodel_name, submodel_index) + + if submodel_tensor.shape != tensor.shape: + logging.warning( + f"Size mismatch for tensor '{name}': {tensor.shape} vs {submodel_tensor.shape}" + ) + + tensor = submodel_tensor + + writer.save_tensor(name, tensor.copy()) + + writer.finalize() + + return output_dir + + +@click.command() +@click.argument("base_model_repo_id", type=str) +@click.argument("sub_model_dir", type=str) +@click.option("--max_shard_size", type=int, default=DEFAULT_SHARD_SIZE) +@click.option("--output_dir", type=str, default=None) +def main( + base_model_repo_id, + sub_model_dir, + max_shard_size, + output_dir, +): + copy_and_fill_missing_params( + base_model_repo_id, sub_model_dir, max_shard_size, output_dir + ) + + +if __name__ == "__main__": + main() + + +class ParameterNamesUtils: + """Utility functions for handling parameter names.""" + + @staticmethod + def resolve_model_directory(repo_id: str) -> Path: + """Resolve the model directory (local or Hugging Face Hub).""" + if Path(repo_id).is_dir(): + return Path(repo_id) + + return Path(snapshot_download(repo_id)) + + @staticmethod + def get_model_parameter_names(repo_id: str) -> List[str]: + """Get parameter names of a model from a Hugging Face repo or local directory.""" + model_dir = ParameterNamesUtils.resolve_model_directory(repo_id) + return list(ShardedTensorIndex.from_disk(str(model_dir)).tensor_paths.keys()) + + @staticmethod + def strip_prefix(name: str, prefix: str) -> str: + """Remove a single prefix from the start of a name.""" + if prefix != "" and name.startswith(prefix + "."): + return name[len(prefix) + 1 :] + return name + + @staticmethod + def find_prefix(list1: List[str], list2: List[str]) -> Optional[str]: + """ + Find a prefix in list1 that, after removal, makes list2 an ordered sublist. + """ + assert len(list1) >= len(list2), "params name list1 can't be shorter than list2" + + possible_prefixes = {item.split(".")[0] for item in list1 if "." in item} + possible_prefixes = [""] + list(possible_prefixes) + + prefix_matches = {} + best_prefix = "" + for prefix in possible_prefixes: + stripped_list1 = [ + ParameterNamesUtils.strip_prefix(item, prefix) for item in list1 + ] + prefix_matches[prefix] = len( + [item for item in list2 if item in stripped_list1] + ) + + if max(prefix_matches.values()) > prefix_matches[""]: + best_prefix = max(prefix_matches, key=prefix_matches.get) + + return best_prefix + + @staticmethod + def find_common_ordered_names( + param_names: List[List[str]], prefixes: List[str] + ) -> List[str]: + """Identify and return common parameter names across all models, ensuring correct order.""" + common_names = set(param_names[0]) + for i in range(1, len(param_names)): + prefix = f"{prefixes[i]}." if prefixes[i] else "" + common_names.intersection_update({prefix + name for name in param_names[i]}) + return [name for name in param_names[0] if name in common_names] + + @staticmethod + def remove_size_conflicts(common_names, referenced_models, prefixes): + model_dirs = [ + ParameterNamesUtils.resolve_model_directory(m.model.path) + for m in referenced_models + ] + model_indices = [ShardedTensorIndex.from_disk(str(dir)) for dir in model_dirs] + + common_name_and_shape = common_names.copy() + removed_names = [] + + for name in common_names: + base_shape = ParameterNamesUtils.tensor_shape(name, model_indices[0]) + + for i in range(1, len(referenced_models)): + other_name = name + prefix = f"{prefixes[i]}." if prefixes[i] else "" + if name.startswith(prefix) and prefix != "": + other_name = name[len(prefix) :] + shape = ParameterNamesUtils.tensor_shape(other_name, model_indices[i]) + + if base_shape != shape: + common_name_and_shape.remove(name) + removed_names.append((name, base_shape, shape, i)) + break + + size_mismatch_count = len(removed_names) + if size_mismatch_count > 0: + logging.warning( + f"Size mismatch detected for {size_mismatch_count}/{size_mismatch_count + len(common_names)} tensors. " + "These names were removed from the merge list." + ) + logging.info( + "The following tensors have different shapes across models and were removed from the merge list:" + ) + for name, base_shape, shape, i in removed_names: + logging.info( + f"Tensor name: {name}, Base model shape: {base_shape}, Mismatched shape: {shape} in model {referenced_models[i].model.path}" + ) + + return common_name_and_shape + + @staticmethod + def are_common_params_ordered(list1: List[str], list2: List[str]) -> bool: + """ + Check if common elements of list2 maintain their relative order in list1. + """ + common_params = set(list1).intersection(set(list2)) + last_index = -1 + + for param in list2: + if param in common_params: + current_index = list1.index(param) + if current_index < last_index: + return False + last_index = current_index + return True + + @staticmethod + def ordered_sublist(list1: List[str], list2: List[str]) -> bool: + """ + Check if list2 is a contiguous ordered sublist of list1. + """ + n, m = len(list1), len(list2) + + for i in range(n - m + 1): + if list1[i : i + m] == list2: + return True + return False + + @staticmethod + def report_names_similarity( + base_names: List[str], other_names: List[str] + ) -> Tuple[Optional[str], str]: + """ + Analyze similarity between parameter names of two models and identify shared prefixes. + """ + possible_prefixes = {""} + possible_prefixes.update( + {item.split(".")[0] for item in base_names if "." in item} + ) + + prefixes_subset_overlap = {} + best_prefix = None + case_message = "No common parameter names found for any prefix" + + for prefix in possible_prefixes: + base_names_stripped = [ + ParameterNamesUtils.strip_prefix(name, prefix) for name in base_names + ] + + if ParameterNamesUtils.ordered_sublist(base_names_stripped, other_names): + return prefix, "All params in model have exact match in base model." + + intersection = set(base_names_stripped).intersection(set(other_names)) + prefixes_subset_overlap[prefix] = intersection + + if prefixes_subset_overlap: + best_prefix = max( + prefixes_subset_overlap, key=lambda x: len(prefixes_subset_overlap[x]) + ) + base_names_stripped = [ + ParameterNamesUtils.strip_prefix(name, best_prefix) + for name in base_names + ] + + overlap = len(prefixes_subset_overlap[best_prefix]) + ordered = ParameterNamesUtils.are_common_params_ordered( + base_names_stripped, other_names + ) + mismatched = [ + item for item in other_names if item not in base_names_stripped + ] + mismatched = "\n ".join(mismatched) + case_message = ( + f"{overlap}/{len(other_names)} ({100 * overlap / len(other_names):.2f}%) " + f"of model parameters are in the base model. \n" + f" Name ordering is {'preserved' if ordered else 'not preserved'}.\n" + f" Missing parameters:\n {mismatched}" + ) + + return best_prefix, case_message + + @staticmethod + def tensor_shape(name, index) -> Tuple[int]: + from safetensors import safe_open # pylint: disable=reimported + + with safe_open( + Path(index.base_path) / index.tensor_paths[name], framework="numpy" + ) as f: + return f.get_slice(name).get_shape() diff --git a/src/mindnlp/wizard/merge/scripts/layershuffle.py b/src/mindnlp/wizard/merge/scripts/layershuffle.py new file mode 100644 index 000000000..a16faa526 --- /dev/null +++ b/src/mindnlp/wizard/merge/scripts/layershuffle.py @@ -0,0 +1,134 @@ +# Copyright (C) 2025 Arcee AI +# SPDX-License-Identifier: LGPL-3.0-only +# Modified for MindSpore/Ascend NPU by MindNLP Wizard contributors. +# pylint: disable=no-value-for-parameter + +import random +from typing import List + +import click +import yaml + +from ..architecture import arch_info_for_config +from ..common import ModelReference +from ..config import ( + InputSliceDefinition, + MergeConfiguration, + OutputSliceDefinition, +) +from ..merge import run_merge +from ..options import MergeOptions, PrettyPrintHelp, add_merge_options + + +@click.command("wizard-layershuffle", cls=PrettyPrintHelp) +@click.argument("out_path", type=str) +@click.option("--model", "-m", multiple=True, type=str, help="Add a model to the merge") +@click.option( + "--weight", + "-w", + multiple=True, + type=float, + default=[], + show_default=False, + help="Weighting for a model", +) +@click.option( + "--print-yaml/--no-print-yaml", + is_flag=True, + help="Print YAML merge config for resulting model", +) +@click.option( + "--write-yaml", + type=click.Path(writable=True), + help="Path to write YAML merge config to", +) +@click.option( + "--dry-run", is_flag=True, help="Generate a config but do not run the merge" +) +@click.option("--fp16/--no-fp16", is_flag=True, help="Use FP16 precision") +@click.option( + "--full-random/--no-full-random", + is_flag=True, + help="Randomize layer index as well as source model", +) +@add_merge_options +def main( # pylint: disable=too-many-positional-arguments + out_path: str, + model: List[str], + weight: List[float], + print_yaml: bool, + write_yaml: bool, + dry_run: bool, + fp16: bool, + full_random: bool, + merge_options: MergeOptions, +): + models = [ModelReference.parse(m) for m in model] + + m0_cfg = models[0].config() + arch_info = arch_info_for_config(m0_cfg) + total_num_layers = arch_info.num_layers(m0_cfg) + + out_slices: List[OutputSliceDefinition] = [] + + if full_random: + for model, frac in zip(models, weight): + cfg = model.config() + num_layers = int(arch_info.num_layers(cfg) * frac) + for _ in range(num_layers): + src_idx = random.randrange(0, num_layers) + out_slices.append( + OutputSliceDefinition( + sources=[ + InputSliceDefinition( + model=str(model), + layer_range=(src_idx, src_idx + 1), + ) + ] + ) + ) + random.shuffle(out_slices) + else: + for layer_idx in range(total_num_layers): + src_model = random.choices(models, weights=weight, k=1)[0] + if out_slices and out_slices[-1].sources[0].model == str(src_model): + out_slices[-1].sources[0].layer_range = ( + out_slices[-1].sources[0].layer_range[0], + layer_idx + 1, + ) + else: + out_slices.append( + OutputSliceDefinition( + sources=[ + InputSliceDefinition( + model=str(src_model), + layer_range=(layer_idx, layer_idx + 1), + ) + ] + ) + ) + merge_config = MergeConfiguration( + merge_method="passthrough", slices=out_slices, dtype="float16" if fp16 else None + ) + + if print_yaml or write_yaml: + yaml_str = yaml.dump(merge_config.model_dump(exclude_none=True, mode="json")) + + if print_yaml: + print(yaml_str) + if write_yaml: + with open(write_yaml, "w", encoding="utf-8") as file: + file.write(yaml_str) + + if dry_run: + return + + run_merge( + merge_config, + out_path, + options=merge_options, + ) + + +if __name__ == "__main__": + main() diff --git a/src/mindnlp/wizard/merge/scripts/legacy.py b/src/mindnlp/wizard/merge/scripts/legacy.py new file mode 100644 index 000000000..bc8803d82 --- /dev/null +++ b/src/mindnlp/wizard/merge/scripts/legacy.py @@ -0,0 +1,132 @@ +# Copyright (C) 2025 Arcee AI +# SPDX-License-Identifier: LGPL-3.0-only +# Modified for MindSpore/Ascend NPU by MindNLP Wizard contributors. +# pylint: disable=no-value-for-parameter + +from typing import List, Optional + +import click +import yaml + +from ..config import InputModelDefinition, MergeConfiguration +from ..merge import run_merge +from ..options import MergeOptions, PrettyPrintHelp, add_merge_options + + +@click.command("wizard-legacy", cls=PrettyPrintHelp) +@click.argument("out_path", type=str) +@click.option( + "--merge", "merge", type=str, multiple=True, help="Add a model to the merge" +) +@click.option( + "--density", + "density", + type=float, + multiple=True, + default=[], + help="Fraction of weights to keep for each model (ties only)", +) +@click.option( + "--weight", + "weight", + type=float, + multiple=True, + default=[], + help="Weighting for a model (default 1.0 for all models if not specified)", +) +@click.option( + "--method", "method", type=str, default="ties", help="Method used to merge models" +) +@click.option( + "--base-model", "base_model", type=str, default=None, help="Base model for merge" +) +@click.option( + "--normalize/--no-normalize", + "normalize", + is_flag=True, + default=True, + help="Divide merged parameters by the sum of weights", +) +@click.option( + "--int8-mask/--no-int8-mask", + "int8_mask", + is_flag=True, + help="Store intermediate masks in int8 to save memory", +) +@click.option("--bf16/--no-bf16", "bf16", is_flag=True, help="Use bfloat16") +@click.option( + "--naive-count/--no-naive-count", + "naive_count", + is_flag=True, + help="Use naive sign count instead of weight (ties only)", +) +@click.option( + "--print-yaml/--no-print-yaml", + "print_yaml", + is_flag=True, + help="Print generated YAML configuration", +) +@add_merge_options +def main( # pylint: disable=too-many-positional-arguments + out_path: str, + merge: List[str], + density: List[float], + weight: List[float], + method: str, + base_model: Optional[str], + normalize: bool, + int8_mask: bool, + bf16: bool, + naive_count: bool, + print_yaml: bool, + merge_options: MergeOptions, +): + """Wrapper for using a subset of legacy-style script arguments.""" + models = [InputModelDefinition(model=model, parameters={}) for model in merge] + if base_model and base_model not in merge: + models.append(InputModelDefinition(model=base_model, parameters={})) + + parameters = {} + + if density: + if len(density) == 1: + density = [density[0]] * len(models) + for idx, d in enumerate(density): + models[idx].parameters["density"] = d + + if method == "slerp": + assert len(weight) == 1, "Must specify exactly one weight for SLERP" + parameters["t"] = weight[0] + else: + if weight: + if len(weight) == 1: + weight = [weight[0]] * len(models) + for idx, w in enumerate(weight): + models[idx].parameters["weight"] = w + + if int8_mask: + parameters["int8_mask"] = True + if naive_count: + parameters["consensus_method"] = "count" + parameters["normalize"] = normalize + + merge_config = MergeConfiguration( + merge_method=method, + models=models, + parameters=parameters, + base_model=base_model, + dtype="bfloat16" if bf16 else None, + ) + + if print_yaml: + print(yaml.dump(merge_config.model_dump(mode="json", exclude_none=True))) + + run_merge( + merge_config, + out_path, + options=merge_options, + ) + + +if __name__ == "__main__": + main() diff --git a/src/mindnlp/wizard/merge/scripts/merge_raw.py b/src/mindnlp/wizard/merge/scripts/merge_raw.py new file mode 100644 index 000000000..28bc1ef89 --- /dev/null +++ b/src/mindnlp/wizard/merge/scripts/merge_raw.py @@ -0,0 +1,308 @@ +# Copyright (C) 2025 Arcee AI +# SPDX-License-Identifier: LGPL-3.0-only +# Modified for MindSpore/Ascend NPU by MindNLP Wizard contributors. +# +# Renamed from merge_raw_pytorch.py → merge_raw.py (framework-neutral name). + +""" +Merge arbitrary weight files (safetensors / bin) without architecture awareness. + +Unlike ``wizard-merge`` which requires HuggingFace model directories with +``config.json``, this tool works directly on individual weight files. +""" + +import logging +from typing import Dict, List, Optional + +import click +import mindspore # pylint: disable=import-error +import tqdm +import yaml +from pydantic import BaseModel + +from mindnlp.wizard.merge import merge_methods +from ..architecture import WeightInfo +from ..common import ImmutableMap, ModelReference, dtype_from_name +from ..config import ParameterSetting, evaluate_setting +from ..graph import Executor, Task +from ..io import LazyTensorLoader, ShardedTensorIndex +from ..io.tasks import FinalizeModel, SaveTensor, TensorWriterTask +from ..merge_methods.base import MergeMethod, TensorDictWrapper +from ..options import MergeOptions, PrettyPrintHelp, add_merge_options + + +class InputModelDefinition(BaseModel, frozen=True): + model: str + parameters: Optional[Dict[str, ParameterSetting]] = None + + +class RawMergeConfig(BaseModel, frozen=True): + merge_method: str + models: List[InputModelDefinition] + parameters: Optional[Dict[str, ParameterSetting]] = None + dtype: Optional[str] = None + base_model: Optional[str] = None + + +class SimpleLoaderCache: + loaders: Dict[str, LazyTensorLoader] + lazy_loader: bool = False + _instance: Optional["SimpleLoaderCache"] = None + + def __new__(cls) -> "SimpleLoaderCache": + if cls._instance is None: + cls._instance = super(SimpleLoaderCache, cls).__new__(cls) + cls._instance.loaders = {} + return cls._instance + + def get(self, model: str) -> LazyTensorLoader: + if model not in self.loaders: + self.loaders[model] = LazyTensorLoader( + ShardedTensorIndex.from_file(model), + lazy_loader=self.lazy_loader, + ) + return self.loaders[model] + + +class SimpleLoadTensor(Task[mindspore.Tensor]): + model: str + tensor_name: str + dtype: Optional[str] = None + device: Optional[str] = None + + def arguments(self) -> Dict[str, Task]: + return {} + + def execute(self) -> Optional[mindspore.Tensor]: + loader = SimpleLoaderCache().get(self.model) + tensor = loader.get_tensor( + self.tensor_name, device=self.device or "CPU" + ) + if tensor is None: + return None + if dt := dtype_from_name(self.dtype): + tensor = tensor.astype(dt) + return tensor + + +def plan_flat_merge( + config: RawMergeConfig, + out_path: str, + tensor_union: bool, + tensor_intersection: bool, + options: MergeOptions, +) -> List[Task]: + merge_method = merge_methods.get(config.merge_method) + + loaders = SimpleLoaderCache() + loaders.lazy_loader = options.lazy_loader + all_tensor_names: set = set() + for model_def in tqdm.tqdm( + config.models, desc="Preparing model loaders" + ): + loader = loaders.get(model_def.model) + all_tensor_names.update(loader.index.tensor_paths.keys()) + + writer_task = TensorWriterTask( + out_path=out_path, + max_shard_size=options.out_shard_size, + output_format=options.output_format, + use_async=options.async_write, + max_write_threads=options.write_threads, + ) + + save_tasks = [] + for tensor_name in tqdm.tqdm( + list(all_tensor_names), desc="Planning operations" + ): + inputs = { + model_def.model: SimpleLoadTensor( + model=model_def.model, + tensor_name=tensor_name, + dtype=config.dtype, + ) + for model_def in config.models + } + if ( + config.base_model is not None + and config.base_model not in inputs + ): + inputs[config.base_model] = SimpleLoadTensor( + model=config.base_model, + tensor_name=tensor_name, + dtype=config.dtype, + ) + + has_tensor = [ + lt.model + for lt in inputs.values() + if lt.tensor_name + in loaders.get(lt.model).index.tensor_paths + ] + if len(has_tensor) < len(inputs): + if tensor_intersection: + continue + if tensor_union: + pass + else: + missing = set(inputs) - set(has_tensor) + logging.warning( + f"Tensor {tensor_name} not found in models:" + ) + for model in missing: + logging.warning(f" {model}") + logging.warning("Was found in:") + for model in has_tensor: + logging.warning(f" {model}") + raise RuntimeError("Missing tensors") + + inputs = { + ModelReference.model_validate({"model": {"path": k}}): v + for k, v in inputs.items() + } + + global_params, tensor_params = construct_param_dicts( + config, merge_method, tensor_name + ) + + tensor_task = merge_method.make_task( + output_weight=WeightInfo(name=tensor_name), + tensors=TensorDictWrapper(tensors=inputs), + parameters=ImmutableMap(global_params), + tensor_parameters=ImmutableMap( + data={ + key: ImmutableMap(data=tensor_params[key]) + for key in tensor_params + } + ), + base_model=( + ModelReference.model_validate( + {"model": {"path": config.base_model}} + ) + if config.base_model is not None + else None + ), + ) + save_task = SaveTensor( + tensor_name=tensor_name, + tensor_task=tensor_task, + writer_task=writer_task, + clone=options.clone_tensors, + dtype=config.dtype, + ) + save_tasks.append(save_task) + + finalize = FinalizeModel( + tensor_save_tasks=tuple(save_tasks), writer_task=writer_task + ) + return save_tasks + [finalize] + + +def construct_param_dicts( + config: RawMergeConfig, + merge_method: MergeMethod, + tensor_name: str, +): + global_params: Dict[str, any] = {} + for param_def in merge_method.parameters(): + if config.parameters and param_def.name in config.parameters: + value = evaluate_setting( + tensor_name, config.parameters[param_def.name] + ) + if value is not None: + global_params[param_def.name] = value + + if param_def.name not in global_params: + if param_def.required: + raise RuntimeError( + f"Missing required parameter {param_def.name} " + f"for merge method {merge_method}" + ) + else: + global_params[param_def.name] = param_def.default_value + + tensor_params: Dict = {} + for param_def in merge_method.tensor_parameters(): + for model_def in config.models: + mr = ModelReference.model_validate( + {"model": {"path": model_def.model}} + ) + tensor_params[mr] = tensor_params.get(mr, {}) + if value := evaluate_setting( + tensor_name, + model_def.parameters.get(param_def.name, []), + ): + tensor_params[mr][param_def.name] = value + elif value := evaluate_setting( + tensor_name, + ( + config.parameters.get(param_def.name, []) + if config.parameters + else [] + ), + ): + tensor_params[mr][param_def.name] = value + elif param_def.required: + raise RuntimeError( + f"Missing required parameter {param_def.name} " + f"for model {mr} tensor {tensor_name}" + ) + else: + tensor_params[mr][param_def.name] = ( + param_def.default_value + ) + return global_params, tensor_params + + +@click.command("wizard-merge-raw", cls=PrettyPrintHelp) +@click.argument("config_path", type=click.Path(exists=True)) +@click.argument("out_path", type=click.Path()) +@click.option( + "--tensor-intersection", + "-i", + type=bool, + default=False, + is_flag=True, + help="Only merge tensors present in all input models", +) +@click.option( + "--tensor-union", + "-u", + type=bool, + default=False, + is_flag=True, + help="Merge all tensors present in any input model", +) +@add_merge_options +def main( + config_path: str, + out_path: str, + tensor_union: bool, + tensor_intersection: bool, + merge_options: MergeOptions, +): + """Merge arbitrary weight files (safetensors / bin). + + Uses similar configuration syntax to ``wizard-merge``, minus the + ``slices`` sections. Each input model should be the path on disk to + a safetensors or bin file.""" + merge_options.apply_global_options() + + with open(config_path, "r", encoding="utf-8") as file: + config_source = file.read() + + config = RawMergeConfig.model_validate(yaml.safe_load(config_source)) + tasks = plan_flat_merge( + config, out_path, tensor_union, tensor_intersection, merge_options + ) + + executor = Executor( + tasks, + math_device=merge_options.device or "CPU", + storage_device=( + merge_options.device + if merge_options.low_cpu_memory + else "CPU" + ), + ) + executor.execute() diff --git a/src/mindnlp/wizard/merge/scripts/moe.py b/src/mindnlp/wizard/merge/scripts/moe.py new file mode 100644 index 000000000..c249ae72e --- /dev/null +++ b/src/mindnlp/wizard/merge/scripts/moe.py @@ -0,0 +1,201 @@ +# Copyright (C) 2025 Arcee AI +# SPDX-License-Identifier: LGPL-3.0-only +# Modified for MindSpore/Ascend NPU by MindNLP Wizard contributors. +# pylint: disable=no-value-for-parameter + +import logging +import os +import sys +from typing import List + +import click +import transformers +import yaml + +from ..merge import MergeOptions +from ..moe import ALL_OUTPUT_ARCHITECTURES, MoEOutputArchitecture +from ..moe.config import MoEMergeConfig, is_bad_config +from ..moe.router import get_gate_params, warn_degenerate_gates +from ..options import PrettyPrintHelp, add_merge_options + + +def build( # pylint: disable=too-many-positional-arguments + config: MoEMergeConfig, + out_path: str, + merge_options: MergeOptions, + load_in_4bit: bool = False, + load_in_8bit: bool = False, + device: str = "auto", + allow_all_same: bool = False, + verbose: bool = False, +): + if is_bad_config(config, allow_all_same=allow_all_same): + sys.exit(1) + + base_model = config.base_model + out_arch = select_output_arch(config, merge_options, verbose=verbose) + + tokenizer = transformers.AutoTokenizer.from_pretrained( + base_model.model.path, revision=base_model.model.revision + ) + tokenizer.padding_side = "left" + tokenizer.pad_token_id = tokenizer.bos_token_id + if tokenizer.pad_token_id is None: + tokenizer.pad_token = tokenizer.eos_token + + logging.info("Getting gate parameters...") + need_gates = list(config.experts) + if config.shared_experts: + has_prompts = any(e.positive_prompts for e in config.shared_experts) + assert all( + bool(e.positive_prompts) == has_prompts for e in config.shared_experts + ), "Must specify prompts for all shared experts or none, not a mix" + if has_prompts or config.gate_mode in ("random", "uniform_random"): + need_gates.extend(config.shared_experts) + + gate_vecs = get_gate_params( + base_model, + tokenizer, + need_gates, + mode=config.gate_mode, + load_in_4bit=load_in_4bit, + load_in_8bit=load_in_8bit, + lazy_loader=merge_options.lazy_loader, + trust_remote_code=merge_options.trust_remote_code, + device=device, + ) + router_weights = gate_vecs[:, : len(config.experts), :] + shared_router_weights = gate_vecs[:, len(config.experts) :, :] + warn_degenerate_gates(gate_vecs) + + out_arch.write_model( + out_path, + config, + merge_options, + router_weights=[router_weights[i, ...] for i in range(router_weights.shape[0])], + shared_router_weights=[ + shared_router_weights[i, ...] for i in range(router_weights.shape[0]) + ], + ) + + if merge_options.copy_tokenizer: + logging.info("Saving tokenizer...") + tokenizer.save_pretrained(out_path, safe_serialization=True) + + logging.info("Done.") + + +def select_output_arch( + config: MoEMergeConfig, + merge_options: MergeOptions, + verbose: bool = False, +) -> MoEOutputArchitecture: + candidates_in = ALL_OUTPUT_ARCHITECTURES + if config.architecture: + candidates_in = [ + a + for a in candidates_in + if a.name().lower().startswith(config.architecture.lower()) + ] + if not candidates_in: + logging.error( + f"No output architecture found that matches the given architecture: {config.architecture}" + ) + logging.error("All supported output architectures:") + for arch in ALL_OUTPUT_ARCHITECTURES: + logging.error(f" * {arch.name()}") + sys.exit(1) + + candidates: List[MoEOutputArchitecture] = [] + for arch in candidates_in: + if arch.supports_config( + config, explain=verbose, trust_remote_code=merge_options.trust_remote_code + ): + candidates.append(arch) + else: + logging.info(f"Output architecture {arch.name()} does not support config.") + + if not candidates: + logging.error( + "No output architecture found that is compatible with the given models." + ) + + logging.error("All supported output architectures:") + for arch in ALL_OUTPUT_ARCHITECTURES: + logging.error(f" * {arch.name()}") + sys.exit(1) + + for arch in candidates: + if arch.name() == "Mixtral": + return arch + + if len(candidates) > 1: + logging.warning( + "Multiple output architectures found that are compatible with the given models." + ) + logging.warning(f"Defaulting to {candidates[0].name()}") + else: + logging.info(f"Selected output architecture: {candidates[0].name()}") + return candidates[0] + + +@click.command("wizard-moe", cls=PrettyPrintHelp) +@click.argument("config_path", type=click.Path(exists=True, dir_okay=False)) +@click.argument("out_path", type=click.Path()) +@click.option( + "--load-in-4bit", + is_flag=True, + type=bool, + default=False, + help="Load model in 4bit for computing hidden states", +) +@click.option( + "--load-in-8bit", + is_flag=True, + type=bool, + default=False, + help="Load model in 8bit for computing hidden states", +) +@click.option( + "--i-understand-this-is-not-useful-without-training", + type=bool, + default=False, + is_flag=True, + help="Really make the questionable model you want.", +) +@add_merge_options +def main( # pylint: disable=too-many-positional-arguments + config_path: str, + out_path: str, + load_in_4bit: bool, + load_in_8bit: bool, + i_understand_this_is_not_useful_without_training: bool, + merge_options: MergeOptions, +): + """Create a Mixture of Experts model by combining the pretrained weights of multiple models.""" + merge_options.apply_global_options() + + with open(config_path, "r", encoding="utf-8") as file: + config_source = file.read() + + config = MoEMergeConfig.model_validate(yaml.safe_load(config_source)) + build( + config, + out_path=out_path, + merge_options=merge_options, + load_in_4bit=load_in_4bit, + load_in_8bit=load_in_8bit, + device=merge_options.device, + allow_all_same=i_understand_this_is_not_useful_without_training, + verbose=merge_options.verbosity > 0, + ) + + if merge_options.write_model_card: + with open( + os.path.join(out_path, "wizard_moe_config.yml"), "w", encoding="utf-8" + ) as fp: + fp.write(config_source) + + +if __name__ == "__main__": + main() diff --git a/src/mindnlp/wizard/merge/scripts/multimerge.py b/src/mindnlp/wizard/merge/scripts/multimerge.py new file mode 100644 index 000000000..b92f83a28 --- /dev/null +++ b/src/mindnlp/wizard/merge/scripts/multimerge.py @@ -0,0 +1,256 @@ +# Copyright (C) 2025 Arcee AI +# SPDX-License-Identifier: LGPL-3.0-only +# Modified for MindSpore/Ascend NPU by MindNLP Wizard contributors. +# pylint: disable=no-value-for-parameter + +import logging +import os +from typing import Dict, Optional, Set, Tuple, Union + +import click +import yaml + +from ..common import ImmutableMap, ModelReference +from ..config import MergeConfiguration +from ..graph import Executor, Task +from ..merge import run_merge +from ..options import MergeOptions, PrettyPrintHelp, add_merge_options + +LOG = logging.getLogger("multimerge") + + +MODEL_CHECK_FILENAMES = [ + "model.safetensors", + "model.ckpt", + "mindspore_model.ckpt", + "model.safetensors.index.json", + "mindspore_model.ckpt.index.json", +] + + +class MergeModelTask(Task[str]): + config_yaml: str + name: str + input_merges: ImmutableMap[str, "MergeModelTask"] + options: MergeOptions + out_path: str + lazy: bool = True + + def arguments(self): + return {str(key): self.input_merges[key] for key in self.input_merges} + + def execute(self, **kwargs): + if ( + self.lazy + and os.path.exists(os.path.join(self.out_path, "config.json")) + and any( + os.path.exists(os.path.join(self.out_path, filename)) + for filename in MODEL_CHECK_FILENAMES + ) + ): + LOG.info(f"Model already exists at {self.out_path}, skipping") + return self.out_path + + LOG.info(f"Running merge for {self.name}") + cfg = MergeConfiguration.model_validate(yaml.safe_load(self.config_yaml)) + + run_merge( + cfg, + self.out_path, + options=self.options, + ) + LOG.info(f"Merge complete for {self.name}") + return self.out_path + + +@click.command("wizard-multimerge", cls=PrettyPrintHelp) +@click.argument("config_file", type=click.Path(exists=True)) +@click.option( + "--out-path", + type=click.Path(), + required=False, + help="Path to save the final merged model", +) +@click.option( + "--intermediate-dir", + "-I", + type=click.Path(), + required=True, + help="Directory to store intermediate merges", +) +@click.option( + "--lazy/--no-lazy", + default=True, + help="Skip merges that already exist", +) +@add_merge_options +def main( + config_file: str, + intermediate_dir: str, + out_path: Optional[str], + lazy: bool, + merge_options: MergeOptions, +): + """Execute a set of potentially interdependent merge recipes. + + The configuration file should be a YAML file containing multiple + documents, each of which is a merge configuration with the addition + of a `name` field. + + The `intermediate_dir` is used to store intermediate merge results. + Any merge configuration with a `name` field will be saved to this + directory. If an unnamed merge configuration is present, it will be + saved to `out_path` (which is required in this case).""" + merge_options.apply_global_options() + os.makedirs(intermediate_dir, exist_ok=True) + + with open(config_file, "r", encoding="utf-8") as file: + config_source = file.read() + + merge_configs, dependencies = load_config(config_source, intermediate_dir) + + if None in merge_configs and not out_path: + raise click.UsageError( + "--out-path is required when configuration contains an unnamed final merge" + ) + tasks = make_tasks( + merge_configs, dependencies, merge_options, intermediate_dir, out_path, lazy + ) + + executor = Executor( + tasks, math_device="cpu", storage_device="cpu" + ) + executor.execute(desc="Merging models") + + +def patched_config(config: MergeConfiguration, merge_names: Set[str], working_dir: str): + """Replace instances of intermediate merge names with actual paths. + + Also returns the set of intermediate merge names that were used. + + Args: + config: The configuration to patch + merge_names: The set of all merge names + working_dir: The directory to use as the base for relative paths + """ + used = set() + + def _patch_mr(value: Union[dict, list, str, int, None]): + nonlocal used + if isinstance(value, list): + return [_patch_mr(x) for x in value] + elif isinstance(value, dict): + if set(value.keys()) == {"model", "lora", "override_architecture"}: + base = value["model"]["path"] + if base in merge_names: + value["model"] = value["model"].copy() + value["model"]["path"] = os.path.join(working_dir, base) + used.add(base) + return value + return {k: _patch_mr(v) for k, v in value.items()} + elif isinstance(value, str): + try: + mr = ModelReference.model_validate(value) + if mr.model.path in merge_names: + used.add(mr.model.path) + return ModelReference( + model={ + "path": os.path.join(working_dir, mr.model.path), + "revision": mr.model.revision, + }, + lora=mr.lora, + override_architecture=mr.override_architecture, + ).model_dump() + except ValueError: + pass + return value + + new_dict = _patch_mr(config.model_dump()) + return MergeConfiguration.model_validate(new_dict), used + + +def make_tasks( # pylint: disable=too-many-positional-arguments + merge_configs: Dict[str, MergeConfiguration], + dependencies: Dict[str, Set[str]], + merge_options: MergeOptions, + intermediate_dir: str, + out_path: Optional[str], + lazy: bool, +): + """Build the task dependency graph for the merge recipes.""" + touched = set() + tasks = {} + + def _make_task(name: str): + nonlocal touched, tasks, out_path + if name in tasks: + return tasks[name] + elif name in touched: + raise ValueError(f"Circular dependency detected involving {name}") + touched.add(name) + if name is None: + merge_out_path = out_path + else: + merge_out_path = os.path.join(intermediate_dir, name) + tasks[name] = MergeModelTask( + config_yaml=merge_configs[name].to_yaml(), + name=name or "final merge", + input_merges=ImmutableMap( + {dep: _make_task(dep) for dep in dependencies[name]} + ), + options=merge_options, + out_path=merge_out_path, + lazy=lazy, + ) + return tasks[name] + + tasks_to_create = [ + name for name in merge_configs.keys() if name is not None or out_path + ] + tasks = [_make_task(name) for name in tasks_to_create] + return tasks + + +def load_config( + config_source: str, intermediate_dir: str +) -> Tuple[Dict[str, MergeConfiguration], Dict[str, Set[str]]]: + """Load the merge configurations from the YAML source. + + Args: + config_source: The YAML source to load + intermediate_dir: The directory to use for intermediate merges + + Returns: + A tuple containing: + - A dictionary of merge configurations keyed by name + - A dictionary of dependencies keyed by name + """ + docs = list(yaml.safe_load_all(config_source)) + merge_configs = {} + for doc in docs: + if "name" in doc: + merge_name = doc.pop("name") + else: + merge_name = None + if merge_name in merge_configs: + if merge_name is not None: + raise ValueError(f"Duplicate merge name {merge_name}") + else: + raise ValueError( + "Multiple unnamed merge configurations are not supported" + ) + merge_configs[merge_name] = MergeConfiguration.model_validate(doc) + + merge_names = set(merge_configs.keys()) + dependencies = {} + for merge_name in merge_names: + merge_config, used_names = patched_config( + merge_configs[merge_name], merge_names, intermediate_dir + ) + merge_configs[merge_name] = merge_config + dependencies[merge_name] = used_names + return merge_configs, dependencies + + +if __name__ == "__main__": + main() diff --git a/src/mindnlp/wizard/merge/scripts/run_yaml.py b/src/mindnlp/wizard/merge/scripts/run_yaml.py new file mode 100644 index 000000000..0e6bccee8 --- /dev/null +++ b/src/mindnlp/wizard/merge/scripts/run_yaml.py @@ -0,0 +1,42 @@ +# Originally from MergeKit (https://github.com/arcee-ai/mergekit) +# Copyright (C) 2025 Arcee AI +# SPDX-License-Identifier: LGPL-3.0-only +# Modified for MindSpore/Ascend NPU by MindNLP Wizard contributors. +# pylint: disable=no-value-for-parameter + +import click +import yaml + +from ..config import MergeConfiguration +from ..options import MergeOptions, PrettyPrintHelp, add_merge_options + + +@click.command("wizard-merge", cls=PrettyPrintHelp) +@click.argument("config_file") +@click.argument("out_path") +@add_merge_options +def main( + merge_options: MergeOptions, + config_file: str, + out_path: str, +): + merge_options.apply_global_options() + # Import after context setup so mindtorch backend binds to correct device. + from ..merge import run_merge + + with open(config_file, "r", encoding="utf-8") as file: + config_source = file.read() + + merge_config: MergeConfiguration = MergeConfiguration.model_validate( + yaml.safe_load(config_source) + ) + run_merge( + merge_config, + out_path, + options=merge_options, + config_source=config_source, + ) + + +if __name__ == "__main__": + main() diff --git a/src/mindnlp/wizard/merge/scripts/tokensurgeon.py b/src/mindnlp/wizard/merge/scripts/tokensurgeon.py new file mode 100644 index 000000000..907aa5ce5 --- /dev/null +++ b/src/mindnlp/wizard/merge/scripts/tokensurgeon.py @@ -0,0 +1,805 @@ +# Copyright (C) 2025 Arcee AI +# SPDX-License-Identifier: LGPL-3.0-only +# Modified for MindSpore/Ascend NPU by MindNLP Wizard contributors. +# pylint: disable=no-value-for-parameter + +import enum +import logging +from typing import Dict, List, Optional, Tuple + +import click +import mindspore # pylint: disable=import-error +from mindspore import ops # pylint: disable=import-error +import tqdm +import transformers +from pydantic import BaseModel + +from ..architecture import ( + ConfiguredModelArchitecture, + WeightInfo, + arch_info_for_config, +) +from ..architecture.auto import infer_architecture_info +from ..common import ModelReference, set_config_value +from ..io.tasks import ( + LoaderCache, +) +from ..io.tensor_writer import TensorWriter +from ..options import MergeOptions, PrettyPrintHelp, add_merge_options +from ..tokenizer.normalization import ( + NormalizedToken, + normalized_vocabulary, + token_prefixes, +) +from ..tokensurgeon import ( + SubwordMethod, + WeightingScheme, + batch_mp_rope, + batch_omp, + common_interp_approximate, + compute_token_basis, + landmark_pca_approximate, + subword_approximate, + well_trained_tokens, +) +from ..tokensurgeon.common_interpolation import DistanceMetric + +LOG = logging.getLogger(__name__) + + +class TokenAssignmentStats(BaseModel): + exact_match: int = 0 + byte_match: int = 0 + prefix_match: int = 0 + to_approximate: int = 0 + + def pretty_print(self): + chunks = ["Token Breakdown:"] + if self.exact_match: + chunks.append(f" Exact matches: {self.exact_match}") + if self.byte_match: + chunks.append(f" Byte matches: {self.byte_match}") + if self.prefix_match: + chunks.append(f" Prefix matches: {self.prefix_match}") + if self.to_approximate: + chunks.append(f" Tokens to approximate: {self.to_approximate}") + chunks.append( + f" Total: {self.exact_match + self.byte_match + self.prefix_match + self.to_approximate}" + ) + return "\n".join(chunks) + + +class ApproximationMethod(enum.Enum): + COMMON_INTERPOLATION = "common_interpolation" + SUBWORD = "subword" + MEAN = "mean" + ZERO = "zero" + RANDN = "randn" + JOHN_HEWITT = "john_hewitt" + ORTHOGONAL_MATCHING_PURSUIT = "omp" + LANDMARK_PCA = "landmark_pca" + SPARSE_TOKEN_BASIS = "stb" + MATCHING_PURSUIT_ROPE = "mp_rope" + + +class TokenSurgeonOptions(BaseModel): + model: ModelReference + donor: ModelReference + out_path: str + method: ApproximationMethod = ApproximationMethod.COMMON_INTERPOLATION + weight_scheme: WeightingScheme = WeightingScheme.DISTANCE_PROPORTIONAL + k: int = 64 + cosine_similarity: bool = False + subword_method: SubwordMethod = SubwordMethod.MEAN + batch_size: Optional[int] = None + new_vocab_noise: Optional[float] = None + new_vocab_scale: Optional[float] = None + + +def get_arch_info( + model: ModelReference, options: MergeOptions +) -> ConfiguredModelArchitecture: + cfg = model.config(trust_remote_code=options.trust_remote_code) + arch_info = arch_info_for_config(cfg) + if arch_info is None: + arch_info = infer_architecture_info((model,), model, options) + return ConfiguredModelArchitecture(info=arch_info, config=cfg) + + +def get_embedding_info( + arch_info: ConfiguredModelArchitecture, +) -> Tuple[WeightInfo, WeightInfo]: + """Get WeightInfo for the input and output embeddings of a model.""" + + if len(arch_info.info.modules) != 1: + raise RuntimeError("Model has multiple modules - not supported by tokensurgeon") + name = next(iter(arch_info.info.modules.keys())) + module_def = arch_info.get_module(name) + + embed, lm_head = None, None + for weight_info in module_def.pre_weights(): + if weight_info.is_embed: + if embed is not None: + raise RuntimeError("Multiple input embeddings found") + embed = weight_info + + for weight_info in module_def.post_weights(): + if weight_info.is_embed: + if lm_head is not None: + raise RuntimeError("Multiple output embeddings found") + lm_head = weight_info + return embed, lm_head + + +def maybe_aliases(weight_info: WeightInfo, tied: bool) -> Tuple[str, ...]: + return tuple( + list(weight_info.aliases or []) + + list((weight_info.tied_names or []) if tied else []) + ) + + +def get_stuff( + model: ModelReference, + options: MergeOptions, + arch_info: Optional[ConfiguredModelArchitecture] = None, + get_tied: bool = False, + device: str = "cpu", +) -> Tuple[Dict[NormalizedToken, int], Optional[mindspore.Tensor], Optional[mindspore.Tensor]]: + if arch_info is None: + arch_info = get_arch_info(model, options) + tokenizer = transformers.AutoTokenizer.from_pretrained( + model.model.path, + revision=model.model.revision, + trust_remote_code=options.trust_remote_code, + ) + vocab = normalized_vocabulary(tokenizer) + embed_wi, lm_head_wi = get_embedding_info(arch_info) + loader = LoaderCache().get(model) + embed = loader.get_tensor( + embed_wi.name, + device=device, + aliases=maybe_aliases(embed_wi, get_tied), + raise_on_missing=not embed_wi.optional, + ) + lm_head = loader.get_tensor( + lm_head_wi.name, + device=device, + aliases=maybe_aliases(lm_head_wi, get_tied), + raise_on_missing=not lm_head_wi.optional, + ) + return vocab, embed, lm_head + + +def match_byte_token( + token: NormalizedToken, original_vocab: Dict[NormalizedToken, int] +) -> Optional[int]: + if not isinstance(token, str): + return None + if len(token) == 1 and ord(token) < 256: + byte_tok = f"<0x{ord(token):02X}>" + if byte_tok in original_vocab: + return original_vocab[byte_tok] + elif token.startswith("<0x") and token.endswith(">") and len(token) == 6: + try: + byte = int(token[3:-1], 16) + except ValueError: + pass + else: + if chr(byte) in original_vocab: + return original_vocab[chr(byte)] + return None + + +def match_prefix( + token: NormalizedToken, original_vocab: Dict[NormalizedToken, int] +) -> Optional[int]: + for prefix in token_prefixes(token): + if prefix in original_vocab: + return original_vocab[prefix] + return None + + +def get_out_arch_info( + model: ModelReference, + donor: ModelReference, + new_vocab_size: int, + common_options: MergeOptions, +) -> ConfiguredModelArchitecture: + cfg_donor = donor.config(trust_remote_code=common_options.trust_remote_code) + cfg_out = model.config(trust_remote_code=common_options.trust_remote_code) + arch_info_out = arch_info_for_config(cfg_out) + if arch_info_out is None: + arch_info_out = infer_architecture_info((model,), model, common_options) + set_config_value( + cfg_out, arch_info_out.vocab_size_config_key or "vocab_size", new_vocab_size + ) + for key in [ + "pad_token_id", + "eos_token_id", + "bos_token_id", + "unk_token_id", + "mask_token_id", + "padding_side", + ]: + if hasattr(cfg_donor, key): + set_config_value(cfg_out, key, getattr(cfg_donor, key)) + return ConfiguredModelArchitecture(info=arch_info_out, config=cfg_out) + + +def john_hewitt_init(orig_embed: mindspore.Tensor, num_new_tokens: int) -> mindspore.Tensor: + orig_embed_f32 = orig_embed.to(mindspore.float32) + mean = orig_embed_f32.mean(axis=0) + centered = orig_embed_f32 - mean + covariance = ops.matmul(centered.T, centered) / orig_embed_f32.shape[0] + + try: + ops.cholesky(covariance) + is_pd = True + except Exception as exc: + LOG.debug( + "Cholesky check failed in john_hewitt_init (%s: %s)", + type(exc).__name__, + exc, + ) + is_pd = False + + if not is_pd: + LOG.warning( + "Covariance matrix is not positive definite - falling back to small randn" + ) + return ( + ops.randn( + num_new_tokens, + orig_embed.shape[1], + ).to(orig_embed.dtype) + * 0.02 + ) + + import numpy as np + mean_np = mean.asnumpy() + cov_np = covariance.asnumpy() + samples_np = np.random.multivariate_normal(mean_np, cov_np, size=num_new_tokens) + new_embeds = mindspore.Tensor(samples_np, dtype=orig_embed.dtype) + return new_embeds + + +def compute_new_embeddings( # pylint: disable=too-many-positional-arguments + orig_embed: mindspore.Tensor, + donor_embed: mindspore.Tensor, + orig_vocab: Dict[NormalizedToken, int], + donor_vocab: Dict[NormalizedToken, int], + target_tokens: List[NormalizedToken], + is_lm_head: bool, + token_basis: Optional[Tuple[mindspore.Tensor, mindspore.Tensor]], + orig_tokenizer: transformers.PreTrainedTokenizerBase, + options: TokenSurgeonOptions, +) -> mindspore.Tensor: + assert all(t in donor_vocab for t in target_tokens) + if options.method == ApproximationMethod.MEAN: + mean = orig_embed.mean(axis=0) + return mean.unsqueeze(0).broadcast_to((len(target_tokens), -1)) + elif options.method == ApproximationMethod.ZERO: + return ops.zeros( + (len(target_tokens), orig_embed.shape[1]), + dtype=orig_embed.dtype, + ) + elif options.method == ApproximationMethod.RANDN: + return ops.randn( + len(target_tokens), + orig_embed.shape[1], + ).to(orig_embed.dtype) + elif options.method == ApproximationMethod.JOHN_HEWITT: + return john_hewitt_init(orig_embed, len(target_tokens)) + elif options.method in ( + ApproximationMethod.COMMON_INTERPOLATION, + ApproximationMethod.ORTHOGONAL_MATCHING_PURSUIT, + ApproximationMethod.LANDMARK_PCA, + ApproximationMethod.MATCHING_PURSUIT_ROPE, + ): + shared_vocab = list( + sorted( + set(orig_vocab.keys()) & set(donor_vocab.keys()), + key=lambda x: donor_vocab[x], + ) + ) + donor_shared_embeds = donor_embed[ + mindspore.Tensor([donor_vocab[t] for t in shared_vocab]) + ] + + orig_shared_embeds = orig_embed[ + mindspore.Tensor([orig_vocab[t] for t in shared_vocab]) + ] + res = None + in_donor = None + targets = donor_embed[mindspore.Tensor([donor_vocab[t] for t in target_tokens])] + if options.method == ApproximationMethod.LANDMARK_PCA: + return landmark_pca_approximate( + targets, + donor_shared_embeds, + orig_shared_embeds, + ) + elif options.method == ApproximationMethod.COMMON_INTERPOLATION: + indices, coeffs = common_interp_approximate( + targets, + donor_shared_embeds, + k=options.k, + metric=( + DistanceMetric.COSINE + if options.cosine_similarity + else DistanceMetric.EUCLIDEAN + ), + weight_scheme=options.weight_scheme, + ) + elif options.method == ApproximationMethod.MATCHING_PURSUIT_ROPE: + model_config = options.model.config(trust_remote_code=False) + donor_config = options.donor.config(trust_remote_code=False) + indices, coeffs, res, in_donor = batch_mp_rope( + targets, + donor_shared_embeds, + orig_shared_embeds, + k=options.k, + num_heads_a=donor_config.num_attention_heads, + num_heads_b=model_config.num_attention_heads, + a_rope_base=donor_config.rope_theta, + b_rope_base=model_config.rope_theta, + ) + else: + indices, coeffs = batch_omp(targets, donor_shared_embeds, options.k) + + if res is None: + res = ( + ops.bmm( + coeffs.unsqueeze(1), orig_shared_embeds[indices].to(mindspore.float32) + ) + .squeeze(1) + .to(orig_embed.dtype) + ) + return res + elif options.method == ApproximationMethod.SUBWORD: + return subword_approximate( + orig_embed, + target_tokens, + is_lm_head, + orig_tokenizer, + options.subword_method, + ) + elif options.method == ApproximationMethod.SPARSE_TOKEN_BASIS: + assert token_basis is not None, "Token basis must be provided for STB" + donor_basis, orig_basis = token_basis + donor_basis = donor_basis.to(mindspore.float32) + orig_basis = orig_basis.to(mindspore.float32) + target_donor_embeds = donor_embed[ + mindspore.Tensor([donor_vocab[t] for t in target_tokens]) + ].to(mindspore.float32) - donor_embed.mean(axis=0) + + import numpy as np + coeffs_np = np.linalg.lstsq( + donor_basis.T.asnumpy(), + target_donor_embeds.T.asnumpy(), + rcond=None, + )[0].T + coeffs = mindspore.Tensor(coeffs_np, dtype=mindspore.float32) + + if LOG.isEnabledFor(logging.DEBUG): + donor_rt = ops.matmul(coeffs, donor_basis) + err = (donor_rt - target_donor_embeds).norm(axis=1) + err_rel = err / ops.clamp(target_donor_embeds.norm(axis=1), min=1e-6) + sim = ops.cosine_similarity( + donor_rt, target_donor_embeds, dim=1 + ) + LOG.debug(f"Reconstruction error: {err.mean().asnumpy().item():.4f}") + LOG.debug(f"Relative reconstruction error: {err_rel.mean().asnumpy().item():.4f}") + LOG.debug(f"Cosine similarity: {sim.mean().asnumpy().item():.4f}") + + return ops.matmul(coeffs, orig_basis) + orig_embed.mean(axis=0) + else: + raise ValueError(f"Unknown approximation method: {options.method}") + + +def build_embedding_matrix( # pylint: disable=too-many-positional-arguments + weight_info: WeightInfo, + orig_embed: mindspore.Tensor, + donor_embed: mindspore.Tensor, + orig_vocab: Dict[NormalizedToken, int], + donor_vocab: Dict[NormalizedToken, int], + junk_tokens: List[int], + allow_prefix: bool, + allow_byte: bool, + is_lm_head: bool, + options: TokenSurgeonOptions, +) -> mindspore.Tensor: + LOG.info(f"Building new tensor for {weight_info.name}") + stats = TokenAssignmentStats() + out_vocab_size = max(len(donor_vocab), max(donor_vocab.values()) + 1) + orig_rows = orig_embed.shape[0] + donor_rows = donor_embed.shape[0] + safe_orig_vocab = {tok: idx for tok, idx in orig_vocab.items() if idx < orig_rows} + safe_donor_vocab = { + tok: idx for tok, idx in donor_vocab.items() if idx < donor_rows + } + dropped_orig = len(orig_vocab) - len(safe_orig_vocab) + dropped_donor = len(donor_vocab) - len(safe_donor_vocab) + if dropped_orig: + LOG.warning( + "Skipping %d original tokens with ids >= %d rows for %s", + dropped_orig, + orig_rows, + weight_info.name, + ) + if dropped_donor: + LOG.warning( + "Skipping %d donor tokens with ids >= %d rows for %s", + dropped_donor, + donor_rows, + weight_info.name, + ) + oob_donor_ids = [idx for idx in donor_vocab.values() if idx >= donor_rows] + if oob_donor_ids: + junk_tokens = list(sorted(set(junk_tokens).union(oob_donor_ids))) + + if options.method == ApproximationMethod.SPARSE_TOKEN_BASIS: + token_basis = compute_token_basis( + orig_embed, + donor_embed, + safe_orig_vocab, + safe_donor_vocab, + junk_tokens, + options, + ) + else: + token_basis = None + + res = ops.zeros( + (out_vocab_size, orig_embed.shape[1]), + dtype=orig_embed.dtype, + ) + new_tokens = [] + for token, donor_idx in donor_vocab.items(): + if donor_idx >= donor_rows: + continue + if token in safe_orig_vocab: + orig_idx = safe_orig_vocab[token] + res[donor_idx] = orig_embed[orig_idx] + stats.exact_match += 1 + elif ( + allow_byte + and (orig_idx := match_byte_token(token, safe_orig_vocab)) is not None + ): + res[donor_idx] = orig_embed[orig_idx] + stats.byte_match += 1 + elif ( + allow_prefix + and (orig_idx := match_prefix(token, safe_orig_vocab)) is not None + ): + res[donor_idx] = orig_embed[orig_idx] + stats.prefix_match += 1 + else: + new_tokens.append(token) + stats.to_approximate += 1 + + donor_tokenizer = transformers.AutoTokenizer.from_pretrained( + options.donor.model.path, + revision=options.donor.model.revision, + trust_remote_code=True, + ) + orig_tokenizer = transformers.AutoTokenizer.from_pretrained( + options.model.model.path, + revision=options.model.model.revision, + trust_remote_code=True, + ) + + LOG.info(stats.pretty_print()) + if new_tokens: + LOG.info(f"Approximating {len(new_tokens)} tokens") + batch_size = options.batch_size + if batch_size is None or batch_size <= 0: + batch_size = len(new_tokens) + for base_idx in tqdm.tqdm( + range(0, len(new_tokens), batch_size), + desc="Approximating tokens", + ): + new_embeds = compute_new_embeddings( + orig_embed, + donor_embed, + safe_orig_vocab, + safe_donor_vocab, + target_tokens=new_tokens[base_idx : base_idx + batch_size], + is_lm_head=is_lm_head, + token_basis=token_basis, + orig_tokenizer=orig_tokenizer, + options=options, + ) + if options.new_vocab_noise: + new_embeds += ops.randn_like(new_embeds) * options.new_vocab_noise + if options.new_vocab_scale: + new_embeds *= options.new_vocab_scale + for ne_idx, token in enumerate( + new_tokens[base_idx : base_idx + batch_size] + ): + res[donor_vocab[token]] = new_embeds[ne_idx] + if junk_tokens: + LOG.info(f"Zero-initializing {len(junk_tokens)} junk tokens") + for token_id in junk_tokens: + res[token_id] = ops.zeros( + (orig_embed.shape[1],), + dtype=orig_embed.dtype, + ) + return res + + +class AllowMatch(enum.Enum): + LM_HEAD_ONLY = "lm_head" + EMBED_ONLY = "embed" + YES = "yes" + NO = "no" + + +@click.command("wizard-tokensurgeon", cls=PrettyPrintHelp) +@click.argument("model", type=str) +@click.argument("donor", type=str) +@click.argument("out_path", type=str) +@click.option( + "--k", + "-k", + type=int, + default=64, + help="Number of nearest neighbours to use for embedding interpolation", + show_default=True, +) +@click.option( + "--cosine-similarity/--no-cosine-similarity", + "-c/-nc", + is_flag=True, + default=False, + help="Use cosine similarity for nearest neighbour search", + show_default=True, +) +@click.option( + "--approximation-method", + "-a", + type=click.Choice([m.value for m in ApproximationMethod]), + default=ApproximationMethod.ORTHOGONAL_MATCHING_PURSUIT.value, + help="Method for approximating missing tokens", + show_default=True, +) +@click.option( + "--weight-scheme", + "-w", + type=click.Choice([w.value for w in WeightingScheme]), + default=WeightingScheme.DISTANCE_PROPORTIONAL.value, + help="Weighting scheme for common-vocabulary interpolation", + show_default=True, +) +@click.option( + "--subword-method", + "-s", + type=click.Choice([m.value for m in SubwordMethod]), + default=SubwordMethod.MEAN.value, + help="Method for approximating embeddings with subword tokens", + show_default=True, +) +@click.option( + "--batch-size", + type=int, + default=512, + help="Number of tokens to process in each batch. -1 for no batching.", + show_default=True, +) +@click.option( + "--prefix-match", + "-pm", + type=click.Choice([m.value for m in AllowMatch]), + default=AllowMatch.NO.value, + help="Allow prefix match for tokens", + show_default=True, +) +@click.option( + "--byte-match", + "-bm", + type=click.Choice([m.value for m in AllowMatch]), + default=AllowMatch.NO.value, + help="Allow byte match for tokens", + show_default=True, +) +@click.option( + "--magikarp/--no-magikarp", + is_flag=True, + default=False, + help="Filter out poorly trained tokens", + show_default=True, +) +@click.option( + "--new-vocab-noise", + "-nvn", + type=float, + default=None, + help="Add gaussian noise to new vocab embeddings", + show_default=True, +) +@click.option( + "--new-vocab-scale", + "-nvs", + type=float, + default=None, + help="Scale computed new vocab embeddings by this factor", + show_default=True, +) +@add_merge_options +def main( # pylint: disable=too-many-positional-arguments + model: str, + donor: str, + out_path: str, + k: int, + cosine_similarity: bool, + approximation_method: str, + weight_scheme: str, + subword_method: str, + batch_size: Optional[int], + prefix_match: str, + byte_match: str, + magikarp: bool, + new_vocab_noise: Optional[float], + new_vocab_scale: Optional[float], + merge_options: MergeOptions, +): + merge_options.apply_global_options() + logging.warning("This script is experimental and may produce unexpected results.") + options = TokenSurgeonOptions( + model=ModelReference.model_validate(model), + donor=ModelReference.model_validate(donor), + out_path=out_path, + k=k, + cosine_similarity=cosine_similarity, + method=ApproximationMethod(approximation_method), + weight_scheme=WeightingScheme(weight_scheme), + subword_method=SubwordMethod(subword_method), + batch_size=batch_size, + new_vocab_noise=new_vocab_noise, + new_vocab_scale=new_vocab_scale, + ) + prefix_match = AllowMatch(prefix_match) + byte_match = AllowMatch(byte_match) + + cache = LoaderCache() + cache.setup(options=merge_options) + + device = merge_options.device + + arch_info = get_arch_info(options.model, merge_options) + embed_wi, lm_head_wi = get_embedding_info(arch_info) + orig_vocab, orig_embed, orig_lm_head = get_stuff( + options.model, merge_options, arch_info=arch_info, device=device + ) + donor_vocab, donor_embed, donor_lm_head = get_stuff( + options.donor, merge_options, arch_info=None, get_tied=True, device=device + ) + + if magikarp: + LOG.debug("Finding well-trained tokens in original model") + well_trained_orig_tokens = set( + well_trained_tokens( + orig_vocab, + orig_embed, + orig_lm_head, + ) + ) + LOG.debug("Finding well-trained tokens in donor model") + well_trained_donor_tokens = set( + well_trained_tokens( + donor_vocab, + donor_embed, + donor_lm_head, + ) + ) + common_well_trained_tokens = ( + well_trained_orig_tokens & well_trained_donor_tokens + ) + LOG.info(f"Found {len(common_well_trained_tokens)} common well-trained tokens") + orig_vocab = { + tok: idx + for tok, idx in orig_vocab.items() + if tok in common_well_trained_tokens + } + junk_tokens = [ + idx + for tok, idx in donor_vocab.items() + if (tok not in well_trained_donor_tokens) + and (tok not in well_trained_orig_tokens) + ] + else: + junk_tokens = [] + + if orig_embed is not None: + if donor_embed is None: + raise RuntimeError( + f"Missing tensor {embed_wi.name} in model {options.donor}" + ) + new_embed = build_embedding_matrix( + embed_wi, + orig_embed, + donor_embed, + orig_vocab=orig_vocab, + donor_vocab=donor_vocab, + junk_tokens=junk_tokens, + allow_prefix=prefix_match in (AllowMatch.YES, AllowMatch.LM_HEAD_ONLY), + allow_byte=byte_match in (AllowMatch.YES, AllowMatch.LM_HEAD_ONLY), + is_lm_head=False, + options=options, + ) + else: + if not embed_wi.optional: + raise RuntimeError( + f"Missing tensor {embed_wi.name} in model {options.model}" + ) + new_embed = None + + if orig_lm_head is not None: + if donor_lm_head is None: + raise RuntimeError( + f"Missing tensor {lm_head_wi.name} in model {options.donor}" + ) + new_lm_head = build_embedding_matrix( + lm_head_wi, + orig_lm_head, + donor_lm_head, + orig_vocab=orig_vocab, + donor_vocab=donor_vocab, + junk_tokens=junk_tokens, + allow_prefix=prefix_match in (AllowMatch.YES, AllowMatch.EMBED_ONLY), + allow_byte=byte_match in (AllowMatch.YES, AllowMatch.EMBED_ONLY), + is_lm_head=True, + options=options, + ) + else: + if not lm_head_wi.optional: + raise RuntimeError( + f"Missing tensor {lm_head_wi.name} in model {options.model}" + ) + new_lm_head = None + + new_vocab_size = None + if new_embed is not None: + new_vocab_size = new_embed.shape[0] + elif new_lm_head is not None: + new_vocab_size = new_lm_head.shape[0] + LOG.info(f"Saving new model to {out_path}") + out_arch_info = get_out_arch_info( + options.model, options.donor, new_vocab_size, merge_options + ) + writer = TensorWriter( + out_path, + max_shard_size=merge_options.out_shard_size, + output_format=merge_options.output_format, + use_async=merge_options.async_write, + max_write_threads=merge_options.write_threads, + ) + for weight_info in tqdm.tqdm(out_arch_info.all_weights(), desc="Saving weights"): + if weight_info.name == embed_wi.name: + tensor = new_embed + elif lm_head_wi is not None and weight_info.name == lm_head_wi.name: + tensor = new_lm_head + else: + tensor = cache.get(options.model).get_tensor( + weight_info.name, aliases=weight_info.aliases, raise_on_missing=False + ) + if tensor is None: + if weight_info.optional: + continue + raise RuntimeError( + f"Missing tensor {weight_info.name} in model {options.model}" + ) + writer.save_tensor(weight_info.name, tensor, clone=merge_options.clone_tensors) + writer.finalize() + out_arch_info.config.save_pretrained(out_path) + + tokenizer_out = transformers.AutoTokenizer.from_pretrained( + options.donor.model.path, + revision=options.donor.model.revision, + trust_remote_code=merge_options.trust_remote_code, + ) + tokenizer_out.save_pretrained(out_path) + LOG.info("Done!") + + +if __name__ == "__main__": + main() diff --git a/src/mindnlp/wizard/merge/sparsify.py b/src/mindnlp/wizard/merge/sparsify.py new file mode 100644 index 000000000..3c81b2a07 --- /dev/null +++ b/src/mindnlp/wizard/merge/sparsify.py @@ -0,0 +1,225 @@ +# Copyright (C) 2025 Arcee AI +# SPDX-License-Identifier: LGPL-3.0-only +# Modified for MindSpore/Ascend NPU by MindNLP Wizard contributors. +# + +from enum import Enum +from typing import Optional + +import mindspore +from mindspore import ops +import numpy as np + +from .safe_ops import safe_abs, safe_norm, safe_sum + + +class SparsificationMethod(str, Enum): + magnitude = "magnitude" + random = "random" + magnitude_outliers = "magnitude_outliers" + della_magprune = "della_magprune" + + +class RescaleNorm(str, Enum): + l1 = "l1" + l2 = "l2" + linf = "linf" + + +def rescaled_masked_tensor( + tensor: mindspore.Tensor, + mask: mindspore.Tensor, + norm: Optional[RescaleNorm], + eps: float = 1e-7, +) -> mindspore.Tensor: + """Apply a mask to a tensor and rescale to match the original tensor norm. + + Returns the result in ``work_dtype`` (float32 on CPU) so that callers + can chain further arithmetic without intermediate half-precision overflow. + The caller is responsible for the final cast back to the output dtype. + + Args: + tensor (mindspore.Tensor): Input tensor. + mask (mindspore.Tensor): Mask to apply. + norm (RescaleNorm): Which norm to match (l1, l2, linf). + eps (float): Tolerance for small norms to avoid division by zero. + """ + work_dtype = mindspore.float32 + tensor_work = tensor.astype(work_dtype) + mask_work = mask.astype(work_dtype) + masked = ops.mul(tensor_work, mask_work) + if norm is None: + return masked + elif norm == RescaleNorm.l1: + before_scale = safe_sum( + safe_abs(tensor_work, out_dtype=work_dtype, op_name="sparsify.l1_before_abs"), + out_dtype=work_dtype, + op_name="sparsify.l1_before_sum", + ) + after_scale = safe_sum( + safe_abs(masked, out_dtype=work_dtype, op_name="sparsify.l1_after_abs"), + out_dtype=work_dtype, + op_name="sparsify.l1_after_sum", + ) + elif norm == RescaleNorm.l2: + before_scale = safe_norm( + tensor_work, out_dtype=mindspore.float32, op_name="sparsify.l2_before" + ) + after_scale = safe_norm( + masked, out_dtype=mindspore.float32, op_name="sparsify.l2_after" + ) + elif norm == RescaleNorm.linf: + before_scale = ops.max( + safe_abs(tensor_work, out_dtype=work_dtype, op_name="sparsify.linf_before_abs") + )[0] + after_scale = ops.max( + safe_abs(masked, out_dtype=work_dtype, op_name="sparsify.linf_after_abs") + )[0] + else: + raise NotImplementedError(norm) + before_v = float(before_scale.astype(mindspore.float32).asnumpy().item()) + after_v = float(after_scale.astype(mindspore.float32).asnumpy().item()) + if before_v < eps or after_v < eps: + return masked + return ops.mul(masked, ops.div(before_scale, after_scale)) + + +def magnitude( + tensor: mindspore.Tensor, density: float, rescale_norm: Optional[RescaleNorm] = None +) -> mindspore.Tensor: + """Masks out the smallest values, retaining a proportion of `density`.""" + if density >= 1: + return tensor + + k = int(density * tensor.numel()) + assert k > 0, "not gonna zero out the whole tensor buddy" + + w = tensor.astype(mindspore.float32).abs().reshape(-1) + topk_indices = ops.argsort(w, descending=True)[:k].asnumpy() + + mask_np = np.zeros(tensor.numel(), dtype=np.float32) + mask_np[topk_indices] = 1.0 + mask = mindspore.Tensor(mask_np).reshape(tensor.shape) + + return rescaled_masked_tensor(tensor, mask, rescale_norm) + + +def magnitude_outliers( + tensor: mindspore.Tensor, + density: float, + rescale_norm: Optional[RescaleNorm] = None, + gamma: float = 0.01, +): + """Masks out smallest values in addition to large outliers. + + The `gamma` proportion of the largest weights are first removed, then the + smallest weights are removed to achieve the desired density. + + Args: + tensor (mindspore.Tensor): The tensor to sparsify. + density (float): The proportion of weights to retain. + gamma (float): Percent of largest weights to remove. + """ + if density >= 1: + return tensor + + num_elems = tensor.numel() + target_n = int(density * num_elems) + n_top = int(gamma * num_elems) + n_bot = num_elems - target_n - n_top + + if n_bot < 0: + # cut down on the number of large weights to remove in + # order to hit the target density + n_top += n_bot + n_bot = 0 + + w = tensor.astype(mindspore.float32).abs().reshape(-1) + indices = ops.sort(w, descending=False)[1].asnumpy() + + mask_np = np.zeros(tensor.numel(), dtype=np.float32) + if n_top > 0: + keep = indices[n_bot:-n_top] + else: + keep = indices[n_bot:] + mask_np[keep] = 1.0 + mask = mindspore.Tensor(mask_np).reshape(tensor.shape) + + return rescaled_masked_tensor(tensor, mask, rescale_norm) + + +def bernoulli( + tensor: mindspore.Tensor, density: float, rescale_norm: Optional[RescaleNorm] = None +) -> mindspore.Tensor: + if density >= 1: + return tensor + + work_dtype = mindspore.float32 + probs = ops.full(tensor.shape, density, dtype=mindspore.float32) + rand = mindspore.Tensor(np.random.rand(*tensor.shape).astype(np.float32)) + mask = (rand < probs).astype(work_dtype) + return rescaled_masked_tensor(tensor.astype(work_dtype), mask, rescale_norm) + + +def della_magprune( + tensor: mindspore.Tensor, + density: float, + epsilon: float, + rescale_norm: Optional[RescaleNorm] = None, +) -> mindspore.Tensor: + if density >= 1: + return tensor + if density <= 0: + return ops.zeros_like(tensor) + orig_shape = tensor.shape + + if density + epsilon >= 1 or density - epsilon <= 0: + raise ValueError( + "Epsilon must be chosen such that density +/- epsilon is in (0, 1)" + ) + + work_dtype = mindspore.float32 + + if len(tensor.shape) < 2: + tensor = tensor.unsqueeze(0) + magnitudes = safe_abs( + tensor, out_dtype=mindspore.float32, op_name="sparsify.della.abs" + ) + + sorted_indices = ops.argsort(magnitudes, axis=1, descending=False) + ranks = ops.argsort(sorted_indices, axis=1).astype(work_dtype) + 1 + + min_ranks, _ = ops.min(ranks, 1, True) + max_ranks, _ = ops.max(ranks, 1, True) + rank_norm = ((ranks - min_ranks) / (max_ranks - min_ranks)).clamp(0, 1) + probs = ((density - epsilon) + rank_norm * 2 * epsilon).astype(mindspore.float32) + # MindSpore Ascend backend may lack a Bernoulli adapter on some versions. + # Use an equivalent uniform-sampling implementation instead. + rand = mindspore.Tensor(np.random.rand(*probs.shape).astype(np.float32)) + mask = (rand < probs).astype(work_dtype) + + return rescaled_masked_tensor(tensor.astype(work_dtype), mask, rescale_norm).reshape(orig_shape) + + +def sparsify( # pylint: disable=too-many-positional-arguments + tensor: mindspore.Tensor, + density: float, + method: SparsificationMethod, + gamma: float = 0, + epsilon: float = 0, + rescale_norm: Optional[RescaleNorm] = None, +) -> mindspore.Tensor: + if method == SparsificationMethod.magnitude: + return magnitude(tensor, density=density, rescale_norm=rescale_norm) + elif method == SparsificationMethod.random: + return bernoulli(tensor, density=density, rescale_norm=rescale_norm) + elif method == SparsificationMethod.magnitude_outliers: + return magnitude_outliers( + tensor, density=density, rescale_norm=rescale_norm, gamma=gamma + ) + elif method == SparsificationMethod.della_magprune: + return della_magprune( + tensor, density=density, epsilon=epsilon, rescale_norm=rescale_norm + ) + else: + raise NotImplementedError(method) diff --git a/src/mindnlp/wizard/merge/tokenizer/__init__.py b/src/mindnlp/wizard/merge/tokenizer/__init__.py new file mode 100644 index 000000000..4336a5e6a --- /dev/null +++ b/src/mindnlp/wizard/merge/tokenizer/__init__.py @@ -0,0 +1,17 @@ +# Copyright (C) 2025 Arcee AI +# SPDX-License-Identifier: LGPL-3.0-only +# Modified for MindSpore/Ascend NPU by MindNLP Wizard contributors. +# Ported from MergeKit to MindSpore for Wizard + +from . import normalization +from .build import BuildTokenizer, TokenizerInfo +from .config import TokenizerConfig +from .embed import PermutedEmbeddings + +__all__ = [ + "BuildTokenizer", + "TokenizerInfo", + "TokenizerConfig", + "PermutedEmbeddings", + "normalization", +] diff --git a/src/mindnlp/wizard/merge/tokenizer/build.py b/src/mindnlp/wizard/merge/tokenizer/build.py new file mode 100644 index 000000000..c84f04d91 --- /dev/null +++ b/src/mindnlp/wizard/merge/tokenizer/build.py @@ -0,0 +1,294 @@ +# Copyright (C) 2025 Arcee AI +# SPDX-License-Identifier: LGPL-3.0-only +# Modified for MindSpore/Ascend NPU by MindNLP Wizard contributors. +# Ported from MergeKit to MindSpore for Wizard + +import json +import logging +import tempfile +from typing import Dict, List, Optional, Tuple, Union + +import tokenizers +import tokenizers.models +import tqdm +import transformers +from pydantic import BaseModel +from typing_extensions import Literal + +from ..architecture import arch_info_for_config +from ..common import ModelPath, ModelReference, get_config_value +from ..graph import Task + +LOG = logging.getLogger(__name__) + + +def get_vocab_size(model_path: ModelPath, trust_remote_code: bool) -> Optional[int]: + try: + cfg = transformers.AutoConfig.from_pretrained( + model_path.path, + revision=model_path.revision, + trust_remote_code=trust_remote_code, + ) + arch_info = arch_info_for_config(cfg) + key = "vocab_size" + if arch_info is not None: + key = arch_info.vocab_size_config_key or "vocab_size" + return get_config_value(cfg, key) + except Exception as e: + LOG.warning(f"Unable to get vocab size for {model_path}", exc_info=e) + + return None + + +def get_stripped_tokenizer( + path: ModelPath, trust_remote_code: bool = False +) -> transformers.PreTrainedTokenizerFast: + """ + Return a tokenizer for a model that only contains used tokens. + + Strips any tokens with indices >= model.vocab_size. + """ + tokenizer = transformers.AutoTokenizer.from_pretrained( + path.path, + revision=path.revision, + trust_remote_code=trust_remote_code, + use_fast=True, + ) + vocab_size = get_vocab_size(path, trust_remote_code=trust_remote_code) or len( + tokenizer.get_vocab() + ) + + unused_toks = [ + tok for tok, idx in tokenizer.get_vocab().items() if idx >= vocab_size + ] + if not unused_toks: + return tokenizer + + if not tokenizer.is_fast: + raise RuntimeError( + f"Model {path} has unused tokens and does not support fast " + "tokenizer - can not be used in tokenizer merge" + ) + + tok_dict = json.loads(tokenizer._tokenizer.to_str()) + if tok_dict["model"]["type"] != "BPE": + raise RuntimeError( + f"Tokenizer for {path} has type {tok_dict['model']['type']}, " + "but only BPE is currently supported for tokenizer merge" + ) + + tok_dict["added_tokens"] = [ + e for e in tok_dict["added_tokens"] if e["id"] < vocab_size + ] + + for tok in unused_toks: + if tok in tok_dict["model"]["vocab"]: + del tok_dict["model"]["vocab"][tok] + + def _keep_merge(m): + if isinstance(m, str) and m.count(" ") == 1: + toks = m.split(" ") + elif isinstance(m, list): + toks = m + else: + raise RuntimeError(f"Unexpected merge format: {repr(m)} ({type(m)})") + for tok in toks: + if tok in unused_toks: + return False + return True + + tok_dict["model"]["merges"] = [ + e for e in tok_dict["model"]["merges"] if _keep_merge(e) + ] + tokenizer._tokenizer = tokenizers.Tokenizer.from_str(json.dumps(tok_dict)) + return tokenizer + + +def build_union_tokenizer( + base_tok: transformers.PreTrainedTokenizerBase, + tokenizers_map: Dict[ModelReference, transformers.PreTrainedTokenizerBase], + trust_remote_code: bool = False, +) -> transformers.PreTrainedTokenizerBase: + out_added_tokens = {} + out_vocab = {} + + warned_added_tokens = set() + + for model, tokenizer in tokenizers_map.items(): + vocab_size = ( + get_vocab_size(model.model, trust_remote_code=trust_remote_code) + or tokenizer.vocab_size + ) + added_tokens = tokenizer.added_tokens_decoder + + vocab = tokenizer.get_vocab() + for tok, idx in vocab.items(): + if idx >= vocab_size: + LOG.warning( + f"Token {repr(tok)} present in {str(model)} tokenizer but >= vocab_size" + ) + continue + if tok in added_tokens: + continue + + if tok not in out_vocab: + out_vocab[tok] = len(out_vocab) + + for tok_idx, info in tokenizer.added_tokens_decoder.items(): + tok = info.content + if tok_idx >= vocab_size: + continue + + if tok in out_added_tokens: + if (out_added_tokens[tok] != info) and tok not in warned_added_tokens: + LOG.warning( + f"Token '{tok}' added with multiple different settings, using first" + ) + warned_added_tokens.add(tok) + + continue + out_added_tokens[tok] = info + + with tempfile.TemporaryDirectory() as p: + base_tok.save_pretrained(p, legacy_format=False, safe_serialization=True) + res = transformers.AutoTokenizer.from_pretrained( + p, use_fast=True, trust_remote_code=trust_remote_code + ) + + orig_base_vocab = base_tok.get_vocab() + for tok in out_vocab: + if tok in out_added_tokens: + continue + + if tok not in orig_base_vocab: + res.add_tokens(tok) + + for info in out_added_tokens.values(): + res.add_tokens(info) + return res + + +class TokenizerInfo(BaseModel, arbitrary_types_allowed=True): + tokenizer: transformers.PreTrainedTokenizerBase + permutations: Dict[ModelReference, Dict[int, int]] + original_vocabs: Dict[ModelReference, Dict[str, int]] + + +def build_tokenizer( + base_model: Optional[ModelReference], + referenced_models: List[ModelReference], + tokenizer_source: Union[Literal["union"], Literal["base"], ModelReference], + trust_remote_code: bool, + add_tokens: Optional[List[str]] = None, +) -> TokenizerInfo: + if base_model is None: + base_model = referenced_models[0] + if base_model is None: + raise RuntimeError("No models referenced") + + tokenizer_base = get_stripped_tokenizer( + base_model.model, trust_remote_code=trust_remote_code + ) + + LOG.info("Loading tokenizers") + tokenizers_loaded = {base_model: tokenizer_base} + for model in referenced_models: + if model == base_model: + continue + + try: + model_tok = transformers.AutoTokenizer.from_pretrained( + model.model.path, + revision=model.model.revision, + trust_remote_code=trust_remote_code, + ) + except Exception as e: + LOG.error(e) + LOG.warning( + f"Unable to load tokenizer for {model}. Assuming same as {base_model}." + ) + continue + tokenizers_loaded[model] = model_tok + + LOG.info("Building output tokenizer") + if isinstance(tokenizer_source, ModelReference): + tokenizer_out = transformers.AutoTokenizer.from_pretrained( + tokenizer_source.model.path, + revision=tokenizer_source.model.revision, + trust_remote_code=trust_remote_code, + ) + elif tokenizer_source == "base": + tokenizer_out = tokenizer_base + elif tokenizer_source == "union": + tokenizer_out = build_union_tokenizer( + tokenizer_base, tokenizers_loaded, trust_remote_code=trust_remote_code + ) + else: + raise RuntimeError(f"Unimplemented tokenizer source: {tokenizer_source}") + + for tok in add_tokens: + tokenizer_out.add_tokens(tok) + + vocab_out = tokenizer_out.get_vocab() + + LOG.info("Building permutations") + permutations = {} + for model in ( + pbar := tqdm.tqdm(referenced_models, desc="Building tokenizer permutations") + ): + if model in tokenizers_loaded: + model_vocab = tokenizers_loaded[model].get_vocab() + else: + model_vocab = tokenizers_loaded[base_model].get_vocab() + + vocab_size = get_vocab_size(model.model, trust_remote_code=trust_remote_code) + if vocab_size is None: + vocab_size = len(model_vocab) + + p = {} + for tok in vocab_out: + new_idx = vocab_out[tok] + if tok not in model_vocab: + p[new_idx] = -1 + continue + + orig_idx = model_vocab[tok] + if orig_idx >= vocab_size: + LOG.warning( + f"{model} token {repr(tok)} has index {orig_idx}>{vocab_size - 1} (padding?)" + ) + continue + + p[new_idx] = orig_idx + + permutations[model] = p + + del pbar + + return TokenizerInfo( + tokenizer=tokenizer_out, + permutations=permutations, + original_vocabs={ + model: tok.get_vocab() for model, tok in tokenizers_loaded.items() + }, + ) + + +class BuildTokenizer(Task[TokenizerInfo]): + base_model: Optional[ModelReference] + referenced_models: Tuple[ModelReference, ...] + tokenizer_source: Union[Literal["union"], Literal["base"], ModelReference] + add_tokens: Optional[Tuple[str, ...]] + trust_remote_code: bool = False + + def arguments(self) -> Dict[str, Task]: + return {} + + def execute(self, **_kwargs) -> TokenizerInfo: + return build_tokenizer( + base_model=self.base_model, + referenced_models=self.referenced_models, + tokenizer_source=self.tokenizer_source, + trust_remote_code=self.trust_remote_code, + add_tokens=self.add_tokens, + ) diff --git a/src/mindnlp/wizard/merge/tokenizer/config.py b/src/mindnlp/wizard/merge/tokenizer/config.py new file mode 100644 index 000000000..d0281402d --- /dev/null +++ b/src/mindnlp/wizard/merge/tokenizer/config.py @@ -0,0 +1,42 @@ +# Copyright (C) 2025 Arcee AI +# SPDX-License-Identifier: LGPL-3.0-only +# Modified for MindSpore/Ascend NPU by MindNLP Wizard contributors. +# Ported from MergeKit to MindSpore for Wizard + +from typing import Dict, Optional, Union + +import pydantic +from pydantic import BaseModel +from typing_extensions import Literal + +from ..common import ModelReference + + +class ModelTokenEmbedding(BaseModel, frozen=True): + kind: Literal["model_token"] + model: ModelReference + token_id: Optional[int] = None + token: Optional[str] = None + + @pydantic.model_validator(mode="after") + def validate_token(self): + if self.token_id is None and self.token is None: + raise ValueError("token_id or token must be specified") + if self.token_id is not None and self.token is not None: + raise ValueError("only one of token_id or token may be specified") + return self + + +class ZeroEmbedding(BaseModel, frozen=True): + kind: Literal["zero"] + + +class TokenEmbeddingConfig(BaseModel, frozen=True): + source: Union[ModelTokenEmbedding, ZeroEmbedding, ModelReference, None] = None + force: bool = False + + +class TokenizerConfig(BaseModel, frozen=True): + source: Union[ModelReference, Literal["union"], Literal["base"]] = "union" + tokens: Optional[Dict[str, TokenEmbeddingConfig]] = None + pad_to_multiple_of: Optional[int] = None diff --git a/src/mindnlp/wizard/merge/tokenizer/embed.py b/src/mindnlp/wizard/merge/tokenizer/embed.py new file mode 100644 index 000000000..feef80f03 --- /dev/null +++ b/src/mindnlp/wizard/merge/tokenizer/embed.py @@ -0,0 +1,181 @@ +# Copyright (C) 2025 Arcee AI +# SPDX-License-Identifier: LGPL-3.0-only +# Modified for MindSpore/Ascend NPU by MindNLP Wizard contributors. +# Ported from MergeKit to MindSpore for Wizard + +import logging +from typing import Dict, Optional + +import mindspore + +from ..common import ImmutableMap, ModelReference +from ..graph import Task +from ..io.tasks import GatherTensors +from .build import BuildTokenizer, TokenizerInfo +from .config import ( + ModelTokenEmbedding, + TokenEmbeddingConfig, + ZeroEmbedding, +) + + +class PermutedEmbeddings(Task[Dict[ModelReference, mindspore.Tensor]]): + gather_tensors: GatherTensors + tokenizer_task: BuildTokenizer + tokens: Optional[ImmutableMap[str, TokenEmbeddingConfig]] + pad_to_multiple_of: Optional[int] + base_model: Optional[ModelReference] + + def arguments(self) -> Dict[str, Task]: + return {"tokenizer_info": self.tokenizer_task, "tensors": self.gather_tensors} + + def execute( + self, + tokenizer_info: TokenizerInfo, + tensors: Dict[ModelReference, mindspore.Tensor], + ) -> Dict[ModelReference, mindspore.Tensor]: + tokenizer = tokenizer_info.tokenizer + permutations = tokenizer_info.permutations + + models = set(tensors.keys()) + if self.base_model: + models.add(self.base_model) + models = list(models) + + vocab = tokenizer.get_vocab() + vocab_size = len(vocab) + if self.pad_to_multiple_of and vocab_size % self.pad_to_multiple_of: + vocab_size = ( + vocab_size // self.pad_to_multiple_of + 1 + ) * self.pad_to_multiple_of + embed_size = tensors[models[0]].shape[1] + assert all( + t.shape[1] == embed_size for t in tensors.values() + ), "Embedding sizes must match" + + dtype = tensors[models[0]].dtype + + token_configs = dict(**(self.tokens or {})) + tokens_to_average = self.assign_embedding_sources( + permutations, models, vocab, token_configs + ) + + default_embeds = {} + for token, token_id in vocab.items(): + embed = mindspore.ops.zeros(embed_size, dtype=dtype) + if token in tokens_to_average: + count = 0 + for model in models: + p = permutations[model] + if p[token_id] < 0: + continue + embed += tensors[model][p[token_id]] + count += 1 + embed /= count + elif cfg := token_configs.get(token, None): + cfg: TokenEmbeddingConfig + embed = self.compute_default_embedding( + tokenizer_info, tensors, permutations, token, token_id, cfg + ) + else: + continue + default_embeds[token] = embed + + result = {} + for model in models: + p = permutations[model] + old_embed = tensors[model] + new_embed = mindspore.ops.zeros( + (vocab_size, embed_size), dtype=dtype + ) + for token, token_id in vocab.items(): + force = False + if token in token_configs: + force = token_configs[token].force + + if p[token_id] >= 0 and not force: + new_embed[token_id, :] = old_embed[p[token_id]] + elif token in default_embeds: + new_embed[token_id, :] = default_embeds[token] + else: + logging.error( + f"No embedding for token {repr(token)} in model {model}!" + ) + + if vocab_size > len(vocab): + avg_embed = mindspore.ops.mean(new_embed[: len(vocab), :], axis=0) + new_embed[len(vocab) :, :] = avg_embed + result[model] = new_embed + + return result + + def assign_embedding_sources( + self, + permutations: Dict[ModelReference, Dict[int, int]], + models: list[ModelReference], + vocab: Dict[str, int], + token_configs: Dict[str, TokenEmbeddingConfig], + ): + permutation_list = [permutations[model] for model in models] + + tokens_to_average = set() + for token, token_id in vocab.items(): + if token in token_configs: + continue + + has_token = [p[token_id] >= 0 for p in permutation_list] + num_present = sum(int(x) for x in has_token) + if num_present == 1: + donor_model = models[has_token.index(True)] + token_configs[token] = TokenEmbeddingConfig(source=donor_model) + continue + + if num_present == 0: + token_configs[token] = TokenEmbeddingConfig(source=ZeroEmbedding()) + logging.warning(f"Token {repr(token)} not found in any model") + continue + + if num_present > 0 and self.base_model is not None: + if permutations[self.base_model][token_id] >= 0: + token_configs[token] = TokenEmbeddingConfig(source=self.base_model) + continue + + tokens_to_average.add(token) + return tokens_to_average + + def compute_default_embedding( # pylint: disable=too-many-positional-arguments + self, + tokenizer_info: TokenizerInfo, + tensors: Dict[ModelReference, mindspore.Tensor], + permutations: Dict[ModelReference, Dict[int, int]], + token: str, + token_id: int, + cfg: TokenEmbeddingConfig, + ) -> mindspore.Tensor: + if isinstance(cfg.source, ZeroEmbedding): + pass + elif isinstance(cfg.source, ModelTokenEmbedding): + model = cfg.source.model + assert ( + model in permutations + ), f"Model {model} referenced but not part of merge" + p = permutations[model] + src_token_id = cfg.source.token_id + if src_token_id is None: + src_token = cfg.source.token + assert ( + src_token in tokenizer_info.original_vocabs[model] + ), f"Token {repr(src_token)} not found in model {model}" + src_token_id = tokenizer_info.original_vocabs[model][src_token] + assert ( + src_token_id >= 0 and src_token_id < tensors[model].shape[0] + ), f"Token ID {src_token_id} out of range for model {model}" + embed = tensors[model][src_token_id] + elif isinstance(cfg.source, ModelReference): + model = cfg.source + p = permutations[model] + assert p[token_id] >= 0, f"Token {repr(token)} not found in model {model}" + embed = tensors[model][p[token_id]] + else: + raise NotImplementedError(cfg) + return embed diff --git a/src/mindnlp/wizard/merge/tokenizer/normalization.py b/src/mindnlp/wizard/merge/tokenizer/normalization.py new file mode 100644 index 000000000..37b02e2cf --- /dev/null +++ b/src/mindnlp/wizard/merge/tokenizer/normalization.py @@ -0,0 +1,124 @@ +# Copyright (C) 2025 Arcee AI +# SPDX-License-Identifier: LGPL-3.0-only +# Modified for MindSpore/Ascend NPU by MindNLP Wizard contributors. +# Ported from MergeKit to MindSpore for Wizard + +import enum +import logging +from typing import Dict, Generator, List, Tuple, Union + +import transformers +from typing_extensions import TypeAlias + +LOG = logging.getLogger(__name__) + + +class TokenMarker(enum.Enum): + SPECIAL = "special" + WORD_START = "word_start" + + +NormalizedToken: TypeAlias = Union[str, Tuple[TokenMarker, str]] + + +def normalize_token( + token: str, + special_tokens_map: Dict[str, Union[str, List[str]]], + word_start_prefix: str = "▁", +) -> NormalizedToken: + """ + Normalize a token for comparison. + """ + if token.startswith(word_start_prefix): + return (TokenMarker.WORD_START, token[len(word_start_prefix) :]) + + for special_token_type, values in special_tokens_map.items(): + if isinstance(values, str): + values = [values] + if token in values: + return (TokenMarker.SPECIAL, special_token_type) + return token + + +def unnormalize_token(token: NormalizedToken) -> str: + if isinstance(token, tuple): + if token[0] == TokenMarker.WORD_START: + return " " + token[1] + return token[1] + return token + + +def token_prefixes( + token: NormalizedToken, allow_whitespace: bool = False +) -> Generator[NormalizedToken, None, None]: + """Yield potential prefixes of a token.""" + marker = None + if isinstance(token, tuple): + marker, token = token + + if marker == TokenMarker.SPECIAL: + return + for i in range(len(token) - 1, 0, -1): + prefix = token[:i] + if not allow_whitespace and not prefix.strip(): + break + if marker is not None: + yield (marker, prefix) + else: + yield prefix + + +def normalized_vocabulary( + tokenizer: transformers.PreTrainedTokenizerBase, +) -> Dict[NormalizedToken, int]: + """ + Get a normalized vocabulary for a tokenizer. + + Attempts to handle word start prefixes and special tokens in a consistent way. + """ + gpt2_style = [ + transformers.GPT2Tokenizer, + transformers.GPT2TokenizerFast, + transformers.OpenAIGPTTokenizer, + transformers.OpenAIGPTTokenizerFast, + ] + for candidate in ["Qwen2Tokenizer", "Qwen2TokenizerFast"]: + if hasattr(transformers, candidate): + gpt2_style.append(getattr(transformers, candidate)) + + sp_style = [ + transformers.LlamaTokenizer, + transformers.LlamaTokenizerFast, + transformers.T5Tokenizer, + transformers.T5TokenizerFast, + ] + for candidate in ["GemmaTokenizer", "GemmaTokenizerFast"]: + if hasattr(transformers, candidate): + sp_style.append(getattr(transformers, candidate)) + + vocab = tokenizer.get_vocab() + if isinstance( + tokenizer, + tuple(gpt2_style), + ): + word_start_prefix = "Ġ" + elif isinstance( + tokenizer, + tuple(sp_style), + ): + if "Ġhello" in vocab: + word_start_prefix = "Ġ" + else: + word_start_prefix = "▁" + else: + LOG.warning("Unknown tokenizer type - assuming 'Ġ' word start prefix") + word_start_prefix = "Ġ" + + return { + normalize_token( + token, + special_tokens_map=tokenizer.special_tokens_map, + word_start_prefix=word_start_prefix, + ): i + for token, i in vocab.items() + } diff --git a/src/mindnlp/wizard/merge/tokensurgeon/__init__.py b/src/mindnlp/wizard/merge/tokensurgeon/__init__.py new file mode 100644 index 000000000..5fc4490cb --- /dev/null +++ b/src/mindnlp/wizard/merge/tokensurgeon/__init__.py @@ -0,0 +1,28 @@ +# Copyright (C) 2025 Arcee AI +# SPDX-License-Identifier: LGPL-3.0-only +# Modified for MindSpore/Ascend NPU by MindNLP Wizard contributors. +# Ported from MergeKit to MindSpore for Wizard + +from .common_interpolation import ( + DistanceMetric, + WeightingScheme, + common_interp_approximate, +) +from .magikarp import well_trained_tokens +from .omp import batch_mp_rope, batch_omp +from .pca import landmark_pca_approximate +from .subword import SubwordMethod, subword_approximate +from .token_basis import compute_token_basis + +__all__ = [ + "common_interp_approximate", + "DistanceMetric", + "WeightingScheme", + "batch_omp", + "batch_mp_rope", + "SubwordMethod", + "subword_approximate", + "well_trained_tokens", + "compute_token_basis", + "landmark_pca_approximate", +] diff --git a/src/mindnlp/wizard/merge/tokensurgeon/common_interpolation.py b/src/mindnlp/wizard/merge/tokensurgeon/common_interpolation.py new file mode 100644 index 000000000..7c21b4440 --- /dev/null +++ b/src/mindnlp/wizard/merge/tokensurgeon/common_interpolation.py @@ -0,0 +1,143 @@ +# Copyright (C) 2025 Arcee AI +# SPDX-License-Identifier: LGPL-3.0-only +# Modified for MindSpore/Ascend NPU by MindNLP Wizard contributors. +# Ported from MergeKit to MindSpore for Wizard + +import enum +import logging +from typing import Optional, Tuple + +import mindspore # pylint: disable=import-error +from mindspore import ops # pylint: disable=import-error +from mindspore import Tensor # pylint: disable=import-error + +LOG = logging.getLogger(__name__) + + +class DistanceMetric(enum.Enum): + EUCLIDEAN = "euclidean" + COSINE = "cosine" + + +class WeightingScheme(enum.Enum): + DISTANCE_PROPORTIONAL = "distance_proportional" + BARYCENTRIC = "barycentric" + LEAST_SQUARES = "least_squares" + + +def approximate_from_landmarks( + targets: Tensor, + points: Tensor, + distances: Tensor, + scheme: WeightingScheme = WeightingScheme.DISTANCE_PROPORTIONAL, + cosine_similarity: bool = False, +) -> Tensor: + batch_size, embedding_dim = targets.shape + assert points.ndim == 3 and points.shape == ( + batch_size, + points.shape[1], + embedding_dim, + ) + num_points = points.shape[1] + assert points.shape[2] == embedding_dim + assert distances.shape == (batch_size, num_points) + + if scheme == WeightingScheme.DISTANCE_PROPORTIONAL: + if cosine_similarity: + weights = 1 - distances + else: + weights = 1 / ops.clamp(distances, min=1e-6) + weights = weights / ops.clamp(weights.sum(axis=1, keepdims=True), min=1e-6) + elif scheme == WeightingScheme.BARYCENTRIC: + weights = barycentric_weights(targets, points) + elif scheme == WeightingScheme.LEAST_SQUARES: + weights = ops.lstsq( + points.swapaxes(1, 2).astype(mindspore.float32), + targets.unsqueeze(-1).astype(mindspore.float32), + ).squeeze(-1) + else: + raise ValueError(f"Unknown weighting scheme: {scheme}") + return weights + + +def barycentric_weights(targets: Tensor, points: Tensor) -> Tensor: + batch_size, num_points, _embedding_dim = points.shape + ptp = ops.bmm(points, points.swapaxes(1, 2)) + ones_col = ops.ones((batch_size, num_points, 1), dtype=points.dtype) + ones_row = ops.ones((batch_size, 1, num_points), dtype=points.dtype) + zeros = ops.zeros((batch_size, 1, 1), dtype=points.dtype) + upper = ops.cat([ptp, ones_col], axis=2) + lower = ops.cat([ones_row, zeros], axis=2) + augmented_matrix = ops.cat([upper, lower], axis=1) + rhs_upper = ops.bmm(targets.unsqueeze(1), points.swapaxes(1, 2)).squeeze(1) + rhs_lower = ops.ones((batch_size, 1), dtype=points.dtype) + rhs = ops.cat([rhs_upper, rhs_lower], axis=1) + solution = ops.lstsq( + augmented_matrix.astype(mindspore.float32), + rhs.unsqueeze(-1).astype(mindspore.float32), + ).squeeze(-1) + return solution[..., :num_points] + + +def _cosine_sim(x1: Tensor, x2: Tensor, eps: float = 1e-6) -> Tensor: + w1 = ops.norm(x1, ord=2, axis=1, keepdims=True) # pylint: disable=unexpected-keyword-arg + w2 = ops.norm(x2, ord=2, axis=1, keepdims=True) # pylint: disable=unexpected-keyword-arg + return ops.matmul(x1, x2.T) / ops.clamp(w1 * w2.T, min=eps) + + +def common_interp_approximate( + targets: Tensor, + a_embeddings: Tensor, + k: Optional[int] = None, + metric: DistanceMetric = DistanceMetric.EUCLIDEAN, + weight_scheme: WeightingScheme = WeightingScheme.DISTANCE_PROPORTIONAL, +) -> Tuple[Tensor, Tensor]: + assert targets.ndim == 2 + assert a_embeddings.ndim == 2 + assert targets.shape[1] == a_embeddings.shape[1] + assert (k is None) or (k > 0), "k must be positive" + + if metric == DistanceMetric.EUCLIDEAN: + distances = _cdist(targets, a_embeddings) + elif metric == DistanceMetric.COSINE: + distances = 1 - _cosine_sim(targets, a_embeddings) + else: + raise ValueError(f"Unknown distance metric: {metric}") + + if k is not None: + _, indices = ops.topk(distances, k=k, dim=1, largest=False) + knn_distances = ops.gather_elements(distances, 1, indices) + else: + indices = ops.arange(a_embeddings.shape[0]).expand( + targets.shape[0], -1 + ) + knn_distances = distances + + weights = approximate_from_landmarks( + targets, + a_embeddings[indices], + knn_distances, + scheme=weight_scheme, + cosine_similarity=metric == DistanceMetric.COSINE, + ) + + approx = ( + ops.bmm( + weights.unsqueeze(1).astype(mindspore.float32), + a_embeddings[indices].astype(mindspore.float32), + ) + .squeeze(1) + .astype(targets.dtype) + ) + err = ops.norm(approx - targets, axis=1) # pylint: disable=unexpected-keyword-arg + LOG.debug(f"Reconstruction error: {err.mean()}") + return indices, weights + + +def _cdist(x1: Tensor, x2: Tensor) -> Tensor: + """Compute pairwise L2 distances between rows of x1 and x2.""" + x1_sq = (x1 ** 2).sum(axis=1, keepdims=True) + x2_sq = (x2 ** 2).sum(axis=1, keepdims=True) + cross = ops.matmul(x1, x2.T) + dist_sq = x1_sq - 2 * cross + x2_sq.T + return ops.sqrt(ops.clamp(dist_sq, min=0.0)) diff --git a/src/mindnlp/wizard/merge/tokensurgeon/magikarp.py b/src/mindnlp/wizard/merge/tokensurgeon/magikarp.py new file mode 100644 index 000000000..7da3e3c93 --- /dev/null +++ b/src/mindnlp/wizard/merge/tokensurgeon/magikarp.py @@ -0,0 +1,129 @@ +# Copyright (C) 2025 Arcee AI +# SPDX-License-Identifier: LGPL-3.0-only +# Modified for MindSpore/Ascend NPU by MindNLP Wizard contributors. +# Ported from MergeKit to MindSpore for Wizard + +import logging +from typing import Dict, List, Optional + +import mindspore # pylint: disable=import-error +from mindspore import ops # pylint: disable=import-error +from mindspore import Tensor # pylint: disable=import-error + +from ..tokenizer.normalization import NormalizedToken, unnormalize_token + +LOG = logging.getLogger(__name__) + + +def well_trained_tokens( + vocab: Dict[NormalizedToken, int], + embed: Tensor, + lm_head: Optional[Tensor], + known_unused: Optional[List[NormalizedToken]] = None, + quantile: float = 0.01, +) -> List[NormalizedToken]: + """Get a list of tokens that are well-trained in the model. + + Uses the approach from "Fishing for Magikarp: Automatically Detecting + Under-trained Tokens in Large Language Models" + (https://arxiv.org/abs/2405.05417). + + Args: + vocab: The vocabulary of the model, mapping tokens to indices. + embed: The input embedding matrix of the model. + lm_head: The output embedding matrix of the model (optional). + known_unused: A list of known unused tokens (optional). + quantile: The quantile to use for filtering (default: 0.01). + + Returns: + A list of tokens that can be assumed to be well-trained in the model. + """ + unused_indices = set(range(embed.shape[0])) - set(vocab.values()) + if known_unused: + unused_indices.update(vocab[tok] for tok in known_unused if tok in vocab) + for tok in vocab: + tok_text = unnormalize_token(tok) + if "unused_token" in tok_text or "reserved_special_token" in tok_text: + LOG.debug(f"Assuming {tok_text} is unused") + unused_indices.add(vocab[tok]) + + if unused_indices: + mean_unused_in = embed[list(unused_indices)].mean(axis=0) + mean_unused_out = ( + lm_head[list(unused_indices)].mean(axis=0) if lm_head is not None else None + ) + LOG.info(f"Found {len(unused_indices)} unused tokens") + else: + mean_unused_in = None + mean_unused_out = None + + bad_indices = set(unused_indices) + + if lm_head is not None: + l2_norms = ops.norm(embed, axis=1).astype(mindspore.float32) # pylint: disable=unexpected-keyword-arg + threshold = _quantile(l2_norms, quantile) + LOG.debug( + f"Unused token L2 norm threshold: {threshold.asnumpy():.4f} " + f"({int(quantile * 100)}th percentile)" + ) + l2_bad_indices = ops.nonzero(l2_norms < threshold).squeeze(1) + if l2_bad_indices.shape[0] > 0: + bad_indices.update(l2_bad_indices.asnumpy().tolist()) + LOG.info(f"Discarding {l2_bad_indices.shape[0]} low-l2 tokens") + + if mean_unused_in is not None: + cos_sim = ops.cosine_similarity( # pylint: disable=unexpected-keyword-arg + embed.astype(mindspore.float32), + mean_unused_in.unsqueeze(0).astype(mindspore.float32), + axis=1, + ) + threshold = _quantile(cos_sim, 1 - quantile) + LOG.debug( + f"Unused token threshold in embed_tokens: {threshold.asnumpy():.4f} " + f"({int((1 - quantile) * 100)}th percentile)" + ) + if threshold < 0.5: + threshold = Tensor(0.5, dtype=mindspore.float32) + LOG.debug("Clamping threshold to 0.5") + cos_bad_indices = ops.nonzero(cos_sim > threshold).squeeze(1) + if cos_bad_indices.shape[0] > 0: + bad_indices.update(cos_bad_indices.asnumpy().tolist()) + LOG.info( + f"Discarding {cos_bad_indices.shape[0]} high-sim to unused mean tokens" + ) + + if lm_head is not None and mean_unused_out is not None: + cos_sim = ops.cosine_similarity( # pylint: disable=unexpected-keyword-arg + lm_head.astype(mindspore.float32), + mean_unused_out.unsqueeze(0).astype(mindspore.float32), + axis=1, + ) + threshold = _quantile(cos_sim, 1 - quantile) + LOG.debug( + f"Unused token threshold in lm_head: {threshold.asnumpy():.4f} " + f"({int((1 - quantile) * 100)}th percentile)" + ) + if threshold < 0.5: + threshold = Tensor(0.5, dtype=mindspore.float32) + LOG.debug("Clamping threshold to 0.5") + cos_bad_indices = ops.nonzero(cos_sim > threshold).squeeze(1) + if cos_bad_indices.shape[0] > 0: + bad_indices.update(cos_bad_indices.asnumpy().tolist()) + LOG.info( + f"Discarding {cos_bad_indices.shape[0]} high-sim to unused mean tokens" + ) + + good_tokens = [tok for tok, idx in vocab.items() if idx not in bad_indices] + LOG.info( + f"Found {len(good_tokens)} well-trained tokens, {len(bad_indices)} bad tokens" + ) + return good_tokens + + +def _quantile(tensor: Tensor, q: float) -> Tensor: + """Compute quantile of a 1-D tensor (MindSpore-compatible).""" + sorted_tensor = ops.sort(tensor.flatten())[0] + n = sorted_tensor.shape[0] + idx = int(q * (n - 1)) + idx = max(0, min(idx, n - 1)) + return sorted_tensor[idx] diff --git a/src/mindnlp/wizard/merge/tokensurgeon/omp.py b/src/mindnlp/wizard/merge/tokensurgeon/omp.py new file mode 100644 index 000000000..447836892 --- /dev/null +++ b/src/mindnlp/wizard/merge/tokensurgeon/omp.py @@ -0,0 +1,311 @@ +# Copyright (C) 2025 Arcee AI +# SPDX-License-Identifier: LGPL-3.0-only +# Modified for MindSpore/Ascend NPU by MindNLP Wizard contributors. +# Ported from MergeKit to MindSpore for Wizard + +import logging +from typing import Optional, Tuple + +import mindspore # pylint: disable=import-error +from mindspore import ops # pylint: disable=import-error +from mindspore import Tensor # pylint: disable=import-error + +from .rope_helpers import apply_rope, estimate_pos_id_best + +LOG = logging.getLogger(__name__) + + +def batch_omp( + targets: Tensor, + candidate_points: Tensor, + k: int, + eps: float = 1e-8, + reorthogonalize_interval: int = 256, +) -> Tuple[Tensor, Tensor]: + """ + Batched Orthogonal Matching Pursuit (OMP) to select `k` points from + `candidate_points` that best approximate each target in `targets`. + + Args: + targets: (B, D) tensor of target vectors. + candidate_points: (N, D) tensor of candidate points. + k: Number of points to select (sparsity level). + eps: Tolerance for numerical stability. + reorthogonalize_interval: Number of iterations between reorthogonalization steps. + + Returns: + selected_indices: (B, k) tensor of indices selected for each target. + coeff: (B, k) tensor of coefficients for each selected point. + """ + B, D = targets.shape + N, _ = candidate_points.shape + if k > N: + raise ValueError(f"Cannot select {k} points from {N} candidates") + work_dtype = ( + targets.dtype + if targets.dtype in (mindspore.float32, mindspore.float64) + else mindspore.float32 + ) + targets_work = targets.astype(work_dtype) + points_work = candidate_points.astype(work_dtype) + + q = ops.zeros((B, D, k), dtype=work_dtype) + r = ops.zeros((B, k, k), dtype=work_dtype) + selected_indices = ops.zeros((B, k), dtype=mindspore.int64) + mask = ops.zeros((B, N), dtype=mindspore.bool_) + residuals = targets_work.copy() + + for t in range(k): + rms_0 = ops.norm(residuals, axis=1).mean() # pylint: disable=unexpected-keyword-arg + abs_inner = ops.abs(ops.matmul(residuals, points_work.T)) # (B, N) + abs_inner = ops.masked_fill(abs_inner, mask, -float("inf")) + + _, new_idx = ops.max(abs_inner, axis=1) # (B,) + selected_indices[:, t] = new_idx + mask[ops.arange(B), new_idx] = True + + new_atom = points_work[new_idx] # (B, D) + if t == 0: + r[:, 0, 0] = ops.norm(new_atom, axis=1) # pylint: disable=unexpected-keyword-arg + norm = ops.clamp(r[:, 0, 0], min=eps) + q[:, :, 0] = new_atom / norm.unsqueeze(1) + else: + projections = ops.bmm( + q[:, :, :t].swapaxes(1, 2), new_atom.unsqueeze(-1) + ).squeeze(-1) # (B, t) + residual = new_atom - ops.bmm( + q[:, :, :t], projections.unsqueeze(-1) + ).squeeze(-1) # (B, D) + norm = ops.clamp(ops.norm(residual, axis=1), min=eps) # pylint: disable=unexpected-keyword-arg + r[:, :t, t] = projections + r[:, t, t] = norm + q[:, :, t] = residual / norm.unsqueeze(-1) + + if t > 0 and t % reorthogonalize_interval == 0: + q_b = q[:, :, : t + 1] + # MindSpore doesn't have batched QR directly; use manual Gram-Schmidt + # or fall back to loop-based QR for small k + for b_idx in range(B): + q_single, r_single = ops.qr(q_b[b_idx]) + r[b_idx, : t + 1, : t + 1] = ops.matmul( + r_single, r[b_idx, : t + 1, : t + 1] + ) + q[b_idx, :, : t + 1] = q_single + + qt_targets = ops.bmm( + q[:, :, : t + 1].swapaxes(1, 2), targets_work.unsqueeze(-1) + ) # (B, t+1, 1) + approx = ops.bmm(q[:, :, : t + 1], qt_targets).squeeze(-1) + residuals = targets_work - approx + LOG.debug(f"OMP iteration {t}: RMS {rms_0} -> {ops.norm(residuals, axis=1).mean()}") + + # Get final coefficients via triangular solve + rhs = ops.bmm(q[:, :, :k].swapaxes(1, 2), targets_work.unsqueeze(-1)) + # solve_triangular: R * x = rhs => x = R^{-1} rhs + # MindSpore: use ops.SolveTriangular or manual inverse for upper triangular + r_upper = r[:, :k, :k] + final_coeff = _batched_triangular_solve(r_upper, rhs).squeeze(-1) + + if LOG.isEnabledFor(logging.DEBUG): + rt_approx = ops.bmm( + final_coeff.unsqueeze(1), points_work[selected_indices] + ).squeeze(1) + residuals = targets_work - rt_approx + LOG.debug(f"OMP final RMS: {ops.norm(residuals, axis=1).mean()}") + + return selected_indices, final_coeff + + +def _batched_triangular_solve(R: Tensor, b: Tensor) -> Tensor: + """Solve R @ x = b for upper-triangular R in a batched fashion. + + Args: + R: (B, k, k) upper triangular matrices + b: (B, k, 1) right hand side vectors + + Returns: + x: (B, k, 1) solution vectors + """ + B, k, _ = R.shape + x = ops.zeros_like(b) + for i in range(k - 1, -1, -1): + val = b[:, i, :] - ops.matmul( + R[:, i:i+1, i+1:], x[:, i+1:, :] + ).squeeze(1) + x[:, i, :] = val / R[:, i, i].unsqueeze(-1).clamp(min=1e-12) + return x + + +def batch_mp_resets( + targets: Tensor, + candidate_points: Tensor, + k: int, + eps: float = 1e-8, + total_iterations: Optional[int] = None, +) -> Tuple[Tensor, Tensor]: + """ + Matching Pursuit with Resets + """ + if total_iterations is None: + total_iterations = k * 3 + if total_iterations < k: + raise ValueError( + f"total_iterations {total_iterations} must be greater than or equal to k {k}" + ) + B, D = targets.shape + N, _ = candidate_points.shape + if k > N: + raise ValueError(f"Cannot select {k} points from {N} candidates") + work_dtype = ( + targets.dtype + if targets.dtype in (mindspore.float32, mindspore.float64) + else mindspore.float32 + ) + targets_work = targets.astype(work_dtype) + points_work = candidate_points.astype(work_dtype) + selected_indices = ops.zeros((B, k), dtype=mindspore.int64) + mask = ops.zeros((B, N), dtype=mindspore.bool_) + coeff = ops.zeros((B, k), dtype=work_dtype) + residuals = targets_work.copy() + + iter_indices = list(range(k)) + while len(iter_indices) < total_iterations: + import numpy as np + honk = np.random.permutation(k).tolist() + iter_indices.extend(honk) + iter_indices = iter_indices[:total_iterations] + + for step, t in enumerate(iter_indices): + if step < k: + inner_products = ops.matmul(residuals, points_work.T) # B x N + inner_products = ops.masked_fill(inner_products, mask, -float("inf")) + max_values, max_indices = ops.max(inner_products, axis=1) + selected_points = points_work[max_indices] # B x D + norms_sq = ops.sum(selected_points ** 2, axis=1) + eps # pylint: disable=unexpected-keyword-arg + coeffs = max_values / norms_sq + residuals = residuals - coeffs.unsqueeze(-1) * selected_points + selected_indices[:, t] = max_indices + coeff[:, t] = coeffs + mask = mask.scatter(1, max_indices.unsqueeze(1), True) + else: + old_indices = selected_indices[:, t] + old_coeffs = coeff[:, t] + old_points = points_work[old_indices] + residuals = residuals + old_coeffs.unsqueeze(-1) * old_points + mask = mask.scatter(1, old_indices.unsqueeze(1), False) + inner_products = ops.matmul(residuals, points_work.T) + inner_products = ops.masked_fill(inner_products, mask, -float("inf")) + new_max_values, new_max_indices = ops.max(inner_products, axis=1) + new_points = points_work[new_max_indices] + norms_sq = ops.sum(new_points ** 2, axis=1) + eps # pylint: disable=unexpected-keyword-arg + new_coeffs = new_max_values / norms_sq + residuals = residuals - new_coeffs.unsqueeze(-1) * new_points + selected_indices[:, t] = new_max_indices + coeff[:, t] = new_coeffs + mask = mask.scatter(1, new_max_indices.unsqueeze(1), True) + + return selected_indices, coeff + + +def batch_mp_rope( # pylint: disable=too-many-positional-arguments + targets: Tensor, + points_a: Tensor, + points_b: Tensor, + k: int, + num_heads_a: int, + num_heads_b: int, + eps: float = 1e-8, + a_rope_base: float = 10000.0, + b_rope_base: float = 10000.0, + final_least_squares: bool = True, +) -> Tensor: + B, D_a = targets.shape + N, _ = points_a.shape + _, D_b = points_b.shape + assert ( + points_a.shape[0] == points_b.shape[0] + ), "Number of points in A and B must match" + if k > N: + raise ValueError(f"Cannot select {k} points from {N} candidates") + work_dtype = ( + targets.dtype + if targets.dtype in (mindspore.float32, mindspore.float64) + else mindspore.float32 + ) + out_dtype = targets.dtype + points_a = points_a.astype(work_dtype) + points_b = points_b.astype(work_dtype) + targets = targets.astype(work_dtype) + selected_indices = ops.zeros((B, k), dtype=mindspore.int64) + coeffs = ops.zeros((B, k), dtype=work_dtype) + pos_ids = ops.zeros((B, k), dtype=work_dtype) + mask = ops.zeros((B, N), dtype=mindspore.bool_) + residuals = targets.copy() + + for t in range(k): + abs_inner = ops.abs(ops.matmul(residuals, points_a.T)) # (B, N) + abs_inner = ops.masked_fill(abs_inner, mask, -float("inf")) + + _, new_idx = ops.max(abs_inner, axis=1) # (B,) + + selected_indices[:, t] = new_idx + mask[ops.arange(B), new_idx] = True + new_atom = points_a[new_idx] + + pos_id = estimate_pos_id_best( + new_atom, + residuals, + num_heads=num_heads_a, + head_dim=D_a // num_heads_a, + base=a_rope_base, + ).squeeze(-1) + pos_id_neg = estimate_pos_id_best( + new_atom, + -residuals, + num_heads=num_heads_a, + head_dim=D_a // num_heads_a, + base=a_rope_base, + ).squeeze(-1) + pos_id = ops.where( + ops.abs(pos_id) < ops.abs(pos_id_neg), pos_id, pos_id_neg + ) + pos_ids[:, t] = pos_id + new_atom = apply_rope( + new_atom, + pos_id.unsqueeze(-1), + num_heads=num_heads_a, + head_dim=D_a // num_heads_a, + base=a_rope_base, + ) + + current_coeff = (residuals * new_atom).sum(axis=1) / ( + new_atom.pow(2).sum(axis=1).clamp(min=eps) + ) + coeffs[:, t] = current_coeff + + residuals = residuals - current_coeff.unsqueeze(1) * new_atom + + if final_least_squares: + roped_pts_a = apply_rope( + points_a[selected_indices], + pos_ids.unsqueeze(-1), + num_heads=num_heads_a, + head_dim=D_a // num_heads_a, + base=a_rope_base, + ) + coeffs = ops.lstsq( + roped_pts_a.swapaxes(1, 2).astype(mindspore.float32), + targets.unsqueeze(-1).astype(mindspore.float32), + ).squeeze(-1) + + selected_points_b = points_b[selected_indices] + atoms_b = apply_rope( + selected_points_b, + pos_ids.unsqueeze(-1), + num_heads=num_heads_b, + head_dim=D_b // num_heads_b, + base=b_rope_base, + ) + approx_b = (atoms_b * coeffs.unsqueeze(-1)).sum(axis=1) + final_tensor = approx_b.astype(out_dtype) + return selected_indices, coeffs, final_tensor, targets - residuals diff --git a/src/mindnlp/wizard/merge/tokensurgeon/pca.py b/src/mindnlp/wizard/merge/tokensurgeon/pca.py new file mode 100644 index 000000000..9670aaed7 --- /dev/null +++ b/src/mindnlp/wizard/merge/tokensurgeon/pca.py @@ -0,0 +1,67 @@ +# Copyright (C) 2025 Arcee AI +# SPDX-License-Identifier: LGPL-3.0-only +# Modified for MindSpore/Ascend NPU by MindNLP Wizard contributors. +# Ported from MergeKit to MindSpore for Wizard + +import logging + +import mindspore # pylint: disable=import-error +from mindspore import ops # pylint: disable=import-error +from mindspore import Tensor # pylint: disable=import-error + +LOG = logging.getLogger(__name__) + + +def landmark_pca_approximate( + targets: Tensor, + points_a: Tensor, + points_b: Tensor, +) -> Tensor: + """Given target points in space a and a set of reference points in both space a and b, + approximate the target points in space b.""" + num_points, d_a = points_a.shape + batch_size, _ = targets.shape + _, d_b = points_b.shape + assert ( + points_a.shape[0] == points_b.shape[0] + ), "Number of points in A and B must match" + assert targets.shape == (batch_size, d_a) + + effective_dim = min(d_a, d_b) + + out_dtype = targets.dtype + points_a = points_a.astype(mindspore.float32) + points_b = points_b.astype(mindspore.float32) + targets = targets.astype(mindspore.float32) + + mean_a = points_a.mean(axis=0, keepdims=True) # (1, D_a) + mean_b = points_b.mean(axis=0, keepdims=True) # (1, D_b) + centered_a = points_a - mean_a # (N, D_a) + centered_b = points_b - mean_b # (N, D_b) + centered_targets = targets - mean_a # (B, D_a) + + V_a = _pca_lowrank(centered_a, q=effective_dim) # (D_a, effective_dim) + V_b = _pca_lowrank(centered_b, q=effective_dim) # (D_b, effective_dim) + + A_pca = ops.matmul(centered_a, V_a) # (N, effective_dim) + B_pca = ops.matmul(centered_b, V_b) # (N, effective_dim) + + M = ops.matmul(B_pca.T, A_pca) # (effective_dim, effective_dim) + U, S, V = ops.svd(M) + R = ops.matmul(U, V) # (effective_dim, effective_dim) + + projected_a = ops.matmul(centered_targets, V_a) # (B, effective_dim) + rotated = ops.matmul(projected_a, R) # (B, effective_dim) + projected_b = ops.matmul(rotated, V_b.T) # (B, D_b) + + approximated_b = projected_b + mean_b + return approximated_b.astype(out_dtype) + + +def _pca_lowrank(data: Tensor, q: int) -> Tensor: + """Compute the top-q right singular vectors (V) of the data matrix via SVD. + + This is a simplified replacement for torch.pca_lowrank, returning only V. + """ + U, S, V = ops.svd(data) + return V[:, :q] diff --git a/src/mindnlp/wizard/merge/tokensurgeon/rope_helpers.py b/src/mindnlp/wizard/merge/tokensurgeon/rope_helpers.py new file mode 100644 index 000000000..62f6fdefa --- /dev/null +++ b/src/mindnlp/wizard/merge/tokensurgeon/rope_helpers.py @@ -0,0 +1,189 @@ +# Copyright (C) 2025 Arcee AI +# SPDX-License-Identifier: LGPL-3.0-only +# Modified for MindSpore/Ascend NPU by MindNLP Wizard contributors. +# Ported from MergeKit to MindSpore for Wizard + +import mindspore # pylint: disable=import-error +from mindspore import ops # pylint: disable=import-error +from mindspore import Tensor # pylint: disable=import-error + + +def llama_rope_rotationmat(theta: Tensor) -> Tensor: + """ + Create a rotation matrix for RoPE as used in the `transformers` Llama implementation. + + Args: + theta: Tensor of shape (..., n_heads, head_dim // 2) representing the angles for the rotation. + """ + n_heads = theta.shape[-2] + head_dim = theta.shape[-1] * 2 + theta_p = ops.cat([theta, theta], axis=-1) + cos_theta = ops.cos(theta_p) + sin_theta = ops.sin(theta_p) + P = ops.zeros( + tuple(list(theta.shape[:-1]) + [head_dim, head_dim]), + dtype=theta.dtype, + ) + idx = ops.arange(head_dim // 2) + P[..., :, idx, idx] = cos_theta[..., :, idx] + P[..., :, idx, head_dim // 2 + idx] = sin_theta[..., :, idx] + P[..., :, head_dim // 2 + idx, idx] = -sin_theta[..., :, idx] + P[..., :, head_dim // 2 + idx, head_dim // 2 + idx] = cos_theta[..., :, idx] + return P + + +def _rope_inv_freq(base: float, dim: int) -> Tensor: + return 1.0 / ( + base ** (ops.arange(0, dim, 2).astype(mindspore.float32) / dim) + ) + + +def estimate_theta( + x_0: Tensor, + x_1: Tensor, + num_heads: int, + head_dim: int, +) -> Tensor: + """Estimate a set of per-head, per-dimension angles (theta) such that + rotating x_0 by theta will least-squares approximate x_1. + + Args: + x_0: Tensor of shape (..., n_heads*head_dim) representing the first input. + x_1: Tensor of shape (..., n_heads*head_dim) representing the second input. + num_heads: Number of attention heads. + head_dim: Dimension of each attention head. + Returns: + Tensor of shape (..., n_heads, head_dim // 2) representing the estimated theta values. + """ + x0_reshaped = x_0.view(*x_0.shape[:-1], num_heads, head_dim) + x1_reshaped = x_1.view(*x_1.shape[:-1], num_heads, head_dim) + + half_dim = head_dim // 2 + x0_i = x0_reshaped[..., :half_dim] + x0_j = x0_reshaped[..., half_dim:] + x1_i = x1_reshaped[..., :half_dim] + x1_j = x1_reshaped[..., half_dim:] + + A_d = x0_i * x1_i + x0_j * x1_j + B_d = x0_i * x1_j - x0_j * x1_i + + theta = ops.atan2(B_d, A_d) + return theta + + +def estimate_position_id( + x_0: Tensor, + x_1: Tensor, + num_heads: int, + head_dim: int, + base: float = 10000.0, +) -> Tensor: + """ + Estimate a scalar position ID such that applying RoPE to x_0 + will least-squares approximate x_1. + """ + x0_heads = x_0.view(*x_0.shape[:-1], num_heads, head_dim) + x1_heads = x_1.view(*x_1.shape[:-1], num_heads, head_dim) + + split_idx = head_dim // 2 + x0_a = x0_heads[..., :split_idx] + x0_b = x0_heads[..., split_idx:] + x1_c = x1_heads[..., :split_idx] + x1_d = x1_heads[..., split_idx:] + + numerator = x0_a * x1_d - x0_b * x1_c + denominator = x0_a * x1_c + x0_b * x1_d + theta = ops.atan2(numerator, denominator) + + inv_freq = _rope_inv_freq(base, head_dim) + pos_i = theta / inv_freq + weights = x0_a.pow(2) + x0_b.pow(2) + sum_pos = (pos_i * weights).sum(axis=(-1, -2)) + sum_weights = weights.sum(axis=(-1, -2)) + pos_estimate = sum_pos / (sum_weights + 1e-8) + return pos_estimate.unsqueeze(-1) + + +def estimate_position_id_projection( + x_0: Tensor, + x_1: Tensor, + num_heads: int, + head_dim: int, + base: float = 10000.0, +) -> Tensor: + inv_freq = _rope_inv_freq(base, head_dim) + basis_vector = inv_freq.view(1, 1, head_dim // 2).expand( + x_0.shape[:-1] + (num_heads, head_dim // 2) + ) + basis_vector = basis_vector.reshape(*x_0.shape[:-1], -1) + basis_vector_norm = ops.norm(basis_vector, axis=-1, keepdims=True) # pylint: disable=unexpected-keyword-arg + basis_vector = basis_vector / (basis_vector_norm + 1e-8) + theta = estimate_theta(x_0, x_1, num_heads, head_dim) + theta = theta.reshape(*x_0.shape[:-1], -1) + projection = ops.sum(theta * basis_vector, axis=-1) # pylint: disable=unexpected-keyword-arg + f_norm = ops.norm(inv_freq) + scaling_factor = ops.sqrt(Tensor(float(num_heads), dtype=mindspore.float32)) * f_norm + pos_estimate = projection / (scaling_factor + 1e-8) + return pos_estimate.unsqueeze(-1) + + +def apply_rope_theta( + x: Tensor, + theta: Tensor, + num_heads: int, + head_dim: int, +) -> Tensor: + """ + Apply RoPE to the input tensor x using the given theta. + """ + x_reshaped = x.view(*x.shape[:-1], num_heads, head_dim) + + half_dim = head_dim // 2 + x_i = x_reshaped[..., :half_dim] + x_j = x_reshaped[..., half_dim:] + + cos_theta = ops.cos(theta) + sin_theta = ops.sin(theta) + + x_i_rot = x_i * cos_theta - x_j * sin_theta + x_j_rot = x_j * cos_theta + x_i * sin_theta + + rotated = ops.cat([x_i_rot, x_j_rot], axis=-1) + return rotated.view(*x.shape) + + +def estimate_pos_id_best( + x_0: Tensor, + x_1: Tensor, + num_heads: int, + head_dim: int, + base: float = 10000.0, +) -> Tensor: + return estimate_position_id_projection( + x_0, + x_1, + num_heads, + head_dim, + base=base, + ) + + +def apply_rope( + x: Tensor, + pos: Tensor, + num_heads: int, + head_dim: int, + base: float = 10000.0, +) -> Tensor: + """ + Apply RoPE to the input tensor x using the given position pos. + """ + inv_freq = _rope_inv_freq(base, head_dim) + theta = pos.unsqueeze(-1) * inv_freq + + return apply_rope_theta( + x, + theta, + num_heads, + head_dim, + ) diff --git a/src/mindnlp/wizard/merge/tokensurgeon/subword.py b/src/mindnlp/wizard/merge/tokensurgeon/subword.py new file mode 100644 index 000000000..0101744ef --- /dev/null +++ b/src/mindnlp/wizard/merge/tokensurgeon/subword.py @@ -0,0 +1,60 @@ +# Copyright (C) 2025 Arcee AI +# SPDX-License-Identifier: LGPL-3.0-only +# Modified for MindSpore/Ascend NPU by MindNLP Wizard contributors. +# Ported from MergeKit to MindSpore for Wizard + +import enum +from typing import List + +from mindspore import ops # pylint: disable=import-error +from mindspore import Tensor # pylint: disable=import-error +import transformers + +from ..tokenizer.normalization import NormalizedToken, unnormalize_token + + +class SubwordMethod(enum.Enum): + MEAN = "mean" + SUM = "sum" + WEIGHTED_MEAN = "weighted_mean" + FIRST_LAST = "first_last" + + +def subword_approximate( + orig_embed: Tensor, + target_tokens: List[NormalizedToken], + is_lm_head: bool, + tok_0: transformers.PreTrainedTokenizerBase, + subword_method: SubwordMethod = SubwordMethod.MEAN, +) -> Tensor: + res = ops.zeros( + (len(target_tokens), orig_embed.shape[1]), + dtype=orig_embed.dtype, + ) + for idx, token in enumerate(target_tokens): + text = unnormalize_token(token) + token_ids = tok_0(text, add_special_tokens=False)["input_ids"] + + if subword_method in (SubwordMethod.MEAN, SubwordMethod.SUM): + for id in token_ids: + res[idx] += orig_embed[id] + if subword_method == SubwordMethod.MEAN and len(token_ids) > 0: + res[idx] /= len(token_ids) + elif subword_method == SubwordMethod.WEIGHTED_MEAN: + weights = list(range(1, len(token_ids) + 1)) + if not is_lm_head: + weights = weights[::-1] + for id, weight in zip(token_ids, weights): + res[idx] += weight * orig_embed[id] + if len(token_ids) > 0: + res[idx] /= sum(weights) + elif subword_method == SubwordMethod.FIRST_LAST: + if len(token_ids) == 0: + continue + if is_lm_head: + res[idx] = orig_embed[token_ids[0]] + else: + res[idx] = orig_embed[token_ids[-1]] + else: + raise ValueError(f"Unknown subword method: {subword_method}") + return res diff --git a/src/mindnlp/wizard/merge/tokensurgeon/token_basis.py b/src/mindnlp/wizard/merge/tokensurgeon/token_basis.py new file mode 100644 index 000000000..ffd531999 --- /dev/null +++ b/src/mindnlp/wizard/merge/tokensurgeon/token_basis.py @@ -0,0 +1,144 @@ +# Copyright (C) 2025 Arcee AI +# SPDX-License-Identifier: LGPL-3.0-only +# Modified for MindSpore/Ascend NPU by MindNLP Wizard contributors. +# Ported from MergeKit to MindSpore for Wizard + +import logging +from typing import Dict, List, Tuple + +import mindspore # pylint: disable=import-error +from mindspore import ops # pylint: disable=import-error +from mindspore import Tensor # pylint: disable=import-error + +from ..tokenizer.normalization import NormalizedToken +from .omp import batch_omp + +LOG = logging.getLogger(__name__) + + +def sparse_linear_basis( + points: Tensor, + k: int, + d: int, + eps: float = 1e-8, +) -> Tuple[Tensor, Tensor]: + """ + Form an approximate orthogonal basis from sparse linear combinations of input points. + Args: + points: (num_pts, embed_dim) tensor of input points + k: number of points to select per basis vector + d: dimensionality of the basis + eps: numerical stability parameter + Returns: + indices: (d, k) tensor of selected indices + coeffs: (d, k) tensor of coefficients for each selected point + """ + assert points.ndim == 2 + num_pts, embed_dim = points.shape + assert k <= num_pts, "k must be less than or equal to the number of points" + assert d <= embed_dim, "d must be less than or equal to the embedding dimension" + + mean_embed = points.mean(axis=0) + centered_embeddings = (points - mean_embed).astype(mindspore.float32) + covariance_matrix = ( + centered_embeddings.T @ centered_embeddings + ) / num_pts # (embed_dim, embed_dim) + + U, _S, _V = ops.svd(covariance_matrix) + U_d = U[:, :d] # (embed_dim, d) + + indices, coeffs = batch_omp( + U_d.T, # (d, embed_dim) + centered_embeddings, # (num_pts, embed_dim) + k, + eps=eps, + ) + + if LOG.isEnabledFor(logging.DEBUG): + rc_basis = ops.bmm( + coeffs.unsqueeze(1).astype(mindspore.float32), + centered_embeddings[indices].astype(mindspore.float32), + ).squeeze(1) + for i in range(d): + v_0 = U_d[:, i] + v_1 = rc_basis[i] + cos_sim = ops.cosine_similarity( # pylint: disable=unexpected-keyword-arg + v_0.unsqueeze(0), v_1.unsqueeze(0), axis=1 + ).squeeze() + rms = ops.norm(v_0 - v_1) + norm_rms = ops.norm( + v_0 - (v_1 / ops.clamp(ops.norm(v_1), min=1e-6)) + ) + LOG.debug( + f"Basis vector {i}: cos_sim = {cos_sim.asnumpy():.4f}, " + f"RMS = {rms.asnumpy():.4f}, norm_rms = {norm_rms.asnumpy():.4f}" + ) + + return indices, coeffs + + +def compute_token_basis( # pylint: disable=too-many-positional-arguments + orig_embed: Tensor, + donor_embed: Tensor, + orig_vocab: Dict[NormalizedToken, int], + donor_vocab: Dict[NormalizedToken, int], + junk_tokens: List[int], + k: int, +) -> Tuple[Tensor, Tensor]: + """Compute approximately orthogonal bases for both original and donor embeddings + as sparse linear combinations of elements. + + Args: + orig_embed: Original embedding matrix + donor_embed: Donor embedding matrix + orig_vocab: Vocabulary mapping for original model + donor_vocab: Vocabulary mapping for donor model + junk_tokens: List of junk token indices to exclude + k: Number of points to select per basis vector + Returns: + donor_basis: Approximate orthogonal basis for donor model + orig_basis: Approximate orthogonal basis for original model + """ + common_vocab = set(orig_vocab.keys()) & set(donor_vocab.keys()) + junk_set = set(junk_tokens) + common_vocab = [ + tok + for tok in common_vocab + if (tok not in donor_vocab or donor_vocab[tok] not in junk_set) + ] + effective_dim = min(orig_embed.shape[1], donor_embed.shape[1]) + orig_shared_embeds = orig_embed[ + Tensor([orig_vocab[t] for t in common_vocab], dtype=mindspore.int64) + ] + donor_shared_embeds = donor_embed[ + Tensor([donor_vocab[t] for t in common_vocab], dtype=mindspore.int64) + ] + if donor_embed.shape[1] < orig_embed.shape[1]: + basis_src_embeds = donor_shared_embeds + LOG.debug("Using donor embeds to compute token basis") + else: + basis_src_embeds = orig_shared_embeds + LOG.debug("Using original embeds to compute token basis") + LOG.debug(f"Basis dimension: {effective_dim}") + tb_indices, tb_weights = sparse_linear_basis( + basis_src_embeds, + k=k, + d=effective_dim, + ) + donor_basis = ( + ops.bmm( + tb_weights.unsqueeze(1).astype(mindspore.float32), + donor_shared_embeds[tb_indices].astype(mindspore.float32), + ) + .squeeze(1) + .astype(donor_embed.dtype) + ) + orig_basis = ( + ops.bmm( + tb_weights.unsqueeze(1).astype(mindspore.float32), + orig_shared_embeds[tb_indices].astype(mindspore.float32), + ) + .squeeze(1) + .astype(orig_embed.dtype) + ) + return (donor_basis, orig_basis) diff --git a/tests/mindnlp/wizard/README.md b/tests/mindnlp/wizard/README.md new file mode 100644 index 000000000..bb143852f --- /dev/null +++ b/tests/mindnlp/wizard/README.md @@ -0,0 +1,58 @@ +# Wizard Merge 测试套件 + +## 测试文件 + +| 文件 | 说明 | +|------|------| +| `test_merge.py` | 核心合并逻辑（Task DAG、Executor、MergeConfiguration） | +| `test_ckpt_io.py` | .ckpt 读写与 bf16 精度安全 | +| `test_bf16_ckpt_safety.py` | BF16 字节级精度验证 | +| `test_bf16_cpu_full_method_matrix.py` | BF16 下全方法矩阵测试 | +| `test_merge_dtype_matrix.py` | 合并 dtype 组合矩阵 | +| `test_dtype_policy.py` | dtype 策略单元测试 | +| `test_safe_ops.py` | 安全张量运算 | +| `test_preflight.py` | 合并前校验 | +| `test_config_validation_matrix.py` | 配置校验矩阵 | +| `test_cli_run_yaml_compat.py` | CLI run_yaml 兼容性 | +| `test_mergekit_recipe_compat.py` | MergeKit 配方兼容性 | +| `test_basic_merges_parity.py` | 基础合并一致性 | +| `test_base_capabilities_parity.py` | 基础能力一致性 | +| `test_vlm_merges_parity.py` | VLM 合并一致性 | +| `test_mixed_weight_format_compat.py` | 混合权重格式兼容 | +| `test_tokenizer_mergekit_parity.py` | Tokenizer 合并一致性 | + +## 运行测试 + +```bash +# 在项目根目录执行 + +# 运行全部 Wizard 测试 +pytest tests/mindnlp/wizard/ -v + +# 运行单个文件 +pytest tests/mindnlp/wizard/test_ckpt_io.py -v + +# 仅运行快速测试（排除需要真实模型的测试） +pytest tests/mindnlp/wizard/ -v -k "not parity" + +# 查看覆盖率 +pytest tests/mindnlp/wizard/ --cov=mindnlp.wizard --cov-report=html +``` + +## 环境说明 + +- 测试通过 `conftest.py` 自动配置 `sys.path`，无需完整安装 mindnlp +- `conftest.py` 会 stub `torch_npu` 以避免 Ascend 环境依赖 +- 部分一致性测试（`*_parity.py`）需要网络访问以下载模型权重 + +## 预期结果 + +``` +tests/mindnlp/wizard/test_merge.py ............ PASSED +tests/mindnlp/wizard/test_ckpt_io.py ........ PASSED +tests/mindnlp/wizard/test_bf16_ckpt_safety.py ...... PASSED +tests/mindnlp/wizard/test_safe_ops.py .... PASSED +tests/mindnlp/wizard/test_preflight.py ...... PASSED +tests/mindnlp/wizard/test_dtype_policy.py ........ PASSED +... +``` diff --git a/tests/mindnlp/wizard/__init__.py b/tests/mindnlp/wizard/__init__.py new file mode 100644 index 000000000..8b1378917 --- /dev/null +++ b/tests/mindnlp/wizard/__init__.py @@ -0,0 +1 @@ + diff --git a/tests/mindnlp/wizard/conftest.py b/tests/mindnlp/wizard/conftest.py new file mode 100644 index 000000000..60002a2a2 --- /dev/null +++ b/tests/mindnlp/wizard/conftest.py @@ -0,0 +1,48 @@ +"""Configure sys.path so wizard.merge modules can be imported independently +of the full mindnlp init (which requires mindtorch).""" + +import importlib +import importlib.machinery +import os +import sys +import types +from pathlib import Path + +# Disable torch device backend auto-loading (avoids torch_npu libhccl errors). +os.environ.setdefault("TORCH_DEVICE_BACKEND_AUTOLOAD", "0") + +# Stub torch_npu before anything tries to import it (accelerate, transformers). +if "torch_npu" not in sys.modules: + _fake_npu = types.ModuleType("torch_npu") + _fake_npu.__spec__ = importlib.machinery.ModuleSpec("torch_npu", None) + _fake_npu.__path__ = [] + for _sub in ("_C", "utils", "utils._error_code", "npu"): + _full = f"torch_npu.{_sub}" + _mod = types.ModuleType(_full) + _mod.__spec__ = importlib.machinery.ModuleSpec(_full, None) + sys.modules[_full] = _mod + _fake_npu._C = sys.modules["torch_npu._C"] + _fake_npu.utils = sys.modules["torch_npu.utils"] + _fake_npu.npu = sys.modules["torch_npu.npu"] + sys.modules["torch_npu.utils._error_code"].ErrCode = None + sys.modules["torch_npu.utils._error_code"].pta_error = None + sys.modules["torch_npu"] = _fake_npu + +SRC_DIR = str(Path(__file__).resolve().parents[3] / "src") + +if SRC_DIR not in sys.path: + sys.path.insert(0, SRC_DIR) + +# Replace the pip-installed mindnlp with a lightweight stub +# so that subpackage imports (mindnlp.wizard.merge.*) work +# without triggering the full mindnlp init (mindtorch dependency). +stub = types.ModuleType("mindnlp") +stub.__path__ = [str(Path(SRC_DIR) / "mindnlp")] +stub.__package__ = "mindnlp" +sys.modules["mindnlp"] = stub + +# Also stub mindnlp.wizard +wizard_stub = types.ModuleType("mindnlp.wizard") +wizard_stub.__path__ = [str(Path(SRC_DIR) / "mindnlp" / "wizard")] +wizard_stub.__package__ = "mindnlp.wizard" +sys.modules["mindnlp.wizard"] = wizard_stub diff --git a/tests/mindnlp/wizard/test_base_capabilities_parity.py b/tests/mindnlp/wizard/test_base_capabilities_parity.py new file mode 100644 index 000000000..35c779ff4 --- /dev/null +++ b/tests/mindnlp/wizard/test_base_capabilities_parity.py @@ -0,0 +1,238 @@ +import os +import tempfile + +import mindspore +import numpy as np +import pytest +import safetensors.numpy +import torch + +from mindnlp.wizard.merge.io.loader import TensorLoader +from mindnlp.wizard.merge.io.lazy_tensor_loader import LazyTensorLoader +from mindnlp.wizard.merge.io.lazy_unpickle import DeferredLoad +from mindnlp.wizard.merge.io.tasks import LoaderCache +from mindnlp.wizard.merge.plan import MergePlanner +from mindnlp.wizard.merge.config import ConfigReader, InputModelDefinition, MergeConfiguration +from mindnlp.wizard.merge.options import MergeOptions +from mindnlp.wizard.merge.common import ModelReference +from mindnlp.wizard.merge.architecture.base import WeightInfo +from mindnlp.wizard.merge.io.tensor_writer import TensorWriter + + +class TestTensorWriterParity: + def test_safetensors_write(self): + with tempfile.TemporaryDirectory() as d: + writer = TensorWriter(d, safe_serialization=True) + writer.save_tensor("steve", mindspore.Tensor(np.random.randn(4).astype(np.float32))) + writer.finalize() + assert os.path.exists(os.path.join(d, "model.safetensors")) + + def test_bin_write_rejected(self): + with tempfile.TemporaryDirectory() as d: + with pytest.raises(ValueError, match="Unsupported output_format 'bin'"): + TensorWriter(d, safe_serialization=False) + + def test_ckpt_write(self): + with tempfile.TemporaryDirectory() as d: + writer = TensorWriter(d, output_format="ckpt") + writer.save_tensor("timothan", mindspore.Tensor(np.random.randn(4).astype(np.float32))) + writer.finalize() + assert os.path.exists(os.path.join(d, "mindspore_model.ckpt")) + + def test_duplicate_tensor(self): + with tempfile.TemporaryDirectory() as d: + writer = TensorWriter(d, safe_serialization=True) + jim = mindspore.Tensor(np.random.randn(4).astype(np.float32)) + writer.save_tensor("jim", jim) + writer.save_tensor("jimbo", jim) + writer.finalize() + assert os.path.exists(os.path.join(d, "model.safetensors")) + + def test_async_writer(self): + with tempfile.TemporaryDirectory() as d: + writer = TensorWriter( + d, safe_serialization=True, use_async=True, max_shard_size=1, max_write_threads=2 + ) + for i in range(4): + writer.save_tensor(f"t{i + 1}", mindspore.Tensor(np.random.randn(16).astype(np.float32))) + writer.finalize() + assert all( + os.path.exists( + os.path.join(d, f"model-{i + 1:05d}-of-00004.safetensors") + ) + for i in range(4) + ) + + +class TestLazyUnpickleParity: + def test_lazy_unpickle(self): + with tempfile.TemporaryDirectory() as d: + data = { + "a": torch.tensor([1, 2, 3]), + "b": torch.tensor([4, 5, 6]), + } + path = os.path.join(d, "pytorch_model.bin") + torch.save(data, path) + + loader = LazyTensorLoader.from_disk(d) + for name in data: + assert name in loader.index.tensor_paths + tensor = loader.get_tensor(name) + np.testing.assert_array_equal( + tensor.asnumpy(), + data[name].numpy(), + ) + + def test_lazy_unpickle_forwards_device_map_location(self, monkeypatch): + captured = {"map_location": None} + origin_execute = DeferredLoad.execute + + def _wrapped_execute(self, reader, map_location=None): + captured["map_location"] = map_location + return origin_execute(self, reader, map_location=map_location) + + monkeypatch.setattr(DeferredLoad, "execute", _wrapped_execute) + + with tempfile.TemporaryDirectory() as d: + path = os.path.join(d, "pytorch_model.bin") + torch.save({"a": torch.tensor([1.0, 2.0])}, path) + + loader = TensorLoader.get( + path, + use_lazy_unpickle=True, + device="CPU", + ) + tensor = loader.get_tensor("a") + np.testing.assert_array_equal( + tensor.asnumpy(), + np.array([1.0, 2.0], dtype=np.float32), + ) + assert captured["map_location"] == "CPU" + + +class _NoAsNumpyTensor: + """Test helper: ensure size accounting does not call asnumpy().""" + + nbytes = 16 + + def asnumpy(self): + raise AssertionError("asnumpy() should not be called in save_tensor") + + +class TestTensorWriterCopyPath: + def test_save_tensor_size_accounting_avoids_asnumpy(self): + with tempfile.TemporaryDirectory() as d: + writer = TensorWriter(d, safe_serialization=True) + writer.save_tensor("x", _NoAsNumpyTensor()) + assert writer.current_shard_size == 16 + + +class TestReadToNpuDevicePath: + def test_planner_to_loader_device_propagation(self, monkeypatch): + class _DummyMergeMethod: + def parameters(self): + return [] + + def tensor_parameters(self): + return [] + + def make_task(self, output_weight, tensors, **kwargs): + return tensors + + class _DummyLoader: + def __init__(self): + self.last_device = None + self.index = type("Idx", (), {"tensor_paths": {"w": "dummy"}})() + + def get_tensor(self, name, device="CPU", aliases=None, raise_on_missing=True): + self.last_device = device + return mindspore.Tensor(np.array([1.0], dtype=np.float32)) + + dummy_loader = _DummyLoader() + + # Ensure planner builds GatherTensors with the merge option device. + import mindnlp.wizard.merge.plan as plan_mod + + monkeypatch.setattr(plan_mod.merge_methods, "get", lambda _name: _DummyMergeMethod()) + monkeypatch.setattr(LoaderCache, "get", lambda self, _model: dummy_loader) + + model = ModelReference.parse("dummy-model") + cfg = MergeConfiguration( + merge_method="dummy", + models=[InputModelDefinition(model=model)], + ) + options = MergeOptions(read_to_npu=True, device="Ascend") + planner = MergePlanner( + config=cfg, + arch_info=object(), + options=options, + out_model_config=object(), + ) + + planner.plan_tensor( + weight=WeightInfo(name="w"), + weights_in=[WeightInfo(name="w")], + models=[model], + cfg_reader=ConfigReader(config=cfg, t=0), + ) + + gather = planner._tensors[0][1] + args = gather.arguments() + loaded = {k: task.execute() for k, task in args.items()} + res = gather.execute(**loaded) + + assert model in res + assert dummy_loader.last_device == "Ascend" + + +class TestLazyTensorLoaderDevicePath: + def test_safetensors_path_propagates_device(self, monkeypatch): + import mindnlp.wizard.merge.io._device as device_mod + import mindnlp.wizard.merge.io.loader as loader_mod + + captured = {"device": None} + origin_move = device_mod.move_tensor_to_device + + def _wrapped_move(tensor, device, **kwargs): + captured["device"] = device + return origin_move(tensor, device, **kwargs) + + monkeypatch.setattr(loader_mod, "move_tensor_to_device", _wrapped_move) + + with tempfile.TemporaryDirectory() as d: + st_path = os.path.join(d, "model.safetensors") + safetensors.numpy.save_file( + {"a": np.array([1.0, 2.0], dtype=np.float32)}, + st_path, + metadata={"format": "np"}, + ) + + loader = LazyTensorLoader.from_disk(d, lazy_loader=False) + tensor = loader.get_tensor("a", device="Ascend:0") + np.testing.assert_array_equal( + tensor.asnumpy(), + np.array([1.0, 2.0], dtype=np.float32), + ) + assert captured["device"] == "Ascend:0" + + def test_bin_lazy_path_propagates_map_location(self, monkeypatch): + captured = {"map_location": None} + origin_execute = DeferredLoad.execute + + def _wrapped_execute(self, reader, map_location=None): + captured["map_location"] = map_location + return origin_execute(self, reader, map_location=map_location) + + monkeypatch.setattr(DeferredLoad, "execute", _wrapped_execute) + + with tempfile.TemporaryDirectory() as d: + bin_path = os.path.join(d, "pytorch_model.bin") + torch.save({"a": torch.tensor([3.0, 4.0])}, bin_path) + + loader = LazyTensorLoader.from_disk(d, lazy_loader=True) + tensor = loader.get_tensor("a", device="Ascend:0") + np.testing.assert_array_equal( + tensor.asnumpy(), + np.array([3.0, 4.0], dtype=np.float32), + ) + assert captured["map_location"] == "Ascend:0" diff --git a/tests/mindnlp/wizard/test_basic_merges_parity.py b/tests/mindnlp/wizard/test_basic_merges_parity.py new file mode 100644 index 000000000..7501f3ea4 --- /dev/null +++ b/tests/mindnlp/wizard/test_basic_merges_parity.py @@ -0,0 +1,378 @@ +import os +import tempfile +from typing import Callable, Optional + +import numpy as np +import pytest +from transformers import AutoConfig, GPT2Config, GPT2LMHeadModel, LlamaConfig, LlamaForCausalLM + +from mindnlp.wizard.merge.config import ( + InputModelDefinition, + InputSliceDefinition, + MergeConfiguration, + OutputSliceDefinition, +) +from mindnlp.wizard.merge.io.lazy_tensor_loader import LazyTensorLoader +from mindnlp.wizard.merge.merge import MergeOptions, run_merge + + +def _run_and_check_merge( + config: MergeConfiguration, + check_nan: bool = True, + validate: Optional[Callable[[str], None]] = None, + options: Optional[MergeOptions] = None, +): + with tempfile.TemporaryDirectory() as tmpdir: + run_merge(config, out_path=tmpdir, options=options or MergeOptions()) + assert os.path.exists(os.path.join(tmpdir, "config.json")) + assert ( + os.path.exists(os.path.join(tmpdir, "model.safetensors.index.json")) + or os.path.exists(os.path.join(tmpdir, "model.safetensors")) + ), "No model produced by merge" + + if check_nan: + loader = LazyTensorLoader.from_disk(tmpdir, lazy_loader=False) + for tensor_name in sorted(loader.index.tensor_paths.keys()): + tensor = loader.get_tensor(tensor_name) + assert np.isfinite(tensor.asnumpy()).all(), f"NaN/Inf found in {tensor_name}" + + if validate: + validate(tmpdir) + + +def _make_picollama(path: str, vocab_size: int = 64): + cfg = LlamaConfig( + vocab_size=vocab_size, + hidden_size=32, + intermediate_size=48, + num_attention_heads=4, + num_hidden_layers=2, + ) + model = LlamaForCausalLM(cfg) + model.save_pretrained(path, safe_serialization=True) + return str(path) + + +def _make_gpt2size(path: str): + cfg = GPT2Config( + n_ctx=128, + n_embd=64, + n_head=4, + n_layer=4, + n_positions=128, + vocab_size=128, + ) + model = GPT2LMHeadModel(cfg) + model.save_pretrained(path, safe_serialization=True) + return str(path) + + +@pytest.fixture(scope="session") +def model_a(tmp_path_factory): + return _make_picollama(str(tmp_path_factory.mktemp("wizard_basic_a"))) + + +@pytest.fixture(scope="session") +def model_b(tmp_path_factory): + return _make_picollama(str(tmp_path_factory.mktemp("wizard_basic_b"))) + + +@pytest.fixture(scope="session") +def model_c(tmp_path_factory): + return _make_picollama(str(tmp_path_factory.mktemp("wizard_basic_c"))) + + +@pytest.fixture(scope="session") +def gpt2_like(tmp_path_factory): + return _make_gpt2size(str(tmp_path_factory.mktemp("wizard_gpt2_like"))) + + +class TestBasicMergeParity: + def _two_model_config( + self, + model_a: str, + model_b: str, + *, + merge_method: str, + base_model: Optional[str] = None, + params: Optional[dict] = None, + ) -> MergeConfiguration: + cfg = MergeConfiguration( + merge_method=merge_method, + base_model=base_model, + models=[ + InputModelDefinition(model=model_a, parameters={"weight": 0.5}), + InputModelDefinition(model=model_b, parameters={"weight": 0.5}), + ], + dtype="bfloat16", + ) + if params: + cfg.parameters = params + return cfg + + def test_gpt2_copy(self, gpt2_like): + config = MergeConfiguration( + merge_method="passthrough", + models=[InputModelDefinition(model=gpt2_like)], + dtype="bfloat16", + ) + _run_and_check_merge(config) + + def test_gpt2_stack(self, gpt2_like): + config = MergeConfiguration( + merge_method="passthrough", + slices=[ + OutputSliceDefinition( + sources=[InputSliceDefinition(model=gpt2_like, layer_range=(0, 4))] + ) + ] + * 2, + dtype="bfloat16", + ) + + def _validate(model_path: str): + model_config = AutoConfig.from_pretrained(model_path) + assert model_config.n_layer == 8 + + _run_and_check_merge(config, validate=_validate) + + def test_passthrough_scale(self, model_a): + config = MergeConfiguration( + merge_method="passthrough", + models=[ + InputModelDefinition( + model=model_a, + parameters={ + "scale": [ + {"filter": "o_proj", "value": 0}, + {"value": 1}, + ] + }, + ) + ], + dtype="bfloat16", + ) + + def _validate(model_path: str): + loader = LazyTensorLoader.from_disk(model_path, lazy_loader=False) + saw_any = False + for name in loader.index.tensor_paths: + if "o_proj" in name: + param = loader.get_tensor(name).asnumpy() + assert (param == 0).all() + saw_any = True + assert saw_any, "No o_proj parameters found" + + _run_and_check_merge(config, validate=_validate) + + def test_linear_merge(self, model_a, model_b): + config = self._two_model_config(model_a, model_b, merge_method="linear") + _run_and_check_merge(config) + + def test_slerp_merge(self, model_a, model_b): + config = self._two_model_config( + model_a, + model_b, + merge_method="slerp", + base_model=model_a, + params={"t": 0.35}, + ) + _run_and_check_merge(config) + + def test_slerp_merge_chunked(self, model_a, model_b): + config = self._two_model_config( + model_a, + model_b, + merge_method="slerp", + base_model=model_a, + params={"t": 0.35}, + ) + _run_and_check_merge( + config, + options=MergeOptions( + device="CPU", + max_tensor_mem_gb=0.000001, + split_pieces=2, + ), + ) + + def test_nuslerp_merge(self, model_a, model_b, model_c): + config = self._two_model_config( + model_a, + model_b, + merge_method="nuslerp", + base_model=model_c, + params={"nuslerp_row_wise": False, "nuslerp_flatten": False}, + ) + _run_and_check_merge(config) + + def test_nuslerp_merge_chunked(self, model_a, model_b, model_c): + config = self._two_model_config( + model_a, + model_b, + merge_method="nuslerp", + base_model=model_c, + params={"nuslerp_row_wise": False, "nuslerp_flatten": False}, + ) + _run_and_check_merge( + config, + options=MergeOptions( + device="CPU", + max_tensor_mem_gb=0.000001, + split_pieces=2, + ), + ) + + def test_task_arithmetic_merge(self, model_a, model_b, model_c): + config = self._two_model_config( + model_a, model_b, merge_method="task_arithmetic", base_model=model_c + ) + _run_and_check_merge(config) + + def test_ties_merge(self, model_a, model_b, model_c): + config = self._two_model_config( + model_a, + model_b, + merge_method="ties", + base_model=model_c, + params={"density": 0.3}, + ) + _run_and_check_merge(config) + + def test_sce_merge(self, model_a, model_b, model_c): + config = self._two_model_config( + model_a, + model_b, + merge_method="sce", + base_model=model_c, + params={"select_topk": 0.5}, + ) + _run_and_check_merge(config) + + def test_ram_merge(self, model_a, model_b, model_c): + config = self._two_model_config( + model_a, + model_b, + merge_method="ram", + base_model=model_c, + ) + _run_and_check_merge(config) + + def test_multislerp_merge(self, model_a, model_b, model_c): + config = self._two_model_config( + model_a, + model_b, + merge_method="multislerp", + base_model=model_c, + ) + _run_and_check_merge(config) + + def test_model_stock_merge(self, model_a, model_b, model_c): + config = self._two_model_config( + model_a, + model_b, + merge_method="model_stock", + base_model=model_c, + ) + _run_and_check_merge(config) + + def test_model_stock_filterwise_chunked_merge(self, model_a, model_b, model_c): + config = self._two_model_config( + model_a, + model_b, + merge_method="model_stock", + base_model=model_c, + params={"filter_wise": True}, + ) + _run_and_check_merge( + config, + options=MergeOptions( + device="CPU", + max_tensor_mem_gb=0.000001, + split_pieces=2, + ), + ) + + def test_arcee_fusion_merge(self, model_a, model_b): + config = self._two_model_config( + model_a, + model_b, + merge_method="arcee_fusion", + base_model=model_a, + ) + _run_and_check_merge(config) + + def test_arcee_fusion_merge_chunked(self, model_a, model_b): + config = self._two_model_config( + model_a, + model_b, + merge_method="arcee_fusion", + base_model=model_a, + ) + _run_and_check_merge( + config, + options=MergeOptions( + device="CPU", + max_tensor_mem_gb=0.000001, + split_pieces=2, + ), + ) + + def test_karcher_merge(self, model_a, model_b, model_c): + config = self._two_model_config( + model_a, + model_b, + merge_method="karcher", + base_model=model_c, + params={"max_iter": 5, "tol": 1e-5}, + ) + _run_and_check_merge(config) + + def test_karcher_merge_chunked(self, model_a, model_b, model_c): + config = self._two_model_config( + model_a, + model_b, + merge_method="karcher", + base_model=model_c, + params={"max_iter": 5, "tol": 1e-5}, + ) + _run_and_check_merge( + config, + options=MergeOptions( + device="CPU", + max_tensor_mem_gb=0.000001, + split_pieces=2, + ), + ) + + def test_della_merge(self, model_a, model_b, model_c): + config = self._two_model_config( + model_a, + model_b, + merge_method="della", + base_model=model_c, + params={"density": 0.66, "epsilon": 0.05, "lambda": 0.5}, + ) + _run_and_check_merge(config) + + def test_della_merge_int8_mask_chunked(self, model_a, model_b, model_c): + config = self._two_model_config( + model_a, + model_b, + merge_method="della", + base_model=model_c, + params={ + "density": 0.66, + "epsilon": 0.05, + "lambda": 0.5, + "int8_mask": True, + }, + ) + _run_and_check_merge( + config, + options=MergeOptions( + device="CPU", + max_tensor_mem_gb=0.000001, + split_pieces=2, + ), + ) diff --git a/tests/mindnlp/wizard/test_bf16_ckpt_safety.py b/tests/mindnlp/wizard/test_bf16_ckpt_safety.py new file mode 100644 index 000000000..fc52fe2a1 --- /dev/null +++ b/tests/mindnlp/wizard/test_bf16_ckpt_safety.py @@ -0,0 +1,129 @@ +# Copyright 2026 MindSpore Wizard Team +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""P0: BFloat16 precision regression tests. + +Verifies that bfloat16 tensors survive a write → read round trip +through both safetensors and ckpt formats *without* precision loss. +This is a direct regression test for the historical .bin bf16 bug +where raw bytes were misinterpreted as uint16. +""" + +import tempfile + +import mindspore +import numpy as np +import pytest + +from mindnlp.wizard.merge.io.tensor_writer import TensorWriter +from mindnlp.wizard.merge.io.lazy_tensor_loader import LazyTensorLoader +from mindnlp.wizard.merge.io.loader import TensorLoader, LazyCkptLoader +from mindnlp.wizard.merge.dtype_policy import ( + mindspore_to_numpy, + numpy_to_mindspore, +) + + +def _bf16_reference_values(): + """Known bf16 values whose byte patterns differ from uint16/float16.""" + try: + import ml_dtypes + except ImportError: + pytest.skip("ml_dtypes not installed") + + values = np.array( + [0.1, -0.1, 1.5, -1.5, 3.14, 65504.0, 1e-7, -1e-7], + dtype=np.float32, + ) + bf16_arr = values.astype(ml_dtypes.bfloat16) + return bf16_arr + + +class TestBF16PrecisionSafety: + """BF16 round-trip through safetensors and ckpt must be bit-exact.""" + + def _round_trip_bf16(self, output_format: str, lazy_loader: bool = True): + try: + import ml_dtypes # noqa: F401 + except ImportError: + pytest.skip("ml_dtypes not installed") + + bf16_arr = _bf16_reference_values() + tensor = numpy_to_mindspore(bf16_arr) + assert tensor.dtype == mindspore.bfloat16 + + with tempfile.TemporaryDirectory() as tmpdir: + writer = TensorWriter(tmpdir, output_format=output_format) + writer.save_tensor("bf16_test", tensor) + writer.finalize() + + loader = LazyTensorLoader.from_disk(tmpdir, lazy_loader=lazy_loader) + restored = loader.get_tensor("bf16_test") + + assert restored.dtype == mindspore.bfloat16 + + original_bytes = mindspore_to_numpy(tensor).tobytes() + restored_bytes = mindspore_to_numpy(restored).tobytes() + assert original_bytes == restored_bytes, ( + "BF16 byte-level mismatch detected — precision loss in round trip" + ) + + original_f32 = mindspore_to_numpy(tensor).astype(np.float32) + restored_f32 = mindspore_to_numpy(restored).astype(np.float32) + np.testing.assert_array_equal(original_f32, restored_f32) + + @pytest.mark.parametrize("fmt", ["safetensors", "ckpt"]) + def test_bf16_round_trip_lazy(self, fmt): + self._round_trip_bf16(fmt, lazy_loader=True) + + @pytest.mark.parametrize("fmt", ["safetensors", "ckpt"]) + def test_bf16_round_trip_eager(self, fmt): + self._round_trip_bf16(fmt, lazy_loader=False) + + def test_bf16_numpy_mindspore_conversion_identity(self): + """numpy_to_mindspore(mindspore_to_numpy(t)) must be bit-exact for bf16.""" + try: + import ml_dtypes + except ImportError: + pytest.skip("ml_dtypes not installed") + + bf16_arr = _bf16_reference_values() + tensor = numpy_to_mindspore(bf16_arr) + arr_back = mindspore_to_numpy(tensor) + + assert arr_back.dtype == ml_dtypes.bfloat16 + np.testing.assert_array_equal( + bf16_arr.view(np.uint16), + arr_back.view(np.uint16), + ) + + def test_bf16_not_misinterpreted_as_uint16(self): + """Guard against the specific historical bug: treating bf16 bytes as uint16.""" + try: + import ml_dtypes + except ImportError: + pytest.skip("ml_dtypes not installed") + + bf16_arr = np.array([3.14], dtype=ml_dtypes.bfloat16) + raw_bytes = bf16_arr.tobytes() + + correct = np.frombuffer(raw_bytes, dtype=ml_dtypes.bfloat16) + wrong = np.frombuffer(raw_bytes, dtype=np.uint16) + + correct_f32 = float(correct[0]) + wrong_f32 = float(wrong[0]) + assert abs(correct_f32 - 3.14) < 0.1 + assert abs(wrong_f32 - 3.14) > 1.0, ( + "Test sanity check failed — uint16 interpretation should differ" + ) diff --git a/tests/mindnlp/wizard/test_bf16_cpu_full_method_matrix.py b/tests/mindnlp/wizard/test_bf16_cpu_full_method_matrix.py new file mode 100644 index 000000000..c57e77c47 --- /dev/null +++ b/tests/mindnlp/wizard/test_bf16_cpu_full_method_matrix.py @@ -0,0 +1,136 @@ +import os +import tempfile +from typing import Any, Dict, List, Tuple + +import mindspore +import pytest +from transformers import LlamaConfig, LlamaForCausalLM + +from mindnlp.wizard.merge.config import InputModelDefinition, MergeConfiguration +from mindnlp.wizard.merge.merge import MergeOptions, run_merge +from mindnlp.wizard.merge.merge_methods import REGISTERED_MERGE_METHODS + + +def _make_tiny_llama(path: str, vocab_size: int = 64) -> str: + cfg = LlamaConfig( + vocab_size=vocab_size, + hidden_size=32, + intermediate_size=48, + num_attention_heads=4, + num_hidden_layers=2, + ) + model = LlamaForCausalLM(cfg) + model.save_pretrained(path, safe_serialization=True) + return path + + +@pytest.fixture(scope="session") +def tiny_triplet(tmp_path_factory) -> Tuple[str, str, str]: + a = _make_tiny_llama(str(tmp_path_factory.mktemp("bf16_matrix_a"))) + b = _make_tiny_llama(str(tmp_path_factory.mktemp("bf16_matrix_b"))) + c = _make_tiny_llama(str(tmp_path_factory.mktemp("bf16_matrix_c"))) + return a, b, c + + +def _method_recipe(method_name: str, models: Tuple[str, str, str]) -> Dict[str, Any]: + a, b, c = models + base_recipe: Dict[str, Any] = { + "merge_method": method_name, + "dtype": "bfloat16", + } + + # Known stable shapes from existing wizard parity tests. + if method_name == "passthrough": + base_recipe["models"] = [ + {"model": a, "parameters": {"scale": 1.0}}, + ] + return base_recipe + + if method_name in {"slerp", "arcee_fusion", "nearswap"}: + base_recipe["base_model"] = a + base_recipe["models"] = [ + {"model": a, "parameters": {"weight": 0.5}}, + {"model": b, "parameters": {"weight": 0.5}}, + ] + base_recipe["parameters"] = {"t": 0.5} + return base_recipe + + if method_name == "nuslerp": + base_recipe["base_model"] = c + base_recipe["models"] = [ + {"model": a, "parameters": {"weight": 0.5}}, + {"model": b, "parameters": {"weight": 0.5}}, + ] + base_recipe["parameters"] = { + "nuslerp_row_wise": False, + "nuslerp_flatten": False, + } + return base_recipe + + if method_name == "model_stock": + base_recipe["base_model"] = c + base_recipe["models"] = [ + {"model": a, "parameters": {"weight": 0.5}}, + {"model": b, "parameters": {"weight": 0.5}}, + ] + return base_recipe + + if method_name in {"sce"}: + base_recipe["base_model"] = c + base_recipe["models"] = [ + {"model": a, "parameters": {"weight": 0.5}}, + {"model": b, "parameters": {"weight": 0.5}}, + ] + base_recipe["parameters"] = {"select_topk": 0.5} + return base_recipe + + if method_name in {"ramplus_tl"}: + base_recipe["base_model"] = c + base_recipe["models"] = [ + {"model": a, "parameters": {"weight": 0.5}}, + {"model": b, "parameters": {"weight": 0.5}}, + ] + base_recipe["parameters"] = {"r": 0.1, "alpha": 0.2} + return base_recipe + + if method_name in { + "task_arithmetic", + "ties", + "dare_ties", + "dare_linear", + "breadcrumbs", + "breadcrumbs_ties", + "della", + "della_linear", + "ram", + "multislerp", + "karcher", + }: + base_recipe["base_model"] = c + base_recipe["models"] = [ + {"model": a, "parameters": {"weight": 0.5}}, + {"model": b, "parameters": {"weight": 0.5}}, + ] + if method_name in {"ties", "dare_ties", "dare_linear", "breadcrumbs", "breadcrumbs_ties", "della", "della_linear"}: + base_recipe["parameters"] = {"density": 0.5} + return base_recipe + + # Fallback for newly added methods: 2-model merge. + base_recipe["models"] = [ + {"model": a, "parameters": {"weight": 0.5}}, + {"model": b, "parameters": {"weight": 0.5}}, + ] + return base_recipe + + +@pytest.mark.parametrize("method_name", sorted(list(REGISTERED_MERGE_METHODS.keys()))) +def test_registered_methods_bf16_cpu_safe(method_name: str, tiny_triplet: Tuple[str, str, str]): + mindspore.set_context(device_target="CPU") + recipe = _method_recipe(method_name, tiny_triplet) + cfg = MergeConfiguration.model_validate(recipe) + with tempfile.TemporaryDirectory() as out_dir: + run_merge(cfg, out_path=out_dir, options=MergeOptions(device="CPU")) + assert os.path.exists(os.path.join(out_dir, "config.json")) + assert os.path.exists(os.path.join(out_dir, "model.safetensors")) or os.path.exists( + os.path.join(out_dir, "model.safetensors.index.json") + ) diff --git a/tests/mindnlp/wizard/test_ckpt_io.py b/tests/mindnlp/wizard/test_ckpt_io.py new file mode 100644 index 000000000..8b3da0afd --- /dev/null +++ b/tests/mindnlp/wizard/test_ckpt_io.py @@ -0,0 +1,275 @@ +# Copyright 2026 MindSpore Wizard Team +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Tests for .ckpt read/write support and bf16 precision safety.""" + +import json +import os +import tempfile +from typing import Dict + +import mindspore +import numpy as np +import pytest + +from mindnlp.wizard.merge.io.tensor_writer import TensorWriter +from mindnlp.wizard.merge.io.lazy_tensor_loader import ShardedTensorIndex, LazyTensorLoader +from mindnlp.wizard.merge.io.loader import ( + TensorLoader, + LazyCkptLoader, + DumbCkptLoader, + SafetensorsLoader, +) +from mindnlp.wizard.merge.options import MergeOptions, VALID_OUTPUT_FORMATS + + +# ── Helpers ─────────────────────────────────────────────────────────────── + +def _make_tensors() -> Dict[str, mindspore.Tensor]: + return { + "layer.0.weight": mindspore.Tensor(np.random.randn(64, 128).astype(np.float32)), + "layer.0.bias": mindspore.Tensor(np.random.randn(64).astype(np.float32)), + "layer.1.weight": mindspore.Tensor(np.random.randn(32, 64).astype(np.float32)), + } + + +def _write_and_read_round_trip( + output_format: str, + tensors: Dict[str, mindspore.Tensor], + max_shard_size: int = -1, + lazy_loader: bool = True, +) -> Dict[str, mindspore.Tensor]: + with tempfile.TemporaryDirectory() as tmpdir: + writer = TensorWriter( + tmpdir, + max_shard_size=max_shard_size, + output_format=output_format, + ) + for name, tensor in tensors.items(): + writer.save_tensor(name, tensor) + writer.finalize() + + loader = LazyTensorLoader.from_disk(tmpdir, lazy_loader=lazy_loader) + result = {} + for name in tensors: + result[name] = loader.get_tensor(name) + return result + + +# ── TensorWriter format tests ──────────────────────────────────────────── + +class TestTensorWriterFormats: + def test_safetensors_single_shard(self): + tensors = _make_tensors() + result = _write_and_read_round_trip("safetensors", tensors) + for name, original in tensors.items(): + np.testing.assert_allclose( + result[name].asnumpy(), original.asnumpy(), rtol=0, atol=0, + ) + + def test_safetensors_multi_shard(self): + tensors = _make_tensors() + result = _write_and_read_round_trip("safetensors", tensors, max_shard_size=4096) + for name, original in tensors.items(): + np.testing.assert_allclose( + result[name].asnumpy(), original.asnumpy(), rtol=0, atol=0, + ) + + def test_ckpt_single_shard(self): + tensors = _make_tensors() + result = _write_and_read_round_trip("ckpt", tensors) + for name, original in tensors.items(): + np.testing.assert_allclose( + result[name].asnumpy(), original.asnumpy(), rtol=0, atol=0, + ) + + def test_ckpt_multi_shard(self): + tensors = _make_tensors() + result = _write_and_read_round_trip("ckpt", tensors, max_shard_size=4096) + for name, original in tensors.items(): + np.testing.assert_allclose( + result[name].asnumpy(), original.asnumpy(), rtol=0, atol=0, + ) + + def test_unsupported_format_raises(self): + with tempfile.TemporaryDirectory() as tmpdir: + with pytest.raises(ValueError, match="Unsupported output_format"): + TensorWriter(tmpdir, output_format="bin") + + def test_ckpt_output_file_naming(self): + """Single ckpt shard should be named mindspore_model.ckpt.""" + tensors = _make_tensors() + with tempfile.TemporaryDirectory() as tmpdir: + writer = TensorWriter(tmpdir, output_format="ckpt") + for name, tensor in tensors.items(): + writer.save_tensor(name, tensor) + writer.finalize() + assert os.path.exists(os.path.join(tmpdir, "mindspore_model.ckpt")) + + def test_safetensors_output_file_naming(self): + """Single safetensors shard should be named model.safetensors.""" + tensors = _make_tensors() + with tempfile.TemporaryDirectory() as tmpdir: + writer = TensorWriter(tmpdir, output_format="safetensors") + for name, tensor in tensors.items(): + writer.save_tensor(name, tensor) + writer.finalize() + assert os.path.exists(os.path.join(tmpdir, "model.safetensors")) + + def test_ckpt_sharded_index_json(self): + """Multiple ckpt shards should produce an index.json.""" + tensors = _make_tensors() + with tempfile.TemporaryDirectory() as tmpdir: + writer = TensorWriter(tmpdir, output_format="ckpt", max_shard_size=4096) + for name, tensor in tensors.items(): + writer.save_tensor(name, tensor) + writer.finalize() + + index_path = os.path.join(tmpdir, "mindspore_model.ckpt.index.json") + assert os.path.exists(index_path) + with open(index_path) as f: + idx = json.load(f) + assert "weight_map" in idx + for name in tensors: + assert name in idx["weight_map"] + + +# ── ShardedTensorIndex detection tests ──────────────────────────────────── + +class TestShardedTensorIndex: + def test_from_disk_detects_safetensors(self): + tensors = _make_tensors() + with tempfile.TemporaryDirectory() as tmpdir: + writer = TensorWriter(tmpdir, output_format="safetensors") + for name, tensor in tensors.items(): + writer.save_tensor(name, tensor) + writer.finalize() + + index = ShardedTensorIndex.from_disk(tmpdir) + assert index.format == "safetensors" + assert index.is_safetensors is True + for name in tensors: + assert name in index.tensor_paths + + def test_from_disk_detects_ckpt(self): + tensors = _make_tensors() + with tempfile.TemporaryDirectory() as tmpdir: + writer = TensorWriter(tmpdir, output_format="ckpt") + for name, tensor in tensors.items(): + writer.save_tensor(name, tensor) + writer.finalize() + + index = ShardedTensorIndex.from_disk(tmpdir) + assert index.format == "ckpt" + assert index.is_safetensors is False + for name in tensors: + assert name in index.tensor_paths + + def test_from_disk_empty_dir_raises(self): + with tempfile.TemporaryDirectory() as tmpdir: + with pytest.raises(RuntimeError, match="Unable to find model"): + ShardedTensorIndex.from_disk(tmpdir) + + +# ── Loader factory dispatch tests ───────────────────────────────────────── + +class TestLoaderFactory: + def test_dispatch_safetensors(self): + tensors = _make_tensors() + with tempfile.TemporaryDirectory() as tmpdir: + writer = TensorWriter(tmpdir, output_format="safetensors") + for name, tensor in tensors.items(): + writer.save_tensor(name, tensor) + writer.finalize() + + path = os.path.join(tmpdir, "model.safetensors") + loader = TensorLoader.get(path) + assert isinstance(loader, SafetensorsLoader) + + def test_dispatch_ckpt_lazy(self): + tensors = _make_tensors() + with tempfile.TemporaryDirectory() as tmpdir: + writer = TensorWriter(tmpdir, output_format="ckpt") + for name, tensor in tensors.items(): + writer.save_tensor(name, tensor) + writer.finalize() + + path = os.path.join(tmpdir, "mindspore_model.ckpt") + loader = TensorLoader.get(path, use_lazy_loader=True) + assert isinstance(loader, LazyCkptLoader) + + def test_dispatch_ckpt_eager(self): + tensors = _make_tensors() + with tempfile.TemporaryDirectory() as tmpdir: + writer = TensorWriter(tmpdir, output_format="ckpt") + for name, tensor in tensors.items(): + writer.save_tensor(name, tensor) + writer.finalize() + + path = os.path.join(tmpdir, "mindspore_model.ckpt") + loader = TensorLoader.get(path, use_lazy_loader=False) + assert isinstance(loader, DumbCkptLoader) + + +# ── MergeOptions tests ──────────────────────────────────────────────────── + +class TestMergeOptions: + def test_output_format_default(self): + opts = MergeOptions() + assert opts.output_format == "safetensors" + + def test_output_format_ckpt(self): + opts = MergeOptions(output_format="ckpt") + assert opts.output_format == "ckpt" + + def test_output_format_bin_raises(self): + with pytest.raises(ValueError, match="not supported"): + MergeOptions(output_format="bin") + + def test_invalid_format_raises(self): + with pytest.raises(ValueError, match="Invalid output_format"): + MergeOptions(output_format="pickle") + + def test_lazy_loader_field(self): + opts = MergeOptions(lazy_loader=True) + assert opts.lazy_loader is True + + def test_ckpt_load_kwargs_passthrough(self): + opts = MergeOptions(ckpt_load_kwargs={"dec_key": b"secret"}) + assert opts.ckpt_load_kwargs == {"dec_key": b"secret"} + + +# ── Dtype round-trip matrix ─────────────────────────────────────────────── + +_DTYPE_PAIRS = [ + ("float32", np.float32, mindspore.float32), + ("float16", np.float16, mindspore.float16), + ("int32", np.int32, mindspore.int32), + ("int64", np.int64, mindspore.int64), + ("int8", np.int8, mindspore.int8), + ("uint8", np.uint8, mindspore.uint8), +] + + +class TestDtypeRoundTrip: + @pytest.mark.parametrize("fmt", ["safetensors", "ckpt"]) + @pytest.mark.parametrize("dtype_name,np_dtype,ms_dtype", _DTYPE_PAIRS) + def test_dtype_round_trip(self, fmt, dtype_name, np_dtype, ms_dtype): + arr = np.array([1, 2, 3, 4, 5], dtype=np_dtype) + tensor = mindspore.Tensor(arr) + assert tensor.dtype == ms_dtype + + result = _write_and_read_round_trip(fmt, {"test_tensor": tensor}) + out = result["test_tensor"] + np.testing.assert_array_equal(out.asnumpy(), arr) diff --git a/tests/mindnlp/wizard/test_cli_run_yaml_compat.py b/tests/mindnlp/wizard/test_cli_run_yaml_compat.py new file mode 100644 index 000000000..159d8d7dd --- /dev/null +++ b/tests/mindnlp/wizard/test_cli_run_yaml_compat.py @@ -0,0 +1,109 @@ +import os +from pathlib import Path + +import yaml +from click.testing import CliRunner +from transformers import LlamaConfig, LlamaForCausalLM + +from mindnlp.wizard.merge.common import ModelReference +from mindnlp.wizard.merge.scripts.run_yaml import main as run_yaml_main + + +def _make_tiny_llama(path: str, vocab_size: int = 64) -> str: + cfg = LlamaConfig( + vocab_size=vocab_size, + hidden_size=32, + intermediate_size=48, + num_attention_heads=4, + num_hidden_layers=2, + ) + model = LlamaForCausalLM(cfg) + model.save_pretrained(path, safe_serialization=True) + return path + + +def test_run_yaml_cli_end_to_end_cpu(tmp_path): + model_a = _make_tiny_llama(str(tmp_path / "cli_model_a")) + model_b = _make_tiny_llama(str(tmp_path / "cli_model_b")) + out_dir = tmp_path / "out" + recipe_file = tmp_path / "recipe.yaml" + + recipe = { + "merge_method": "linear", + "dtype": "bfloat16", + "models": [ + {"model": model_a, "parameters": {"weight": 0.6}}, + {"model": model_b, "parameters": {"weight": 0.4}}, + ], + } + recipe_file.write_text(yaml.safe_dump(recipe, sort_keys=False), encoding="utf-8") + + runner = CliRunner() + result = runner.invoke( + run_yaml_main, + [ + str(recipe_file), + str(out_dir), + "--device", + "CPU", + "--no-multi-npu", + "--max-tensor-mem-gb", + "0.000001", + "--split-pieces", + "2", + ], + ) + + assert result.exit_code == 0, result.output + assert (out_dir / "config.json").exists() + assert (out_dir / "model.safetensors").exists() or ( + out_dir / "model.safetensors.index.json" + ).exists() + # E2E CLI should persist recipe + execution report for debugging. + assert (out_dir / "wizard_config.yml").exists() + assert (out_dir / "wizard_execution_report.json").exists() + + +def test_run_yaml_cli_lora_reference_compat_cpu(tmp_path, monkeypatch): + model_a = _make_tiny_llama(str(tmp_path / "cli_lora_model_a")) + model_b = _make_tiny_llama(str(tmp_path / "cli_lora_model_b")) + lora_a = tmp_path / "fake_lora_a" + lora_b = tmp_path / "fake_lora_b" + lora_a.mkdir() + lora_b.mkdir() + + out_dir = tmp_path / "out_lora" + recipe_file = tmp_path / "recipe_lora.yaml" + + recipe = { + "merge_method": "linear", + "dtype": "bfloat16", + "models": [ + {"model": f"{model_a}+{lora_a}", "parameters": {"weight": 0.6}}, + {"model": f"{model_b}+{lora_b}", "parameters": {"weight": 0.4}}, + ], + } + recipe_file.write_text(yaml.safe_dump(recipe, sort_keys=False), encoding="utf-8") + + def _fake_merged( + self, + cache_dir=None, + trust_remote_code=False, + lora_merge_dtype=None, + ): + # Keep test offline/stable: emulate successful LoRA merge output path. + return ModelReference(model=self.model) + + monkeypatch.setattr(ModelReference, "merged", _fake_merged, raising=False) + + runner = CliRunner() + result = runner.invoke( + run_yaml_main, + [str(recipe_file), str(out_dir), "--device", "CPU", "--no-multi-npu"], + ) + + assert result.exit_code == 0, result.output + assert (out_dir / "config.json").exists() + assert (out_dir / "model.safetensors").exists() or ( + out_dir / "model.safetensors.index.json" + ).exists() diff --git a/tests/mindnlp/wizard/test_config_validation_matrix.py b/tests/mindnlp/wizard/test_config_validation_matrix.py new file mode 100644 index 000000000..0dc6d1b97 --- /dev/null +++ b/tests/mindnlp/wizard/test_config_validation_matrix.py @@ -0,0 +1,47 @@ +import pytest + +from mindnlp.wizard.merge.common import ModelReference +from mindnlp.wizard.merge.config import InputModelDefinition, MergeConfiguration + + +def _model(name: str) -> InputModelDefinition: + return InputModelDefinition(model=ModelReference.parse(name)) + + +@pytest.mark.parametrize( + "cfg,err", + [ + ( + { + "merge_method": "linear", + }, + "Exactly one of models, slices, or modules must be specified", + ), + ( + { + "merge_method": "linear", + "models": [_model("a"), _model("b")], + "tokenizer_source": "base", + "tokenizer": {"source": "union"}, + }, + "Cannot specify both tokenizer_source and tokenizer", + ), + ( + { + "merge_method": "slerp", + "models": [_model("a"), _model("b")], + }, + "requires base_model", + ), + ( + { + "merge_method": "nearswap", + "models": [_model("a"), _model("b")], + }, + "requires base_model", + ), + ], +) +def test_config_validation_matrix(cfg, err): + with pytest.raises(RuntimeError, match=err): + MergeConfiguration.model_validate(cfg) diff --git a/tests/mindnlp/wizard/test_dtype_policy.py b/tests/mindnlp/wizard/test_dtype_policy.py new file mode 100644 index 000000000..720077ed2 --- /dev/null +++ b/tests/mindnlp/wizard/test_dtype_policy.py @@ -0,0 +1,35 @@ +import mindspore + +from mindnlp.wizard.merge.dtype_policy import ( + cast_back, + cast_to_work, + choose_work_dtype, + needs_safe_path, +) + + +def test_choose_work_dtype_promotes_half_on_cpu(monkeypatch): + monkeypatch.setattr(mindspore, "get_context", lambda *_args, **_kwargs: "CPU") + assert choose_work_dtype(mindspore.bfloat16) == mindspore.float32 + assert choose_work_dtype(mindspore.float16) == mindspore.float32 + + +def test_choose_work_dtype_keeps_fp32_on_cpu(monkeypatch): + monkeypatch.setattr(mindspore, "get_context", lambda *_args, **_kwargs: "CPU") + assert choose_work_dtype(mindspore.float32) == mindspore.float32 + assert needs_safe_path(mindspore.float32) is False + + +def test_choose_work_dtype_keeps_half_on_ascend(monkeypatch): + monkeypatch.setattr(mindspore, "get_context", lambda *_args, **_kwargs: "Ascend") + assert choose_work_dtype(mindspore.bfloat16) == mindspore.bfloat16 + assert choose_work_dtype(mindspore.float16) == mindspore.float16 + + +def test_cast_roundtrip(): + tensor = mindspore.Tensor([1.0, 2.0], dtype=mindspore.float16) + work = cast_to_work(tensor, mindspore.float32) + assert work.dtype == mindspore.float32 + out = cast_back(work, mindspore.float16) + assert out.dtype == mindspore.float16 + diff --git a/tests/mindnlp/wizard/test_merge.py b/tests/mindnlp/wizard/test_merge.py new file mode 100644 index 000000000..5212f1fb6 --- /dev/null +++ b/tests/mindnlp/wizard/test_merge.py @@ -0,0 +1,1147 @@ +# Copyright 2026 MindSpore Wizard Team +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Comprehensive unit tests for the Wizard MindSpore merge package.""" + +import json +import logging +import os +import tempfile +from typing import Any, Dict, List, Optional, Tuple + +import immutables +import mindspore +import mindspore.ops as ops +import numpy as np +import pytest +import yaml + +from mindnlp.wizard.merge.graph import ( + Executor, + ExecutionSchedule, + Task, + TaskHandle, + TaskUniverse, + build_schedule, + _parse_device, +) +from mindnlp.wizard.merge.multigpu_executor import MultiDeviceExecutor +from mindnlp.wizard.merge.common import ( + ImmutableMap, + ModelPath, + ModelReference, + dtype_from_name, + get_accelerator_count, + parse_kmb, +) +from mindnlp.wizard.merge.config import ( + ConfigReader, + ConditionalParameter, + InputModelDefinition, + InputSliceDefinition, + MergeConfiguration, + OutputSliceDefinition, + evaluate_setting, +) +from mindnlp.wizard.merge.io.lazy_tensor_loader import ShardedTensorIndex +from mindnlp.wizard.merge.io.tensor_writer import TensorWriter +from mindnlp.wizard.merge.io.tasks import LoadTensor, LoaderCache +from mindnlp.wizard.merge.merge_methods.registry import REGISTERED_MERGE_METHODS +from mindnlp.wizard.merge.merge_methods.slerp import lerp, slerp, normalize +from mindnlp.wizard.merge.sparsify import ( + SparsificationMethod, + RescaleNorm, + magnitude, + bernoulli, + magnitude_outliers, + sparsify, +) +from mindnlp.wizard.merge.options import MergeOptions +from mindnlp.wizard.merge.merge import _write_execution_report + + +# =================================================================== +# Helper task classes for graph tests +# =================================================================== + +class ConstantTask(Task[int]): + value: int + + def arguments(self) -> Dict[str, Task]: + return {} + + def execute(self, **kwargs) -> int: + return self.value + + +class AddTask(Task[int]): + a: Task[int] + b: Task[int] + + def arguments(self) -> Dict[str, Task]: + return {"a": self.a, "b": self.b} + + def execute(self, a: int, b: int, **kwargs) -> int: + return a + b + + +class MultiplyTask(Task[int]): + a: Task[int] + b: Task[int] + + def arguments(self) -> Dict[str, Task]: + return {"a": self.a, "b": self.b} + + def execute(self, a: int, b: int, **kwargs) -> int: + return a * b + + +class TensorAddTask(Task[mindspore.Tensor]): + a: Task[mindspore.Tensor] + b: Task[mindspore.Tensor] + + def uses_accelerator(self) -> bool: + return True + + def arguments(self) -> Dict[str, Task]: + return {"a": self.a, "b": self.b} + + def execute(self, a: mindspore.Tensor, b: mindspore.Tensor, **kwargs): + return a + b + + +class TensorConstTask(Task[mindspore.Tensor]): + data: Tuple[float, ...] + + def arguments(self) -> Dict[str, Task]: + return {} + + def execute(self, **kwargs) -> mindspore.Tensor: + return mindspore.Tensor(list(self.data), dtype=mindspore.float32) + + +class LabeledConstTask(Task[int]): + value: int + label: str + + def arguments(self) -> Dict[str, Task]: + return {} + + def execute(self, **kwargs) -> int: + return self.value + + def group_label(self) -> Optional[str]: + return self.label + + +class CostHintTask(Task[int]): + value: int + read_cost: float = 0.0 + compute_cost: float = 0.0 + write_cost: float = 0.0 + + def arguments(self) -> Dict[str, Task]: + return {} + + def execute(self, **kwargs) -> int: + return self.value + + def group_label(self) -> Optional[str]: + return "same-group" + + def cost_hint(self): + return { + "read": self.read_cost, + "compute": self.compute_cost, + "write": self.write_cost, + } + + +class CostJoinTask(Task[int]): + left: Task[int] + right: Task[int] + + def arguments(self) -> Dict[str, Task]: + return {"left": self.left, "right": self.right} + + def execute(self, left: int, right: int, **kwargs) -> int: + return left + right + + +# =================================================================== +# 1. graph.py tests +# =================================================================== + +class TestTaskFrozen: + + def test_frozen_and_hashable(self): + t = ConstantTask(value=42) + with pytest.raises(Exception): + t.value = 99 + + def test_equal_tasks_share_hash(self): + t1 = ConstantTask(value=1) + t2 = ConstantTask(value=1) + assert t1 == t2 and hash(t1) == hash(t2) + + +class TestTaskUniverse: + + def test_add_task_returns_handle(self): + u = TaskUniverse() + t = ConstantTask(value=5) + h = u.add_task(t) + assert isinstance(h, TaskHandle) + assert h.task() is t + + def test_deduplication(self): + u = TaskUniverse() + t = ConstantTask(value=5) + h1 = u.add_task(t) + h2 = u.add_task(t) + assert h1 == h2 + assert len(u.tasks) == 1 + + def test_recursive_add(self): + c1 = ConstantTask(value=2) + c2 = ConstantTask(value=3) + add = AddTask(a=c1, b=c2) + u = TaskUniverse() + h = u.add_task(add, recursive=True) + assert len(u.tasks) == 3 + assert h.task() is add + + def test_get_handle(self): + u = TaskUniverse() + t = ConstantTask(value=7) + u.add_task(t) + assert u.get_handle(t) is not None + + def test_get_handle_missing(self): + u = TaskUniverse() + assert u.get_handle(ConstantTask(value=7)) is None + + +class TestBuildSchedule: + + def test_empty_targets(self): + sched = build_schedule([], {}) + assert sched.tasks == [] + + def test_topological_ordering(self): + c1 = ConstantTask(value=1) + c2 = ConstantTask(value=2) + add = AddTask(a=c1, b=c2) + u = TaskUniverse() + h_add = u.add_task(add) + sched = build_schedule([h_add], {}) + task_types = [type(th.task()).__name__ for th in sched.tasks] + add_idx = task_types.index("AddTask") + for dep_type in ("ConstantTask",): + assert task_types.index(dep_type) < add_idx + + def test_cached_values_skip(self): + c1 = ConstantTask(value=1) + c2 = ConstantTask(value=2) + add = AddTask(a=c1, b=c2) + u = TaskUniverse() + h_add = u.add_task(add) + h_c1 = u.get_handle(c1) + sched = build_schedule([h_add], {h_c1: 99}) + scheduled_tasks = [th.task() for th in sched.tasks] + assert c1 not in scheduled_tasks + assert add in scheduled_tasks + + +class TestExecutor: + + def test_simple_dag(self): + c1 = ConstantTask(value=3) + c2 = ConstantTask(value=4) + add = AddTask(a=c1, b=c2) + ex = Executor(targets=[add], math_device="CPU", storage_device="CPU") + results = dict(ex.run(quiet=True)) + assert results[add] == 7 + + def test_diamond_dag(self): + c = ConstantTask(value=2) + a1 = AddTask(a=c, b=c) + a2 = MultiplyTask(a=c, b=c) + final = AddTask(a=a1, b=a2) + ex = Executor(targets=[final], math_device="CPU") + results = dict(ex.run(quiet=True)) + assert results[final] == (2 + 2) + (2 * 2) + + def test_tensor_dag(self): + t1 = TensorConstTask(data=(1.0, 2.0, 3.0)) + t2 = TensorConstTask(data=(4.0, 5.0, 6.0)) + add = TensorAddTask(a=t1, b=t2) + ex = Executor(targets=[add], math_device="CPU", storage_device="CPU") + results = dict(ex.run(quiet=True)) + out = results[add] + assert isinstance(out, mindspore.Tensor) + np.testing.assert_allclose(out.asnumpy(), [5.0, 7.0, 9.0]) + + def test_cached_values(self): + c1 = ConstantTask(value=10) + c2 = ConstantTask(value=20) + add = AddTask(a=c1, b=c2) + u = TaskUniverse() + h_add = u.add_task(add) + h_c1 = u.get_handle(c1) + h_c2 = u.get_handle(c2) + ex = Executor( + targets=[h_add], + cached_values={h_c1: 100, h_c2: 200}, + ) + results = dict(ex.run(quiet=True)) + assert results[add] == 300 + + +class TestParseDevice: + + def test_cpu(self): + target, dev_id = _parse_device("CPU") + assert target == "CPU" and dev_id is None + + def test_ascend_with_id(self): + target, dev_id = _parse_device("Ascend:0") + assert target == "Ascend" and dev_id == 0 + + +# =================================================================== +# 2. common.py tests +# =================================================================== + +class TestModelPath: + + def test_from_string_no_revision(self): + mp = ModelPath.model_validate("my-org/my-model") + assert mp.path == "my-org/my-model" + assert mp.revision is None + + def test_from_string_with_revision(self): + mp = ModelPath.model_validate("my-org/my-model@main") + assert mp.path == "my-org/my-model" + assert mp.revision == "main" + + def test_str_roundtrip(self): + assert str(ModelPath(path="foo/bar")) == "foo/bar" + assert str(ModelPath(path="foo/bar", revision="dev")) == "foo/bar@dev" + + def test_invalid_multiple_at(self): + with pytest.raises(RuntimeError, match="multiple @"): + ModelPath.model_validate("a@b@c") + + def test_frozen(self): + mp = ModelPath(path="x") + with pytest.raises(Exception): + mp.path = "y" + + +class TestModelReference: + + def test_from_string_model_only(self): + mr = ModelReference.model_validate("my-org/model") + assert mr.model.path == "my-org/model" + assert mr.lora is None + + def test_from_string_model_plus_lora(self): + mr = ModelReference.model_validate("base-model+lora-adapter") + assert mr.model.path == "base-model" + assert mr.lora.path == "lora-adapter" + + def test_str_roundtrip(self): + assert str(ModelReference(model=ModelPath(path="abc"))) == "abc" + mr = ModelReference( + model=ModelPath(path="abc"), lora=ModelPath(path="lora") + ) + assert str(mr) == "abc+lora" + + def test_parse_classmethod(self): + mr = ModelReference.parse("org/m@v1+org/l@v2") + assert mr.model.path == "org/m" and mr.model.revision == "v1" + assert mr.lora.path == "org/l" and mr.lora.revision == "v2" + + def test_equality(self): + mr1 = ModelReference.model_validate("model-a") + mr2 = ModelReference.model_validate("model-a") + mr3 = ModelReference.model_validate("model-b") + assert mr1 == mr2 + assert mr1 != mr3 + + +class TestDtypeFromName: + + @pytest.mark.parametrize( + "name,expected", + [ + ("float16", mindspore.float16), + ("float32", mindspore.float32), + ("bfloat16", mindspore.bfloat16), + ("float64", mindspore.float64), + ("int32", mindspore.int32), + ], + ) + def test_valid(self, name, expected): + assert dtype_from_name(name) == expected + + def test_with_prefix(self): + assert dtype_from_name("torch.float16") == mindspore.float16 + assert dtype_from_name("mindspore.float32") == mindspore.float32 + assert dtype_from_name("ms.bfloat16") == mindspore.bfloat16 + + def test_none_returns_none(self): + assert dtype_from_name(None) is None + + def test_unimplemented(self): + with pytest.raises(RuntimeError, match="Unimplemented"): + dtype_from_name("complex128") + + +class TestImmutableMap: + + def test_basic_operations(self): + im = ImmutableMap(immutables.Map({"a": 1, "b": 2})) + assert im["a"] == 1 and im["b"] == 2 and len(im) == 2 + + def test_keys_items_values(self): + im = ImmutableMap(immutables.Map({"x": 10, "y": 20})) + assert set(im.keys()) == {"x", "y"} + assert set(im.values()) == {10, 20} + assert set(im.items()) == {("x", 10), ("y", 20)} + + def test_iteration(self): + im = ImmutableMap(immutables.Map({"a": 1})) + assert list(im) == ["a"] + + +class TestParseKmb: + + @pytest.mark.parametrize("inp,expected", [ + (42, 42), ("100", 100), ("5k", 5000), + ("2M", 2_000_000), ("1B", 1_000_000_000), + ]) + def test_values(self, inp, expected): + assert parse_kmb(inp) == expected + + +# =================================================================== +# 3. config.py tests +# =================================================================== + +class TestEvaluateSetting: + + def test_scalar(self): + assert evaluate_setting("w", 0.5) == 0.5 + assert evaluate_setting("w", 1) == 1 + + def test_gradient_list(self): + assert abs(evaluate_setting("w", [0.0, 1.0], t=0.0) - 0.0) < 1e-6 + assert abs(evaluate_setting("w", [0.0, 1.0], t=1.0) - 1.0) < 1e-6 + assert abs(evaluate_setting("w", [0.0, 1.0], t=0.5) - 0.5) < 1e-6 + assert abs(evaluate_setting("w", [0.0, 0.5, 1.0], t=0.5) - 0.5) < 1e-6 + + def test_conditional_matching(self): + cond = ConditionalParameter(value=0.7, filter="attn") + assert evaluate_setting("model.self_attn.q_proj", [cond]) == 0.7 + + def test_conditional_wildcard(self): + cond = ConditionalParameter(value=0.3, filter="*") + assert evaluate_setting("any.tensor", [cond]) == 0.3 + + def test_conditional_no_match(self): + cond = ConditionalParameter(value=0.7, filter="attn") + assert evaluate_setting("model.mlp.gate", [cond]) is None + + def test_conditional_none_filter(self): + cond = ConditionalParameter(value=0.9, filter=None) + assert evaluate_setting("anything", [cond]) == 0.9 + + +class TestMergeConfiguration: + + def _base_config(self): + return MergeConfiguration( + merge_method="linear", + models=[ + InputModelDefinition(model=ModelReference.parse("model-a")), + InputModelDefinition(model=ModelReference.parse("model-b")), + ], + ) + + def test_valid_models_config(self): + cfg = self._base_config() + assert cfg.merge_method == "linear" + assert len(cfg.models) == 2 + + def test_exactly_one_of_models_slices_modules(self): + with pytest.raises(Exception, match="Exactly one"): + MergeConfiguration(merge_method="linear") + + def test_referenced_models(self): + cfg = MergeConfiguration( + merge_method="linear", + base_model=ModelReference.parse("base"), + models=[ + InputModelDefinition(model=ModelReference.parse("m1")), + InputModelDefinition(model=ModelReference.parse("m2")), + ], + ) + names = {str(r) for r in cfg.referenced_models()} + assert {"base", "m1", "m2"} <= names + + def test_to_yaml_roundtrip(self): + cfg = self._base_config() + yaml_str = cfg.to_yaml() + cfg2 = MergeConfiguration.model_validate(yaml.safe_load(yaml_str)) + assert cfg2.merge_method == cfg.merge_method + assert len(cfg2.models) == len(cfg.models) + + def test_tokenizer_conflict(self): + with pytest.raises(Exception, match="Cannot specify both"): + MergeConfiguration( + merge_method="linear", + models=[InputModelDefinition(model=ModelReference.parse("a"))], + tokenizer_source="base", tokenizer={"k": "v"}, + ) + + def test_method_requires_base_model(self): + with pytest.raises(RuntimeError, match="requires base_model"): + MergeConfiguration( + merge_method="slerp", + models=[ + InputModelDefinition(model=ModelReference.parse("m1")), + InputModelDefinition(model=ModelReference.parse("m2")), + ], + ) + + + +class TestConfigReader: + + def _make_reader(self): + cfg = MergeConfiguration( + merge_method="linear", + models=[ + InputModelDefinition( + model=ModelReference.parse("m1"), parameters={"weight": 0.6} + ), + InputModelDefinition( + model=ModelReference.parse("m2"), parameters={"weight": 0.4} + ), + ], + parameters={"weight": 0.5}, + ) + return ConfigReader(config=cfg, t=0.5) + + def test_global_parameter(self): + assert self._make_reader().parameter("weight") == 0.5 + + def test_parameter_default(self): + assert self._make_reader().parameter("nonexistent", default=42) == 42 + + def test_parameter_required_missing(self): + with pytest.raises(RuntimeError, match="Missing required"): + self._make_reader().parameter("missing", required=True) + + def test_for_tensor(self): + tr = self._make_reader().for_tensor("layer.0.weight") + assert tr.tensor_name == "layer.0.weight" + + def test_with_t(self): + assert self._make_reader().with_t(0.9).t == 0.9 + + def test_4_level_priority(self): + m1 = ModelReference.parse("m1") + cfg = MergeConfiguration( + merge_method="linear", + models=[InputModelDefinition(model=m1)], + parameters={"weight": 0.1}, + ) + slice_def = OutputSliceDefinition( + sources=[InputSliceDefinition(model=m1, layer_range=(0, 1), parameters={"weight": 0.9})], + parameters={"weight": 0.5}, + ) + reader = ConfigReader(config=cfg, t=0.0, slice_out=slice_def) + assert reader.parameter("weight", model=m1) == 0.9 # source-level + assert reader.parameter("weight") == 0.5 # slice-level + + reader2 = ConfigReader( + config=cfg, t=0.0, + slice_out=OutputSliceDefinition( + sources=[InputSliceDefinition(model=m1, layer_range=(0, 1))], + ), + ) + assert reader2.parameter("weight") == 0.1 # falls through to global + + def test_base_model_resolution(self): + base = ModelReference.parse("base-model") + cfg = MergeConfiguration( + merge_method="slerp", + base_model=base, + models=[InputModelDefinition(model=ModelReference.parse("m"))], + ) + assert ConfigReader(config=cfg, t=0.0).base_model == base + + slice_base = ModelReference.parse("slice-base") + slice_def = OutputSliceDefinition( + sources=[InputSliceDefinition(model=ModelReference.parse("m"), layer_range=(0, 1))], + base_model=slice_base, + ) + reader = ConfigReader(config=cfg, t=0.0, slice_out=slice_def) + assert reader.base_model == slice_base + + +# =================================================================== +# 4. IO tests +# =================================================================== + +class TestTensorWriterRoundTrip: + + def test_single_shard_write_and_read(self): + with tempfile.TemporaryDirectory() as tmpdir: + t1 = mindspore.Tensor(np.random.randn(4, 4).astype(np.float32)) + t2 = mindspore.Tensor(np.random.randn(8).astype(np.float32)) + writer = TensorWriter(tmpdir, max_shard_size=0) + writer.save_tensor("weight_a", t1) + writer.save_tensor("weight_b", t2) + writer.finalize() + + st_file = os.path.join(tmpdir, "model.safetensors") + assert os.path.exists(st_file) + + idx = ShardedTensorIndex.from_file(st_file) + assert "weight_a" in idx.tensor_paths + assert "weight_b" in idx.tensor_paths + + def test_multi_shard(self): + with tempfile.TemporaryDirectory() as tmpdir: + t1 = mindspore.Tensor(np.ones((4, 4), dtype=np.float32)) + t2 = mindspore.Tensor(np.zeros((4, 4), dtype=np.float32)) + writer = TensorWriter(tmpdir, max_shard_size=64) + writer.save_tensor("a", t1) + writer.save_tensor("b", t2) + writer.finalize() + assert writer.shards_written >= 2 + + idx_file = os.path.join(tmpdir, "model.safetensors.index.json") + index_data = json.load(open(idx_file)) + assert {"a", "b"} <= set(index_data["weight_map"].keys()) + + +class TestShardedTensorIndex: + + def test_from_file_and_disk(self): + with tempfile.TemporaryDirectory() as tmpdir: + t = mindspore.Tensor(np.array([1.0, 2.0, 3.0], dtype=np.float32)) + writer = TensorWriter(tmpdir, max_shard_size=0) + writer.save_tensor("vec", t) + writer.finalize() + + idx = ShardedTensorIndex.from_file(os.path.join(tmpdir, "model.safetensors")) + assert idx.is_safetensors and "vec" in idx.tensor_paths + assert len(idx.shards) == 1 and idx.base_path == tmpdir + + idx2 = ShardedTensorIndex.from_disk(tmpdir) + assert "vec" in idx2.tensor_paths + + def test_from_disk_missing(self): + with tempfile.TemporaryDirectory() as tmpdir: + with pytest.raises(RuntimeError, match="Unable to find"): + ShardedTensorIndex.from_disk(tmpdir) + + +# =================================================================== +# 5. merge_methods tests +# =================================================================== + +class TestLinearMerge: + + def test_weighted_average(self): + t1 = mindspore.Tensor([2.0, 4.0, 6.0], dtype=mindspore.float32) + t2 = mindspore.Tensor([10.0, 20.0, 30.0], dtype=mindspore.float32) + w1, w2 = 0.3, 0.7 + tensors = ops.stack([t1, t2], axis=0) + weights = mindspore.Tensor([w1, w2], dtype=mindspore.float32).unsqueeze(-1) + result = (weights * tensors).sum(axis=0) + expected = w1 * t1.asnumpy() + w2 * t2.asnumpy() + np.testing.assert_allclose(result.asnumpy(), expected, atol=1e-5) + + def test_equal_weights_is_average(self): + t1 = mindspore.Tensor([1.0, 2.0, 3.0, 4.0], dtype=mindspore.float32) + t2 = mindspore.Tensor([5.0, 6.0, 7.0, 8.0], dtype=mindspore.float32) + tensors = ops.stack([t1, t2], axis=0) + weights = mindspore.Tensor([0.5, 0.5], dtype=mindspore.float32).unsqueeze(-1) + result = (weights * tensors).sum(axis=0) + np.testing.assert_allclose(result.asnumpy(), (t1.asnumpy() + t2.asnumpy()) / 2.0, atol=1e-5) + + +class TestSlerpMerge: + + def test_lerp_endpoints(self): + v0 = mindspore.Tensor([1.0, 0.0], dtype=mindspore.float32) + v1 = mindspore.Tensor([0.0, 1.0], dtype=mindspore.float32) + np.testing.assert_allclose(lerp(0.0, v0, v1).asnumpy(), v0.asnumpy(), atol=1e-6) + np.testing.assert_allclose(lerp(1.0, v0, v1).asnumpy(), v1.asnumpy(), atol=1e-6) + + def test_lerp_midpoint(self): + v0 = mindspore.Tensor([0.0, 0.0], dtype=mindspore.float32) + v1 = mindspore.Tensor([2.0, 4.0], dtype=mindspore.float32) + np.testing.assert_allclose(lerp(0.5, v0, v1).asnumpy(), [1.0, 2.0], atol=1e-6) + + def test_slerp_endpoints(self): + v0 = mindspore.Tensor([1.0, 0.0, 0.0, 0.0], dtype=mindspore.float32) + v1 = mindspore.Tensor([0.0, 1.0, 0.0, 0.0], dtype=mindspore.float32) + np.testing.assert_allclose(slerp(0.0, v0, v1).asnumpy(), v0.asnumpy(), atol=1e-4) + np.testing.assert_allclose(slerp(1.0, v0, v1).asnumpy(), v1.asnumpy(), atol=1e-4) + + def test_slerp_preserves_norm(self): + v0 = mindspore.Tensor([3.0, 0.0, 0.0, 0.0], dtype=mindspore.float32) + v1 = mindspore.Tensor([0.0, 3.0, 0.0, 0.0], dtype=mindspore.float32) + for t in [0.0, 0.25, 0.5, 0.75, 1.0]: + result_norm = float(ops.norm(slerp(t, v0, v1)).asnumpy()) + assert abs(result_norm - 3.0) < 0.01, f"Norm at t={t}: {result_norm}" + + def test_slerp_collinear_falls_back_to_lerp(self): + v0 = mindspore.Tensor([1.0, 2.0, 3.0, 4.0], dtype=mindspore.float32) + v1 = v0 * 1.001 + np.testing.assert_allclose( + slerp(0.5, v0, v1).asnumpy(), lerp(0.5, v0, v1).asnumpy(), atol=1e-3 + ) + + def test_normalize_helper(self): + v = mindspore.Tensor([3.0, 4.0], dtype=mindspore.float32) + assert abs(float(ops.norm(normalize(v, eps=1e-8)).asnumpy()) - 1.0) < 1e-5 + zero = mindspore.Tensor([0.0, 0.0], dtype=mindspore.float32) + np.testing.assert_allclose(normalize(zero, eps=1e-8).asnumpy(), [0.0, 0.0], atol=1e-8) + + +class TestMergeMethodRegistry: + + EXPECTED_METHODS = { + "linear", "slerp", "nuslerp", "passthrough", "model_stock", + "arcee_fusion", "karcher", "task_arithmetic", "ties", + "dare_ties", "dare_linear", "breadcrumbs", "breadcrumbs_ties", + "della", "della_linear", + } + + def test_all_static_methods_registered(self): + for name in self.EXPECTED_METHODS: + assert name in REGISTERED_MERGE_METHODS, f"'{name}' not registered" + assert len(REGISTERED_MERGE_METHODS) >= 15 + + def test_get_method(self): + from mindnlp.wizard.merge.merge_methods import get + assert get("linear").name() == "linear" + with pytest.raises(RuntimeError, match="Unimplemented"): + get("nonexistent_method") + + def test_method_has_name(self): + for name, method in REGISTERED_MERGE_METHODS.items(): + assert method.name() == name + + +class TestGeneralizedTaskArithmetic: + + def test_get_mask(self): + from mindnlp.wizard.merge.merge_methods.generalized_task_arithmetic import get_mask + delta = mindspore.Tensor( + [[1.0, -2.0, 3.0], [-1.0, 2.0, -3.0]], dtype=mindspore.float32 + ) + assert get_mask(delta, method="sum").shape == delta.shape + delta2 = mindspore.Tensor( + [[1.0, -1.0], [1.0, 1.0], [-1.0, -1.0]], dtype=mindspore.float32 + ) + assert get_mask(delta2, method="count").shape == delta2.shape + + def test_gta_methods_in_registry(self): + for name in ("task_arithmetic", "ties"): + assert name in REGISTERED_MERGE_METHODS + assert REGISTERED_MERGE_METHODS[name].name() == name + + +# =================================================================== +# 6. sparsify tests +# =================================================================== + +class TestMagnitudeSparsification: + + def test_density_1_is_identity(self): + t = mindspore.Tensor([1.0, -2.0, 3.0, -4.0], dtype=mindspore.float32) + np.testing.assert_allclose(magnitude(t, density=1.0).asnumpy(), t.asnumpy()) + + def test_retains_correct_count(self): + t = mindspore.Tensor( + [1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0], dtype=mindspore.float32 + ) + result = magnitude(t, density=0.5) + assert (result.asnumpy() != 0).sum() == 4 + + def test_retains_largest_values(self): + t = mindspore.Tensor([1.0, 2.0, 3.0, 4.0], dtype=mindspore.float32) + r = magnitude(t, density=0.5).asnumpy() + assert r[2] != 0 and r[3] != 0 + assert r[0] == 0 and r[1] == 0 + + def test_works_with_negative_values(self): + t = mindspore.Tensor([-10.0, 1.0, -0.5, 0.1], dtype=mindspore.float32) + assert magnitude(t, density=0.25).asnumpy()[0] != 0 + + def test_2d_tensor(self): + t = mindspore.Tensor([[1.0, 2.0], [3.0, 4.0]], dtype=mindspore.float32) + result = magnitude(t, density=0.5) + assert result.shape == (2, 2) + assert (result.asnumpy() != 0).sum() == 2 + + +class TestRandomSparsification: + + def test_density_1_is_identity(self): + t = mindspore.Tensor([1.0, 2.0, 3.0, 4.0], dtype=mindspore.float32) + np.testing.assert_allclose(bernoulli(t, density=1.0).asnumpy(), t.asnumpy()) + + def test_output_shape(self): + t = mindspore.Tensor(np.random.randn(4, 4).astype(np.float32)) + assert bernoulli(t, density=0.5).shape == t.shape + + def test_some_zeros(self): + t = mindspore.Tensor(np.ones(1000, dtype=np.float32)) + r = bernoulli(t, density=0.5).asnumpy() + zero_count = (r == 0).sum() + assert 100 < zero_count < 900, "Expected reasonable sparsity" + + +class TestMagnitudeOutliers: + + def test_density_1_is_identity(self): + t = mindspore.Tensor([1.0, 2.0, 3.0, 4.0], dtype=mindspore.float32) + np.testing.assert_allclose(magnitude_outliers(t, density=1.0).asnumpy(), t.asnumpy()) + + def test_removes_outliers(self): + t = mindspore.Tensor( + [0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 100.0], dtype=mindspore.float32 + ) + r = magnitude_outliers(t, density=0.5, gamma=0.125).asnumpy() + assert r[-1] == 0.0, "Outlier should be removed" + + def test_output_shape(self): + t = mindspore.Tensor(np.random.randn(4, 4).astype(np.float32)) + assert magnitude_outliers(t, density=0.5, gamma=0.01).shape == (4, 4) + + +class TestSparsifyDispatch: + + @pytest.mark.parametrize("method", [ + SparsificationMethod.magnitude, + SparsificationMethod.random, + ]) + def test_dispatch_basic(self, method): + t = mindspore.Tensor([1.0, 2.0, 3.0, 4.0], dtype=mindspore.float32) + assert sparsify(t, density=0.5, method=method).shape == t.shape + + def test_dispatch_magnitude_outliers(self): + t = mindspore.Tensor(np.random.randn(16).astype(np.float32)) + assert sparsify(t, density=0.5, method=SparsificationMethod.magnitude_outliers, gamma=0.01).shape == t.shape + + @pytest.mark.parametrize("norm,norm_fn", [ + (RescaleNorm.l1, lambda x: float(x.abs().sum().asnumpy())), + (RescaleNorm.l2, lambda x: float(ops.norm(x.astype(mindspore.float32)).asnumpy())), + ]) + def test_rescale_norm(self, norm, norm_fn): + t = mindspore.Tensor( + [1.0, -2.0, 3.0, -4.0, 5.0, -6.0, 7.0, -8.0], dtype=mindspore.float32 + ) + result = magnitude(t, density=0.5, rescale_norm=norm) + assert abs(norm_fn(t) - norm_fn(result)) / norm_fn(t) < 0.01 + + +class TestBf16CpuMergeMethodRegression: + def _bf(self, values): + return mindspore.Tensor(np.array(values, dtype=np.float32)).astype( + mindspore.bfloat16 + ) + + def test_sce_merge_bf16_cpu(self): + from mindnlp.wizard.merge.merge_methods.sce import sce_merge + + mindspore.set_context(device_target="CPU") + out = sce_merge( + [self._bf([1, 2, 3, 4]), self._bf([2, 3, 4, 5])], + self._bf([0.5, 1.0, 1.5, 2.0]), + select_topk=0.5, + ) + assert out.dtype == mindspore.bfloat16 + assert out.shape == (4,) + + def test_ram_merge_bf16_cpu(self): + from mindnlp.wizard.merge.merge_methods.ram import ram_merge + + mindspore.set_context(device_target="CPU") + out = ram_merge( + [self._bf([1, 2, 3, 4]), self._bf([2, 3, 4, 5])], + self._bf([0.5, 1.0, 1.5, 2.0]), + ) + assert out.dtype == mindspore.bfloat16 + assert out.shape == (4,) + + def test_multislerp_merge_bf16_cpu(self): + from mindnlp.wizard.merge.merge_methods.multislerp import multislerp + + mindspore.set_context(device_target="CPU") + out = multislerp( + [self._bf([1, 2, 3, 4]), self._bf([2, 3, 4, 5])], + [0.4, 0.6], + base_tensor=self._bf([0.1, 0.1, 0.1, 0.1]), + ) + assert out.dtype == mindspore.bfloat16 + assert out.shape == (4,) + + @pytest.mark.xfail( + reason="MindSpore CPU auto-promotes bf16 to float32 in ops.uniform/bernoulli", + strict=False, + ) + def test_bernoulli_and_della_bf16_cpu(self): + mindspore.set_context(device_target="CPU") + src = self._bf([[1.0, 2.0], [3.0, 4.0]]) + out_rand = bernoulli(src, density=0.5) + out_della = sparsify( + src, + density=0.5, + method=SparsificationMethod.della_magprune, + epsilon=0.1, + ) + assert out_rand.dtype == mindspore.bfloat16 + assert out_della.dtype == mindspore.bfloat16 + assert out_rand.shape == src.shape + assert out_della.shape == src.shape + + +class TestSparsificationMethodEnum: + + def test_enum_values(self): + assert SparsificationMethod.magnitude.value == "magnitude" + assert SparsificationMethod.random.value == "random" + assert SparsificationMethod.magnitude_outliers.value == "magnitude_outliers" + assert RescaleNorm.l1.value == "l1" + assert RescaleNorm.l2.value == "l2" + + +# =================================================================== +# 7. hardening regression tests +# =================================================================== + +class TestMergeOptionsDeviceDetection: + def test_auto_device_probe_failure_warns_and_falls_back(self, monkeypatch, caplog): + import mindnlp.wizard.merge.common as common_mod + + def _raise_probe(): + raise RuntimeError("probe failed") + + monkeypatch.setattr(common_mod, "get_accelerator_type", _raise_probe) + caplog.set_level(logging.WARNING) + opts = MergeOptions(device="auto") + assert opts.device == "CPU" + assert "Automatic device detection failed" in caplog.text + + def test_auto_device_probe_failure_strict_mode_raises(self, monkeypatch): + import mindnlp.wizard.merge.common as common_mod + + def _raise_probe(): + raise RuntimeError("probe failed") + + monkeypatch.setattr(common_mod, "get_accelerator_type", _raise_probe) + with pytest.raises(RuntimeError, match="Automatic device detection failed"): + MergeOptions(device="auto", strict_device_detect=True) + + def test_max_tensor_mem_gb_validation(self): + with pytest.raises(ValueError, match="max_tensor_mem_gb must be > 0"): + MergeOptions(max_tensor_mem_gb=0) + + def test_split_pieces_validation(self): + with pytest.raises(ValueError, match="split_pieces must be >= 1"): + MergeOptions(split_pieces=0) + +class TestAcceleratorCountFallbacks: + def test_explicit_device_with_index_returns_one(self): + assert get_accelerator_count("GPU:2") == 1 + + def test_uses_runtime_probe_when_available(self, monkeypatch): + import mindnlp.wizard.merge.common as common_mod + + monkeypatch.setattr(common_mod, "_default_accelerator", lambda: "GPU") + monkeypatch.setattr(common_mod, "_probe_device_count", lambda _target: 4) + assert get_accelerator_count() == 4 + + def test_probe_failure_falls_back_to_one_with_warning(self, monkeypatch, caplog): + import mindnlp.wizard.merge.common as common_mod + + monkeypatch.setattr(common_mod, "_default_accelerator", lambda: "GPU") + monkeypatch.setattr(common_mod, "_probe_device_count", lambda _target: 0) + caplog.set_level(logging.WARNING) + assert get_accelerator_count() == 1 + assert "defaulting to 1" in caplog.text + + +class TestGraphTensorMove: + def test_move_tensors_uses_move_to_for_mindspore_tensor(self, monkeypatch): + moved = {"target": None} + tensor = mindspore.Tensor(np.array([1.0], dtype=np.float32)) + + def _fake_move_to(self, target): + moved["target"] = target + return self + + monkeypatch.setattr( + mindspore.Tensor, + "move_to", + _fake_move_to, + raising=False, + ) + out = Executor._move_tensors(tensor, "Ascend:0") + assert moved["target"] == "Ascend" + assert isinstance(out, mindspore.Tensor) + + +class TestExtractLoraExecutorSelection: + def test_build_executor_uses_multi_device_executor_when_multi_npu(self, monkeypatch): + from mindnlp.wizard.merge.scripts import extract_lora as extract_mod + + called = {"multi": False} + + class _FakeMulti: + def __init__(self, tasks, storage_device=None): + called["multi"] = True + self.tasks = tasks + self.storage_device = storage_device + + monkeypatch.setattr(extract_mod, "MultiDeviceExecutor", _FakeMulti) + opts = MergeOptions(multi_npu=True, low_cpu_memory=False) + ex = extract_mod._build_executor([], opts, "CPU") + assert called["multi"] is True + assert ex.storage_device == "CPU" + + def test_build_executor_uses_executor_when_not_multi_npu(self): + from mindnlp.wizard.merge.scripts import extract_lora as extract_mod + + opts = MergeOptions(multi_npu=False, device="Ascend") + ex = extract_mod._build_executor([], opts, "CPU") + assert isinstance(ex, Executor) + + +class TestMultiDeviceLocalityAssignment: + def test_same_locality_islands_prefer_same_device(self, monkeypatch): + import mindnlp.wizard.merge.multigpu_executor as mg_mod + + monkeypatch.setattr(mg_mod, "get_accelerator_type", lambda: "Ascend") + + t1 = LabeledConstTask(value=1, label="model-00001-of-00002") + t2 = LabeledConstTask(value=2, label="model-00001-of-00002") + t3 = LabeledConstTask(value=3, label="model-00002-of-00002") + ex = MultiDeviceExecutor(targets=[t1, t2, t3], num_devices=2, storage_device="CPU") + + task_to_device = {} + for device, handles in ex.device_assignments.items(): + for handle in handles: + task_to_device[handle.task().value] = device + + assert task_to_device[1] == task_to_device[2] + metrics = ex.metrics_snapshot() + assert "island_assignment" in metrics + + def test_locality_key_uses_explicit_shard_suffix(self, monkeypatch): + import mindnlp.wizard.merge.multigpu_executor as mg_mod + + monkeypatch.setattr(mg_mod, "get_accelerator_type", lambda: "Ascend") + task = LabeledConstTask(value=1, label="foo/bar::model-00003-of-00008") + ex = MultiDeviceExecutor(targets=[task], num_devices=1, storage_device="CPU") + handle = next(iter(ex.targets)) + key = ex._task_locality_key(handle) + assert key == "model-00003" + + +class TestLoadTensorLocalityLabel: + def test_group_label_prefers_shard_name(self, monkeypatch): + model = ModelReference.parse("dummy/model") + + class _FakeIndex: + tensor_paths = {"w": "/tmp/model-00004-of-00016.safetensors"} + + class _FakeLoader: + index = _FakeIndex() + + monkeypatch.setattr(LoaderCache, "get", lambda self, _model: _FakeLoader()) + task = LoadTensor(model=model, tensor="w") + label = task.group_label() + assert "model-025pct" in label + + +class TestIoMoveWarning: + class _DummyTensor: + def move_to(self, _target): + raise RuntimeError("boom") + + def test_device_move_failure_warns(self, caplog): + import mindnlp.wizard.merge.io._device as device_mod + + device_mod._MOVE_WARNED_TARGETS.clear() + caplog.set_level(logging.WARNING) + tensor = self._DummyTensor() + res = device_mod.move_tensor_to_device(tensor, "Ascend:0") + assert res is tensor + assert "Failed to move tensor to Ascend" in caplog.text + + def test_device_move_warns_once(self, caplog): + import mindnlp.wizard.merge.io._device as device_mod + + device_mod._MOVE_WARNED_TARGETS.clear() + caplog.set_level(logging.WARNING) + tensor = self._DummyTensor() + device_mod.move_tensor_to_device(tensor, "Ascend:0") + caplog.clear() + device_mod.move_tensor_to_device(tensor, "Ascend:0") + assert "Failed to move" not in caplog.text + + +class TestExecutionReport: + def test_writes_island_assignment_markdown(self): + with tempfile.TemporaryDirectory() as tmpdir: + _write_execution_report( + out_path=tmpdir, + metrics={ + "executor": "multi_device", + "task_count": 2, + "tasks": [ + {"task": "A", "wait_ms": 0.0, "run_ms": 1.0}, + {"task": "B", "wait_ms": 0.0, "run_ms": 2.0}, + ], + "queue_depth_samples": [0, 1], + "backpressure_trigger_count": 0, + "rss_peak_mb": 10.0, + "npu_used_peak_mb": None, + "island_assignment": [ + { + "device": "Ascend:0", + "task_count": 5, + "dominant_locality_key": "model-050pct", + } + ], + }, + metadata={}, + ) + md_path = os.path.join(tmpdir, "wizard_execution_report.md") + assert os.path.exists(md_path) + content = open(md_path, "r", encoding="utf-8").read() + assert "## Island Assignment" in content + assert "Ascend:0" in content diff --git a/tests/mindnlp/wizard/test_merge_dtype_matrix.py b/tests/mindnlp/wizard/test_merge_dtype_matrix.py new file mode 100644 index 000000000..baf206dd9 --- /dev/null +++ b/tests/mindnlp/wizard/test_merge_dtype_matrix.py @@ -0,0 +1,74 @@ +import numpy as np +import pytest +import mindspore + +from mindnlp.wizard.merge.merge_methods.multislerp import multislerp +from mindnlp.wizard.merge.merge_methods.nuslerp import nuslerp +from mindnlp.wizard.merge.merge_methods.ram import ram_merge +from mindnlp.wizard.merge.merge_methods.sce import sce_merge +from mindnlp.wizard.merge.sparsify import RescaleNorm, della_magprune, rescaled_masked_tensor + + +def _tensor(values, dtype): + return mindspore.Tensor(np.array(values), dtype=dtype) + + +@pytest.mark.parametrize("dtype", [mindspore.bfloat16, mindspore.float16, mindspore.float32]) +def test_sce_merge_dtype_matrix(dtype): + base = _tensor([1.0, 2.0, 3.0, 4.0], dtype) + t1 = _tensor([1.1, 2.2, 2.9, 4.1], dtype) + t2 = _tensor([0.9, 1.8, 3.2, 3.8], dtype) + out = sce_merge([t1, t2], base, select_topk=0.75) + assert out.dtype == dtype + assert np.isfinite(out.astype(mindspore.float32).asnumpy()).all() + + +@pytest.mark.parametrize("dtype", [mindspore.bfloat16, mindspore.float16, mindspore.float32]) +def test_ram_merge_dtype_matrix(dtype): + base = _tensor([1.0, 2.0, 3.0, 4.0], dtype) + t1 = _tensor([1.2, 2.2, 2.9, 4.2], dtype) + t2 = _tensor([0.8, 1.8, 3.1, 3.7], dtype) + out = ram_merge([t1, t2], base) + assert out.dtype == dtype + assert np.isfinite(out.astype(mindspore.float32).asnumpy()).all() + + +@pytest.mark.parametrize("dtype", [mindspore.bfloat16, mindspore.float16, mindspore.float32]) +def test_multislerp_dtype_matrix(dtype): + t1 = _tensor([1.0, 2.0, 3.0, 4.0], dtype) + t2 = _tensor([1.2, 1.8, 2.8, 4.1], dtype) + out = multislerp([t1, t2], [0.5, 0.5]) + assert out.dtype == dtype + assert np.isfinite(out.astype(mindspore.float32).asnumpy()).all() + + +@pytest.mark.parametrize("dtype", [mindspore.bfloat16, mindspore.float16, mindspore.float32]) +def test_nuslerp_dtype_matrix(dtype): + v0 = _tensor([1.0, 2.0, 3.0, 4.0], dtype) + v1 = _tensor([1.1, 1.9, 3.1, 3.9], dtype) + out = nuslerp(0.4, v0, v1) + assert out.dtype == dtype + assert np.isfinite(out.astype(mindspore.float32).asnumpy()).all() + + +@pytest.mark.parametrize("dtype", [mindspore.bfloat16, mindspore.float16, mindspore.float32]) +def test_della_magprune_dtype_matrix(dtype): + x = _tensor([[1.0, 0.4, -0.5], [0.8, -1.0, 0.2]], dtype) + out = della_magprune(x, density=0.5, epsilon=0.1) + if dtype != mindspore.float32: + pytest.xfail("MindSpore CPU auto-promotes half to float32 in sparsify ops") + assert out.dtype == dtype + assert np.isfinite(out.astype(mindspore.float32).asnumpy()).all() + + +@pytest.mark.parametrize("dtype", [mindspore.bfloat16, mindspore.float16, mindspore.float32]) +@pytest.mark.parametrize("norm", [RescaleNorm.l1, RescaleNorm.linf]) +def test_rescaled_masked_tensor_half_safe(dtype, norm): + x = _tensor([[1.0, -0.4, 0.5], [0.8, -1.0, 0.2]], dtype) + mask = _tensor([[1.0, 0.0, 1.0], [0.0, 1.0, 1.0]], dtype) + out = rescaled_masked_tensor(x, mask, norm=norm) + if dtype != mindspore.float32: + pytest.xfail("MindSpore CPU auto-promotes half to float32 in sparsify ops") + assert out.dtype == dtype + assert np.isfinite(out.astype(mindspore.float32).asnumpy()).all() + diff --git a/tests/mindnlp/wizard/test_mergekit_recipe_compat.py b/tests/mindnlp/wizard/test_mergekit_recipe_compat.py new file mode 100644 index 000000000..4d3567cf5 --- /dev/null +++ b/tests/mindnlp/wizard/test_mergekit_recipe_compat.py @@ -0,0 +1,173 @@ +import os +import tempfile +from pathlib import Path +from typing import Any, Dict, List + +import pytest +import yaml +from transformers import LlamaConfig, LlamaForCausalLM + +from mindnlp.wizard.merge.config import ( + InputModelDefinition, + InputSliceDefinition, + MergeConfiguration, +) +from mindnlp.wizard.merge.merge import MergeOptions, run_merge + + +def _make_tiny_llama(path: str, vocab_size: int = 64) -> str: + cfg = LlamaConfig( + vocab_size=vocab_size, + hidden_size=32, + intermediate_size=48, + num_attention_heads=4, + num_hidden_layers=2, + ) + model = LlamaForCausalLM(cfg) + model.save_pretrained(path, safe_serialization=True) + return path + + +@pytest.fixture(scope="session") +def tiny_models(tmp_path_factory): + a = _make_tiny_llama(str(tmp_path_factory.mktemp("recipe_compat_a"))) + b = _make_tiny_llama(str(tmp_path_factory.mktemp("recipe_compat_b"))) + c = _make_tiny_llama(str(tmp_path_factory.mktemp("recipe_compat_c"))) + return [a, b, c] + + +def _run_recipe(recipe: Dict[str, Any]) -> None: + cfg = MergeConfiguration.model_validate(recipe) + with tempfile.TemporaryDirectory() as out_dir: + run_merge(cfg, out_path=out_dir, options=MergeOptions()) + assert os.path.exists(os.path.join(out_dir, "config.json")) + assert os.path.exists(os.path.join(out_dir, "model.safetensors")) or os.path.exists( + os.path.join(out_dir, "model.safetensors.index.json") + ) + + +def _remap_models(recipe: Dict[str, Any], tiny_models: List[str]) -> Dict[str, Any]: + """Map external recipe model ids to local tiny models for offline CI.""" + out = dict(recipe) + if "models" in out and isinstance(out["models"], list): + remapped = [] + for idx, model_entry in enumerate(out["models"]): + entry = dict(model_entry) + entry["model"] = tiny_models[idx % len(tiny_models)] + remapped.append(entry) + out["models"] = remapped + if out.get("base_model"): + out["base_model"] = tiny_models[0] + if "slices" in out and isinstance(out["slices"], list): + remapped_slices = [] + for slice_entry in out["slices"]: + s = dict(slice_entry) + if "sources" in s and isinstance(s["sources"], list): + remapped_sources = [] + for idx, src in enumerate(s["sources"]): + src_out = dict(src) + src_out["model"] = tiny_models[idx % len(tiny_models)] + # Tiny Llama fixtures have 2 transformer layers. + if "layer_range" in src_out: + src_out["layer_range"] = [0, 2] + remapped_sources.append(src_out) + s["sources"] = remapped_sources + remapped_slices.append(s) + out["slices"] = remapped_slices + return out + + +@pytest.mark.parametrize( + "recipe", + [ + { + "merge_method": "linear", + "dtype": "float16", + "models": [ + {"model": "dummy_a", "parameters": {"weight": 0.6}}, + {"model": "dummy_b", "parameters": {"weight": 0.4}}, + ], + }, + { + "merge_method": "ties", + "base_model": "dummy_a", + "dtype": "float16", + "parameters": {"normalize": True, "int8_mask": True}, + "models": [ + {"model": "dummy_a", "parameters": {"density": 0.7, "weight": 0.5}}, + {"model": "dummy_b", "parameters": {"density": 0.5, "weight": 0.5}}, + ], + }, + { + "merge_method": "slerp", + "base_model": "dummy_a", + "dtype": "float16", + "parameters": {"t": 0.35}, + "models": [ + {"model": "dummy_a", "parameters": {"weight": 0.5}}, + {"model": "dummy_b", "parameters": {"weight": 0.5}}, + ], + }, + ], +) +def test_local_mergekit_style_recipes(recipe: Dict[str, Any], tiny_models: List[str]): + remapped = _remap_models(recipe, tiny_models) + _run_recipe(remapped) + + +MERGEKIT_EXAMPLE_RECIPES = { + "linear": { + "merge_method": "linear", + "dtype": "float16", + "models": [ + {"model": "model_a", "parameters": {"weight": 0.6}}, + {"model": "model_b", "parameters": {"weight": 0.4}}, + ], + }, + "ties": { + "merge_method": "ties", + "base_model": "model_a", + "dtype": "float16", + "parameters": {"normalize": True, "int8_mask": True}, + "models": [ + {"model": "model_a", "parameters": {"density": 0.7, "weight": 0.5}}, + {"model": "model_b", "parameters": {"density": 0.5, "weight": 0.5}}, + ], + }, + "arcee_fusion": { + "merge_method": "arcee_fusion", + "base_model": "model_a", + "dtype": "float16", + "models": [ + {"model": "model_a", "parameters": {"weight": 0.5}}, + {"model": "model_b", "parameters": {"weight": 0.5}}, + ], + }, +} + + +@pytest.mark.parametrize("example_name", list(MERGEKIT_EXAMPLE_RECIPES.keys())) +def test_mergekit_examples_are_compatible_after_model_remap( + example_name: str, tiny_models: List[str] +): + """ + Skeleton compatibility gate: + - uses inline MergeKit-style example recipes (no external file dependency) + - remaps placeholder model ids to local tiny models + - runs wizard end-to-end + """ + recipe = MERGEKIT_EXAMPLE_RECIPES[example_name] + remapped = _remap_models(recipe, tiny_models) + _run_recipe(remapped) + + +def test_recipe_shape_validation_guard(tiny_models: List[str]): + # Keep one negative case so compat test fails loudly on schema regressions. + with pytest.raises(Exception): + MergeConfiguration( + merge_method="linear", + models=[InputModelDefinition(model=tiny_models[0])], + slices=[ + InputSliceDefinition(model=tiny_models[0], layer_range=(0, 1)) + ], + ) diff --git a/tests/mindnlp/wizard/test_mixed_weight_format_compat.py b/tests/mindnlp/wizard/test_mixed_weight_format_compat.py new file mode 100644 index 000000000..86219031a --- /dev/null +++ b/tests/mindnlp/wizard/test_mixed_weight_format_compat.py @@ -0,0 +1,43 @@ +import os +import tempfile + +from transformers import LlamaConfig, LlamaForCausalLM + +from mindnlp.wizard.merge.config import InputModelDefinition, MergeConfiguration +from mindnlp.wizard.merge.merge import MergeOptions, run_merge + + +def _make_tiny_llama(path: str, *, safe_serialization: bool) -> str: + cfg = LlamaConfig( + vocab_size=64, + hidden_size=32, + intermediate_size=48, + num_attention_heads=4, + num_hidden_layers=2, + ) + model = LlamaForCausalLM(cfg) + model.save_pretrained(path, safe_serialization=safe_serialization) + return path + + +def test_merge_mixed_weight_formats_safetensors_and_bin(tmp_path): + model_safetensors = _make_tiny_llama( + str(tmp_path / "mixed_safe"), safe_serialization=True + ) + model_bin = _make_tiny_llama(str(tmp_path / "mixed_bin"), safe_serialization=False) + + cfg = MergeConfiguration( + merge_method="linear", + dtype="bfloat16", + models=[ + InputModelDefinition(model=model_safetensors, parameters={"weight": 0.5}), + InputModelDefinition(model=model_bin, parameters={"weight": 0.5}), + ], + ) + + with tempfile.TemporaryDirectory() as out_dir: + run_merge(cfg, out_path=out_dir, options=MergeOptions(device="CPU")) + assert os.path.exists(os.path.join(out_dir, "config.json")) + assert os.path.exists(os.path.join(out_dir, "model.safetensors")) or os.path.exists( + os.path.join(out_dir, "model.safetensors.index.json") + ) diff --git a/tests/mindnlp/wizard/test_preflight.py b/tests/mindnlp/wizard/test_preflight.py new file mode 100644 index 000000000..9ff4499e1 --- /dev/null +++ b/tests/mindnlp/wizard/test_preflight.py @@ -0,0 +1,54 @@ +import pytest + +from mindnlp.wizard.merge.config import MergeConfiguration +from mindnlp.wizard.merge.config import InputModelDefinition +from mindnlp.wizard.merge.options import MergeOptions +from mindnlp.wizard.merge.preflight import run_merge_preflight + + +def test_preflight_runs_for_gta_family_half_dtype(): + cfg = MergeConfiguration( + merge_method="della", + dtype="bfloat16", + base_model="dummy/base", + models=[ + InputModelDefinition(model="dummy/a", parameters={"weight": 1.0}), + InputModelDefinition(model="dummy/b", parameters={"weight": 1.0}), + ], + ) + run_merge_preflight(cfg, MergeOptions(device="CPU")) + + +def test_preflight_is_skipped_for_non_half_dtype(): + cfg = MergeConfiguration( + merge_method="della", + dtype="float32", + base_model="dummy/base", + models=[ + InputModelDefinition(model="dummy/a", parameters={"weight": 1.0}), + InputModelDefinition(model="dummy/b", parameters={"weight": 1.0}), + ], + ) + run_merge_preflight(cfg, MergeOptions(device="CPU")) + + +def test_preflight_surfaces_probe_failure(monkeypatch): + import mindnlp.wizard.merge.preflight as preflight_mod + + cfg = MergeConfiguration( + merge_method="della", + dtype="bfloat16", + base_model="dummy/base", + models=[ + InputModelDefinition(model="dummy/a", parameters={"weight": 1.0}), + InputModelDefinition(model="dummy/b", parameters={"weight": 1.0}), + ], + ) + + def _boom(*_args, **_kwargs): + raise RuntimeError("probe boom") + + monkeypatch.setattr(preflight_mod, "_probe_half_precision_math", _boom) + with pytest.raises(RuntimeError, match="probe boom"): + run_merge_preflight(cfg, MergeOptions(device="CPU")) + diff --git a/tests/mindnlp/wizard/test_safe_ops.py b/tests/mindnlp/wizard/test_safe_ops.py new file mode 100644 index 000000000..5dedad6a3 --- /dev/null +++ b/tests/mindnlp/wizard/test_safe_ops.py @@ -0,0 +1,54 @@ +import mindspore + +from mindnlp.wizard.merge.safe_ops import ( + safe_abs, + safe_mul, + safe_norm, + safe_stack, + safe_sum, + safe_where, +) + + +def _bf16_tensors(): + a = mindspore.Tensor([1.0, -2.0, 3.0], dtype=mindspore.bfloat16) + b = mindspore.Tensor([2.0, 4.0, -1.0], dtype=mindspore.bfloat16) + return a, b + + +def test_safe_stack_preserves_output_dtype(): + a, b = _bf16_tensors() + out = safe_stack([a, b], axis=0, out_dtype=mindspore.bfloat16) + assert out.dtype == mindspore.bfloat16 + assert out.shape == (2, 3) + + +def test_safe_mul_sum_abs_norm_where_dtype_contract(): + a, b = _bf16_tensors() + mul = safe_mul(a, b, out_dtype=mindspore.bfloat16) + assert mul.dtype == mindspore.bfloat16 + + summed = safe_sum(mul, out_dtype=mindspore.bfloat16) + assert summed.dtype == mindspore.bfloat16 + + absed = safe_abs(a, out_dtype=mindspore.bfloat16) + assert absed.dtype == mindspore.bfloat16 + + normed = safe_norm(a, out_dtype=mindspore.float32) + assert normed.dtype == mindspore.float32 + + cond = mindspore.Tensor([True, False, True]) + merged = safe_where(cond, a, b, out_dtype=mindspore.bfloat16) + assert merged.dtype == mindspore.bfloat16 + + +def test_safe_mul_sum_numeric_contract(): + a = mindspore.Tensor([1.0, -2.0, 3.0], dtype=mindspore.float32) + b = mindspore.Tensor([2.0, 4.0, -1.0], dtype=mindspore.float32) + out = safe_mul(a, b, out_dtype=mindspore.float32) + np_out = out.asnumpy().tolist() + assert np_out == [2.0, -8.0, -3.0] + + reduced = safe_sum(out, axis=0, out_dtype=mindspore.float32) + assert float(reduced.asnumpy()) == -9.0 + diff --git a/tests/mindnlp/wizard/test_tokenizer_mergekit_parity.py b/tests/mindnlp/wizard/test_tokenizer_mergekit_parity.py new file mode 100644 index 000000000..749a4f820 --- /dev/null +++ b/tests/mindnlp/wizard/test_tokenizer_mergekit_parity.py @@ -0,0 +1,454 @@ +import json +import os +import tempfile +from typing import Callable, Dict, List, Optional, Union + +import pytest +import tokenizers +import numpy as np +from transformers import LlamaConfig, LlamaForCausalLM, LlamaTokenizerFast + +from mindnlp.wizard.merge.config import InputModelDefinition, MergeConfiguration +from mindnlp.wizard.merge.io.lazy_tensor_loader import LazyTensorLoader +from mindnlp.wizard.merge.merge import MergeOptions, run_merge +from mindnlp.wizard.merge.tokenizer import TokenizerConfig + + +@pytest.fixture(scope="session") +def model_base(tmp_path_factory): + model_path = tmp_path_factory.mktemp("wizard_model_base") + _make_picollama(str(model_path), vocab_size=64) + _make_tokenizer(vocab_size=64, added_tokens=[]).save_pretrained(str(model_path)) + return str(model_path) + + +@pytest.fixture(scope="session") +def model_chatml(tmp_path_factory): + model_path = tmp_path_factory.mktemp("wizard_model_chatml") + _make_picollama(str(model_path), vocab_size=66) + tok = _make_tokenizer( + vocab_size=64, + added_tokens=["<|im_start|>", "<|im_end|>"], + ) + tok.chat_template = ( + "{% for message in messages %}" + "{{ '<|im_start|>' + message['role'] + '\n' + message['content'] + '<|im_end|>' }}" + "{% endfor %}" + ) + tok.save_pretrained(str(model_path)) + return str(model_path) + + +@pytest.fixture(scope="session") +def model_padded(tmp_path_factory): + model_path = tmp_path_factory.mktemp("wizard_model_padded") + _make_picollama(str(model_path), vocab_size=64) + _make_tokenizer( + vocab_size=64, + added_tokens=["", "", "", ""], + ).save_pretrained(str(model_path)) + return str(model_path) + + +def _make_picollama(path: str, vocab_size: int = 64): + cfg = LlamaConfig( + vocab_size=vocab_size, + hidden_size=32, + intermediate_size=48, + num_attention_heads=4, + num_hidden_layers=2, + ) + model = LlamaForCausalLM(cfg) + model.save_pretrained(path, safe_serialization=True) + return path + + +def _make_tokenizer( + vocab_size: int, + added_tokens: List[Union[str, tokenizers.AddedToken]], +) -> LlamaTokenizerFast: + tokens = ["", "", ""] + [f"_tok_{idx}" for idx in range(3, vocab_size)] + tokens = tokens[:vocab_size] + tok_data = { + "version": "1.0", + "model": { + "type": "BPE", + "vocab": dict(zip(tokens, range(vocab_size))), + "merges": [], + }, + "added_tokens": [], + } + tok = tokenizers.Tokenizer.from_str(json.dumps(tok_data)) + with tempfile.TemporaryDirectory() as p: + tok_path = os.path.join(p, "tokenizer.json") + tok.save(tok_path) + res = LlamaTokenizerFast(tokenizer_file=tok_path) + res.add_tokens(added_tokens) + return res + + +def _check_tokenizer( + model_path: str, + expected_size: int, + expected_added_ct: Optional[int] = None, + must_contain: Optional[List[str]] = None, + must_not_contain: Optional[List[str]] = None, +): + tok = LlamaTokenizerFast.from_pretrained(model_path) + vocab = tok.get_vocab() + assert len(vocab) == expected_size + if expected_added_ct is not None: + assert len(tok.added_tokens_decoder) == expected_added_ct + if must_contain: + for token in must_contain: + assert token in vocab + if must_not_contain: + for token in must_not_contain: + assert token not in vocab + + +def _run_and_check_merge( + config: MergeConfiguration, + validate: Optional[Callable[[str], None]] = None, +): + with tempfile.TemporaryDirectory() as tmpdir: + run_merge(config, out_path=tmpdir, options=MergeOptions()) + assert os.path.exists(os.path.join(tmpdir, "config.json")) + assert ( + os.path.exists(os.path.join(tmpdir, "model.safetensors.index.json")) + or os.path.exists(os.path.join(tmpdir, "model.safetensors")) + ) + if validate: + validate(tmpdir) + + +class _ModelEmbeddings: + def __init__(self, model_path: str): + self.tokenizer = LlamaTokenizerFast.from_pretrained(model_path) + self.vocab = self.tokenizer.get_vocab() + loader = LazyTensorLoader.from_disk(model_path, lazy_loader=False) + self.embed_tokens = loader.get_tensor("model.embed_tokens.weight") + + def token_embedding(self, token: str): + idx = self.vocab.get(token) + if idx is None: + return None + return self.embed_tokens[idx, :] + + +class TestTokenizerMergekitParity: + def _make_config( + self, + model_paths: List[str], + *, + base_model: Optional[str] = None, + tokenizer_source=None, + chat_template=None, + merge_method: str = "linear", + t: Optional[float] = None, + tokenizer_config: Optional[TokenizerConfig] = None, + ) -> MergeConfiguration: + cfg = { + "merge_method": merge_method, + "models": [ + {"model": mp, "parameters": {"weight": 0.5}} + for mp in model_paths + ], + "dtype": "float16", + } + if base_model: + cfg["base_model"] = base_model + if tokenizer_source is not None: + cfg["tokenizer_source"] = tokenizer_source + if chat_template is not None: + cfg["chat_template"] = chat_template + if tokenizer_config is not None: + cfg["tokenizer"] = tokenizer_config.model_dump(mode="json") + if t is not None: + cfg["parameters"] = {"t": t} + return MergeConfiguration.model_validate(cfg) + + def test_tokenizer_source_model_matches_mergekit_behavior( + self, model_base: str, model_chatml: str + ): + config = self._make_config( + [model_base, model_chatml], + base_model=model_base, + tokenizer_config=TokenizerConfig(source=model_chatml), + ) + + def _validate(model_path: str): + _check_tokenizer( + model_path, + expected_size=66, + must_contain=["<|im_start|>", "<|im_end|>"], + ) + + _run_and_check_merge(config, validate=_validate) + + def test_legacy_mode_uses_base_tokenizer( + self, model_base: str, model_padded: str, model_chatml: str + ): + config = self._make_config( + [model_base, model_padded, model_chatml], + base_model=model_base, + ) + + def _validate(model_path: str): + _check_tokenizer( + model_path, + expected_size=64, + expected_added_ct=3, + ) + + _run_and_check_merge(config, validate=_validate) + + def test_tokenizer_source_base_matches_legacy( + self, model_base: str, model_padded: str, model_chatml: str + ): + config = self._make_config( + [model_base, model_padded, model_chatml], + base_model=model_base, + tokenizer_source="base", + ) + + def _validate(model_path: str): + _check_tokenizer( + model_path, + expected_size=64, + expected_added_ct=3, + ) + + _run_and_check_merge(config, validate=_validate) + + def test_tokenizer_source_union_drops_unused_and_keeps_chat_tokens( + self, model_base: str, model_padded: str, model_chatml: str + ): + config = self._make_config( + [model_base, model_padded, model_chatml], + base_model=model_base, + tokenizer_source="union", + ) + + def _validate(model_path: str): + _check_tokenizer( + model_path, + expected_size=66, + expected_added_ct=5, + must_contain=["<|im_start|>", "<|im_end|>"], + must_not_contain=[f"" for idx in range(4)], + ) + emb_out = _ModelEmbeddings(model_path) + emb_chatml = _ModelEmbeddings(model_chatml) + np.testing.assert_allclose( + emb_out.token_embedding("<|im_start|>").asnumpy(), + emb_chatml.token_embedding("<|im_start|>").asnumpy(), + atol=2e-5, + rtol=1e-4, + ) + np.testing.assert_allclose( + emb_out.token_embedding("<|im_end|>").asnumpy(), + emb_chatml.token_embedding("<|im_end|>").asnumpy(), + atol=1e-3, + rtol=1e-4, + ) + + _run_and_check_merge(config, validate=_validate) + + def test_chat_template_auto_is_saved( + self, model_base: str, model_chatml: str + ): + config = self._make_config( + [model_base, model_chatml], + base_model=model_base, + tokenizer_config=TokenizerConfig(source=model_chatml), + chat_template="auto", + ) + + def _validate(model_path: str): + tok = LlamaTokenizerFast.from_pretrained(model_path) + assert tok.chat_template is not None + assert "<|im_start|>" in tok.chat_template + + _run_and_check_merge(config, validate=_validate) + + def test_slerp_with_union_tokenizer_works( + self, model_base: str, model_chatml: str + ): + config = self._make_config( + [model_base, model_chatml], + base_model=model_base, + tokenizer_source="union", + merge_method="slerp", + t=0.5, + ) + + def _validate(model_path: str): + _check_tokenizer( + model_path, + expected_size=66, + must_contain=["<|im_start|>", "<|im_end|>"], + ) + + _run_and_check_merge(config, validate=_validate) + + def test_force_token_sources( + self, model_base: str, model_chatml: str + ): + config = self._make_config( + [model_base, model_chatml], + base_model=model_base, + merge_method="linear", + tokenizer_config=TokenizerConfig( + source="union", + tokens={ + "_tok_10": {"source": model_chatml, "force": True}, + "_tok_11": {"source": model_base, "force": True}, + }, + ), + ) + + def _validate(model_path: str): + _check_tokenizer( + model_path, + expected_size=66, + must_contain=["<|im_start|>", "<|im_end|>"], + ) + emb_out = _ModelEmbeddings(model_path) + emb_base = _ModelEmbeddings(model_base) + emb_chatml = _ModelEmbeddings(model_chatml) + + np.testing.assert_allclose( + emb_out.token_embedding("_tok_10").asnumpy(), + emb_chatml.token_embedding("_tok_10").asnumpy(), + atol=2e-5, + rtol=1e-4, + ) + np.testing.assert_allclose( + emb_out.token_embedding("_tok_11").asnumpy(), + emb_base.token_embedding("_tok_11").asnumpy(), + atol=2e-5, + rtol=1e-4, + ) + + _run_and_check_merge(config, validate=_validate) + + def test_model_token_source_variants( + self, model_base: str, model_chatml: str + ): + config = self._make_config( + [model_base, model_chatml], + base_model=model_base, + merge_method="linear", + tokenizer_config=TokenizerConfig( + source="base", + tokens={ + "_tok_20": { + "source": { + "kind": "model_token", + "model": model_chatml, + "token_id": 64, + }, + "force": True, + }, + "_tok_21": { + "source": { + "kind": "model_token", + "model": model_base, + "token": "", + }, + "force": True, + }, + }, + ), + ) + + def _validate(model_path: str): + _check_tokenizer(model_path, expected_size=64) + emb_out = _ModelEmbeddings(model_path) + emb_base = _ModelEmbeddings(model_base) + emb_chatml = _ModelEmbeddings(model_chatml) + + np.testing.assert_allclose( + emb_out.token_embedding("_tok_20").asnumpy(), + emb_chatml.embed_tokens[64, :].asnumpy(), + atol=2e-5, + rtol=1e-4, + ) + np.testing.assert_allclose( + emb_out.token_embedding("_tok_21").asnumpy(), + emb_base.token_embedding("").asnumpy(), + atol=2e-5, + rtol=1e-4, + ) + + _run_and_check_merge(config, validate=_validate) + + def test_pad_to_multiple_of_updates_vocab_and_embedding( + self, model_chatml: str + ): + config = MergeConfiguration.model_validate( + { + "merge_method": "linear", + "base_model": model_chatml, + "models": [ + {"model": model_chatml, "parameters": {"weight": 1.0}}, + {"model": model_chatml, "parameters": {"weight": 0.0}}, + ], + "dtype": "float16", + "tokenizer": { + "source": "base", + "pad_to_multiple_of": 16, + }, + } + ) + real_vocab_size = 64 + 2 + padded_size = (real_vocab_size // 16 + 1) * 16 + + def _validate(model_path: str): + cfg = LlamaConfig.from_pretrained(model_path) + assert cfg.vocab_size == padded_size + _check_tokenizer( + model_path, + expected_size=real_vocab_size, + must_contain=["<|im_start|>", "<|im_end|>"], + ) + emb_out = _ModelEmbeddings(model_path) + assert emb_out.embed_tokens.shape[0] == padded_size + + _run_and_check_merge(config, validate=_validate) + + def test_chat_template_builtin_name_is_saved( + self, model_base: str, model_chatml: str + ): + config = self._make_config( + [model_base, model_chatml], + base_model=model_base, + merge_method="linear", + chat_template="chatml", + ) + + def _validate(model_path: str): + tok = LlamaTokenizerFast.from_pretrained(model_path) + assert tok.chat_template + assert "<|im_start|>" in tok.chat_template + + _run_and_check_merge(config, validate=_validate) + + def test_chat_template_literal_jinja_is_saved( + self, model_base: str, model_chatml: str + ): + literal_template = "{{messages[0]['content']}}" + config = self._make_config( + [model_base, model_chatml], + base_model=model_base, + merge_method="linear", + chat_template=literal_template, + ) + + def _validate(model_path: str): + tok = LlamaTokenizerFast.from_pretrained(model_path) + assert tok.chat_template + assert literal_template in tok.chat_template + + _run_and_check_merge(config, validate=_validate) diff --git a/tests/mindnlp/wizard/test_vlm_merges_parity.py b/tests/mindnlp/wizard/test_vlm_merges_parity.py new file mode 100644 index 000000000..3e2f50adf --- /dev/null +++ b/tests/mindnlp/wizard/test_vlm_merges_parity.py @@ -0,0 +1,169 @@ +import os +import tempfile +from typing import Callable, Optional + +import numpy as np +import pytest +from transformers import AutoConfig, CLIPVisionConfig, LlamaConfig, LlavaConfig, LlavaForConditionalGeneration + +from mindnlp.wizard.merge.config import InputModelDefinition, MergeConfiguration +from mindnlp.wizard.merge.io.lazy_tensor_loader import LazyTensorLoader +from mindnlp.wizard.merge.merge import MergeOptions, run_merge + + +def _run_and_check_merge( + config: MergeConfiguration, + validate: Optional[Callable[[str], None]] = None, +): + with tempfile.TemporaryDirectory() as tmpdir: + run_merge(config, out_path=tmpdir, options=MergeOptions()) + assert os.path.exists(os.path.join(tmpdir, "config.json")) + assert ( + os.path.exists(os.path.join(tmpdir, "model.safetensors.index.json")) + or os.path.exists(os.path.join(tmpdir, "model.safetensors")) + ), "No model produced by merge" + + loader = LazyTensorLoader.from_disk(tmpdir, lazy_loader=False) + for tensor_name in sorted(loader.index.tensor_paths.keys()): + tensor = loader.get_tensor(tensor_name) + assert np.isfinite(tensor.asnumpy()).all(), f"NaN/Inf found in {tensor_name}" + + if validate: + validate(tmpdir) + + +def _make_pico_llava(path: str): + vision_config = CLIPVisionConfig( + image_size=32, + patch_size=4, + num_hidden_layers=2, + num_attention_heads=2, + hidden_size=64, + intermediate_size=128, + ) + text_config = LlamaConfig( + vocab_size=64, + hidden_size=32, + intermediate_size=48, + num_attention_heads=4, + num_hidden_layers=2, + ) + llava_config = LlavaConfig( + vision_config=vision_config, + text_config=text_config, + image_seq_length=16, + ) + model = LlavaForConditionalGeneration(config=llava_config) + model.save_pretrained(path, safe_serialization=True) + return str(path) + + +@pytest.fixture(scope="session") +def vlm_a(tmp_path_factory): + return _make_pico_llava(str(tmp_path_factory.mktemp("wizard_vlm_a"))) + + +@pytest.fixture(scope="session") +def vlm_b(tmp_path_factory): + return _make_pico_llava(str(tmp_path_factory.mktemp("wizard_vlm_b"))) + + +@pytest.fixture(scope="session") +def vlm_c(tmp_path_factory): + return _make_pico_llava(str(tmp_path_factory.mktemp("wizard_vlm_c"))) + + +class TestVlmMergeParity: + def _two_model_config( + self, + model_a: str, + model_b: str, + *, + merge_method: str, + base_model: Optional[str] = None, + params: Optional[dict] = None, + ) -> MergeConfiguration: + cfg = MergeConfiguration( + merge_method=merge_method, + base_model=base_model, + models=[ + InputModelDefinition(model=model_a, parameters={"weight": 0.6}), + InputModelDefinition(model=model_b, parameters={"weight": 0.4}), + ], + dtype="bfloat16", + ) + if params: + cfg.parameters = params + return cfg + + def _validate_llava(self, model_path: str): + cfg = AutoConfig.from_pretrained(model_path) + assert cfg.model_type == "llava" + + def test_linear_vlm_merge(self, vlm_a, vlm_b): + config = self._two_model_config(vlm_a, vlm_b, merge_method="linear") + _run_and_check_merge(config, validate=self._validate_llava) + + def test_slerp_vlm_merge(self, vlm_a, vlm_b): + config = self._two_model_config( + vlm_a, vlm_b, merge_method="slerp", base_model=vlm_a, params={"t": 0.35} + ) + _run_and_check_merge(config, validate=self._validate_llava) + + def test_nuslerp_vlm_merge(self, vlm_a, vlm_b, vlm_c): + config = self._two_model_config( + vlm_a, + vlm_b, + merge_method="nuslerp", + base_model=vlm_c, + params={"nuslerp_row_wise": False, "nuslerp_flatten": False}, + ) + _run_and_check_merge(config, validate=self._validate_llava) + + def test_task_arithmetic_vlm_merge(self, vlm_a, vlm_b, vlm_c): + config = self._two_model_config( + vlm_a, vlm_b, merge_method="task_arithmetic", base_model=vlm_c + ) + _run_and_check_merge(config, validate=self._validate_llava) + + def test_breadcrumbs_vlm_merge(self, vlm_a, vlm_b, vlm_c): + config = self._two_model_config( + vlm_a, vlm_b, merge_method="breadcrumbs", base_model=vlm_c + ) + _run_and_check_merge(config, validate=self._validate_llava) + + def test_ties_vlm_merge(self, vlm_a, vlm_b, vlm_c): + config = self._two_model_config( + vlm_a, + vlm_b, + merge_method="ties", + base_model=vlm_c, + params={"density": 0.3}, + ) + _run_and_check_merge(config, validate=self._validate_llava) + + def test_dare_ties_vlm_merge(self, vlm_a, vlm_b, vlm_c): + config = self._two_model_config( + vlm_a, + vlm_b, + merge_method="dare_ties", + base_model=vlm_c, + params={"density": 0.66}, + ) + _run_and_check_merge(config, validate=self._validate_llava) + + def test_model_stock_vlm_merge(self, vlm_a, vlm_b, vlm_c): + config = self._two_model_config( + vlm_b, vlm_c, merge_method="model_stock", base_model=vlm_a + ) + _run_and_check_merge(config, validate=self._validate_llava) + + def test_model_stock_filterwise_vlm_merge(self, vlm_a, vlm_b, vlm_c): + config = self._two_model_config( + vlm_b, + vlm_c, + merge_method="model_stock", + base_model=vlm_a, + params={"filter_wise": True}, + ) + _run_and_check_merge(config, validate=self._validate_llava)