聊聊我在小红书搞推荐算法工具链的一些踩坑经历

早上八点的成都，雾气还没散干净，我已经坐在工位上泡好了第三杯美式。说实话，作为一个在小红书干了两年推荐算法的工程师，每天的生活节奏其实挺舒服的——至少表面上看是这样。但只有我们自己知道，推荐系统这个东西，就像一台永远停不下来的跑步机，你稍一松懈，线上指标就给你掉下来看。

今天想跟大家聊聊我这一年多在工具链建设上的一些实践和踩坑经历。为啥要聊这个？因为上周五晚上，我又被一个线上问题搞到了十一点半。搞完之后我坐在空荡荡的办公室里，突然觉得，有些坑真的不该再踩第二次了。

事情的起因

事情要从去年Q3说起。当时我们团队接了个任务：要搭一套新的模型实验平台，支撑算法同学快速迭代推荐模型。听起来挺高大上的对吧？但实际上，我们当时的现状是这样的——

模型训练在A集群，特征存储在B集群，模型服务部署在C集群，中间还夹杂着各种离线数据处理任务。每次算法同学想跑个实验，得手动在好几个系统之间倒腾数据、配置参数、提交任务。整个流程走下来，一个实验从想法到上线，最快也得两三天。

"这也太慢了吧？"新来的算法同学如是说。

确实太慢了。而且更可怕的是，这种手工操作的方式，出错的概率极高。我记得特别清楚，有一次一个同学把特征配置的版本号写错了，导致线上模型用了一套完全不对的特征，DAU直接掉了0.3%。那天下午，整个团队的气氛可以用"凝重"来形容。

所以，搞一套靠谱的工具链，把整个流程自动化、标准化，就成了我们的当务之急。

工具链的整体设计

先说下我们最终搞出来的东西长什么样。整体架构大概分这么几层：

数据层：负责特征存储、样本管理、数据版本控制。这块我们用了Hive做离线存储，Redis做在线特征缓存，中间加了一层自研的数据同步组件。

模型管理层：负责模型的训练、评估、版本管理。这里我们深度集成了Hugging Face的Transformers库，做了一套统一的模型训练框架。

实验平台层：负责实验的配置、调度、监控。算法同学只需要在Web界面上配好参数，点一下"提交"，后面的事情就全自动了。

服务层：负责模型的在线推理。这块用的是我们自研的推理引擎，基于Triton Inference Server做的二次开发。

听起来好像挺清晰的，但实际上，从设计到落地，中间踩的坑能写一本书。

Hugging Face集成那些事儿

先重点聊聊Hugging Face这块，因为这是我觉得踩坑最多、也最有价值的部分。

我们选择Hugging Face作为模型训练的基础框架，原因很简单：它的生态太完善了。从预训练模型到Tokenizer，从数据加载到训练循环，基本上你能想到的东西，它都给你封装好了。而且社区活跃度极高，遇到问题基本都能在GitHub Issues里找到答案。

但是，"封装得好"也意味着"黑盒得多"。当我们试图把它集成到自己的训练平台里时，问题就来了。

坑一：分布式训练的显存管理

我们用的训练集群是8卡A100，跑一些大的推荐模型时，需要用到DeepSpeed的ZeRO-3优化。理论上，Hugging Face的Trainer API对DeepSpeed的支持已经很好了，但实际上，当模型参数量超过一定阈值时，显存管理就会出各种幺蛾子。

我记得有一次，训练一个大概3B参数的模型，跑到第2000步的时候，突然OOM了。一开始我以为是batch size设太大了，调小之后还是OOM。后来仔细排查才发现，是DeepSpeed在某个特定的通信节点上，没有正确释放梯度的显存。

这个问题折腾了我整整三天。最后怎么解决的呢？我魔改了Hugging Face的Trainer源码，在每一步训练结束后，手动调用了一次torch.cuda.empty_cache()，并且在DeepSpeed的配置里加了一个自定义的gradient checkpointing策略。

# 自定义的Trainer，解决显存泄漏问题
class MemoryOptimizedTrainer(Trainer):
    def training_step(self, model, inputs):
        loss = super().training_step(model, inputs)
        
        # 每100步手动清理一次显存碎片
        if self.state.global_step % 100 == 0:
            torch.cuda.empty_cache()
            
        return loss

    def _save_checkpoint(self, model, trial, metrics=None):
        # 自定义checkpoint保存逻辑，避免保存时显存峰值过高
        with torch.cuda.amp.autocast(enabled=False):
            super()._save_checkpoint(model, trial, metrics)

这段代码看起来简单，但背后的调试过程真的是一言难尽。我当时对着nvidia-smi的输出看了不下几百遍，眼睛都快看瞎了。

坑二：Tokenizer的坑

这个坑可能很多人都会遇到。我们有一个场景，需要对用户的评论文本做embedding，然后拼接到推荐模型里。一开始，我们直接用Hugging Face的AutoTokenizer加载了一个预训练的Tokenizer。

结果线上效果很差。后来分析才发现，我们的评论文本里充斥着大量的emoji、网络用语、甚至是一些火星文。而预训练Tokenizer的vocab里根本没有这些token，导致大量文本被切成了UNK，信息损失严重。

解决方案是自己训练一个Tokenizer。我们用了Hugging Face的tokenizers库，从训练数据里采样了100万条评论，训练了一个BPETokenizer，vocab大小设成了50000。

from tokenizers import Tokenizer, models, pre_tokenizers, trainers

# 初始化BPE模型
tokenizer = Tokenizer(models.BPE(unk_token="<unk>"))
tokenizer.pre_tokenizer = pre_tokenizers.ByteLevel(add_prefix_space=False)

# 配置训练器
trainer = trainers.BpeTrainer(
    vocab_size=50000,
    special_tokens=["<pad>", "<unk>", "<s>", "</s>"],
    min_frequency=2,
    show_progress=True
)

# 训练
files = ["comments_sample_1.txt", "comments_sample_2.txt"]
tokenizer.train(files, trainer)

# 保存
tokenizer.save("xiaohongshu_bpe_tokenizer.json")

训练完之后，线上的文本embedding质量明显提升了一个档次。这个经历让我深刻认识到：工具再好，也得根据业务场景做适配。直接拿来就用，大概率会翻车。

坑三：模型导出的兼容性

这个问题比较隐蔽。我们用Hugging Face训练完模型后，需要导出成TorchScript或者ONNX格式，部署到线上推理服务。但Hugging Face的model.save_pretrained()导出的模型，在某些自定义层上，转换到TorchScript时会报错。

具体来说，我们在模型里用了一个自定义的Attention层，里面有一些动态shape的操作。这些操作在PyTorch的动态图模式下没问题，但trace成TorchScript时就挂了。

# 有问题的自定义Attention层
class CustomAttention(nn.Module):
    def forward(self, x, mask=None):
        # 这里的动态shape操作在TorchScript trace时会出问题
        batch_size, seq_len = x.shape[:2]
        attn_weights = torch.matmul(x, x.transpose(-1, -2))
        
        # 动态mask处理
        if mask is not None:
            mask = mask.unsqueeze(1).expand(-1, seq_len, -1)
            attn_weights = attn_weights.masked_fill(mask == 0, -1e9)
        
        attn_weights = F.softmax(attn_weights, dim=-1)
        return torch.matmul(attn_weights, x)

最后的解决方案是把这部分逻辑重写了，用torch.jit.script代替torch.jit.trace，并且把动态shape的操作改成静态的。

# 修改后的版本，兼容TorchScript
class CustomAttention(nn.Module):
    def __init__(self, max_seq_len: int = 512):
        super().__init__()
        self.max_seq_len = max_seq_len
        # 预分配mask buffer
        self.register_buffer(
            "causal_mask", 
            torch.triu(torch.ones(max_seq_len, max_seq_len), diagonal=1).bool()
        )
    
    @torch.jit.export
    def forward(self, x: Tensor, mask: Optional[Tensor] = None) -> Tensor:
        attn_weights = torch.matmul(x, x.transpose(-1, -2))
        
        if mask is not None:
            # 使用预分配的mask，避免动态shape
            seq_len = x.size(1)
            static_mask = self.causal_mask[:seq_len, :seq_len]
            attn_weights = attn_weights.masked_fill(static_mask, -1e9)
        
        attn_weights = F.softmax(attn_weights, dim=-1)
        return torch.matmul(attn_weights, x)

改完之后，模型导出终于顺畅了。但这个坑让我学到一个教训：在设计模型架构的时候，就要考虑到部署的需求。不能只管训练爽，不管部署死活。

实验平台的工具化实践

说完了模型训练这块，再聊聊实验平台的工具化建设。

我们的核心思路是：一切皆配置，一切皆版本。

特征配置管理

特征管理是推荐系统里最复杂的部分之一。一个推荐模型可能用到几百个特征，这些特征的来源、计算逻辑、更新频率各不相同。如果没有一套好的管理工具，简直是灾难。

我们搞了一个特征注册中心，每个特征都有唯一的ID，关联了它的计算逻辑、数据源、更新频率等信息。算法同学在配置实验时，只需要选择需要的特征ID，系统会自动生成特征处理的pipeline。

# 特征配置示例
feature_config:
  - feature_id: "user_age_bucket"
    feature_type: "dense"
    source: "user_profile"
    transform:
      type: "bucketize"
      boundaries: [18, 25, 30, 35, 40, 50]
    update_freq: "daily"
    
  - feature_id: "item_click_rate_7d"
    feature_type: "dense"
    source: "item_stats"
    transform:
      type: "log_normalize"
    update_freq: "hourly"
    
  - feature_id: "user_item_interaction_seq"
    feature_type: "sequence"
    source: "user_behavior_log"
    transform:
      type: "embedding_lookup"
      embedding_dim: 64
      max_seq_len: 50
    update_freq: "realtime"

这套东西上线之后，特征配置的错误率直接降了一个数量级。之前靠人肉配置，经常会出现特征ID写错、版本号搞混的情况。现在全部走系统，配置即代码，还有版本控制和回滚能力。

实验调度系统

实验调度这块，我们借鉴了Kubernetes的设计理念，搞了一套轻量级的任务调度系统。核心概念就三个：Experiment（实验）、Trial（试验）、Worker（执行器）。

一个Experiment可以包含多个Trial，每个Trial是一组具体的超参配置。Worker负责执行具体的训练任务。调度系统会根据集群的资源情况，自动分配Trial到合适的Worker上执行。

概念	说明	类比
Experiment	一次完整的实验，包含目标和约束	一个项目
Trial	一组超参配置，是实验的最小执行单元	一个任务
Worker	执行训练任务的节点	一个工人
ResourceQuota	资源配额，限制实验能用的最大资源	预算

这套调度系统上线后，实验的吞吐量提升了大概3倍。之前算法同学提交一个实验，可能要排队等半天。现在系统会自动调度，资源利用率也上去了。

安全意识的觉醒

说到这里，我不得不提一个让我印象深刻的安全事件。

大概是今年年初的时候，我们有一个算法同学，为了图方便，把模型的训练数据直接上传到了Hugging Face的公共Hub上。这些数据里包含了大量用户的真实行为数据，虽然做了一些脱敏处理，但还是有泄露风险。

这件事被我们的安全团队发现后，整个部门都被通报批评了。虽然最后没有造成实际的泄露，但这件事给我们敲响了警钟。

从那以后，我们在工具链里加了很多安全相关的约束：

数据层面：所有训练数据必须经过脱敏处理，并且有严格的访问权限控制。我们搞了一个数据脱敏工具，会自动识别和替换敏感信息。

# 数据脱敏工具的核心逻辑
class DataAnonymizer:
    def __init__(self):
        self.phone_pattern = re.compile(r'1[3-9]\d{9}')
        self.email_pattern = re.compile(r'[\w\.-]+@[\w\.-]+\.\w+')
        self.id_card_pattern = re.compile(r'\d{17}[\dXx]')
    
    def anonymize(self, text: str) -> str:
        # 替换手机号
        text = self.phone_pattern.sub('[PHONE]', text)
        # 替换邮箱
        text = self.email_pattern.sub('[EMAIL]', text)
        # 替换身份证号
        text = self.id_card_pattern.sub('[ID_CARD]', text)
        return text

模型层面：模型文件在上传到内部模型仓库之前，必须经过安全扫描，确保没有包含敏感信息。我们还会定期检查Hugging Face的公共Hub，确保没有内部模型被误传上去。

代码层面：所有的配置文件中，禁止硬编码任何密钥和敏感信息。我们接入了公司的密钥管理系统，所有的密钥都通过环境变量注入。

说实话，这些事情在开发的时候确实会觉得麻烦，会拖慢一些进度。但经历过那次安全事件之后，我真心觉得：安全不是可选项，而是必选项。 你省下的每一分钟，都可能在未来变成一颗定时炸弹。

一些心得体会

写了这么多，最后分享几点我个人的心得体会吧。

第一，工具是为业务服务的，不要为了用工具而用工具。 我们选择Hugging Face，不是因为它最牛，而是因为它最适合我们的场景。如果有一天它不适合了，我们也会毫不犹豫地换掉它。技术选型永远要回到业务需求上来。

第二，自动化是减少错误的最好方式。 人能犯的错误，机器不会犯。只要是可以标准化的流程，就应该尽量自动化。我们工具链上线后，最大的收益不是效率提升，而是错误率的大幅下降。

第三，安全意识要贯穿始终。 不要觉得安全是安全团队的事，每个工程师都应该有安全意识。特别是在处理用户数据的时候，一定要慎之又慎。

第四，文档和代码一样重要。 我们工具链的文档，花了差不多两个月时间才写完。虽然写的时候很痛苦，但后来新人入职的时候，看文档就能快速上手，省了大量的沟通成本。

好了，今天就聊到这里。窗外的成都已经华灯初上了，我也该收拾收拾下班了。对了，今晚得早点回去，明天还要早起review代码呢。

如果大家对推荐系统的工具链建设有什么想法或者问题，欢迎在评论区交流。毕竟，独乐乐不如众乐乐嘛。

---

P.S. 如果这篇文章对你有帮助，点个赞再走呗。你们的认可是我继续写下去的动力。