基于Daft的百万级文本嵌入实现

火山引擎LAS团队最近干了一件挺漂亮的事：他们用Daft框架处理百万级文本Embedding，把GPU利用率压到了接近100%。没错，就是那种“跑满”的状态——不是偶尔99%，而是持续高位运行。这事儿听起来简单，真做起来门道不少。今天这篇就从头拆解，从数据读取到向量入库，把整个流程摆出来聊聊。

在计算密集型任务里，GPU就是核心生产力。利用率不高，等于白花钱还耽误事。所以工程师们绞尽脑汁想让GPU“满负荷运转”，尤其是面对海量数据时，怎么避免GPU空转，怎么把并行能力发挥到极致，往往就是优化的胜负手。

这次我们用Qwen3-Embedding-0.6B模型对百万级文本做向量化，最终做到了GPU几乎全程100%利用率。不仅省时间，更重要的是验证了一套可复用的方法论。下面的内容会用一个具体样例，把关键步骤和调优细节掰开揉碎。另外，在写这篇文章的时候，我们其实还发现了另一条路，能把处理速度再提3倍——不是靠压榨GPU，而是换了一套思路。那个要留到下一篇再讲，今天先把基础打扎实。

Daft任务描述

在Daft的日常开发中，我们注意到一个高频需求：用户想找个好用的工具，把海量文档快速塞进向量数据库，支撑下游检索。于是我们基于Daft搭了一个样例作业，覆盖了从数据读取到入库的完整链路：

• 从S3读取百万级文档数据
• 用spaCy把每个文档切分成句子
• 用开源顶流模型Qwen3-Embedding-0.6B给每句话生成Embedding向量
• 结果写入TurboPuffer向量数据库，供后续业务检索

准备工作

开始之前，先创建一个虚拟环境，virtualenv、uv、conda都行。这里拿uv举例：

uv venv --seed --python 3.11
source .venv/bin/activate

然后装依赖：

pip install "daft[ray]" turbopuffer torch sentence-transformers spacy accelerate transformers
python -m spacy download en_core_web_sm

导入依赖并配置常量

接下来写代码，先引入必要的库，同时设定一些工作负载参数。这些参数可以根据集群规模、数据大小、GPU可用资源灵活调整。这里用en_core_web_sm做句子切分，用intfloat/multilingual-e5-large-instruct生成Embedding向量——当然你也可以换成其他开源模型或商用模型。

import torch
import daft
from daft import col

NUM_GPU_NODES = 8                     # GPU nodes in your cluster
NLP_MODEL_NAME = "en_core_web_sm"     # spaCy model for sentence detection
CHUNKING_PARALLELISM = 8              # Parallel chunking processes per node
EMBEDDING_MODEL_NAME = "Qwen/Qwen3-Embedding-0.6B"  # Text embedding model
ENCODING_DIM = 1024                   # Embedding dimensions
BATCH_SIZE = 512                      # Records per embedding batch
SENTENCE_TRANSFORMER_BATCH_SIZE = 16  # GPU batch size for embeddings

步骤 1：文本切块

给文本生成Embedding之前，通常得先做切分。文档天然有层级结构：文档→章节→段落→句子→单词→字符。不同层级对应不同切分策略：

• 语句级分块：最常用，尤其文档结构不清晰时
• 段落级分块：适合RAG（检索增强生成）场景
• 章节级分块：适合结构清晰的长文档
• 固定长度分块：实现简单，但可能破坏语义边界，适合自定义逻辑

具体怎么选？看业务场景：文档结构未知或内容类型多样，优先用语句级分块；打算做RAG系统、需要跨语句上下文，优先用段落级分块；数据主要是推文、信息、代码等非标准结构，可以考虑自定义分割。

这个样例我们选语句级分块，用spaCy来处理。相比简单按标点分割，spaCy的句子边界检测更健壮，能更好地应对边缘情况。

# Define the return type for chunked text
# Here we'll keep both the chunked text and the chunk ID which
# we'll later use for creating IDs for the sentences
chunked_type = daft.DataType.list(
    daft.DataType.struct({
        "text": daft.DataType.string(),
        "chunk_id": daft.DataType.int32()
    })
)

@daft.udf(
    return_dtype=chunked_type,
    concurrency=NUM_GPU_NODES * CHUNKING_PARALLELISM,
    batch_size=BATCH_SIZE // CHUNKING_PARALLELISM // 2
)
class ChunkingUDF:
    def __init__(self):
        import spacy
        self.nlp = spacy.load(NLP_MODEL_NAME)

    def __call__(self, text_col):
        results = []
        for text in text_col:
            doc = self.nlp(text)
            sentence_texts = [
                {"text": sentence.text, "chunk_id": i}
                for i, sentence in enumerate(doc.sents)
            ]
            results.append(sentence_texts)
        return results

这里用Daft的Class UDF实现了ChunkingUDF。关键点：

• 初始化时加载spaCy模型，每个UDF实例只加载一次
• 批量处理文档数据（text_col），减少调用开销
• 返回切分后的文本列表，每个元素包含文本内容和对应的chunk ID
• 支持分布式运行——示例中启动64个实例并行处理（NUM_GPU_NODES * CHUNKING_PARALLELISM）

步骤 2：GPU加速Embedding向量生成

Embedding质量取决于模型选择。选模型时可以关注这几个维度：

• MTEB Leaderboard：参考排行榜，看不同模型在不同任务上的排名
• 特定任务性能：比如语义搜索、聚类，不同模型擅长的领域不同
• 多语言支持：需不需要处理多语言文本？单语言场景应该关注该语言的专项排行榜

几个主流模型的特点：

• Qwen3-Embedding-0.6B：性价比很高，同级别开源模型里性能最优，本示例用的就是它
• all-MiniLM-L6-v2：Sentence Transformer文档默认模型，教程里常见
• gemini-embedding-001：目前MTEB上的顶级多语言模型，需要Gemini API
• Seed1.6-Embedding：中国MTEB排行榜的顶级模型，需要Volcengine API

选定模型后，改一下EMBEDDING_MODEL_NAME变量就能切换。接下来实现Embedding逻辑：

# Define the return type for embeddings
embedding_type = daft.DataType.embedding(daft.DataType.float32(), ENCODING_DIM)

@daft.udf(
    return_dtype=embedding_type,
    concurrency=NUM_GPU_NODES,
    num_gpus=1,
    batch_size=BATCH_SIZE
)
class EncodingUDF:
    def __init__(self):
        from sentence_transformers import SentenceTransformer
        device = 'cuda' if torch.cuda.is_a vailable() else 'cpu'
        self.model = SentenceTransformer(EMBEDDING_MODEL_NAME, device=device)
        self.model.compile()

    def __call__(self, text_col):
        embeddings = self.model.encode(
            text_col.to_pylist(),
            batch_size=SENTENCE_TRANSFORMER_BATCH_SIZE,
            convert_to_tensor=True,
            torch_dtype=torch.bfloat16,
        )
        return embeddings.cpu().numpy()

和ChunkingUDF类似，EncodingUDF也走自定义UDF模式：

• 初始化时检测GPU资源，有的话就把模型加载到GPU内存
• 用bfloat16精度降低内存占用，同时批量处理文本（SENTENCE_TRANSFORMER_BATCH_SIZE=128）优化GPU利用率
• 最终返回Numpy数组格式的Embedding向量

步骤 3：配置分布式处理

核心逻辑写完后，先在小批量数据上本地测试跑一遍。真正面对百万级数据时，通常得把任务提交到分布式集群。以下配置基于8个g5.2xlarge节点（每个节点含一块A10G GPU）的Ray集群：

# Configure Daft to use Ray to schedule work on different worker nodes
daft.context.set_runner_ray()

# Configure S3 access for reading data
daft.set_planning_config(
    default_io_config=daft.io.IOConfig(
        s3=daft.io.S3Config.from_env()
    )
)

步骤 4：执行完整流程

下面是端到端的完整管道：

(
    daft.read_parquet("s3://desmond-demo/text-embedding-dataset.parquet")
    .with_column("sentences", ChunkingUDF(col("text")))
    .explode("sentences")
    .with_column("text", col("sentences")["text"])
    .with_column("chunk_id", col("sentences")["chunk_id"])
    .exclude("sentences")
    .with_column("embedding", EncodingUDF(col("text")))
    .with_column(
        "id",
        col("url").str.right(50) + "-" + col("chunk_id").cast(daft.DataType.string())
    )
    .select("id", "url", "language", "source", "text", "embedding")
    .write_turbopuffer(
        namespace="desmond-scale-experiment6",
        region="aws-us-west-2",
        id_column="id",
        vector_column="embedding",
        distance_metric="cosine_distance"
    )
)

流程拆解一下：

• 读取数据：从S3加载Parquet文件，高效读入文档内容
• 文本分块：ChunkingUDF对text列切分句子，新增sentences列
• 展开列表：sentences列是list类型，用explode展开后每个元素变成Struct，包含文本和chunk ID
• 提取字段：从sentences中提取text和chunk_id，然后移除sentences列
• 生成嵌入：EncodingUDF为每个句子生成Embedding向量，存到embedding列
• 创建唯一ID：组合url和chunk_id，生成每个向量的唯一标识
• 筛选列：只保留需要的字段
• 写入TurboPuffer：把数据写进向量数据库

在集群上跑起来之后，应该能看到网络I/O、CPU和GPU资源都被充分利用，GPU负载持续高位。还可以通过调整批处理大小和并发参数，进一步优化GPU利用率。

自定义处理逻辑

这个样例可以按业务需求灵活调整：

• 调整批大小：增大SENTENCE_TRANSFORMER_BATCH_SIZE可以提升吞吐量，同时提高GPU内存使用率
• 扩展节点：根据集群规模调整NUM_GPU_NODES和CHUNKING_PARALLELISM
• 更换模型：修改EMBEDDING_MODEL_NAME就能切换Sentence Transformers模型
• 自定义分块：改写ChunkingUDF实现其他分块策略
• 替换向量库：Daft也原生支持写入其他向量数据库，比如LanceDB

性能优化实践

如果想高效处理百万级文本，下面几点值得留意：

• GPU内存：任务运行期间建议配合监控，关注GPU内存使用。如果分配失败或超过模型最大序列长度，就减小SENTENCE_TRANSFORMER_BATCH_SIZE
• 模型加载：UDF初始化时会加载一次模型，首次可能需要从远程下载。如果带宽受限，加载时间会很长。可以提前把模型缓存到本地；容器化场景则提前下载并挂载到容器中
• 模型量化：对精度要求不高的场景，用bfloat16/float16做量化，能降低内存占用，提升吞吐量

下一代计划

通过这个示例，我们已经把GPU利用率推到了接近100%的水平——很多工程师追求的“圣杯”。但Daft的探索远未停止。我们正在试验一条新路径：通过自定义GPU流水线，用vLLM替换Sentence Transformers，让整体处理速度再提升3倍。方案还在打磨中，等完善了会在后续博客里继续分享——怎么突破“峰值利用率”的瓶颈，让吞吐量再上一个台阶。

你是一名 AI 行业编辑，请围绕下面这条热点输出一份资讯解读：
热点：基于Daft的百万级文本嵌入实现要求：
1. 先用一句话解释这条热点在讲什么
2. 再总结它为什么重要
3. 说明会影响哪些 AI 产品或内容方向
4. 最后给出 3 个适合资讯站使用的标题