floret-torch：PyTorch 开源实现，GPU 训练 fastText 风格词向量，比 CPU 快约 20 倍

floret-torch：fastText 风格的词向量，在 GPU 上训练

floret-torch 在 GPU 上训练 fastText 风格的词向量，在 4 核笔记本上以相同质量比 CPU floret 快约 20 倍，且向量可直接加载到 spaCy 中。其背后的赌注是：训练主要是在大表上进行 gather 和 scatter 操作，因此受内存带宽限制，而 GPU 的内存带宽远超 CPU。

fastText 是 Facebook AI Research 的词向量库，仅能在 CPU 线程上训练。floret-torch 的灵感来自亚马逊的 SageMaker BlazingText，后者作为托管 AWS 服务在 GPU 上训练 fastText 风格的子词向量，并覆盖相同的词向量用例：大型语料的嵌入，以及由字符 n-gram 构建的罕见、拼写错误或未见过单词的向量。与 BlazingText 不同，它是开源的，可在你自己的 GPU 上运行。它在 PyTorch 中重新实现了 fastText 的训练循环，采用相同的算法和更新顺序，并训练 fastText 的 floret 变体。

floret 是带有紧凑子词表的 fastText：每个字符 n-gram 被哈希到固定表 bucket 行中的 hashCount 行（布隆过滤器风格），因此任何单词，即使是未见的单词，也能从其 n-gram 获得向量。floret-torch 写出与 CPU floret 相同的 .floret 和 .vec 文件。

当 CPU 训练是慢速环节时使用它：大型语料，或只有两个 CPU 核心但配有 GPU 的 Colab 会话。

结果

text8（1700 万 token 的英文维基百科），cbow，dim=300 bucket=50000 hashCount=2 minn=4 maxn=5 neg=5。 质量是 WordSim-353（WS353）上的 Spearman ρ，这是一个词相似度基准。

在三个 epoch 时，向量与 CPU floret 匹配（0.4740 对 0.4738）。训练比 4 核笔记本快约 20 倍，在 T4 上约每秒 100 万训练对。

940MX 的结果是个例外。它以其击败的 CPU 一半的内存带宽获胜（16 GB/s 对约 34 GB/s 双通道 DDR4）；原因见 the idea。

哈希与 CPU floret 字节级精确一致，因此一个单词在两者中映射到相同的表行。tests/test_parity.py 针对一个实时 floret 模型检查了这一点：MurmurHash3 种子、每个 128 位哈希拆分为四个 32 位键、hashCount 截断、n-gram 提取和桶取模。

安装

需要 Python >=3.9 和 PyTorch（任何构建；CUDA 仅用于 GPU 路径）。torch 故意不作为声明的依赖项，因此 pip 永远不会拉取第二个数 GB 的副本。

pip install git+https://github.com/Fazel94/floret-torch

从克隆安装，复用已安装的 torch：

python -m venv --system-site-packages .venv
.venv/bin/pip install -e '.[bench]'
.venv/bin/python -m pytest tests/test_parity.py -q      # hashing parity gate

添加 weasel 额外项（'.[bench,weasel]'）以用于下面的波斯语管道。

用法

获取 text8，然后训练：

mkdir -p data
curl -sL -o data/text8.zip http://mattmahoney.net/dc/text8.zip
unzip -o -q -d data data/text8.zip

python -m floret_torch.train \
  --input data/text8 --output out/vectors \
  --model cbow --dim 300 --minn 4 --maxn 5 --hashCount 2 --bucket 50000 \
  --epoch 3 --lr 0.05 --batch 8192 --dtype fp32 --device cuda

写出 out/vectors.floret 和 out/vectors.vec。用以下方式加载到 spaCy：

python -m spacy init vectors en out/vectors.floret out/pipeline --mode floret

仅实现了 floret 布局。--mode fasttext（普通 fastText 的独立单词表和 n-gram 表）会引发 NotImplementedError；请使用 CPU fastText 或 floret。

文本分类（fastText 的 supervised 模式，如 BlazingText 中那样）尚未实现；正在进行中。在此之前，向量可作为 spaCy textcat 组件或任何其他分类器的特征。

fp16

--dtype fp16 将两个表存储为 fp16，但以 fp32 计算每个 logit 和梯度。这种拆分是故意的：单次更新约为 1e-5，而权重约为 3e-3，fp16 在该大小下的最小步长（其 ulp）约为 4e-6。以 fp16 存储是安全的；一旦 lr 衰减，以 fp16 累加会停止学习。

预计约 10–15%，不会更多；这个负载受限于表查找，而不是算术。在 T4 上实测：matmul 5.02→0.55ms（9×），但 embedding_bag 前向+反向仅 28.27→24.29ms（1.16×），而 embedding_bag 才是热点路径。在 sm_50 GPU 上，fp16 比 fp32 更慢（没有 tensor cores，且 sm_53 以下没有原生成对 fp16 运算）。

思路

训练词向量主要是 gather 和 scatter，而不是数学运算。每一步从大表中读取几百个分散的行（一个词的子词行、真实词及其负样本的输出行），做几次点积，然后把小的更新加回这些相同的行。在 PyTorch 中，这就是 embedding_bag（gather 和 sum）和 index_add_（scatter-add）。每移动一字节的算术量如此之少，瓶颈是内存带宽，而不是计算。这就是为什么 fp16 更快的算术几乎帮不上忙（见 fp16）。

GPU 的内存带宽远超 CPU：T4 有 320 GB/s，而双通道 DDR4 笔记本约为 34 GB/s，数据中心 GPU 则有数 TB/s。它们还能同时保持数千次内存读取在途，而 CPU 线程会在每次缓存未命中时停顿。这些读取又小又随机，所以第二点也很重要。这很可能就是为什么即使 940MX 的峰值带宽低于笔记本 CPU，也能胜出的原因。

它如何获得速度

没有自定义 CUDA kernel：Triton 和 torch.compile 需要计算能力 7.0（sm_70）或更新的 GPU，我的开发机上的 940MX 是 5.0（sm_50），而且它没有 nvcc。两项改动完成了工作。

按批次词去重。一个词的输入向量是其子词行的和（它的“袋”），而一个词在批次中每次出现都有相同的袋。因此每个袋只 gather 一次并被索引。这是精确的，不是近似：表在一次前向传播内是恒定的。在 text8 上，65536 个上下文 token 坍缩为约 3900 个唯一词。

手写梯度。cbow 和 skip-gram 都用负采样训练（代码中为 SGNS）：将真实输出词与几个随机词进行打分。使用 autograd 时，这在输出侧每批次花费 139ms，外加 33ms 的稠密优化器扫描，因为 autograd 每一步都会分配并清零稠密的 (rows, dim) 梯度缓冲区，然后更新全部 121k 行。SGNS 梯度有闭式解且只触及少数行，因此直接计算并用 index_add_ 进行 scatter，顺序与 fastText 的 Model::update 相同。对于隐藏向量 h（平均后的输入袋）、输出词向量 w_o 和标签 1（真实词）或 0（负样本）：

alpha  = lr * (label - sigmoid(w_o . h))
grad  += alpha * w_o        # uses the PRE-update w_o
w_o   += alpha * h
W_r   += grad               # every row r of the input bag

真实代码在 floret_torch/model.py 中。在完整 text8 上实测效果：235.7s → 178.5s。

两个正确性陷阱

输入行梯度绝不能除以袋大小。fastText 仅在监督模式下才除以。跳过它看起来像 bug，但它是关键支撑：一个 cbow 窗口平均约 240 个子词行，所以除以袋大小 L 会使梯度小约 240 倍，无法启动学习，而 w_o 从零开始。使用 1/L 时，模型会静默地产生看似合理的向量，但得分接近零（WS353 −0.12，每对余弦相似度均为 1.00）。--input-grad normalized 为实验选择该变体；不要用它训练。默认的 fasttext 是正确的。

批次大小是一个正确性旋钮。未除的更新为每次出现添加一个完整梯度，因此一个词在批次中出现 k 次就会移动 k 倍远，过大的批次会直接发散。将 --batch 保持在 4096–16384（默认 8192）。--max-grad-norm（默认 1.0）限制每批次的行移动量，并将发散转化为缓慢学习；如果损失无论如何都变为非有限值，训练会以可操作的消息中止。

布局

floret_torch/     vocab (hashing) · data (binarize) · model (SGNS) ·
                  trainer (batching) · export (.vec/.floret) · train (CLI)
tests/            hashing + word-vector parity against CPU floret
bench/            cpu_baseline · eval (WS353 + neighbours) · sweep
colab/            T4 notebook, plain-command version, bundle script
colab/farsi_weasel/   Persian Wikipedia → floret → spaCy pipeline

Farsi / Weasel 流水线

Weasel 是 spaCy 的工作流运行器：一个 project.yml 列出命名步骤，weasel run 执行它们。colab/farsi_weasel/ 改编自 explosion 的 floret_wiki_oscar_vectors 项目，该项目下载 Wikipedia 以及 OSCAR（一个大型网络爬取语料库），对它们进行分词， 在 CPU 上训练 floret 并打包一个 spaCy 模型。为了在 Colab 上运行，它改变了四件事：

• 向量使用 floret-torch 在 GPU 上训练（train-vectors-gpu）。CPU floret 仍 可作为 train-vectors-cpu 用于比较。
• OSCAR 被舍弃：unshuffled_deduplicated_fa 有访问限制，未认证会失败。 语料库仅为波斯语 Wikipedia，即通过 WikiExtractor --no-templates 处理的 fawiki 转储，并用 cleanup-corpus 在分词完成后删除转储和提取的 JSONL。
• 分词使用 nlp.pipe 在进程内运行，而不是为 16 核 机器配置的多进程池。
• 路径指向 /content，而不是 /scratch，进程数适合约 2 个 vCPU。

weasel run all colab/farsi_weasel   # Wikipedia → GPU vectors → spaCy package
weasel run cpu colab/farsi_weasel   # CPU floret instead; no packaging

致谢

算法和文件格式遵循 explosion/floret 和 fastText。GPU 方法受 Amazon SageMaker BlazingText 启发。Weasel 项目源自 explosion/projects pipelines/floret_wiki_oscar_vectors。

引用

如果您在工作中使用 floret-torch，请引用：

@software{fazeli2026florettorch,
  author  = {Fazeli, Mohammad},
  title   = {floret-torch: fastText-style word vectors, trained on the GPU},
  year    = {2026},
  version = {0.1.0},
  url     = {https://github.com/Fazel94/floret-torch},
  note    = {Contact: kiyarash@nlogn.ir}
}

许可证

MIT，参见 LICENSE。floret、fastText 和 explosion/projects 也是 MIT。