大概念模型

句子表示空间中的语言建模

[博客] [论文]

本仓库提供了大概念模型（LCM）的官方实现与实验。

LCM 在一种显式的高层语义表示上运行，我们将其称为“概念”。概念具有语言和模态无关性，代表更高层次的思想。在本工作中，一个概念对应于一句话，我们使用 SONAR 嵌入空间，该空间支持文本中的 200 种语言以及语音中的 57 种语言。受支持的语言列表请参见 此处。

方法

LCM 是一个在概念空间中进行自回归句子预测的序列到序列模型。我们探索了多种方法：

• 均方误差回归（此代码中的 base_lcm）。
• 基于扩散的生成变体（本次发布包含 two_tower_diffusion_lcm）。
• 在量化 SONAR 空间中运行的模型（即将推出）。

这些探索均采用 16 亿参数的模型，并以约 1.3 万亿个 token 的数据进行训练。本仓库包含了重现 16 亿参数 MSE LCM 和双塔扩散 LCM 训练及微调的脚本。具体说明请参见 下方。

安装

使用 UV

LCM 仓库依赖于 fairseq2。如果你的系统已安装 uv，可以通过运行以下命令来创建一个包含所有必要包的虚拟环境：

uv sync --extra cpu --extra eval --extra data

你也可以使用 uv run 来在正确的环境中运行演示命令。

请注意，我们仅提供 CPU 依赖项的要求；若需使用 GPU 支持，则需要选择适合你系统的 PyTorch 和 fairseq2 版本。例如，对于 PyTorch 2.5.1 和 CUDA 1.21，你可以执行如下操作：

uv pip install torch==2.5.1 --extra-index-url https://download.pytorch.org/whl/cu121 --upgrade
uv pip install fairseq2==v0.3.0rc1 --pre --extra-index-url  https://fair.pkg.atmeta.com/fairseq2/whl/rc/pt2.5.1/cu121 --upgrade

请参考 fairseq2 版本 以了解可能的版本。需要注意的是，LCM 目前依赖于 fairseq2 0.3.0 rc1 发布候选版本。

使用 pip

使用 pip 安装时，命令非常相似，但你需要自行管理环境，并确保先手动安装 fairseq2。例如，对于 CPU 安装：

pip install --upgrade pip
pip install fairseq2==v0.3.0rc1 --pre --extra-index-url  https://fair.pkg.atmeta.com/fairseq2/whl/rc/pt2.5.1/cpu
pip install -e ".[data,eval]"

如果 fairseq2 没有为你使用的机器提供构建版本，请查阅该项目的 README 文件，以在当地编译构建。

使用

[!NOTE] 如果使用 uv，请在所有命令前加上 uv run，以使用默认在 .venv 中创建的环境，例如： uv run torchrun --standalone。 或者，你也可以通过 source .venv/bin/activate 一次性激活环境。

数据准备

LCM 可以使用按句子分割并用 SONAR 嵌入的文本数据进行训练和评估。我们提供了一个可用于准备此类训练数据的示例处理流程，你可以通过以下命令运行它：

 uv run --extra data scripts/prepare_wikipedia.py /output/dir/for/the/data

该流程展示了如何从 Hugging Face 获取数据，并使用 SONAR 和 SaT 进行处理。请查看文件以获取更多关于处理自有数据的详细信息。虽然该脚本提供了从 Hugging Face 拉取数据的示例，但我们还提供了 API 来处理 jsonl、parquet 和 CSV 文件。

数据卡片

下文所述的训练器依赖于用于配置数据集的数据卡片。这些数据卡片是 YAML 文件，其中包含指向数据文件（本地或 S3 上）的指针以及有关其模式的信息。我们在 lcm/datacards/datacards.yaml 中提供了一些示例数据卡片。当你处理完一些数据后，可以更新数据卡片中的路径。

拟合归一化器

要为给定的加权数据集混合物拟合一个新的嵌入空间归一化器，可以使用以下命令：

python scripts/fit_embedding_normalizer.py --ds dataset1:4 dataset2:1 dataset3:10 --save_path "path/to/new/normalizer.pt" --max_nb_samples 1000000

其中，dataset1、dataset2、dataset3 是如上所示数据卡片中声明的数据集名称，而 (4, 1, 10) 则是它们各自的相对权重。 生成的归一化器随后可以按照 lcm/cards/sonar_normalizer.yaml 中所示的方式声明为模型，并在所有模型训练配置中引用。

预训练模型

基础 MSE LCM

要训练一个 MSE LCM，我们可以使用以下命令之一：

选项 1. 使用 SLURM 和 submitit 通过 stopes 的启动器进行训练：

python -m lcm.train \
    +pretrain=mse \
    ++trainer.output_dir="checkpoints/mse_lcm" \
    ++trainer.experiment_name=training_mse_lcm \

使用此命令，我们将提交一个名为 training_mse_lcm 的 SLURM 作业，并按照配方的要求执行，在本例中为：

requirements:
  nodes: 4
  tasks_per_node: 8
  gpus_per_node: 8
  cpus_per_task: 32
  mem_gb: 0
  timeout_min: 10000

您可以通过以下方式覆盖作业的要求，例如超时限制和启动器的 SLURM 分区：

python -m lcm.train \
    +pretrain=mse \
    ++trainer.output_dir="checkpoints/mse_lcm" \
    ++trainer.experiment_name=training_mse_lcm \
    ++trainer.requirements.timeout_min=100 \
    ++trainer.requirements.cpus_per_task=8 \
    ++launcher.partition=$partition_name

选项 2. 使用 torchrun 在本地进行训练（例如仅使用 2 张 GPU），并采用较小的批大小（覆盖 ++trainer.data_loading_config.max_tokens=1000）：

CUDA_VISIBLE_DEVICES=0,1 torchrun --standalone --nnodes=1 --nproc-per-node=2 \
    -m lcm.train launcher=standalone \
    +pretrain=mse \
    ++trainer.data_loading_config.max_tokens=1000 \
    ++trainer.output_dir="checkpoints/mse_lcm" \
    +trainer.use_submitit=false \

[!重要] 由于我们更改了配方所需的 GPU 数量，因此这将无法复现论文中的实验设置。

检查点目录 checkpoints/mse_lcm 的结构如下：

.
├── checkpoints
│   ├── step_2000
│   ├── ...
│   └── step_250000
├── config_logs
├── executor_logs
├── model_card.yaml
├── tb   # tensorboard 日志
└── wandb  # W&B 日志

请注意，除非安装了 wandb，否则将跳过 W&B 日志记录。 您可以使用 uv pip install wandb 来安装 wandb。 可以通过覆盖配方中的 Hydra 配置值来更改 W&B 参数：

++trainer.wandb_project=$project_name
++trainer.wandb_run_name=$run_name

双塔扩散 LCM

与基础 MSE LCM 类似，我们也可以按照 ./recipes/train/pretrain/two_tower.yaml 中的配方提交训练作业：

python -m lcm.train \
    +pretrain=two_tower \
    ++trainer.output_dir="checkpoints/two_tower_lcm" \
    ++trainer.experiment_name=training_two_tower_lcm \

[!提示] 要了解训练配方中的不同组成部分，请查看 此 README。

微调模型

要对先前预训练的双塔扩散 LCM 进行监督数据上的微调，请按照以下步骤操作：

步骤 1. 将预训练的检查点注册为 fairseq2 资产。

您可以使用带有卡片 checkpoints/two_tower_lcm/model_card.yaml 的最终检查点，或任何经过特定训练步数后的检查点，例如 checkpoints/two_tower_lcm/checkpoints/step_2000/model_card.yaml。要注册选定的检查点，将自动生成的 YAML 文件复制到 ./lcm/cards/mycards.yaml，并重命名模型以替换默认的 on_the_fly_lcm。./lcm/cards/mycards.yaml 将如下所示：

__source__: inproc
 checkpoint: file://path_to/large_concept_model/checkpoints/two_tower_lcm/checkpoints/step_2000/model.pt
 model_arch: two_tower_diffusion_lcm_1_6B
 model_family: two_tower_diffusion_lcm
 name: my_pretrained_two_tower

有关如何管理 fairseq2 资产的更多信息，请参阅 文档。

步骤 2. 启动指向要微调模型的微调作业，在本例中为 my_pretrained_two_tower：

CUDA_VISIBLE_DEVICES=0,1 torchrun --standalone --nnodes=1 --nproc-per-node=2 \
    -m lcm.train launcher=standalone \
    +finetune=two_tower \
    ++trainer.output_dir="checkpoints/finetune_two_tower_lcm" \
    ++trainer.data_loading_config.max_tokens=1000 \
    +trainer.use_submitit=false \
    ++trainer.model_config_or_name=my_pretrained_two_tower

或者

python -m lcm.train \
    +finetune=two_tower \
    ++trainer.output_dir="checkpoints/finetune_two_tower_lcm" \
    ++trainer.experiment_name=finetune_two_tower_lcm \
    ++trainer.model_config_or_name=my_pretrained_two_tower

同样，要微调一个 MSE LCM，可以按照相同的预训练检查点注册说明，并使用相应的配方（./recipes/train/finetune/mse.yaml）提交微调作业：

python -m lcm.train \
    +finetune=mse \
    ++trainer.output_dir="checkpoints/finetune_mse_lcm" \
    ++trainer.experiment_name=finetune_mse_lcm \
    ++trainer.model_config_or_name=my_pretrained_mse_lcm

评估模型

[!NOTE] 对于高级评估（例如不同任务的基准测试、与大语言模型的结果比较等），请查看评估文档。

步骤 0. 下载用于评估 ROUGE 所需的 NLTK 数据：

python -m nltk.downloader punkt_tab

步骤 1. 可以通过指向模型的 model_card YAML 文件，或者在将其注册为 fairseq2 资产之后（与我们注册 my_pretrained_two_tower 的方式相同），生成并评分模型的输出：

model_card=./checkpoints/finetune_two_tower_lcm/checkpoints/step_1000/model_card.yaml
OUTPUT_DIR=evaluation_outputs/two_tower

torchrun --standalone --nnodes=1 --nproc-per-node=1 -m lcm.evaluation  \
  --predictor two_tower_diffusion_lcm  \
  --show_progress true \
  --data_loading.max_samples 100 \
  --model_card ${model_card} \
  --launcher standalone \
  --dataset.source_suffix_text '[MODEL]:' \
  --tasks finetuning_data_lcm.validation \
   --task_args '{"max_gen_len": 10, "eos_config": {"text": "End of text."}}' \
  --data_loading.batch_size 4  --generator_batch_size 4 \
  --dump_dir ${OUTPUT_DIR} \
  --inference_timesteps 40 \
  --initial_noise_scale 0.6 \
  --guidance_scale 3 \
  --guidance_rescale 0.7

其中，在示例中我们仅评估 100 个样本（--data_loading.max_samples 100），并将模型输出长度限制为 10 句话（--task_args '{"max_gen_len": 10}'）。

输出将被转储到 ./evaluation_outputs/two_tower 目录下，其结构如下：

.
├── metadata.jsonl
├── metrics.eval.jsonl
├── raw_results
├── results
└── tb

其中 metrics.eval.jsonl 包含语料级别的评分。

要评估 MSE LCM，我们使用相应的预测器 (base_lcm) 并按以下方式评估：

model_card=./checkpoints/finetune_mse_lcm/checkpoints/step_1000/model_card.yaml
OUTPUT_DIR=evaluation_outputs/mse_lcm

torchrun --standalone --nnodes=1 --nproc-per-node=1 -m lcm.evaluation  \
  --predictor base_lcm --sample_latent_variable False \
  --show_progress true \
  --data_loading.max_samples 100 \
  --model_card ${model_card} \
  --launcher standalone \
  --dataset.source_suffix_text '[MODEL]:' \
  --tasks finetuning_data_lcm.validation \
   --task_args '{"max_gen_len": 10, "eos_config": {"text": "End of text."}}' \
  --data_loading.batch_size 4  --generator_batch_size 4 \
  --dump_dir ${OUTPUT_DIR} \

请注意，在此示例中，我们仅展示了如何在相同的微调数据集（验证集）上评估 LCM。若要在下游任务中进行评估，并与大语言模型的结果进行比较，请参阅评估文档。

贡献

有关如何参与贡献，请参阅CONTRIBUTING文件。

引用

如果您使用本代码库，请引用以下内容：

@article{lcm2024,
  author = {{LCM 团队}, Lo\"{i}c Barrault, Paul-Ambroise Duquenne, Maha Elbayad, Artyom Kozhevnikov, Belen Alastruey, Pierre Andrews, Mariano Coria, Guillaume Couairon, Marta R. Costa-juss\`{a}, David Dale, Hady Elsahar, Kevin Heffernan, Jo\~{a}o Maria Janeiro, Tuan Tran, Christophe Ropers, Eduardo Sánchez, Robin San Roman, Alexandre Mourachko, Safiyyah Saleem, Holger Schwenk},
  title = {{大型概念模型}：基于句子表示空间的语言建模},
  publisher = {arXiv},
  year = {2024},
  url = {https://arxiv.org/abs/2412.08821},
}

许可证

本代码以 MIT 许可证发布（详见LICENSE）。

Large Concept Model (LCM) 快速上手指南

Large Concept Models (LCM) 是 Meta AI 开源的一种新型语言模型架构。它不在传统的 token 空间操作，而是在显式的“概念”（Concept）空间中进行建模。在本实现中，一个概念对应一个句子，利用 SONAR 嵌入空间，支持多达 200 种语言的文本和 57 种语言的语音处理。

环境准备

系统要求

• 操作系统: Linux (推荐 Ubuntu 20.04+) 或 macOS
• Python: 3.10 或更高版本
• GPU (可选): 如需加速训练或推理，需安装 CUDA 驱动。本仓库默认提供 CPU 依赖配置，GPU 用户需手动指定 PyTorch 和 fairseq2 的 CUDA 版本。
• 磁盘空间: 建议预留至少 50GB 空间用于存放模型权重、数据集及中间缓存。

前置依赖

• 包管理工具: 推荐使用 uv (极速 Python 包管理器)，也可使用 pip。
• 核心框架: 项目强依赖于 fairseq2 (当前版本需 v0.3.0rc1) 和 torch。

注意：国内开发者若遇到网络问题，建议在配置 pip/uv 源时使用清华或阿里镜像（见下文安装步骤中的注释）。

---

安装步骤

方式一：使用 UV（推荐）

uv 能自动管理虚拟环境并解析依赖，速度极快。

1. 安装 uv (如果尚未安装):

   curl -LsSf https://astral.sh/uv/install.sh | sh
   

2. 克隆仓库并同步环境:

   git clone https://github.com/facebookresearch/large_concept_model.git
   cd large_concept_model
   
   # 安装 CPU 版本依赖（包含数据处理和评估工具）
   # 国内用户可添加 --index-url 指定镜像源，例如：--index-url https://pypi.tuna.tsinghua.edu.cn/simple
   uv sync --extra cpu --extra eval --extra data
   

3. 配置 GPU 支持 (可选): 如果需要 GPU 加速，需手动升级 torch 和 fairseq2 到对应的 CUDA 版本。以下以 CUDA 12.1 为例：

   # 替换为适合你系统的 CUDA 版本索引地址
   uv pip install torch==2.5.1 --extra-index-url https://download.pytorch.org/whl/cu121 --upgrade
   uv pip install fairseq2==v0.3.0rc1 --pre --extra-index-url https://fair.pkg.atmeta.com/fairseq2/whl/rc/pt2.5.1/cu121 --upgrade
   

方式二：使用 Pip

如果你习惯使用原生 pip，需自行管理虚拟环境。

python -m venv .venv
source .venv/bin/activate

pip install --upgrade pip
# 安装 fairseq2 (CPU 版本示例)
pip install fairseq2==v0.3.0rc1 --pre --extra-index-url https://fair.pkg.atmeta.com/fairseq2/whl/rc/pt2.5.1/cpu
# 安装 LCM 及其额外依赖
pip install -e ".[data,eval]"

---

基本使用

以下流程展示如何准备数据、预训练模型以及进行简单的推理评估。所有命令若在 uv 环境下运行，请在命令前加 uv run。

1. 准备数据

LCM 需要将被分割为句子的文本数据转换为 SONAR 嵌入向量。以下脚本演示了如何从 HuggingFace 拉取维基百科数据并进行处理。

# 将处理后的数据保存到 /output/dir/for/the/data
uv run --extra data scripts/prepare_wikipedia.py /output/dir/for/the/data

提示：处理完成后，你需要更新 lcm/datacards/datacards.yaml 文件，将路径指向你新生成的数据目录。

2. 预训练模型 (Base MSE LCM)

你可以选择使用 SLURM 集群或在本地单机多卡上进行训练。

本地单机训练示例 (使用 2 张 GPU): 为了快速验证，我们减小 batch size 并在本地运行：

CUDA_VISIBLE_DEVICES=0,1 torchrun --standalone --nnodes=1 --nproc-per-node=2 \
    -m lcm.train launcher=standalone \
    +pretrain=mse \
    ++trainer.data_loading_config.max_tokens=1000 \
    ++trainer.output_dir="checkpoints/mse_lcm" \
    +trainer.use_submitit=false

训练完成后，检查点将保存在 checkpoints/mse_lcm 目录下。

3. 微调模型 (Finetuning)

在预训练基础上，可以使用监督数据进行微调。首先需要注册预训练模型卡片。

步骤 1: 注册模型 复制生成的 model_card.yaml 到 ./lcm/cards/mycards.yaml，修改内容如下（假设使用 two-tower 模型的某个检查点）：

__source__: inproc
checkpoint: file://<绝对路径>/large_concept_model/checkpoints/two_tower_lcm/checkpoints/step_2000/model.pt
model_arch: two_tower_diffusion_lcm_1_6B
model_family: two_tower_diffusion_lcm
name: my_pretrained_two_tower

步骤 2: 启动微调

CUDA_VISIBLE_DEVICES=0,1 torchrun --standalone --nnodes=1 --nproc-per-node=2 \
    -m lcm.train launcher=standalone \
    +finetune=two_tower \
    ++trainer.output_dir="checkpoints/finetune_two_tower_lcm" \
    ++trainer.data_loading_config.max_tokens=1000 \
    +trainer.use_submitit=false \
    ++trainer.model_config_or_name=my_pretrained_two_tower

4. 模型评估与推理

在评估前，需下载 NLTK 数据以支持 ROUGE 指标计算：

python -m nltk.downloader punkt_tab

运行评估脚本生成结果并计算指标：

model_card=./checkpoints/finetune_two_tower_lcm/checkpoints/step_1000/model_card.yaml
OUTPUT_DIR=evaluation_outputs/two_tower

torchrun --standalone --nnodes=1 --nproc-per-node=1 -m lcm.evaluation \
  --predictor two_tower_diffusion_lcm \
  --show_progress true \
  --data_loading.max_samples 100 \
  --model_card ${model_card} \
  --launcher standalone \
  --dataset.source_suffix_text '[MODEL]:' \
  --tasks finetuning_data_lcm.validation \
  --task_args '{"max_gen_len": 10, "eos_config": {"text": "End of text."}}' \
  --data_loading.batch_size 4 --generator_batch_size 4 \
  --dump_dir ${OUTPUT_DIR} \
  --inference_timesteps 40 \
  --initial_noise_scale 0.6 \
  --guidance_scale 3 \
  --guidance_rescale 0.7

评估结果（包括指标分数和生成文本）将输出至 evaluation_outputs/two_tower 目录。