You need to enable JavaScript to run this app.
导航
AWQ量化 -- LlaMA3.1-8B 在 MLP 的最佳实践
最近更新时间:2024.09.14 15:37:22首次发布时间:2024.09.14 15:37:22

机器学习平台是一套服务于专业算法工程师,整合云原生的工具+算力 (GPU、CPU云服务器),进行一站式AI算法开发和迭代的平台。本方案将为您介绍如何在 MLP 对 Meta-Llama-3.1-8B 的开源模型进行AWQ量化。

使用前提

本方案以对 Meta-Llama-3.1-8B 模型进行AWQ量化为例,在开始执行操作前,请确认您已经完成以下准备工作:

步骤一:环境配置
  1. 进入开发机。

    1. 登陆 账号登录-火山引擎
    2. 在左侧导航栏单击开发机,在开发机列表页面中单击待操作的开发机名称,进入对应开发机内。
    3. 通过 SSH 远程登录开发机,具体操作详见通过SSH远程连接开发机--机器学习平台-火山引擎
  2. 配置volc configure。具体操作详见使用文档--机器学习平台-火山引擎

# 配置AK/SK,以及环境的Region
volc configure
  1. 进行Huggingface 环境配置。

    export HF_TOKEN=<User Specific Huggingface Access Token>
    export HF_ENDPOINT=[https://hf-mirror.com](https://hf-mirror.com)
    pip3 install -U huggingface_hub
    

    参数说明如下:

    参数描述
    hf_token具有llama使用权限的Huggingface账号的hf_token。
    hf_endpoint镜像地址。

    说明:

    1. --token/--hf_username/--hf_token等参数非必须,仅为部分Repo有license限制,需登录申请许可。如果必要,请在 https://huggingface.co/settings/tokens 获取token后进行下载。
    2. 推荐您将模型文件、训练数据集的下载、存放地址都设置在 TOS/VEPFS/SCFS 等共享存储内,之后在进行训练、推理等跨节点的任务时,可以更加轻松地挂载到不同的机器上。
步骤二:准备 Llama-3.1-8B 模型
  1. 在Terminal中执行以下命令,通过镜像源下载 Llama-3.1-8B 模型,并打开挂载的存储桶路径。
cd /vemlp-demo-models
mkdir model-llama
cd model-llama
huggingface-cli download --token $HF_TOKEN --resume-download meta-llama/Meta-Llama-3.1-8B-Instruct --local-dir $ORIGINAL_MODEL_PATH

其中运行 ./hfd.sh 文件,需要传入的参数列表说明如下:

参数描述
hf_token具有llama使用权限的Huggingface账号的hf_token。
original_model_path基础模型的本地存储路径。
步骤三:准备数据集

在Terminal执行以下命令,下载 MLP 已准备好的数据集。

cd /vemlp-demo-models
mkdir datasets-llama-quant
huggingface-cli download --repo-type dataset --resume-download carlosejimenez/wikitext__wikitext-2-raw-v1 --local-dir $CALIB_DATASET_PATH

参数说明如下:

参数描述
calib_dataset_path用来校验量化结果的数据集的路径。
步骤四:在开发机中执行量化

4.1 创建量化脚本

  1. 在Terminal中执行以下命令,创建量化脚本 awq-quant.py
cd /vemlp-demo-models
touch awq-quant.py
  1. 编辑脚本内容,替换为需要配置的信息。
import os
from awq import AutoAWQForCausalLM
from transformers import AutoTokenizer
from datasets import load_dataset

# 配置模型、数据集地址
os.environ['BASE_MODEL_PATH'] = '/vemlp-demo-models/model-llama/Meta-Llama-3.1-8B'
os.environ['CALIB_DATASET_PATH'] = '/vemlp-demo-models/datasets-llama-quant/'
os.environ['QUANT_MODEL_PATH'] = '/vemlp-demo-models/model-quant-awq'

model_path = os.getenv('BASE_MODEL_PATH')
dataset_path = os.getenv('CALIB_DATASET_PATH')
quant_path = os.getenv('QUANT_MODEL_PATH')

# 加载基础模型及tokenizer
tokenizer = AutoTokenizer.from_pretrained(model_path)
model = AutoAWQForCausalLM.from_pretrained(model_path, device_map="auto", safetensors=True)

# 加载校准数据集,并准备数据以进行校准。
data = load_dataset(dataset_path, split="train")
dataset = [text for text in data["text"] if text.strip() != '' and len(text.split(' ')) > 20]

# 设置AWQ量化的超参数,目前AutoAWQ仅支持4bit量化,需设置w_bit为4
quant_config = { "zero_point": True, "q_group_size": 128, "w_bit": 4, "version": "GEMM" }
model.quantize(tokenizer, quant_config=quant_config, calib_data=dataset)

model.save_quantized(quant_path, safetensors=True, shard_size="4GB")
tokenizer.save_pretrained(quant_path)

4.2 创建Shell脚本并执行

  1. 在Terminal中执行以下命令,创建shell 脚本 awq-quant.sh。请确保awq-quant.sh 与 awq-quant.py 两个文件在同一层级目录。
cd /vemlp-demo-models/
touch awq-quant.sh
  1. 编辑脚本内容,替换为需要配置的信息。
set  -ex

pip install autoawq==0.2.6
python3 awq-quant.py
echo "Finished Model Quantization Task"

# 如果用A100且GPU driver是470,可以运行下述命令
pip install autoawq==0.2.1 torch==2.1.0+cu118
pip install autoawq==0.2.1
pip install torch==2.1.0+cu118  --index-url [https://download.pytorch.org/whl](https://download.pytorch.org/whl)
python3 awq-quant.py
echo "Finished Model Quantization Task"
  1. 执行脚本。
cd /vemlp-demo-models/
bash ./awq-quant.sh

执行完成后,输出路径文件如下。

4.3 量化后模型对比

  1. 在Terminal中执行以下命令,对量化前的基础模型和量化模型进行对比。
du -sh $BASE_MODEL_PATH $QUANT_MODEL_PATH

对比结果如下:可见在模型容量上,经过 AWQ 4bit 量化后模型大小约为量化前模型的1/3。

  1. 在Terminal中执行以下命令,对量化前后的config.json文件进行对比。
diff $QUANT_MODEL_PATH/config.json $BASE_MODEL_PATH/config.json

对比结果如下:可见量化后模型多了关于量化的信息 -- quantization_config。


步骤五:在自定义任务中执行量化

方式一:控制台

  1. 进入自定义任务控制台

  2. 配置环境

    1. 在自定义镜像处配置镜像。
    vemlp-cn-beijing.cr.volces.com/preset-images/pytorch:2.1.0
    
    1. 填写入口命令。请注意tos在开发机上挂载的路径与自定义任务中挂载的路径不一致,在入口命令填写数据集路径环境变量中填写模型路径时,需确保为自定义任务中挂载的路径。
    cat > awq-quant.py <<EOF
    import os
    from awq import AutoAWQForCausalLM
    from transformers import AutoTokenizer
    from datasets import load_dataset
    
    # 模型、数据集地址配置
    os.environ['BASE_MODEL_PATH'] = '/vemlp-demo-models/model-llama/Meta-Llama-3.1-8B'
    os.environ['CALIB_DATASET_PATH'] = '/vemlp-demo-models/datasets-llama-quant/'
    os.environ['QUANT_MODEL_PATH'] = '/vemlp-demo-models/model-quant-awq'
    
    model_path = os.getenv('BASE_MODEL_PATH')
    dataset_path = os.getenv('CALIB_DATASET_PATH')
    quant_path = os.getenv('QUANT_MODEL_PATH')
    
    tokenizer = AutoTokenizer.from_pretrained(model_path)
    model = AutoAWQForCausalLM.from_pretrained(model_path, device_map="auto", safetensors=True)
    
    data = load_dataset(dataset_path, split="train")
    dataset = [text for text in data["text"] if text.strip() != '' and len(text.split(' ')) > 20]
    
    # 设置量化超参数
    quant_config = { "zero_point": True, "q_group_size": 128, "w_bit": 4, "version": "GEMM" }
    model.quantize(tokenizer, quant_config=quant_config, calib_data=dataset)
    
    model.save_quantized(quant_path, safetensors=True, shard_size="4GB")
    tokenizer.save_pretrained(quant_path)
    EOF
    
    pip install autoawq==0.2.6
    python3 awq-quant.py
    

  3. 配置资源

    1. 选择合适的资源,最佳实践所用机器为单机8卡ml.pni2。
    2. 挂载TOS和CFS。需要注意:TOS在开发机上挂载的路径与自定义任务中挂载的路径不一致,在环境变量中填写模型路径时需确保为自定义任务中挂载的路径。

方式二:命令行

  1. 在Terminal中执行以下命令,编辑 awq-quant.py。将开发机中的云盘访问路径变更为自定义任务中的访问路径。
import os
  from awq import AutoAWQForCausalLM
  from transformers import AutoTokenizer
  from datasets import load_dataset

  # 配置模型、数据集地址
  os.environ['BASE_MODEL_PATH'] = '/data/model-llama/Meta-Llama-3.1-8B'
  os.environ['CALIB_DATASET_PATH'] = '/data/datasets-llama-quant/'
  os.environ['QUANT_MODEL_PATH'] = '/data/model-quant-awq'

  model_path = os.getenv('BASE_MODEL_PATH')
	dataset_path = os.getenv('CALIB_DATASET_PATH')
	quant_path = os.getenv('QUANT_MODEL_PATH')

	# 加载基础模型及tokenizer
	tokenizer = AutoTokenizer.from_pretrained(model_path)
	model = AutoAWQForCausalLM.from_pretrained(model_path, device_map="auto", safetensors=True)

	# 加载校准数据集,并准备数据以进行校准。
	data = load_dataset(dataset_path, split="train")
	dataset = [text for text in data["text"] if text.strip() != '' and len(text.split(' ')) > 20]

	# 设置AWQ量化的超参数,目前AutoAWQ仅支持4bit量化,需设置w_bit为4
	quant_config = { "zero_point": True, "q_group_size": 128, "w_bit": 4, "version": "GEMM" }
	model.quantize(tokenizer, quant_config=quant_config, calib_data=dataset)

	model.save_quantized(quant_path, safetensors=True, shard_size="4GB")
	tokenizer.save_pretrained(quant_path)
  1. 在Terminal中执行以下命令,创建一个用于执行量化的自定义任务的 quantization-task.yaml 文件。
cd /vemlp-demo-models/
touch quantization-task.yaml
  1. 编辑量化任务提交的yaml文件。
TaskName:  "Quant-Meta-Llama-3-1-8B-Instruct-AWQ-INT4"
Description:  "Use AWQ to quantize Meta-Llama-3.1-8B-Instruct in INT4 format"
Entrypoint: |
    cd  /data && bash ./awq-quant.sh
Tags: []
Envs: []

ResourceQueueID:  "q-**************"            #请替换为你自己的 ResourceQueueID
Framework:  "PyTorchDDP"
TaskRoleSpecs:                                  #请替换为你自己的 TaskRoleSpecs
        - RoleName:  "worker"
            RoleReplicas: 1
            Flavor:  "custom"
            ResourceSpec:
                    Family:  "ml.xni3c"
                    CPU: 15.000
                    Memory: 70.000
                    GPUNum: 1
            
ActiveDeadlineSeconds: 3600
EnableTensorBoard: False

Storages:                                        #请替换为你自己的 Storages
        - Type:  "Tos"
            MountPath:  "/data"
            Bucket:  "vemlp-demo-models"
            FsName": "testmodel"
            MetaCacheExpiryMinutes: "-1"

ImageUrl: "vemlp-cn-beijing.cr.volces.com/preset-images/pytorch:2.1.0"
CacheType: "Cloudfs"

RetryOptions:
        EnableRetry: False
        MaxRetryTimes: 5
        IntervalSeconds: 120
        PolicySets: []
  1. 在Terminal中执行以下命令,进行自定义任务提交。
volc ml_task submit --conf=/vemlp-demo-models/awq-quantization-task.yaml
  1. 量化模型前后对比查看可参考上文。