对比学习工具包

contrastors 是一个对比学习工具包，使研究人员和工程师能够高效地训练和评估对比模型。

特性

• 基于 Flash Attention 构建，实现快速高效的训练
• 支持多 GPU 训练
• GradCache 支持，在内存受限的环境中使用大批次进行训练
• Huggingface 支持，方便加载常用模型（Pythia/GPTNeoX、BERT 等）
• 掩码语言模型（MLM）预训练
• 套娃表示学习，支持灵活的嵌入尺寸
• CLIP 和 LiT 风格的对比学习
• 支持加载流行的 ViT 模型（例如 timm）

研究成果

• Nomic Embed：训练可复现的长上下文文本嵌入模型 由 Zach Nussbaum、Jack Morris、Andriy Mulyar 和 Brandon Duderstadt 共同撰写
• Nomic Embed Vision：扩展潜在空间 由 Zach Nussbaum、Brandon Duderstadt 和 Andriy Mulyar 共同撰写
• 训练稀疏专家混合文本嵌入模型 由 Zach Nussbaum 和 Brandon Duderstadt 共同撰写

快速入门与环境要求

contrastors 库依赖于 Flash Attention 仓库中的自定义内核。要设置您的环境，请按照以下步骤操作。

请确保您已安装 CUDA 11.8 或更高版本。您可以通过运行 nvcc --version 来检查，或者如果您已经安装了 PyTorch，可以运行 python -c "import torch; print(torch.version.cuda)"。

创建并激活 Python 虚拟环境：

python3 -m venv env
source env/bin/activate

安装 PyTorch。请参阅 PyTorch 文档，获取适用于您系统的具体安装说明（例如，截至 2023 年 12 月 12 日，默认支持的 CUDA 版本为 12.1）。

pip3 install torch torchvision torchaudio

安装 wheel、packaging 和 ninja，以加速 Flash Attention 的构建过程：

pip install wheel packaging ninja setuptools

安装 Flash Attention 及其自定义内核：

pip install --no-cache-dir flash-attn --no-build-isolation git+https://github.com/HazyResearch/flash-attention.git#subdirectory=csrc/rotary git+https://github.com/HazyResearch/flash-attention.git#subdirectory=csrc/layer_norm git+https://github.com/HazyResearch/flash-attention.git#subdirectory=csrc/fused_dense_lib git+https://github.com/HazyResearch/flash-attention.git#subdirectory=csrc/xentropy

安装其余依赖项及该包：

pip install -e .

数据访问

我们通过 nomic 包提供对 nomic-embed-text-v1 数据集的访问权限。要访问数据，您需要先在 atlas.nomic.ai 上注册账号并登录 nomic 包。首先下载 nomic Python 客户端，并运行以下命令：

pip install nomic
nomic login # 按照提示登录
python -c "from nomic import atlas; print(atlas._get_datastream_credentials(name='contrastors'))"

这将打印出您的访问密钥。您可以使用 aws configure 或设置 AWS_ACCESS_KEY_ID 和 AWS_SECRET_ACCESS_KEY 环境变量来配置这些密钥。

如果您尚未安装 AWS CLI，可以从 这里 下载。

要验证您的访问权限，可以运行以下命令列出存储桶内容：

aws s3 ls --endpoint-url=https://9fa58365a1a3d032127970d0bd9a1290.r2.cloudflarestorage.com/ s3://contrastive
aws s3 ls --endpoint-url=https://9fa58365a1a3d032127970d0bd9a1290.r2.cloudflarestorage.com/ s3://contrastive-index-filtered

您应该能够看到存储桶中的内容并下载数据。

如果您打算使用我们的数据和 contrastors 仓库进行训练，还需要为 Cloudflare R2 设置 fsspec 支持。为此，创建文件 ~/.config/fsspec/s3.json，内容如下：

{
  "s3": {
    "client_kwargs": {
      "endpoint_url": "https://9fa58365a1a3d032127970d0bd9a1290.r2.cloudflarestorage.com/",
      "aws_access_key_id": <ACCESS_KEY_ID>,
      "aws_secret_access_key": <SECRET_KEY_ID>
    }
  }
}

Nomic 数据格式

我们的文本数据存储在 gzip 压缩的 jsonl 文件中，同时附带 counts.json 文件和 offsets.json.gz 文件。

counts.json 文件是一个字典，将文件名映射到该文件中的样本数量。offsets.json.gz 文件也是一个字典，将文件名映射到另一个字典，其中每个键是样本的索引，值则是该样本在文件中开始和结束的字节偏移量。这样做是为了支持从 R2 流式传输数据，尤其是在数据量超过缓冲区大小时。

以下是一个小型的数据集配置示例：

datasets:
  - name: "paq"
    bucket: "s3://contrastive-index-filtered/paq_full/shard-{00000..00538}.jsonl.gz"
    query_prefix: "search_query"
    document_prefix: "search_document"
    objective: 
        type: "paired"
        columns: ["query", "document"]

objective 定义是成对还是三元组目标。在两种情况下，columns 字段都定义了每个样本所使用的列。

训练 nomic-embed-text-v1

掩码语言模型预训练

要从头开始训练您自己的 BERT 模型（包含所有优化），请运行：

cd src/contrastors
deepspeed --num_gpus=8 train.py --config=configs/train/mlm.yaml --deepspeed_config=configs/deepspeed/ds_config.json --dtype=bf16

对比预训练与微调

要启动实验，请运行：

cd src/contrastors
torchrun --nproc-per-node=8 train.py --config=configs/train/contrastive_pretrain.yaml --dtype=bf16

这将基于全部约 2 亿个样本训练一个 BERT 模型。如需更改数据集，可以修改 data_args.input_shards。

要微调 nomic-bert-embed-v1-unsupervised，请将配置更新为 configs/train/contrastive_finetune.yaml。

生成您自己的数据

要为流程中的任何步骤生成您自己的数据，可以使用 scripts/text 中提供的脚本。

更多信息请参阅 scripts/text 中的 README。

训练 nomic-embed-vision-v1.5

要对齐视觉模型，您需要整理一个大型的图像-文本数据集。更多详细信息请参见 这里。

要将 nomic-embed-vision-v1.5 与 nomic-embed-text-v1.5 对齐，可以运行以下命令：

deepspeed  train.py --deepspeed_config=configs/deepspeed/image_text.json --config=configs/train/nomic_embed_vision_v1.5.yaml --dtype=bf16

预训练模型

我们在以下位置提供了 Nomic Embed 的预训练模型：

• nomic-embed-text-v2-moe
• nomic-embed-text-v2-moe-unsupervised
• nomic-embed-text-v1
• nomic-embed-vision-v1
• nomic-embed-text-v1.5
• nomic-embed-vision-v1.5
• nomic-embed-text-v1-ablated
• nomic-embed-text-v1-unsupervised
• nomic-bert-2048
• nomic-xlm-2048

加入 Nomic 社区

• Nomic: https://nomic.ai
• Discord: https://discord.gg/myY5YDR8z8
• Twitter: https://twitter.com/nomic_ai

许可证

本代码采用 Apache 2.0 许可证授权。各模型的具体许可证信息请参阅相应的模型卡片。

致谢

我们感谢 Tri Dao 在 Flash Attention 和自定义内核方面的工作，正是这些技术使本项目得以实现；感谢 OpenCLIP 团队提供的优秀仓库，我们的许多工作都基于此；同时也感谢 Hugging Face 团队在 Transformers 库方面的杰出贡献。

引用

如果您认为该模型、数据集或训练代码有用，请引用我们的工作：

@misc{nussbaum2024nomic,
      title={Nomic Embed: 训练可复现的长上下文文本嵌入模型}, 
      author={Zach Nussbaum、John X. Morris、Brandon Duderstadt 和 Andriy Mulyar},
      year={2024},
      eprint={2402.01613},
      archivePrefix={arXiv},
      primaryClass={cs.CL}
}
@misc{nussbaum2024nomicembedvisionexpanding,
      title={Nomic Embed Vision: 扩展潜在空间}, 
      author={Zach Nussbaum、Brandon Duderstadt 和 Andriy Mulyar},
      year={2024},
      eprint={2406.18587},
      archivePrefix={arXiv},
      primaryClass={cs.CV},
      url={https://arxiv.org/abs/2406.18587}, 
}
@misc{nussbaum2025trainingsparsemixtureexperts,
      title={训练稀疏专家混合文本嵌入模型}, 
      author={Zach Nussbaum 和 Brandon Duderstadt},
      year={2025},
      eprint={2502.07972},
      archivePrefix={arXiv},
      primaryClass={cs.CL},
      url={https://arxiv.org/abs/2502.07972}, 
}

Contrastors 快速上手指南

contrastors 是一个基于 Flash Attention 构建的高效对比学习工具包，专为研究人员和工程师设计，支持大规模文本和视觉嵌入模型的训练与评估。

环境准备

在开始之前，请确保您的系统满足以下要求：

• 操作系统: Linux (推荐)
• GPU: NVIDIA GPU (支持多卡训练)
• CUDA: 版本需为 11.8 或更高
  • 检查命令：nvcc --version 或 python -c "import torch; print(torch.version.cuda)"
• Python: 建议 Python 3.8+
• 编译器: 需安装 ninja 以加速编译过程

注意：由于该库依赖自定义 CUDA 内核，Windows 用户可能需要通过 WSL2 运行或自行解决编译环境问题。

安装步骤

请严格按照以下顺序执行命令，以确保依赖项正确编译和链接。

1. 创建并激活虚拟环境

python3 -m venv env
source env/bin/activate

2. 安装 PyTorch

根据您的 CUDA 版本安装对应的 PyTorch。以下为通用安装命令（具体版本请参考 PyTorch 官网）：

pip3 install torch torchvision torchaudio

3. 安装构建工具

安装 wheel, packaging, ninja 等工具以加速 Flash Attention 的编译：

pip install wheel packaging ninja setuptools

4. 安装 Flash Attention 及自定义内核

这是最关键的一步，需要安装主库以及特定的子模块（rotary, layer_norm 等）：

pip install --no-cache-dir flash-attn --no-build-isolation git+https://github.com/HazyResearch/flash-attention.git#subdirectory=csrc/rotary git+https://github.com/HazyResearch/flash-attention.git#subdirectory=csrc/layer_norm git+https://github.com/HazyResearch/flash-attention.git#subdirectory=csrc/fused_dense_lib git+https://github.com/HazyResearch/flash-attention.git#subdirectory=csrc/xentropy

5. 安装 Contrastors

克隆仓库并安装本地包：

git clone https://github.com/nomic-ai/contrastors.git
cd contrastors
pip install -e .

基本使用

数据访问配置（可选）

如果您需要使用 Nomic 提供的预处理数据集（如 nomic-embed-text-v1），需配置 AWS S3 兼容凭证以访问 Cloudflare R2 存储桶。

1. 安装 nomic 客户端并登录获取密钥：

   pip install nomic
   nomic login
   python -c "from nomic import atlas; print(atlas._get_datastream_credentials(name='contrastors'))"
   

2. 配置 fsspec：创建文件 ~/.config/fsspec/s3.json 并填入上一步获取的密钥：

   {
     "s3": {
       "client_kwargs": {
         "endpoint_url": "https://9fa58365a1a3d032127970d0bd9a1290.r2.cloudflarestorage.com/",
         "aws_access_key_id": "<YOUR_ACCESS_KEY_ID>",
         "aws_secret_access_key": "<YOUR_SECRET_ACCESS_KEY>"
       }
     }
   }
   

训练示例

1. 掩码语言建模 (MLM) 预训练

从头训练一个优化的 BERT 模型（示例使用 8 张 GPU 和 bf16 精度）：

cd src/contrastors
deepspeed --num_gpus=8 train.py --config=configs/train/mlm.yaml --deepspeed_config=configs/deepspeed/ds_config.json --dtype=bf16

2. 对比学习预训练

启动对比学习实验（支持 GradCache 以大批次训练）：

cd src/contrastors
torchrun --nproc-per-node=8 train.py --config=configs/train/contrastive_pretrain.yaml --dtype=bf16

• 修改数据集：编辑配置文件中的 data_args.input_shards 字段。
• 微调模型：将 --config 参数更改为 configs/train/contrastive_finetune.yaml 即可进行微调。

3. 视觉 - 文本对齐训练

训练视觉嵌入模型（如 nomic-embed-vision-v1.5）：

deepspeed train.py --deepspeed_config=configs/deepspeed/image_text.json --config=configs/train/nomic_embed_vision_v1.5.yaml --dtype=bf16

加载预训练模型

您可以直接从 Hugging Face 加载已训练好的模型进行推理或进一步微调：

• 文本模型: nomic-ai/nomic-embed-text-v1.5, nomic-ai/nomic-embed-text-v2-moe
• 视觉模型: nomic-ai/nomic-embed-vision-v1.5
• 基础模型: nomic-ai/nomic-bert-2048

使用 Hugging Face transformers 库即可正常加载这些模型。