[feat] Add shortgpt (#10983)

* feat(shortgpt): add new script and polish existing code (typos, paths) * fix: add shortgpt script, fix formatting and imports
2025-08-22 14:02:21 +08:00
parent c83684fe78
commit 1a58555b56
9 changed files with 543 additions and 66 deletions
--- a/slm/pipelines/examples/contrastive_training/README.md
+++ b/slm/pipelines/examples/contrastive_training/README.md
@@ -1,17 +1,23 @@
 # 向量检索模型训练

-推荐安装 gpu 版本的[PaddlePaddle](https://www.paddlepaddle.org.cn/install/quick?docurl=/documentation/docs/zh/install/conda/linux-conda.html)，以 cuda12.3的 paddle 为例，安装命令如下：
+推荐安装 gpu 版本的[PaddlePaddle](https://www.paddlepaddle.org.cn/install/quick?docurl=/documentation/docs/zh/install/conda/linux-conda.html)，以 cuda11.8的 paddle 为例，安装命令如下：

 ```
-conda install nccl -c conda-forge
-conda install paddlepaddle-gpu==3.0.0rc1 -i https://www.paddlepaddle.org.cn/packages/stable/cu123/ -c conda-forge
-```
-安装其他依赖：
-```
-pip install git+https://github.com/PaddlePaddle/PaddleNLP.git@develop
-pip install -r requirements.txt
+# 创建一个名为 paddle_env 的新环境，并激活
+conda create --name paddle_env python=3.10
+conda activate paddle_env
+
+# 安装 paddlenlp develop版本 
+pip install --pre --upgrade paddlenlp -f https://www.paddlepaddle.org.cn/whl/paddlenlp.html
+
+# 安装 paddlepaddle-gpu nightly版本
+pip install --pre paddlepaddle-gpu -i https://www.paddlepaddle.org.cn/packages/nightly/cu118/
+
+#安装其他依赖：
+pip install -r slm/pipelines/examples/contrastive_training/requirements.txt
 ```

+
 下载 DuReader-Retrieval 中文数据集：
 ```
 cd data
@@ -45,7 +51,7 @@ python train.py --do_train \
 python -m paddle.distributed.launch --gpus "0,1,2,3" train.py --do_train \
              --model_name_or_path rocketqa-zh-base-query-encoder \
              --output_dir rocketqa-zh-base-query-encoder-duretrieval \
-              --train_data ./data/dual.train.json \
+              --train_data ./data/dureader_dual.train.jsonl \
              --overwrite_output_dir \
              --fine_tune_type sft \
              --sentence_pooling_method cls \
@@ -132,7 +138,7 @@ python evaluation/benchmarks.py --model_type bert \
    --passage_max_length 512 \
 ```
 可配置参数包括：
- `model_type`: 模型的类似，可选 bert 或 roberta 等等
+- `model_type`: 模型的类型，可选 bert 或 roberta 等等
 - `query_model`: query 向量模型的路径
 - `passage_model`: passage 向量模型的路径
 - `query_max_length`: query 的最大长度
@@ -182,7 +188,7 @@ python -u evaluation/eval_mteb.py \
 - `kv_cache_reuse`: 量化加载时，是否仅预分配首层 kv_cache 并重复利用，0 表示不复用，1 表示复用，默认为 0，此策略可降低量化加载时显存占用


-# MTEB 评估
+## MTEB 评估
 [MTEB](https://github.com/embeddings-benchmark/mteb)
 是一个大规模文本嵌入评测基准，包含了丰富的向量检索评估任务和数据集。
 本仓库主要面向其中的英文检索任务（Retrieval），并额外支持针对 MSMARCO-Title 的评估。
@@ -264,6 +270,39 @@ MTEB-Retrieval 数据集, NDCG@10分数：
 | LLARA-passage               |  52.48   |  47.51  |    26.13     |        37.26         |  44.12  | 81.09  |  43.98   |  69.17   |  45.49  |  37.07   | 61.76  |     82.29      |  17.30  |  76.07  |   36.73    |   81.30   |


+## 压缩
+
+### 模型删层
+模型剪枝脚本 `shortgpt_prune.py`，用于评估并移除大语言模型中重要性较低的层，以生成一个更小、更高效的模型。该脚本采用“块影响”度量来计算层的重要性，并直接在内存中完成剪枝和保存，流程高效。
+
+#### 使用方法
+
+通过以下命令执行剪枝脚本。可指定原始模型、输出路径、要剪枝的层数以及模型中transformer层的路径。
+
+以repllama-v1-7b-lora-passage为例：
+```bash
+python shortgpt_prune.py \
+    --model_name_or_path castorini/repllama-v1-7b-lora-passage \
+    --output_model_path ./pruned-repllama-v1-7b-lora-passage \
+    --n_prune_layers 6 \
+    --layers_path "llama.layers"
+```
+
+以NV-Embed-v1为例：
+```bash
+python shortgpt_prune.py \
+    --model_name_or_path nvidia/NV-Embed-v1 \
+    --output_model_path /pruned-NV-Embed-v1_pruned_26 \
+    --n_prune_layers 6 \
+    --layers_path "layers"
+```
+可配置参数包括：
+- `--model_name_or_path`: 原始模型的名称或本地路径。
+- `--output_model_path`: 剪枝后模型的保存路径。
+- `--n_prune_layers`: 希望移除的层数。脚本会自动找出最不重要的N层。
+- `--layers_path`: 模型对象中指向transformer层列表的点分隔路径（例如repllama为`"llama.layers"`, llama为`"model.layers"`）。
+
+可用output_model_path路径中的模型跑评估[评估部分的代码](#评估)

 ## Reference

@@ -284,3 +323,5 @@ MTEB-Retrieval 数据集, NDCG@10分数：
 [8] Yingqi Qu, Yuchen Ding, Jing Liu, Kai Liu, Ruiyang Ren, Wayne Xin Zhao, Daxiang Dong, Hua Wu, Haifeng Wang: RocketQA: An Optimized Training Approach to Dense Passage Retrieval for Open-Domain Question Answering. NAACL 2021

 [9] Ruiyang Ren, Yingqi Qu, Jing Liu, Wayne Xin Zhao, Qiaoqiao She, Hua Wu, Haifeng Wang, Ji-Rong Wen: RocketQAv2: A Joint Training Method for Dense Passage Retrieval and Passage Re-ranking. EMNLP 2021
+
+[10] Xin Men, Mingyu Xu, Qingyu Zhang, Bingning Wang, Hongyu Lin, Yaojie Lu, Xianpei Han, Weipeng Chen: Shortgpt: Layers in large language models are more redundant than you expect. ACL Findings 2025
--- a/slm/pipelines/examples/contrastive_training/evaluation/benchmarks.py
+++ b/slm/pipelines/examples/contrastive_training/evaluation/benchmarks.py
@@ -21,7 +21,7 @@ from typing import Dict, List, cast

 from datasets import load_dataset
 from mteb.abstasks import AbsTaskRetrieval
-from prediction import Eval_modle
+from prediction import Eval_model

 csv.field_size_limit(500 * 1024 * 1024)

@@ -51,7 +51,7 @@ class PaddleModel:
        pooling_mode="mean_tokens",
        **kwargs,
    ):
-        self.query_model = Eval_modle(
+        self.query_model = Eval_model(
            model=query_model,
            max_seq_len=max_seq_len,
            batch_size=batch_size,
--- a/slm/pipelines/examples/contrastive_training/evaluation/eval_mteb.py
+++ b/slm/pipelines/examples/contrastive_training/evaluation/eval_mteb.py
@@ -29,7 +29,7 @@ from paddlenlp.transformers import AutoTokenizer, BiEncoderModel, NVEncodeModel
 class MSMARCOTITLE(AbsTaskRetrieval):
    metadata = TaskMetadata(
        dataset={
-            "corpus_path": "Tevatron/msmarco-passage-corpus",
+            "corpus_path": "Tevatron/msmarco-passage-corpus-new",
            "path": "mteb/msmarco",
            "revision": "c5a29a104738b98a9e76336939199e264163d4a0",
        },
@@ -53,6 +53,9 @@ class MSMARCOTITLE(AbsTaskRetrieval):
        bibtex_citation=None,
        n_samples=None,
        avg_character_length=None,
+        modalities=["text"],
+        sample_creation="created",
+        descriptive_stats={},
    )

    def load_data(self, **kwargs):
--- a/slm/pipelines/examples/contrastive_training/evaluation/eval_mteb.sh
+++ b/slm/pipelines/examples/contrastive_training/evaluation/eval_mteb.sh
@@ -12,18 +12,43 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.

+#!/bin/bash

-for task in "ArguAna" "ClimateFEVER" "DBPedia" "FEVER" "FiQA2018" "HotpotQA" "MSMARCO" "NFCorpus" "NQ" "QuoraRetrieval" "SCIDOCS" "SciFact" "Touche2020" "TRECCOVID" "CQADupstackAndroidRetrieval" "CQADupstackEnglishRetrieval" "CQADupstackGamingRetrieval" "CQADupstackGisRetrieval" "CQADupstackMathematicaRetrieval" "CQADupstackPhysicsRetrieval" "CQADupstackProgrammersRetrieval" "CQADupstackStatsRetrieval" "CQADupstackTexRetrieval" "CQADupstackUnixRetrieval" "CQADupstackWebmastersRetrieval" "CQADupstackWordpressRetrieval" "MSMARCOTITLE"
-do
+# --- Script Configuration ---
+# Exit immediately if a command exits with a non-zero status.
+set -e

-       # 1. RocketQA V1
-       python3.10 -u eval_mteb.py \
-              --corpus_model_name_or_path rocketqa-en-base-v1/passage_model \
-              --query_model_name_or_path rocketqa-en-base-v1/query_model \
+# Define the list of all tasks (datasets) to be evaluated.
+# TASKS=(
+#     "ArguAna" "ClimateFEVER" "DBPedia" "FEVER" "FiQA2018" "HotpotQA" "MSMARCO" "NFCorpus" "NQ" "QuoraRetrieval" 
+#     "SCIDOCS" "SciFact" "Touche2020" "TRECCOVID" "CQADupstackAndroidRetrieval" "CQADupstackEnglishRetrieval" 
+#     "CQADupstackGamingRetrieval" "CQADupstackGisRetrieval" "CQADupstackMathematicaRetrieval" "CQADupstackPhysicsRetrieval" 
+#     "CQADupstackProgrammersRetrieval" "CQADupstackStatsRetrieval" "CQADupstackTexRetrieval" "CQADupstackUnixRetrieval" 
+#     "CQADupstackWebmastersRetrieval" "CQADupstackWordpressRetrieval" "MSMARCOTITLE"
+# )
+
+TASKS=("ArguAna" "SCIDOCS" "FEVER")
+
+
+# You can uncomment the models you wish to evaluate.
+# MODELS_TO_RUN=("RocketQA-V1" "RocketQA-V2" "BGE" "RepLLaMA" "NV-Embed-v1" "BGE-EN-ICL" "LLARA-passage")
+MODELS_TO_RUN=("BGE") 
+
+
+# ===================================================================================
+# 🚀 1. RocketQA V1
+# ===================================================================================
+if [[ " ${MODELS_TO_RUN[*]} " =~ " RocketQA-V1 " ]]; then
+    echo "===== Running Evaluation for Model: RocketQA V1 ====="
+    for task in "${TASKS[@]}"; do
+        echo "--- Task: $task ---"
+        python3.10 -u evaluation/eval_mteb.py \
+              --corpus_model_name_or_path rocketqa-v1-marco-para-encoder \
+              --query_model_name_or_path rocketqa-v1-marco-query-encoder \
              --model_flag RocketQA-V1 \
              --output_folder en_results/rocketqa-en-base-v1 \
              --task_name "$task" \
-              --task_split $(if [[ "$task" == *"MSMARCO"* ]]; then echo "dev"; else echo "test"; fi) \
+              --task_split $([[ "$task" == *"MSMARCO"* ]] && echo "dev" || echo "test") \
              --query_instruction "" \
              --document_instruction "" \
              --max_seq_length 512 \
@@ -31,15 +56,24 @@ do
              --dtype "float32" \
              --padding_side right \
              --pooling_method "cls"
+    done
+fi

-       # 2. RocketQA V2     
-       python3.10 -u eval_mteb.py \
-              --corpus_model_name_or_path rocketqa-en-base-v2/passage_model \
-              --query_model_name_or_path rocketqa-en-base-v2/query_model \
+
+# ===================================================================================
+# 🚀 2. RocketQA V2
+# ===================================================================================
+if [[ " ${MODELS_TO_RUN[*]} " =~ " RocketQA-V2 " ]]; then
+    echo "===== Running Evaluation for Model: RocketQA V2 ====="
+    for task in "${TASKS[@]}"; do
+        echo "--- Task: $task ---"
+        python3.10 -u evaluation/eval_mteb.py \
+              --corpus_model_name_or_path rocketqav2-en-marco-para-encoder \
+              --query_model_name_or_path rocketqav2-en-marco-query-encoder \
              --model_flag RocketQA-V2 \
              --output_folder en_results/rocketqa-en-base-v2 \
              --task_name "$task" \
-              --task_split $(if [[ "$task" == *"MSMARCO"* ]]; then echo "dev"; else echo "test"; fi) \
+              --task_split $([[ "$task" == *"MSMARCO"* ]] && echo "dev" || echo "test") \
              --query_instruction "" \
              --document_instruction "" \
              --max_seq_length 512 \
@@ -47,13 +81,22 @@ do
              --dtype "float32" \
              --padding_side right \
              --pooling_method "cls"
+    done
+fi

-       # 3. BGE
-       python3.10 eval_mteb.py \
+
+# ===================================================================================
+# 🎯 3. BGE (BAAI/bge-large-en-v1.5)
+# ===================================================================================
+if [[ " ${MODELS_TO_RUN[*]} " =~ " BGE " ]]; then
+    echo "===== Running Evaluation for Model: BGE (bge-large-en-v1.5) ====="
+    for task in "${TASKS[@]}"; do
+        echo "--- Task: $task ---"
+        python3.10 evaluation/eval_mteb.py \
              --base_model_name_or_path BAAI/bge-large-en-v1.5 \
-              --output_folder en_results/bge-large-en-v1.5 \
+              --output_folder en_results/bge-large-en-v1.5_2 \
              --task_name "$task" \
-              --task_split $(if [[ "$task" == *"MSMARCO"* ]]; then echo "dev"; else echo "test"; fi) \
+              --task_split $([[ "$task" == *"MSMARCO"* ]] && echo "dev" || echo "test") \
              --document_instruction 'Represent this sentence for searching relevant passages: ' \
              --pooling_method mean \
              --max_seq_length 512 \
@@ -61,13 +104,22 @@ do
              --padding_side right \
              --add_bos_token 0 \
              --add_eos_token 0 
+    done
+fi

-       # 4. RepLLaMA
-       python3.10 eval_mteb.py \
+
+# ===================================================================================
+# 🦙 4. RepLLaMA
+# ===================================================================================
+if [[ " ${MODELS_TO_RUN[*]} " =~ " RepLLaMA " ]]; then
+    echo "===== Running Evaluation for Model: RepLLaMA ====="
+    for task in "${TASKS[@]}"; do
+        echo "--- Task: $task ---"
+        python3.10 evaluation/eval_mteb.py \
              --base_model_name_or_path castorini/repllama-v1-7b-lora-passage \
              --output_folder en_results/repllama-v1-7b-lora-passage \
              --task_name "$task" \
-              --task_split $(if [[ "$task" == *"MSMARCO"* ]]; then echo "dev"; else echo "test"; fi) \
+              --task_split $([[ "$task" == *"MSMARCO"* ]] && echo "dev" || echo "test") \
              --query_instruction 'query: ' \
              --document_instruction 'passage: ' \
              --pooling_method last \
@@ -76,22 +128,40 @@ do
              --padding_side right \
              --add_bos_token 0 \
              --add_eos_token 1
+    done
+fi

-       # 5. NV-Embed-v1
-       python3.10 eval_mteb.py \
+
+# ===================================================================================
+#  Nvidia 5. NV-Embed-v1
+# ===================================================================================
+if [[ " ${MODELS_TO_RUN[*]} " =~ " NV-Embed-v1 " ]]; then
+    echo "===== Running Evaluation for Model: NV-Embed-v1 ====="
+    for task in "${TASKS[@]}"; do
+        echo "--- Task: $task ---"
+        python3.10 evaluation/eval_mteb.py \
              --base_model_name_or_path nvidia/NV-Embed-v1 \
              --output_folder en_results/nv-embed-v1 \
              --query_instruction "Given a claim, find documents that refute the claim" \
              --task_name "$task" \
-              --task_split $(if [[ "$task" == *"MSMARCO"* ]]; then echo "dev"; else echo "test"; fi) \
+              --task_split $([[ "$task" == *"MSMARCO"* ]] && echo "dev" || echo "test") \
              --eval_batch_size 8
+    done
+fi

-       # 6. BGE-EN-ICL
-       python3.10 eval_mteb.py \
+
+# ===================================================================================
+# 🎯 6. BGE-EN-ICL
+# ===================================================================================
+if [[ " ${MODELS_TO_RUN[*]} " =~ " BGE-EN-ICL " ]]; then
+    echo "===== Running Evaluation for Model: BGE-EN-ICL ====="
+    for task in "${TASKS[@]}"; do
+        echo "--- Task: $task ---"
+        python3.10 evaluation/eval_mteb.py \
              --base_model_name_or_path BAAI/bge-en-icl \
              --output_folder en_results/bge-en-icl \
              --task_name "$task" \
-              --task_split $(if [[ "$task" == *"MSMARCO"* ]]; then echo "dev"; else echo "test"; fi) \
+              --task_split $([[ "$task" == *"MSMARCO"* ]] && echo "dev" || echo "test") \
              --query_instruction $'<instruct> Given a scientific claim, retrieve documents that support or refute the claim.\n<query>' \
              --max_seq_length 512 \
              --eval_batch_size 32 \
@@ -99,18 +169,31 @@ do
              --padding_side left \
              --add_bos_token 1 \
              --add_eos_token 1
+    done
+fi

-       # 7. LLARA-passage
-       python3.10 eval_mteb.py \
+
+# ===================================================================================
+# 🦙 7. LLARA-passage
+# ===================================================================================
+if [[ " ${MODELS_TO_RUN[*]} " =~ " LLARA-passage " ]]; then
+    echo "===== Running Evaluation for Model: LLARA-passage ====="
+    for task in "${TASKS[@]}"; do
+        echo "--- Task: $task ---"
+        python3.10 evaluation/eval_mteb.py \
              --base_model_name_or_path BAAI/LLARA-passage \
              --output_folder en_results/llara-passage \
              --task_name "$task" \
-              --task_split $(if [[ "$task" == *"MSMARCO"* ]]; then echo "dev"; else echo "test"; fi) \
+              --task_split $([[ "$task" == *"MSMARCO"* ]] && echo "dev" || echo "test") \
              --eval_batch_size 8 \
              --pooling_method last_8 \
              --model_flag llara \
              --add_bos_token 1 \
              --add_eos_token 0 \
              --max_seq_length 532
+    done
+fi

-done
+
+
+echo "All specified evaluations are complete."
--- a/slm/pipelines/examples/contrastive_training/evaluation/prediction.py
+++ b/slm/pipelines/examples/contrastive_training/evaluation/prediction.py
@@ -11,20 +11,14 @@
 # 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.
-import os
-import sys
-
 import numpy as np
 import paddle

 from paddlenlp.data import DataCollatorWithPadding
-from paddlenlp.transformers import AutoTokenizer
-
-sys.path.append(os.path.abspath("."))
-from models.modeling import BiEncoderModel
+from paddlenlp.transformers import AutoTokenizer, BiEncoderModel


-class Eval_modle:
+class Eval_model:
    def __init__(
        self,
        model: str = None,
@@ -45,10 +39,8 @@ class Eval_modle:
        Construct the inference model for the predictor.
        """
        if self.model_type in ["bert", "roberta", "ernie"]:
-            self._model = BiEncoderModel.from_pretrained(
-                model_name_or_path=self.model,
-                normalized=True,
-                sentence_pooling_method="cls",
+            self._model = BiEncoderModel(
+                model_name_or_path=self.model, normalized=True, sentence_pooling_method="cls", dtype="float32"
            )
            print(f"loading checkpoints {self.model}")
        else:
@@ -138,7 +130,7 @@ class Eval_modle:
        with paddle.no_grad():
            for batch_inputs in inputs["batches"]:
                batch_inputs = self._collator(batch_inputs)
-                token_embeddings = self._model.encode(batch_inputs)
+                token_embeddings = self._model.encode(batch_inputs, self._model.model)
                all_feats.append(token_embeddings.detach().cpu().numpy())
            return all_feats

--- a/slm/pipelines/examples/contrastive_training/requirements.txt
+++ b/slm/pipelines/examples/contrastive_training/requirements.txt
@@ -1,2 +1,4 @@
 mteb==1.12.80
 typer==0.9.0
+h11==0.16.0
+requests[socks]
--- a/slm/pipelines/examples/contrastive_training/shortgpt_prune.py
+++ b/slm/pipelines/examples/contrastive_training/shortgpt_prune.py
@@ -0,0 +1,355 @@
+# Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
+#
+# 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.
+
+import argparse
+import json
+import math
+import os
+import re
+from collections import OrderedDict
+from typing import List
+
+import paddle
+from datasets import load_dataset
+from paddle.io import DataLoader
+from tqdm import tqdm
+
+from paddlenlp.transformers import AutoModelForCausalLM, AutoTokenizer, NVEncodeModel
+
+
+# =====================================================================================
+# 1. block_influence
+# =====================================================================================
+def block_influence(
+    input_hidden_state: paddle.Tensor,
+    output_hidden_state: paddle.Tensor,
+    angular: bool = False,
+) -> paddle.Tensor:
+    """
+    Calculates block influence between input and output hidden states.
+    """
+    _, _, d = input_hidden_state.shape
+    input_hidden_state = paddle.reshape(input_hidden_state, [-1, d])
+    output_hidden_state = paddle.reshape(output_hidden_state, [-1, d])
+
+    norm_input = paddle.norm(input_hidden_state, p=2, axis=-1, keepdim=True)
+    norm_output = paddle.norm(output_hidden_state, p=2, axis=-1, keepdim=True)
+
+    sim = paddle.matmul(input_hidden_state, output_hidden_state, transpose_y=True) / (norm_input * norm_output)
+    sim = paddle.diag(sim).astype("float32").nan_to_num(nan=0.5)
+
+    if angular:
+        return paddle.acos(sim) / math.pi
+    return 1 - sim
+
+
+# =====================================================================================
+# 2. ShortGPT
+# =====================================================================================
+class ShortGPT:
+    """
+    A class to evaluate layer importance in LLMs using PaddlePaddle.
+    """
+
+    def __init__(self, model_name: str, layers_path: str):
+        print(f"Loading tokenizer for '{model_name}'...")
+        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
+        self.tokenizer.pad_token = self.tokenizer.eos_token
+
+        print(f"Loading model '{model_name}' with PaddlePaddle backend...")
+        if "NV-Embed" in model_name:
+            self.model = NVEncodeModel.from_pretrained(
+                model_name, tokenizer_path=model_name, query_instruction="", document_instruction=""
+            )
+        else:
+            self.model = AutoModelForCausalLM.from_pretrained(model_name, dtype=paddle.float16)
+
+        self.model.eval()
+        print("Model loaded successfully for importance evaluation.")
+
+        try:
+            path_parts = layers_path.split(".")  # e.g., 'llama.layers' -> ['llama', 'layers']
+
+            self.base_model_for_call = self.model
+            # 遍历路径中除了最后 'layers' 之外的部分 (e.g., 'llama')
+            for part in path_parts[:-1]:
+                self.base_model_for_call = getattr(self.base_model_for_call, part)
+
+            # 从基础模型中获取 'layers' 列表
+            self.layers = getattr(self.base_model_for_call, path_parts[-1])
+            print(f"Successfully located base model for evaluation call: {type(self.base_model_for_call)}")
+            print(f"Successfully located {len(self.layers)} layers.")
+
+        except AttributeError:
+            raise AttributeError(f"Could not find layers at path '{layers_path}' in the model architecture.")
+
+        self.importances = [0.0 for _ in self.layers]
+
+    def compute_bi(self, hiddens: List[paddle.Tensor]):
+        """
+        Computes and accumulates block influence scores from hidden states.
+        """
+        n = 1
+        for i in range(len(hiddens) - n):
+            layer_index = i
+            if layer_index < len(self.importances):
+                in_hidden = hiddens[i]
+                out_hidden = hiddens[i + n]
+                self.importances[layer_index] += block_influence(in_hidden, out_hidden).sum().item()
+
+    @paddle.no_grad()
+    def eval_importance(self, prompts: List[str], model_name: str, stride: int = 256):
+        """
+        Evaluates the importance of model layers on given prompts.
+        """
+        prompt_tokens = self.tokenizer(prompts, padding=True, return_attention_mask=True, return_tensors="pd")
+        input_ids = prompt_tokens.input_ids
+        attn_mask = prompt_tokens.attention_mask
+
+        max_prompt_len = max(len(t) for t in input_ids)
+
+        for start in range(0, max_prompt_len, stride):
+            seq_ids = (attn_mask.sum(axis=-1) > start).nonzero().squeeze()
+            seq_ids = seq_ids.unsqueeze(0) if seq_ids.ndim == 0 else seq_ids
+
+            if seq_ids.shape[0] == 0:
+                continue
+
+            inputs = input_ids[seq_ids, start : start + stride]
+            attn = attn_mask[seq_ids, start : start + stride]
+
+            if "NV-Embed" in model_name:
+                outputs = self.base_model_for_call.m_forward(
+                    input_ids=inputs, attention_mask=attn, output_hidden_states=True, return_dict=True
+                )
+            else:
+                outputs = self.base_model_for_call(
+                    input_ids=inputs, attention_mask=attn, output_hidden_states=True, return_dict=True
+                )
+
+            if outputs.hidden_states:
+                self.compute_bi(outputs.hidden_states)
+
+
+def load_model_weights(model_folder_path: str) -> OrderedDict:
+    print(f"Attempting to load model weights from FOLDER: '{model_folder_path}'...")
+
+    # 1. Ensure the path is a valid directory
+    if not os.path.isdir(model_folder_path):
+        raise NotADirectoryError(f"The provided path is not a valid directory: '{model_folder_path}'")
+
+    state_dict = OrderedDict()
+    index_path = os.path.join(model_folder_path, "model_state.pdparams.index.json")
+
+    # 2. Check for the presence of a sharded index file
+    if os.path.isfile(index_path):
+        # Case A: Sharded model format detected (index file found)
+        print("Sharded model format detected (index file found).")
+        with open(index_path, "r", encoding="utf-8") as f:
+            index_data = json.load(f)
+
+        shard_files = sorted(list(set(index_data["weight_map"].values())))
+        print(f"Found {len(shard_files)} shard(s).")
+
+        for shard_file in shard_files:
+            shard_path = os.path.join(model_folder_path, shard_file)
+            if not os.path.exists(shard_path):
+                raise FileNotFoundError(f"Shard file '{shard_file}' listed in index not found at '{shard_path}'")
+
+            print(f"  > Loading shard: {shard_file}")
+            shard_state_dict = paddle.load(shard_path, return_numpy=True)
+            state_dict.update(shard_state_dict)
+            del shard_state_dict
+        print("All weight shards loaded successfully.")
+
+    else:
+        # Case B: No index file found; look for a single .pdparams file
+        print("No index file found. Searching for a single .pdparams file inside the folder...")
+        pdparams_files = [f for f in os.listdir(model_folder_path) if f.endswith(".pdparams")]
+
+        if len(pdparams_files) == 1:
+            # Found exactly one .pdparams file
+            single_file_path = os.path.join(model_folder_path, pdparams_files[0])
+            print(f"  > Loading single parameters file: {pdparams_files[0]}")
+            state_dict = paddle.load(single_file_path, return_numpy=True)
+            print("Single weight file loaded successfully.")
+        elif len(pdparams_files) > 1:
+            raise ValueError(
+                f"Ambiguous model files. Multiple .pdparams files found in '{model_folder_path}' "
+                "but no 'model_state.pdparams.index.json' to specify order."
+            )
+        else:  # len(pdparams_files) == 0
+            raise FileNotFoundError(f"No .pdparams files found in the directory '{model_folder_path}'.")
+
+    return state_dict
+
+
+# =====================================================================================
+# 3. Prune and Save
+# =====================================================================================
+def prune_and_save_model_in_memory(
+    model,
+    tokenizer,
+    new_model_path,
+    layers_to_delete,
+    layers_path_str,
+):
+    """
+    Prunes and saves a model directly from the in-memory model object.
+    """
+    print("=" * 50)
+    print("PART 2: Starting In-Memory Model Pruning and Saving")
+    print("=" * 50)
+    os.makedirs(new_model_path, exist_ok=True)
+
+    # Step 1: Get state_dict directly from the in-memory model
+    print("Getting state_dict directly from the in-memory model...")
+    state_dict = model.state_dict()
+
+    # Step 2: Iterate, filter, and rename weights
+    print("Processing weights: removing specified layers and re-indexing...")
+    escaped_layers_path = layers_path_str.replace(".", r"\.")
+    layer_pattern = re.compile(rf"^{escaped_layers_path}\.(\d+)\.")
+    new_state_dict = OrderedDict()
+
+    for key, value in state_dict.items():
+        match = layer_pattern.match(key)
+        if not match:
+            new_state_dict[key] = value
+            continue
+
+        layer_idx = int(match.group(1))
+        if layer_idx in layers_to_delete:
+            continue
+
+        num_layers_deleted_before = sum(1 for deleted_idx in layers_to_delete if deleted_idx < layer_idx)
+        new_layer_idx = layer_idx - num_layers_deleted_before
+        old_prefix = f"{layers_path_str}.{layer_idx}."
+        new_prefix = f"{layers_path_str}.{new_layer_idx}."
+        new_key = key.replace(old_prefix, new_prefix, 1)
+        new_state_dict[new_key] = value
+
+    print(f"Processing complete. Removed {len(layers_to_delete)} layer(s): {sorted(list(layers_to_delete))}.")
+
+    # Step 3: Get and modify the configuration from the model object
+    print("Updating configuration file...")
+    config = model.config.to_dict()
+
+    # Fix: Convert non-serializable paddle data types to strings
+    for key, value in config.items():
+        if type(value).__name__ == "DataType":
+            config[key] = str(value).split(".")[-1]
+
+    if "num_hidden_layers" in config:
+        original_num_layers = config["num_hidden_layers"]
+        new_num_layers = original_num_layers - len(layers_to_delete)
+        config["num_hidden_layers"] = new_num_layers
+        print(f"  - Number of layers changed from {original_num_layers} to {new_num_layers}.")
+
+    new_config_path = os.path.join(new_model_path, "config.json")
+    with open(new_config_path, "w", encoding="utf-8") as f:
+        json.dump(config, f, indent=2)
+    print(f"New config saved to '{new_config_path}'.")
+
+    # Step 4: Save the new weights and tokenizer
+    print("Saving pruned weights...")
+    new_weights_path = os.path.join(new_model_path, "model_state.pdparams")
+    paddle.save(new_state_dict, new_weights_path)
+    print(f"Pruned weights saved to '{new_weights_path}'.")
+
+    print("Saving tokenizer files...")
+    tokenizer.save_pretrained(new_model_path)
+    print(f"Tokenizer files saved to '{new_model_path}'.")
+
+    print("\n🎉 Pruning process completed successfully!")
+    print(f"Pruned model has been saved to '{new_model_path}'")
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="Calculate layer importance, prune, and save a new PaddlePaddle model."
+    )
+    parser.add_argument(
+        "--model_name_or_path",
+        type=str,
+        required=True,
+        help="Path or HuggingFace name of the source PaddlePaddle model.",
+    )
+    parser.add_argument(
+        "--output_model_path", type=str, required=True, help="Path to save the new, pruned model directory."
+    )
+    parser.add_argument(
+        "--layers_path", type=str, required=True, help="Dot-separated path to the layers list (e.g., 'llama.layers')."
+    )
+    parser.add_argument(
+        "--n_prune_layers", type=int, required=True, help="The number of layers to identify and prune."
+    )
+    parser.add_argument(
+        "--dataset_name",
+        type=str,
+        default="emozilla/pg19",
+        help="Name of the Hugging Face dataset for calibration. Default: 'emozilla/pg19'.",
+    )
+    parser.add_argument(
+        "--dataset_split",
+        type=str,
+        default="validation",
+        help="The split of the dataset to use. Default: 'validation'.",
+    )
+    args = parser.parse_args()
+
+    # --- PART 1: Calculate Layer Importance ---
+    print("=" * 50)
+    print("PART 1: Calculating Layer Importance")
+    print("=" * 50)
+    print(f"Loading '{args.dataset_split}' split from '{args.dataset_name}' dataset for calibration...")
+    try:
+        data = load_dataset(args.dataset_name, split=args.dataset_split)
+    except Exception as e:
+        print(f"Failed to load dataset. Error: {e}")
+        print(
+            "Please ensure the dataset name and split are correct and you have internet access for Hugging Face datasets."
+        )
+        return
+
+    dataloader = DataLoader(data, batch_size=1, shuffle=False)
+
+    short_model = ShortGPT(model_name=args.model_name_or_path, layers_path=args.layers_path)
+
+    for batch in tqdm(dataloader, desc="Evaluating Layer Importance"):
+        if "text" not in batch:
+            raise ValueError("Dataset must contain a 'text' column.")
+        prompts = batch["text"]
+        short_model.eval_importance(prompts=prompts, model_name=args.model_name_or_path, stride=256)
+
+    prune_order = sorted(range(len(short_model.importances)), key=lambda i: short_model.importances[i])
+    layers_to_delete = set(prune_order[: args.n_prune_layers])
+
+    print("\n--- Importance Calculation Complete ---")
+    print(f"Calculated importances: {[f'{v:.2f}' for v in short_model.importances]}")
+    print(f"Pruning order (least to most important): {prune_order}")
+    print(f"Will delete the {args.n_prune_layers} least important layers: {sorted(list(layers_to_delete))}")
+
+    # --- PART 2: Perform In-Memory Pruning and Saving ---
+    prune_and_save_model_in_memory(
+        model=short_model.model,
+        tokenizer=short_model.tokenizer,
+        new_model_path=args.output_model_path,
+        layers_to_delete=layers_to_delete,
+        layers_path_str=args.layers_path,
+    )
+
+
+if __name__ == "__main__":
+    main()
--- a/slm/pipelines/examples/contrastive_training/train.py
+++ b/slm/pipelines/examples/contrastive_training/train.py
@@ -109,6 +109,7 @@ def main():
            use_inbatch_neg=training_args.use_inbatch_neg,
            matryoshka_dims=training_args.matryoshka_dims if training_args.use_matryoshka else None,
            matryoshka_loss_weights=training_args.matryoshka_loss_weights if training_args.use_matryoshka else None,
+            dtype=dtype,
        )

    if training_args.fix_position_embedding:
@@ -129,13 +130,13 @@ def main():
    if training_args.fine_tune_type == "lora":
        if any([x in model_args.model_name_or_path for x in ["llama", "baichuan", "NV-Embed"]]):
            target_modules = [
-                ".*q_proj.*",
-                ".*k_proj.*",
-                ".*v_proj.*",
-                ".*o_proj.*",
-                ".*down_proj.*",
-                ".*up_proj.*",
-                ".*gate_proj.*",
+                ".*q_proj$",
+                ".*k_proj$",
+                ".*v_proj$",
+                ".*o_proj$",
+                ".*down_proj$",
+                ".*up_proj$",
+                ".*gate_proj$",
            ]
        else:
            raise ValueError("need to specify the target modules for LoRA fine-tuning.")