集群-分支-训练-合并 (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管道中使用的检查点。