Logo
Project Icon

cbtm

无监督领域发现技术扩展专家语言模型

c-BTM语言模型聚类专家模型评估Github开源项目

Cluster-Branch-Train-Merge (c-BTM)项目提出了一种通过无监督领域发现技术扩展专家语言模型的方法。该项目包含完整的训练评估流程，涵盖数据准备、聚类、模型训练和评估等步骤。c-BTM旨在提升大规模语言模型在特定领域的性能，为自然语言处理研究提供新思路。

GithubGithubHuggingfaceHuggingface论文论文
介绍相关项目

集群-分支-训练-合并 (c-BTM)

论文《通过无监督领域发现扩展专家语言模型》的代码

本仓库是metaseq的分支。

引用

如果您使用此代码，请考虑引用我们的工作:

@article{cbtm,
 author = {Suchin Gururangan and Margaret Li and Mike Lewis and Weijia Shi and Tim Althoff and Noah A. Smith and Luke Zettlemoyer},
 title = {通过无监督领域发现扩展专家语言模型},
 journal={arXiv预印本 arXiv:2303.14177},
 year = {2023}
}

设置环境

克隆此仓库

git clone https://github.com/kernelmachine/cbtm.git
cd cbtm

创建新的conda环境(推荐)

conda create -n cbtm python=3.9
conda activate cbtm

或者，我们提供了一个environment.yml文件;可用于创建一个包含python 3.9和各种依赖项的conda环境。这将需要几分钟时间，可能不适用于所有版本的conda。

conda env create -n cbtm -f environment.yml
conda activate cbtm

安装PyTorch

我们使用Cuda 11.3编译的torch测试了此代码。

pip3 install torch==1.10.1+cu113  -f https://download.pytorch.org/whl/cu113/torch_stable.html

安装Megatron

确保在此步骤中有GPU和可见的CUDA。

cd ..
git clone --branch fairseq_v2 https://github.com/ngoyal2707/Megatron-LM.git
cd Megatron-LM
pip3 install six regex numpy
pip3 install -e .

安装fairscale

cd ..
git clone https://github.com/facebookresearch/fairscale.git
cd fairscale
git checkout prefetch_fsdp_params_simple
pip3 install -e .

安装balanced-kmeans

cd ..
git clone https://github.com/kernelmachine/balanced-kmeans.git
cd balanced-kmeans
pip3 install -e .

(可选) 安装Apex

Apex可能与所有GPU不兼容。特别是，如果您在前向传播过程中发现CUDA不支持您的模型，您可能需要尝试卸载Apex并重试。

cd ..
git clone https://github.com/NVIDIA/apex.git
cd apex
git checkout e2083df5eb96643c61613b9df48dd4eea6b07690

根据您的硬件，在运行下一个pip安装之前，您可能需要注释掉setup.py中的101-107行。

pip3 install -v --no-cache-dir --global-option="--cpp_ext" --global-option="--cuda_ext" --global-option="--deprecated_fused_adam" --global-option="--xentropy" --global-option="--fast_multihead_attn" ./

安装c-BTM库

构建c-BTM库。

cd ../cbtm
pip3 install -e .

修复setuptools版本

上述torch版本与setuptools>=59.6.0之间存在已知的兼容性问题，因此我们将此包回滚:

pip3 install setuptools==59.5.0

c-BTM 训练与评估

步骤0: 设置数据、模型和目录

为简单起见，我们将在本教程中使用以下环境变量。您可以根据需要设置这些变量。

export CBTM_DIR=~/cbtm; 
export DATA_DIR=${CBTM_DIR}/data;
export SERIALIZATION_DIR=${CBTM_DIR}/experiments;
export KMEANS_DIR=${CBTM_DIR}/clusterers;
export CLUSTERS_DIR=${CBTM_DIR}/clusters;
export VOCAB_DIR=$CBTM_DIR/vocab
export PRETRAINED_MODELS_DIR=${CBTM_DIR}/pretrained_models;
mkdir -p ${CBTM_DIR} ${DATA_DIR} ${SERIALIZATION_DIR} ${KMEANS_DIR} ${CLUSTERS_DIR} ${VOCAB_DIR} ${PRETRAINED_MODELS_DIR};

配置cbtm_constants.py

接下来，使该仓库工作所需的常量位于metaseq/cbtm_constants.py中。根据您的本地环境修改这些常量。

确保metaseq/cbtm_constants.py中的变量与您上面设置为环境变量的路径一致。

下载词汇文件和种子模型

我们使用GPT-2词汇:

wget -O $VOCAB_DIR/gpt2-vocab.json http://s3.wasabisys.com/c4-example/vocab/gpt2-vocab.json
wget -O $VOCAB_DIR/gpt2-merges.txt http://s3.wasabisys.com/c4-example/vocab/gpt2-merges.txt

下载OPT-1.3B和OPT-6.7B检查点，我们将其用作种子模型。如果您不打算使用6.7B模型，可以跳过最后2个wget命令，这些命令需要一些时间才能完成:

mkdir -p ${PRETRAINED_MODELS_DIR}/opt/1.3b/
mkdir -p ${PRETRAINED_MODELS_DIR}/opt/6.7b/
wget -qO- dl.fbaipublicfiles.com/cbtm/opt_models/1.3B/sharded_for_ddp.tgz | tar xvz  --strip-components 6 -C ${PRETRAINED_MODELS_DIR}/opt/1.3b/
wget -qO- dl.fbaipublicfiles.com/cbtm/opt_models/6.7B/sharded_for_ddp_part_0.tgz | tar xvz  --strip-components 6 -C ${PRETRAINED_MODELS_DIR}/opt/6.7b/
wget -qO- dl.fbaipublicfiles.com/cbtm/opt_models/6.7B/sharded_for_ddp_part_1.tgz | tar xvz --strip-components 6 -C ${PRETRAINED_MODELS_DIR}/opt/6.7b/

下载数据

我们提供了一些C4样本数据供您开始使用。我们的模型只需要(分片的)按行分隔的jsonl文件，分为训练和验证数据。如果您想在自己的数据上进行训练，只需按照示例中的整体数据布局即可。

mkdir -p ${DATA_DIR}/c4_example/
wget -qO- http://s3.wasabisys.com/c4-example/c4_example.tar.gz | tar xvz -C ${DATA_DIR}/

这个示例数据集是C4的单个分片和验证数据集的小样本。

您可以从以下链接的Huggingface数据集下载完整的C4数据集: https://huggingface.co/datasets/c4。请记住，该数据集非常大，并以`json.gz`文件的形式提供。我们的代码期望原始jsonl文件采用示例目录中的结构，因此请确保您有足够的空间(总共约1.4TB的未压缩数据)。

Metaseq期望数据采用以下格式:

{"text": "这是一个文档", "id": 0}
{"text": "这是另一个文档", "id": 1}

步骤1: 训练聚类器

此命令在C4训练数据的单个分片上训练一个平衡的k-means聚类器。这里我们使用k=8，并给出一个包含名为C4.jsonl文件的文件夹作为参数，如上所述。

确保在这里有一个GPU可用，以加速训练!

NUM_CLUSTERS=8;
DATASET=c4_example;
python -m metaseq.scripts.train_clusterer \
--data-dir ${DATA_DIR}/${DATASET}/train/00000 \
--num-clusters ${NUM_CLUSTERS} \
--balanced \
--output-dir ${KMEANS_DIR}/${DATASET}/${NUM_CLUSTERS}/

这将在${KMEANS_DIR}/${DATASET}/${NUM_CLUSTERS}/创建tfidf.pkl(一个tf-idf嵌入器)和kmeans.pkl(一个kmeans聚类器)pickle文件。

步骤2: 聚类数据

此代码使用您训练的聚类器对数据集的文档进行聚类。如果您可以将其作为slurm作业并行化，这将显著加快速度。

我们将通过submitit自动为您完成此操作，只需提供您的slurm账户和分区(可以作为程序的标志，也可以在metaseq/cbtm_constants.py中提供)

如果您无法访问slurm，可以使用--run local标志在本地聚类数据，但可能需要一些时间!

DATASET=c4_example;
NUM_CLUSTERS=8;

# 聚类训练数据
python -m metaseq.scripts.cluster \
--job-dir ${CBTM_DIR}/cluster_logs \
--data-dir ${DATA_DIR}/${DATASET} \
--path-to-clusterer ${KMEANS_DIR}/${DATASET}/${NUM_CLUSTERS}/ \
--num-clusters ${NUM_CLUSTERS} \
--output-prefix ${CLUSTERS_DIR}/${DATASET}/${NUM_CLUSTERS}/ \
--split train \
--run slurm;
# 验证集聚类数据
python -m metaseq.scripts.cluster \
--job-dir ${CBTM_DIR}/cluster_logs \
--data-dir ${DATA_DIR}/${DATASET} \
--path-to-clusterer ${KMEANS_DIR}/${DATASET}/${NUM_CLUSTERS}/ \
--num-clusters ${NUM_CLUSTERS} \
--output-prefix ${CLUSTERS_DIR}/${DATASET}/${NUM_CLUSTERS}/ \
--split valid/C4_small \
--run slurm;

这些聚类作业的日志会出现在 ${CBTM_DIR}/cluster_logs 中。

这些作业完成后，打开 ${CLUSTERS_DIR} 中的文件，例如 ${CLUSTERS_DIR}/${DATASET}/${NUM_CLUSTERS}/train/00000/C4.jsonl。您应该看到如下行:

{"sp_id":"\/gscratch\/zlab\/sg01\/data\/c4_example\/train\/00000\/C4.jsonl|0","cluster":5}
{"sp_id":"\/gscratch\/zlab\/sg01\/data\/c4_example\/train\/00000\/C4.jsonl|1","cluster":3}
{"sp_id":"\/gscratch\/zlab\/sg01\/data\/c4_example\/train\/00000\/C4.jsonl|2","cluster":2}

字段 sp_id 表示文件中的一行（即文档），字段 cluster 表示其预测的聚类。

步骤 3: 训练模型

现在我们将使用聚类数据来训练专家模型。您需要至少同时使用 4 个 GPU 来训练每个模型。

本教程使用我们的 train_cbtm 脚本，该脚本与 SLURM 交互。默认情况下，训练作业将为每个 GPU 请求 8 个 CPU 和 48G 内存。这些资源参数可以在 train_cbtm.sh 脚本中自定义。

我们还提供了一个示例 sbatch 脚本（如果需要的话），位于 metaseq/scripts/example_sbatch.sh。您可能需要编辑这个示例 sbatch 命令，以包含在您的系统上使其正常工作所需的任何额外 slurm 参数。

训练专家模型

以下命令将使用每个 4 个 GPU 训练 8 个专家模型，每个模型训练 50 步（增加到 10000 步以复现我们的论文结果）。

DATASET=c4_example;
NUM_CLUSTERS=8;
NUM_GPUS_PER_EXPERT=4;

if [[ "$NUM_GPUS_PER_EXPERT" -lt 8 ]]; then
    NUM_NODES=1
    NUM_GPUS_PER_NODE=$NUM_GPUS
else 
    NUM_NODES=$((NUM_GPUS/8))
    NUM_GPUS_PER_NODE=8
fi
python -m metaseq.scripts.train_cbtm \
   --model-size 1.3b \
   --run slurm   \
   --path-to-clusters-dir $CLUSTERS_DIR/${DATASET}/$NUM_CLUSTERS/ \
   --num-clusters $NUM_CLUSTERS  \
   --num-nodes $NUM_NODES \
   --num-gpus $NUM_GPUS_PER_NODE \
   --data-name ${DATASET}  \
   --path-to-data $DATA_DIR/${DATASET} \
   --learning-rate 2e-4 \
   --max-steps 50 \
   --valid-subset valid/C4_small \
   --train-subset train

要训练特定的聚类，您可以添加标志 --train-cluster 1,3,5

要在本地调试，将 run 标志更改为 --run local。

此命令将输出检查点和日志到 ${SERIALIZATION_DIR}/8_clusters/。

密集训练

以下命令将使用 4 个 GPU 训练一个密集模型，训练 50 步（增加到 10000 步以复现我们的论文结果）。

DATASET=c4_example;
NUM_GPUS_PER_EXPERT=4;

if [[ "$NUM_GPUS_PER_EXPERT" -lt 8 ]]; then
    NUM_NODES=1
    NUM_GPUS_PER_NODE=$NUM_GPUS
else 
    NUM_NODES=$((NUM_GPUS/8))
    NUM_GPUS_PER_NODE=8
fi
python -m metaseq.scripts.train_cbtm \
   --num-clusters 1 \
   --model-size 1.3b \
   --run slurm \
   --num-nodes $NUM_NODES \
   --num-gpus $NUM_GPUS_PER_NODE \
   --data-name ${DATASET}  \
   --path-to-data $DATA_DIR/$DATASET  \
   --learning-rate 2e-4 \
   --max-steps 50 \
   --valid-subset valid/C4_small \
   --train-subset train

要在本地调试，将 run 标志更改为 --run local。

此命令将输出检查点到 ${SERIALIZATION_DIR}/1_clusters/。

评估

要评估您的模型，首先将分片合并到一个单一的检查点文件中。

以下脚本依赖于 gnu-parallel 包。您可能需要修改 metaseq/scripts/consolidate_fsdp_shards.sh 的最后一行，以指向您的 gnu-parallel 路径。

NUM_CLUSTERS=8;
NUM_GPUS_PER_EXPERT=4;
bash metaseq/scripts/consolidate_fsdp_shards.sh ${SERIALIZATION_DIR}/${NUM_CLUSTERS}_clusters/ "*ngpu${NUM_GPUS_PER_EXPERT}"

这将在每个模型的文件夹中创建一个 consolidated.pt 检查点。

现在检查点已准备好进行评估。要在 slurm 上启动：

export NUM_CLUSTERS=8;
# 我们需要的 GPU 数量与聚类数量相同
export NUM_GPUS=${NUM_CLUSTERS};
export DATASET=c4_example;
export EVAL_DIR=${SERIALIZATION_DIR}/${NUM_CLUSTERS}_clusters/eval

mkdir -p ${EVAL_DIR};

# 获取模型检查点
CONSOLIDATED_MODEL_PATHS=;
# 此函数获取所有模型检查点目录并按聚类 ID 排序
# 如果需要，修改文件夹模式以匹配您的模型目录名称
FOLDER_PATTERN="cbtm\.c4_example\.*ngpu4"
mapfile -t MODEL_FOLDERS < <(find ${SERIALIZATION_DIR}/${NUM_CLUSTERS}_clusters/ -type d -name $FOLDER_PATTERN  -name "*\.cluster*" -printf "%f|%p\n" | sort -t "|" -k1,1 -t "|" -k2,2 | cut -d "|" -f 2)
for folder in "${MODEL_FOLDERS[@]}"; do
    # 检查模型文件夹中是否有任何 consolidated.pt 文件
    if test -f "${folder}/consolidated.pt"; then
        CONSOLIDATED_MODEL_PATHS+="${folder}/consolidated.pt ";
    fi;
done
# 将检查点连接成逗号分隔的字符串的函数
function join { local IFS=","; echo "$*"; }

# 这些模型路径应该按聚类 ID 排序！
JOINED_MODEL_PATHS=$(join ${CONSOLIDATED_MODEL_PATHS[@]})

python -m metaseq_cli.eval_cbtm \
    --data-dir ${DATA_DIR}/${DATASET} \
    --data-subset valid/C4_small \
    --path-to-clusterer ${KMEANS_DIR}/${DATASET}/${NUM_CLUSTERS}/ \
    --model-paths $(join ${CONSOLIDATED_MODEL_PATHS[@]}) \
    --job-dir ${EVAL_DIR} \
    --temperature 0.1 \
    --max-valid-steps 200 \
    --ensemble-type clustering \
    --submitit

您可以在 ${EVAL_DIR} 查看 slurm 作业的日志。

要在本地启动，请删除上述命令中的 --submitit 标志。但请确保您有 $NUM_CLUSTERS 个可见的 GPU！

这将输出困惑度结果到 ${EVAL_DIR}/result.json。

使用相同的命令评估您的密集模型，只需将环境变量更改为 NUM_CLUSTERS=1。

在我们的机器上，本教程中的 8 聚类模型得到的 ppl: 17.86，1 聚类模型得到的 ppl: 18.56。（请注意，这些困惑度仅用于可重现性目的；将步数增加到 10000 以复现我们论文中的运行结果）。

通过稀疏升级训练 MoE 基线

请参阅我们的 fairseq fork 以获取训练稀疏升级 MoE 基线的说明和代码。

开源预训练模型

下载聚类器和嵌入器

要获取论文中的预训练聚类器和嵌入器，您可以运行：

# c4 或 s2orc
export DATASET=c4
# 2, 4, 8, 16, 32, 64, 128
export NUM_CLUSTERS=4
mkdir -p ${KMEANS_DIR}/${DATASET}/${NUM_CLUSTERS}/
wget -O ${KMEANS_DIR}/${DATASET}/${NUM_CLUSTERS}/tfidf.pkl https://dl.fbaipublicfiles.com/cbtm/clusterers/${DATASET}/${NUM_CLUSTERS}/tfidf.pkl
wget -O ${KMEANS_DIR}/${DATASET}/${NUM_CLUSTERS}/kmeans.pkl https://dl.fbaipublicfiles.com/cbtm/clusterers/${DATASET}/${NUM_CLUSTERS}/kmeans.pkl

下载语言模型

要访问我们在论文中训练的语言模型，您可以运行以下命令：

# c4 或 s2orc
export DATASET=c4
# opt1.3b 或 opt6.7b
export MODEL_ARCH=opt1.3b
# 4, 8, 16, 32, 64, 128, 256, 512, 1024
export GPUS_PER_EXPERT=4
# 2, 4, 8, 16, 32, 64, 128
export NUM_CLUSTERS=64
# 从0到$NUM_CLUSTERS-1的任意数
export CLUSTER_NUMBER=28
FOLDER=${DATASET}.${MODEL_ARCH}.${NUM_CLUSTERS}_clusters.cluster${CLUSTER_NUMBER}.ngpu${GPUS_PER_EXPERT}/
mkdir -p ${PRETRAINED_MODELS_DIR}/cbtm/$FOLDER
wget -O ${PRETRAINED_MODELS_DIR}/cbtm/$FOLDER/consolidated.pt https://dl.fbaipublicfiles.com/cbtm/cbtm_models/$DATASET/${MODEL_ARCH}/${NUM_CLUSTERS}_clusters/ngpu${GPUS_PER_EXPERT}/${CLUSTER_NUMBER}/consolidated.pt

要获取n集群c-BTM模型的所有n个专家，可以使用for循环：

# c4 或 s2orc
export DATASET=c4
# opt1.3b 或 opt6.7b
export MODEL_ARCH=opt1.3b
# 4, 8, 16, 32, 64, 128, 256, 512, 1024
export GPUS_PER_EXPERT=4
# 2, 4, 8, 16, 32, 64, 128
export NUM_CLUSTERS=4
for CLUSTER_NUMBER in $(seq 0 $((NUM_CLUSTERS-1))); do 
    FOLDER=${DATASET}.${MODEL_ARCH}.${NUM_CLUSTERS}_clusters.cluster${CLUSTER_NUMBER}.ngpu${GPUS_PER_EXPERT}/
    mkdir -p ${PRETRAINED_MODELS_DIR}/cbtm/$FOLDER
    wget -O ${PRETRAINED_MODELS_DIR}/cbtm/$FOLDER/consolidated.pt https://dl.fbaipublicfiles.com/cbtm/cbtm_models/$DATASET/${MODEL_ARCH}/${NUM_CLUSTERS}_clusters/ngpu${GPUS_PER_EXPERT}/${CLUSTER_NUMBER}/consolidated.pt
done

这里可以下载论文核心结果部分中展示的所有模型。最高预算（168B标记）的1.3B参数模型有${NUM_CLUSTERS} * $GPUS_PER_EXPERT = 1024，结果为(${NUM_CLUSTERS}, $GPUS_PER_EXPERT) = (1, 1024), (2, 512) ... (128, 8)。较小计算预算模型的设置可以通过按比例缩小乘积找到。例如，对于42B标记的1.3B参数模型，我们有${NUM_CLUSTERS} * $GPUS_PER_EXPERT = 256，结果为(${NUM_CLUSTERS}, $GPUS_PER_EXPERT) = (1, 256), (2, 128) ... (64, 4)。

这些模型可以使用上面"集群BTM训练和评估"下给出的命令进行评估。然而，要使用模型并行继续训练这些专家中的任何一个，你需要重新分片它们，查看metaseq库以了解更多关于如何做到这一点的信息。

下载模型时的速率限制详情，请参见此README。

从metaseq转换为Huggingface

我们提供了一个脚本，用于将所有metaseq模型转换为与Huggingface transformers兼容的检查点。

首先，你需要从源代码安装huggingface transformers，并确保使用consolidate_fsdp_shards.sh先合并你的metaseq检查点（见上文）。

然后运行：

PATH_TO_TRANSFORMERS=/path/to/installed/transformers/library/
# 包含consolidated.pt检查点的目录路径
INPUT_DIR=/path/to/metaseq/model/dir
OUTPUT_DIR=my_huggingface_model

bash metaseq/scripts/convert_hf.sh $PATH_TO_TRANSFORMERS $INPUT_DIR $OUTPUT_DIR

这将在$OUTPUT_DIR中输出一个可以在任何huggingface transformers管道中使用的检查点。

相关项目
项目侧边栏1项目侧边栏2
推荐项目
Project Cover

豆包MarsCode

豆包 MarsCode 是一款革命性的编程助手，通过AI技术提供代码补全、单测生成、代码解释和智能问答等功能，支持100+编程语言，与主流编辑器无缝集成，显著提升开发效率和代码质量。

Project Cover

AI写歌

Suno AI是一个革命性的AI音乐创作平台，能在短短30秒内帮助用户创作出一首完整的歌曲。无论是寻找创作灵感还是需要快速制作音乐，Suno AI都是音乐爱好者和专业人士的理想选择。

Project Cover

有言AI

有言平台提供一站式AIGC视频创作解决方案，通过智能技术简化视频制作流程。无论是企业宣传还是个人分享，有言都能帮助用户快速、轻松地制作出专业级别的视频内容。

Project Cover

Kimi

Kimi AI助手提供多语言对话支持，能够阅读和理解用户上传的文件内容，解析网页信息，并结合搜索结果为用户提供详尽的答案。无论是日常咨询还是专业问题，Kimi都能以友好、专业的方式提供帮助。

Project Cover

阿里绘蛙

绘蛙是阿里巴巴集团推出的革命性AI电商营销平台。利用尖端人工智能技术，为商家提供一键生成商品图和营销文案的服务，显著提升内容创作效率和营销效果。适用于淘宝、天猫等电商平台，让商品第一时间被种草。

Project Cover

吐司

探索Tensor.Art平台的独特AI模型，免费访问各种图像生成与AI训练工具，从Stable Diffusion等基础模型开始，轻松实现创新图像生成。体验前沿的AI技术，推动个人和企业的创新发展。

Project Cover

SubCat字幕猫

SubCat字幕猫APP是一款创新的视频播放器，它将改变您观看视频的方式！SubCat结合了先进的人工智能技术，为您提供即时视频字幕翻译，无论是本地视频还是网络流媒体，让您轻松享受各种语言的内容。

Project Cover

美间AI

美间AI创意设计平台，利用前沿AI技术，为设计师和营销人员提供一站式设计解决方案。从智能海报到3D效果图，再到文案生成，美间让创意设计更简单、更高效。

Project Cover

稿定AI

稿定设计 是一个多功能的在线设计和创意平台，提供广泛的设计工具和资源，以满足不同用户的需求。从专业的图形设计师到普通用户，无论是进行图片处理、智能抠图、H5页面制作还是视频剪辑，稿定设计都能提供简单、高效的解决方案。该平台以其用户友好的界面和强大的功能集合，帮助用户轻松实现创意设计。

使用协议隐私政策广告服务
投诉举报邮箱: service@vectorlightyear.com
@2024 懂AI·鲁ICP备2024100362号-6·鲁公网安备37021002001498号