Post

供应链场景下百万级 SKU 异构分布式强化学习系统：RL-Infra 工程实践全解析

Posted
By jelech
84 min read
供应链场景下百万级 SKU 异构分布式强化学习系统：RL-Infra 工程实践全解析

前言：为什么供应链需要强化学习，以及为什么它如此之难

供应链的补货决策表面上是一个预测问题——预测未来需求，然后计算安全库存和补货量。但现实远比这复杂。每一个SKU的库存水平、在途量、供应商交期波动、促销计划、季节因素之间存在高度耦合。一个SKU的缺货可能导致替代品的需求暴增，连锁反应横跨整个品类。传统的运筹优化方法在面对数百万SKU、数十个仓、数百个供应商的组合爆炸时，要么求解时间不可接受，要么不得不做大量简化假设而失去精度。

强化学习（Reinforcement Learning，RL）的引入，是为了让每个SKU拥有自己的决策智能体（agent），通过与仿真环境的大量交互学习最优补货策略。每个agent观察自己的实时特征（库存水位、预测需求、在途量、历史销量等），输出补货量决策。训练目标变为了最大化长期累积收益——在满足服务水平的同时最小化库存持有成本。

这个设想听起来优雅，但工程落地是另一回事。当SKU数量达到100万，每个agent需要与仿真环境交互数百步来完成一个episode，仿真本身是纯CPU密集计算（涉及库存逻辑、约束校验、多级联动），而模型训练又需要GPU的并行算力——一个典型的CPU-GPU异构计算场景就此诞生。这篇文章不会花太多篇幅讨论算法本身（PPO、GAE这些在教材里都能找到），而是聚焦于把这个系统跑起来、跑稳定、跑快、跑得可控的全部工程细节。

一、系统全局架构概览

整个系统在架构上采用了”Go做数据面，Python做计算面”的异构设计。这完全就是现实约束的产物。

仿真器是用Go实现的。供应链仿真需要处理大量的状态机逻辑——库存扣减、在途到货、过期淘汰、多仓调拨——这些逻辑天然适合用强类型、高并发、低GC停顿的语言来写。一个仿真实例在处理一个SKU的一个episode（通常60~365步，每步代表一天）时，涉及的逻辑分支非常密集，但几乎没有浮点矩阵运算。Go的goroutine在这种场景下的表现远优于Python的线程模型，而CGo的开销又使得把仿真逻辑嵌入Python进程不太现实。

训练器使用Python + PyTorch，这没什么好争论的。整个深度学习生态都在Python上，DDP、NCCL、CUDA kernel这些基础设施的可用性决定了learner端必须是Python。

两者之间通过gRPC通信。选择gRPC而非REST或者自定义TCP协议，核心原因有三：Protocol Buffers提供的强类型schema能避免大量运行时数据格式错误；HTTP/2的多路复用在高频小包通信场景下效率显著优于HTTP/1.1；streaming RPC为权重同步提供了天然的推送机制。

系统的角色划分如下：

Learner集群：4张GPU卡组成一个DDP组，运行PPO算法的前向推理和反向传播。Learner是整个系统的”大脑”，它接收collector汇总的trajectory数据，完成梯度计算和参数更新，然后把新权重推送出去。

Collector集群：8台collector机器，每台连接8个仿真程序实例，总共64个仿真worker。Collector的职责是拿到最新的模型权重，用它来驱动仿真环境产生trajectory数据，然后把这些数据打包发回给learner。它同时承担了动作推理的职责——在collector端用CPU做模型推理（模型只有20MB，CPU推理完全可行），避免了每一步仿真都要和GPU端做一次网络往返。

权重同步服务：基于Redis的pub/sub机制，带版本号的权重分发系统。Learner训练完一轮就把新权重写入Redis，所有collector通过订阅channel得到通知，拉取最新权重。当然用Redis有个前提条件，你所同步的模型权重是一个小到带宽可以接受的。

监控体系：Prometheus采集各组件的metrics，Grafana做可视化。训练指标（loss、reward、entropy）通过TensorBoard或MLflow记录。

整个数据流是一个闭环：Collector采集 → Trajectory传输 → Learner训练 → 权重同步 → Collector更新权重 → 继续采集。在PPO的on-policy约束下，这个闭环的每一步都不能出错，否则数据的分布就会和当前策略不匹配，导致训练不收敛。

二、仿真层——系统瓶颈所在

为什么仿真是瓶颈

100万个SKU，这不是我们计算的上限，而是不同品类尽量合并后的结果。每个episode执行100步仿真，每步仿真涉及：读取当前库存状态、应用补货决策、模拟当天销量扣减、处理在途到货、计算缺货和过期、更新状态、计算reward。这些操作全部是标量/小向量运算，无法利用GPU的SIMD并行性。粗略计算：100万 × 100步 × 每步约50μs = 5000秒的纯CPU计算。即使用64个worker并行，单个episode的仿真时间也需要约78秒。这就是为什么64个batch在单次episode中同时计算100万SKU——通过batch并行来分摊这个开销。

仿真服务的Go实现

仿真环境被封装为一个gRPC服务，对外暴露Reset、Step、GetInfo、Close四个RPC接口，严格遵循OpenAI Gym的交互范式。每个仿真实例内部维护一个EnvironmentManager，管理多个并行环境的生命周期。

1
2
3
4
5
6

type Environment interface {
    Reset() (map[string][]float32, error)
    Step(action map[string][]float32) (map[string][]float32, float32, bool, map[string]string, error)
    GetInfo() map[string]string
    Close() error
}

之所以observation和action使用map[string][]float32而非简单的[]float32，是因为我们计算的商品量太大，希望能一次性推送多个action，让仿真程序能并行计算更多的数据。另一个重要原因就是一些商品并不是独立的，他们会互相影响。

64个Worker的编排

8台collector机器，每台启动8个仿真进程。这些仿真进程不是直接由collector fork出来的子进程，而是独立的gRPC服务。Collector通过配置文件知道自己要连接哪些仿真服务的地址（例如同一台机器上的:50051到:50058）。这种进程隔离的设计有一个重要的好处：仿真进程崩溃不会导致collector崩溃。供应链仿真逻辑经常需要迭代修改（加一种新的约束、调整过期逻辑），修改后的仿真进程可以独立重启，不影响collector的运行。

每台collector机器的内存规划也是一个需要认真对待的问题。每个仿真进程需要加载该batch对应的SKU的历史数据和参数配置。100万SKU分64个worker，每个worker负责约1.5万个SKU。每个SKU的特征数据（历史销量、预测、参数等）大约占25KB，一个worker的数据量在300MB左右，加上Go runtime自身的开销，每个仿真进程的内存消耗在200~500MB左右。一台collector机器8个进程，峰值内存在4GB左右，这在现代服务器上不是问题。但需要注意的是，Go的GC在内存使用量大时可能导致较长的STW停顿。实际部署中我们将GOGC设置为200（默认100），以减少GC频率，用内存换延迟。

仿真数据的分配策略

100万SKU如何分配到64个worker上？最朴素的做法是均分，但这忽略了SKU之间计算复杂度的差异。一个高周转SKU（日销数千件）的仿真步计算量远大于一个长尾SKU（月销几件），因为前者涉及更频繁的库存变动和更复杂的补货触发逻辑。如果均匀分配，某些worker会显著慢于其他worker，整个episode的完成时间由最慢的worker决定（木桶效应）。

解决方案是在训练开始前做一次profiling pass：对所有SKU的仿真做一次试运行，记录每个SKU的平均单步耗时，然后用贪心算法（类似多机调度问题的近似解法）将SKU分配到各个worker，使得各worker的总计算时间尽量均衡。这个profiling只需在SKU集合变化或仿真逻辑修改后重做一次，开销可以接受。

实际操作中还有一种更灵活的方式：将SKU按计算复杂度排序后，用round-robin方式交替分配，这样即使个体差异较大，每个worker分到的”重”和”轻”SKU也大致均衡。我们在实验中发现，round-robin方式在不需要额外profiling的情况下，能将最慢worker与最快worker的时间差控制在15%以内，对于大多数场景已经足够。

三、Collector层——CPU端的推理与数据搬运

Collector的双重角色

在异构架构中，collector承担了两个关键职责：

**第一，驱动仿真。**Collector从仿真服务获取当前observation，将observation送入本地持有的策略模型做推理，得到action，再将action送回仿真服务执行一步。这个循环在每个episode中重复数百次。

**第二，数据整理。**一个完整的trajectory包含每一步的(state, action, reward, next_state, done, log_prob, value)，collector需要把这些数据按正确的格式组装成RolloutBuffer，用于PPO的训练。

1
2
3
4
5
6
7
8

type Trajectory struct {
    Observations []map[string][]float32
    Actions      []map[string][]float32
    Rewards      []float32
    Dones        []bool
    Values       []float32
    LogProbs     []float32
}

在Collector端做推理——一个关键的架构决策

一个常见的替代方案是让collector只做数据收集，把observation发给learner端做推理，拿回action再送给仿真。这种”远程推理”模式在很多RL框架中被采用（比如早期的RLlib），但对于我们的场景来说是不可接受的。

**原因在于延迟。**一个episode有100步，每一步都需要一次推理。如果推理在远端GPU上做，那每步需要一次网络往返，即使在同一个数据中心内（RTT约0.5ms），100步就是50ms。但真正的问题不在单次延迟，而在排队效应。64个worker同时发送推理请求，GPU端需要做batched inference来提高吞吐，这意味着请求需要等待凑批。凑批的窗口时间（通常几十毫秒）乘以100步，延迟会膨胀到不可接受的程度。

更根本的矛盾是：仿真是严格串行的（第N步的action依赖第N-1步的state），无法做流水线化。每一步都必须等推理结果返回后才能执行下一步仿真，这使得网络延迟成为不可掩盖的串行瓶颈。

所以我们选择在collector端做CPU推理。模型大小约20MB（一个两层隐藏层、每层256个神经元的Actor-Critic网络），在CPU上的单次前向推理时间在0.1~0.5ms之间，远快于网络往返。64个worker各自独立做推理，天然并行，无需凑批。代价是每台collector机器需要加载一份模型权重到内存，但20MB乘以8个进程也不过160MB，完全可以承受。当机器数量增多或者模型增大后，Redis就需要替换为另一个中间存储介质来承接存储（比如S3），而Redis只作为广播信号降低网络通信协议复杂度。

但这引出了另一个问题：模型权重的同步。

权重同步机制

在PPO这样的on-policy算法中，collector用来采集数据的策略必须和learner正在训练的策略一致。如果collector用旧版本的权重采集了数据，这些数据对当前策略来说就是off-policy的，直接用来更新会引入偏差。

权重同步的设计基于Redis的pub/sub机制。整个流程如下：

  • Learner完成一轮参数更新后，将新权重序列化为字节流，连同版本号一起写入Redis。权重数据使用rl:weights:{model_key}作为key，版本号使用rl:version:{model_key}。

  • Learner同时向rl:updates:{model_key} channel发布一条消息，payload就是新的版本号。

  • 每个collector都有一个后台goroutine在订阅这个channel。收到消息后，它比较收到的版本号和自己当前持有的版本号，如果更新，就从Redis拉取新权重并加载。

  • 所有操作通过Redis Pipeline批量执行，减少网络往返。

1
2
3
4
5
6
7
8
9

func (r *RedisSyncer) SetWeights(ctx context.Context, key string, data *WeightData) error {
    jsonData, _ := json.Marshal(data)
    pipe := r.client.Pipeline()
    pipe.Set(ctx, r.weightKey(key), jsonData, 0)
    pipe.Set(ctx, r.versionKey(key), data.Version, 0)
    pipe.Publish(ctx, r.channelKey(key), data.Version)
    _, err = pipe.Exec(ctx)
    return err
}

一个值得讨论的工程细节是权重的序列化格式。最初的实现直接使用JSON序列化PyTorch的state_dict（通过torch.save保存到BytesIO再base64编码），但20MB的模型权重在JSON化之后体积膨胀到约55MB（base64编码本身有33%的膨胀，加上JSON的结构开销），在Redis中的读写和网络传输都产生了不必要的延迟。后来改用Protocol Buffers的bytes字段直接传输torch.save的二进制输出，体积缩小到22MB左右，传输时间从约180ms降低到约60ms。对于每个episode做一次的权重同步来说，60ms是可以接受的。

权重同步的时序问题

一个微妙的问题是：在PPO的强 on-policy要求下，collector采集数据时使用的权重版本必须和learner当前训练使用的权重版本完全一致。但由于网络延迟和处理时间的存在，权重同步不可能是瞬时的。

我们的解决方案是引入”同步屏障”（sync barrier）。训练流程被组织为严格的同步轮次：

  • Learner通知所有collector开始采集（附带当前权重版本号V）

  • 每个collector检查自己持有的权重版本是否为V，如果不是，等待直到获取到版本V的权重

  • 所有collector使用版本V的权重完成采集

  • 所有collector将trajectory数据上传到learner

  • Learner汇总数据，执行K个epoch的PPO更新，产生版本V+1的权重

  • Learner将版本V+1的权重推送到Redis

  • 回到步骤1

这种严格同步的方式虽然牺牲了一些并行度（collector在等待权重更新时是空闲的），但保证了数据的on-policy特性。在实际测量中，权重同步的等待时间约占总训练时间的3~5%，是可以接受的。

这里也调研过能否让on-policy没有那么“强”。一种优化手段是”双缓冲”：collector在使用当前权重V采集数据的同时，后台goroutine预先下载权重V+1（如果learner已经训练完并推送了的话）。这样当采集完成后，新权重可能已经准备好了，等待时间趋近于零。但在严格on-policy场景下，这个优化只在learner训练速度快于collector采集速度时才有效。对于我们的场景（仿真是瓶颈，learner通常先完成），这个优化效果有限。当然如果PPO能切换到近on-policy的模型，就能进一步提高GPU和 CPU利用率。

四、Learner层——GPU训练的核心

DDP + NCCL：4卡并行训练

Learner采用PyTorch的DistributedDataParallel（DDP），后端通信使用NCCL。DDP的工作原理在此不再赘述，但有几个在我们的场景下需要特别注意的地方。

首先是初始化。DDP要求所有参与训练的进程在启动时执行dist.init_process_group。我们使用torchrun（PyTorch的弹性分布式启动工具）来管理进程组的创建：

1
2
3
4
5
6
7
8

class DistributedTrainingContext(TrainingContext):
    def _setup(self):
        backend = "nccl" if torch.cuda.is_available() else "gloo"
        dist.init_process_group(backend)
        self.rank = dist.get_rank()
        if torch.cuda.is_available():
            self._device = torch.device(f"cuda:{self.rank}")
            torch.cuda.set_device(self._device)

NCCL之所以是必选项而非GLOO，是因为NCCL专门为NVIDIA GPU间的集合通信优化（AllReduce、Broadcast等），在多GPU场景下的带宽利用率远高于GLOO。对于4张GPU（假设V100或A100），NCCL能充分利用NVLink或PCIe的带宽做梯度同步，AllReduce的开销几乎可以被计算所掩盖。

PPO训练的具体实现

PPO的核心是裁剪的策略梯度目标函数，加上价值函数损失和熵正则化：

```plain text L = L_clip + c1 * L_value - c2 * H[π]

1
2
3
4
5
6
7
8
9
10
11
12
13

其中`L_clip`使用importance sampling ratio的裁剪版本来避免策略更新过大，`L_value`使用均方误差拟合状态价值函数，`H[π]`是策略的熵用于鼓励探索。

在我们的实现中，一次PPO更新包含K=10个epoch的小批量梯度下降。这里有一个容易被忽视的性能细节：在每个epoch开始时，不应重新从CPU内存加载数据到GPU，而应该在第一个epoch时就把所有数据搬到GPU显存并保持在那里。100万SKU × 100步 × 每步状态维度（假设50维）的数据量约为20GB的float32张量。4张GPU意味着每张GPU分到5GB数据，对于32GB显存的V100来说绑定但可行，对于80GB的A100则游刃有余。

```python
def update_policy(self, policy, optimizer, buffer, context):
    states, actions, returns, advantages, old_log_probs = self.prepare_data(buffer, context, policy)
    for _ in range(self.epochs):
        optimizer.zero_grad(set_to_none=True)
        loss, metrics = self.compute_loss(states, actions, returns, advantages, old_log_probs, policy)
        loss.backward()
        optimizer.step()

zero_grad(set_to_none=True)是一个值得注意的优化：将梯度置为None而非零，可以避免一次不必要的memset操作，在参数量较大时能节省几个百分点的训练时间。

Advantage的归一化问题

在PPO中，advantage的归一化（减均值除标准差）对训练稳定性至关重要。但在DDP场景下，每张GPU只持有全局数据的1/4，局部计算的均值和标准差可能与全局值存在偏差。正确的做法是先通过dist.all_reduce汇总所有GPU上的advantage的总和与平方总和，计算全局均值和标准差，然后在各GPU上分别做归一化。如果跳过这一步，直接在各GPU上做局部归一化，训练仍然能跑，但收敛速度会变慢，最终性能也会打折扣。

这是一个在小规模实验中很难发现的问题——当只有一张GPU或数据量较小时，局部归一化和全局归一化的差异微乎其微，但当数据量达到百万级且分布在多张GPU上时，差异就会显现。我们在最初的上线过程中观察到reward曲线在某个值附近震荡不收敛，经过排查才发现是advantage归一化不一致导致的。

梯度裁剪与数值稳定性

在大batch训练（100万个transition）中，梯度的方差非常低（大数定律），这是好事，但也意味着如果某个SKU的仿真数据出现异常（比如reward突然变得极大或极小），对应的梯度可能在batch中被稀释掉而不会被发现。更危险的情况是当异常数据的比例达到一定程度时，梯度突然变大，一步更新就把模型带飞。

为此我们做了两层防护：

  • 数据层面：在trajectory进入buffer之前做范围检查，reward超过合理范围的transition被标记并记录到异常日志。这种数据清洗必须在collector端完成，因为到了learner端数据已经被混在一起，很难追溯来源。

  • 梯度层面：使用torch.nn.utils.clip_grad_norm_对全局梯度范数做裁剪，阈值设为0.5。在DDP模式下，梯度裁剪必须在AllReduce之后、optimizer.step之前执行，否则每张GPU裁剪自己的局部梯度会导致全局梯度方向不一致。

五、数据流与通信——gRPC、Protocol Buffers与序列化性能

为什么选择gRPC

系统中的跨进程通信主要发生在三个地方：

  • Collector ↔ 仿真服务：高频、低延迟、小payload（每步一个observation + action）

  • Collector ↔ Learner：低频、高吞吐、大payload（一个episode的所有trajectory数据）

  • Learner ↔ Redis ↔ Collector：低频、中等payload（模型权重）

gRPC的Protocol Buffers序列化在这三种场景下都表现良好。对于第一种场景，protobuf的编码/解码速度远快于JSON；对于第二种场景，protobuf的紧凑二进制格式减少了网络传输量；对于第三种场景，gRPC的streaming RPC使得权重更新可以主动推送而非轮询。

proto定义如下（节选）：

1
2
3
4
5
6
7
8
9
10
11
12
13

service CollectorService {
  rpc Collect(CollectRequest) returns (CollectResponse);
  rpc GetBuffer(GetBufferRequest) returns (GetBufferResponse);
  rpc UpdateWeights(UpdateWeightsRequest) returns (UpdateWeightsResponse);
  rpc GetStats(GetStatsRequest) returns (GetStatsResponse);
}

service WeightSyncService {
  rpc SetWeights(SetWeightsRequest) returns (SetWeightsResponse);
  rpc GetWeights(GetWeightsRequest) returns (GetWeightsResponse);
  rpc Subscribe(SubscribeRequest) returns (stream WeightUpdate);
  rpc GetVersion(GetVersionRequest) returns (GetVersionResponse);
}

Subscribe接口使用server-side streaming，collector打开一个长连接，learner每次更新权重后通过stream推送通知。这比轮询机制（collector定时查询最新版本号）节省了大量无效请求，也降低了Redis的读压力。

Trajectory数据的传输优化

一个episode产生的trajectory数据量的估算：100万SKU × 100步 × 每步约50个float32（state维度） × 4字节 = 20GB。这个数据量不可能一次性通过gRPC传输——gRPC默认的消息大小限制是4MB，即使调大限制，单个消息几十GB也会导致内存压力和超时风险。

解决方案是分片传输。每个collector负责的SKU子集（约15625个SKU）的trajectory数据约300MB，这个量级仍然偏大。进一步将每个collector的数据按mini-batch分片，每个分片包含约2000个SKU的数据（约40MB），分多次RPC调用传输。Learner端维护一个汇聚缓冲区，等所有分片到齐后才开始训练。

分片传输引入了一个复杂性：部分失败处理。如果某个collector的某个分片传输失败了（网络抖动、超时），应该重传该分片还是丢弃整个collector的数据？在on-policy的PPO中，数据的完整性直接影响策略梯度的估计精度，但完美主义的全量重传会拖慢整个训练循环。我们的做法是：设定一个可容忍的数据丢失比例（默认5%），如果丢失量在阈值内，用已有数据训练并在日志中记录丢失情况；如果超过阈值，重新执行本轮采集。实际运行中，网络引起的数据丢失几乎不会超过1%。

gRPC的连接管理

在长期运行的训练任务中（可能持续数天），gRPC连接的稳定性是一个不容忽视的问题。我们遇到过几类故障：

连接静默断开：TCP连接在中间设备（防火墙、负载均衡器）上被超时关闭，但客户端和服务端都不知道。下次RPC调用时才发现连接已死，导致超时错误。解决方案是启用gRPC的keepalive机制：客户端每30秒发送一次keepalive ping，服务端配置允许接收这些ping。同时设置keepalive_timeout，如果ping在5秒内没有收到pong，就认为连接已死并重建。

负载不均：当使用容器编排工具（如K8s）管理collector时，Service的负载均衡对gRPC的长连接不友好。gRPC默认维持长连接，新创建的连接会被路由到新的Pod，但旧连接仍然指向旧Pod，导致负载不均。我们的解决方案是在客户端侧实现连接池轮换：每个collector客户端维护一个连接池，定期（每5分钟）关闭最旧的连接并创建新连接，让K8s Service有机会重新做负载均衡。

消息体过大：gRPC默认的最大消息大小是4MB。在传输trajectory数据或模型权重时，需要在客户端和服务端都配置max_send_message_length和max_receive_message_length。这个配置必须在两端同时设置，否则会出现单端能发但对端拒收的情况，错误信息还不太明显（通常是RESOURCE_EXHAUSTED状态码），排查起来比较费时间。

六、RolloutBuffer——PPO数据管道的核心组件

RolloutBuffer是连接collector和learner的数据结构，它的设计直接影响训练效率。

在PPO中，RolloutBuffer需要存储每一步的transition，包括(state, action, reward, done, log_prob, value)。在所有步骤完成后，需要计算每一步的回报（return）和优势（advantage）。回报的计算涉及折扣累积奖励（discounted cumulative reward），通常使用GAE（Generalized Advantage Estimation）来平衡偏差和方差。

一个容易犯的错误是在多个episode的trajectory数据中错误地跨episode做了折扣计算。当一个trajectory中包含多个episode（中间有done=True的转折点）时，必须在done处截断折扣累积——因为done意味着环境重置，后续的reward不应影响之前的advantage估计。

1
2
3
4
5
6
7
8
9

def compute_returns(self, gamma: float) -> list[float]:
    returns = []
    discounted_sum = 0
    for t in reversed(self.storage):
        if t.done:
            discounted_sum = 0
        discounted_sum = t.reward + gamma * discounted_sum
        returns.insert(0, discounted_sum)
    return returns

内存管理的考量

100万SKU × 100步 × 6个字段 × 每个字段约50维 = 大约120GB的float32数据。这个数据量远超单台机器的内存容量。所以数据不会在任何单点完整汇聚，而是以分布式的方式存在：每个collector持有自己负责的SKU子集的buffer，learner端也只加载当前mini-batch的数据到GPU显存。

在DDP训练中，数据的分配策略是：将所有collector汇总的数据按SKU均匀分配到4张GPU上。每张GPU在一个PPO epoch中需要处理25万SKU × 100步的数据。为了避免GPU显存不足，单个epoch内部还会进一步分成若干mini-batch。mini-batch的大小是一个需要调优的超参数：太小会导致梯度估计的方差大、GPU利用率低；太大会导致显存不足或梯度过于平滑而不利于探索。我们的经验值是每个mini-batch包含4096~8192个transition。

Buffer的清理时机

PPO是on-policy算法，每一轮训练完成后，buffer中的数据就不再有用（因为策略已经更新了，旧数据不再反映当前策略）。所以buffer必须在每轮训练后清空。这个”清空”操作看似简单，但如果实现不当，可能导致内存泄漏。Python的引用计数GC在tensor对象形成循环引用时不能及时回收，而PyTorch的CUDA tensor还涉及GPU显存的释放。正确的做法是显式地del掉所有tensor引用，然后调用torch.cuda.empty_cache()释放GPU端的缓存（注意：empty_cache不会释放正在使用的显存，只是释放PyTorch内存分配器缓存的空闲块）。

1
2
3
4
5

def train(self):
    for ep in range(1, self.cfg.training.total_episodes + 1):
        ep_reward = self._collect()
        metrics = self.learner.learn(self.buffer)
        self.buffer.clear()

七、模型架构与推理效率

Actor-Critic网络设计

供应链补货决策的模型使用Actor-Critic架构，actor和critic共享一个backbone，然后分别接各自的head：

1
2
3
4
5
6
7
8
9
10
11

class ActorCritic(nn.Module):
    def __init__(self, obs_dim: int, act_dim: int, hidden: int):
        super().__init__()
        self.shared = nn.Sequential(
            nn.Linear(obs_dim, hidden),
            nn.ReLU(),
            nn.Linear(hidden, hidden),
            nn.ReLU(),
        )
        self.policy_head = nn.Linear(hidden, act_dim)
        self.value_head = nn.Linear(hidden, 1)

共享backbone的设计在训练早期能让actor和critic共同学习有用的特征表示，加速收敛。但在训练后期，两者的学习目标存在冲突——actor需要特征对动作区分度高，critic需要特征对状态价值预测准确——共享backbone可能成为瓶颈。我们的做法是从共享开始训练，在loss曲线出现plateau时切换到独立backbone。这个切换不需要从头训练：将共享backbone的权重复制两份分别给actor和critic，然后继续训练。

模型的大小约20MB，对应参数量约500万。对于供应链场景这个规模是合理的——状态空间的维度不高（通常50~200维），动作空间是离散的（补货量的离散化区间），模型不需要太大的容量。用更大的模型在我们的实验中并没有显著提升，反而增加了推理延迟和通信开销。

DDP包装的注意事项

当模型被DistributedDataParallel包装后，原始模型被放在model.module属性下。在保存checkpoint或提取推理用的模型时，必须正确地解包装：

1
2
3
4

class ModelFactory:
    @staticmethod
    def get_model_for_saving(model: nn.Module, is_ddp: bool) -> nn.Module:
        return model.module if is_ddp else model

如果直接保存DDP wrapper，checkpoint中会包含module.前缀的参数名（如module.shared.0.weight），在不使用DDP的推理环境中加载时会因为key不匹配而报错。这是一个DDP新手经常踩的坑，但对于有经验的工程师来说更大的风险在于：在训练过程中不小心绕过DDP wrapper直接调用了model.module的forward方法（比如在evaluation时），导致梯度同步失效而自己不知道，表现为训练缓慢但不会报错。

八、监控体系——你无法优化你看不见的东西

多层次监控架构

一个百万SKU规模的RL训练系统，如果没有完善的监控，就像盲人开飞机。我们建立了三个层次的监控：

系统层监控（Prometheus + Grafana）：CPU使用率、内存使用率、GPU利用率、GPU显存使用率、网络IO、磁盘IO。这些指标通过node_exporter（系统）和nvidia-smi（GPU）采集到Prometheus，在Grafana上做时序展示。关键告警规则包括：GPU利用率持续低于50%（可能说明数据加载是瓶颈）、某台collector的CPU使用率异常高（可能SKU分配不均）、Redis内存使用率超过80%（可能权重数据没被及时清理）。

应用层监控（自定义Prometheus metrics）：每个组件暴露自己的业务指标。

Collector端的关键指标：

  • collector_episodes_total：已完成的episode总数

  • collector_steps_per_second：每秒完成的仿真步数

  • collector_episode_duration_seconds：单个episode的耗时（histogram，用于发现长尾延迟）

  • collector_weight_sync_lag_seconds：权重同步的延迟（从learner更新到collector收到的时间差）

  • collector_buffer_size：当前buffer中的trajectory数量

Learner端的关键指标：

  • learner_training_step_duration_seconds：单次训练步骤的耗时

  • learner_gpu_memory_allocated_bytes：GPU显存使用量

  • learner_gradient_norm：梯度范数（用于检测梯度爆炸/消失）

  • learner_weight_version：当前权重版本号

训练质量监控（TensorBoard / MLflow / Wandb）：

1
2
3
4
5
6
7

class TensorboardLogger:
    def __init__(self, log_dir: str):
        self.writer = SummaryWriter(log_dir)

    def log(self, metrics: dict, step: int = None):
        for k, v in metrics.items():
            self.writer.add_scalar(k, v, step)

训练质量指标包括：episode_reward的均值和分布（不仅看均值，还要看分位数——如果中位数在涨但P95在跌，说明策略在大多数SKU上改善了但在某些SKU上恶化了）、policy_loss、value_loss、entropy（entropy持续下降到零说明策略坍缩到确定性策略，可能过早收敛到局部最优）、KL散度（当前策略与旧策略之间的KL divergence，PPO通过裁剪来隐式控制，但监控它有助于判断学习率是否合适）。

训练曲线的解读——一些非显而易见的经验

Reward震荡不收敛：最常见的原因不是超参数不对，而是数据管道有bug——比如observation的某些字段没有正确归一化、advantage计算跨了episode边界、或者collector使用的权重版本和learner不一致。我们的调试经验是：先在小规模（1000个SKU，单GPU，单collector）上验证训练能收敛，确认算法正确后再逐步放大规模。

Value loss先降后升：这通常说明bootstrap value（用于GAE计算的V(s’)）的估计偏差在累积。可能的原因是环境的reward scale发生了变化（比如某批新上线的SKU的reward量级和已有SKU差异很大），导致value network需要重新学习。解决方案是对reward做running normalization。

Entropy坍缩：策略的entropy快速降为零，意味着agent对所有状态都给出了确定性的动作。在供应链场景下这可能不是坏事（如果策略确实找到了最优补货量），但更常见的情况是策略过早收敛到了”永远不补货”或”永远补满”这样的退化策略。增大entropy系数或者使用entropy调度（training前半段entropy系数高，后半段逐渐降低）可以缓解。

九、版本管理与实验追踪

模型版本管理

供应链RL系统的模型版本管理比传统ML复杂得多，因为”模型”不只是一个权重文件——它还绑定了仿真环境的参数、特征工程的逻辑、PPO的超参数、SKU的分配方案等。任何一个要素的变化都可能导致模型行为的改变。

我们的版本管理策略分为三个层次：

权重版本（训练内部）：使用单调递增的整数版本号，每次learner完成一轮PPO更新就递增。这个版本号主要用于权重同步的一致性检查，在一次训练任务的生命周期内有效。

Checkpoint版本（训练之间）：每隔固定的episode保存一次checkpoint，命名为policy_ep{episode_number}.pth。Checkpoint包含模型权重、优化器状态、当前episode数，以便从中断处恢复训练。同时保存训练配置文件的快照，确保可以准确复现。

1
2
3
4
5
6
7
8

class Checkpoint:
    def save(self, policy, episode):
        path = os.path.join(self.save_dir, f"policy_ep{episode}.pth")
        torch.save(policy.state_dict(), path)

    def close(self, policy):
        path = os.path.join(self.save_dir, "policy_final.pth")
        torch.save(policy.state_dict(), path)

发布版本（上线推理）：当一个训练任务的最终模型通过离线评估（使用holdout的SKU集合做仿真回测），才会被标记为可发布版本，推送到模型仓库（我们使用MLflow Model Registry）。每个发布版本包含：模型权重、特征schema、推理代码、仿真环境版本hash、训练日志链接。发布版本使用语义化版本号（如v2.3.1），major版本号变更表示模型架构变化，minor版本号变更表示训练数据或超参数变化，patch版本号变更表示bug修复后的重训练。

实验追踪

在调优RL系统时，需要同时跟踪大量实验。每次实验可能改变的变量包括但不限于：

  • 模型架构（隐藏层大小、层数、激活函数）

  • PPO超参数（学习率、clip_ratio、epoch数、mini-batch大小）

  • 仿真参数（reward函数的设计、状态空间的构成）

  • 基础设施参数（collector数量、batch大小、权重同步频率）

我们使用MLflow来管理实验。每次训练启动时，创建一个MLflow run，记录所有配置参数、代码版本（git commit hash）、环境信息（Python版本、PyTorch版本、CUDA版本）。训练过程中实时记录metrics，训练结束后记录最终模型和评估结果。

一个实际问题是：RL的实验结果方差很大。同样的配置跑两次可能得到差距显著的结果，因为初始权重的随机性、仿真环境中随机因素（需求的随机波动）都会影响训练轨迹。所以每组配置至少需要跑3次取平均才能做有意义的比较，这进一步增加了实验管理的复杂度。我们的做法是在MLflow中使用parent-child run的机制：一组配置作为parent run，同一配置的多次重复作为child run，分析时在parent level做聚合。

十、推理服务——从训练到线上的最后一公里

推理链路设计

训练完成后，模型需要在线上环境中为每个SKU生成补货决策。推理链路和训练时的collector端推理有本质区别：训练时推理只需要和仿真环境交互，数据是内存中的虚拟数据；线上推理需要接入实时的特征数据，输出的决策会直接影响真实的供应链操作。

线上推理采用batch inference模式：每天定时触发一次，对所有SKU做一次推理，生成次日的补货建议。推理过程如下：

  • 特征服务从数据仓库和实时数据源（Kafka/Flink）拉取最新的SKU特征

  • 特征按模型的输入schema组装成batch tensor

  • 模型加载最新的发布版本权重

  • batch inference，得到每个SKU的动作（补货量等级或连续补货量）

  • 后处理：将模型输出映射到实际的补货量（考虑起订量、箱规等约束）

  • 结果写入决策表，供下游系统执行

因为是batch inference而非实时推理，对延迟的要求不高（分钟级即可），所以不需要像推荐系统那样部署GPU推理服务。一台8核CPU的机器，加载20MB的模型，对100万SKU的推理时间在10秒以内（100万 × 0.01ms/SKU）。

但这种简单性是以每天只决策一次为代价的。如果未来需要更高频的决策（比如仓内实时调拨），推理链路需要改为在线服务模式，那时就需要考虑模型热更新、请求排队、批处理推理等问题。

模型灰度发布

新模型上线不会一次性替换旧模型。我们采用灰度发布策略：

  • A/B测试：选择一组对照SKU（按品类、仓、周转率分层抽样），将这些SKU的决策切换到新模型，其余SKU继续使用旧模型。运行2~4周后，比较两组SKU的服务水平（满足率、缺货率）和库存成本（平均库存天数、滞销率）。

  • 逐步放量：如果A/B测试的结果显著优于旧模型，逐步扩大新模型的覆盖范围：10% → 30% → 50% → 100%。每个阶段运行1周以上，确认没有异常后再继续放量。

  • 快速回滚：如果任何阶段发现服务水平下降超过阈值，5分钟内回滚到旧模型。这要求模型切换是配置驱动的（改一个配置项指向旧版本模型），而不是代码发布驱动的。

线上推理的监控

推理结果的监控不仅仅是检查推理服务是否正常运行，更重要的是检查模型的输出是否”合理”。我们设计了以下监控指标：

  • 动作分布漂移：将线上推理的动作分布与训练时的动作分布做比较（使用KL散度或JS散度）。如果漂移超过阈值，可能说明线上特征的分布发生了变化（数据漂移），模型需要重新训练。

  • 异常动作检测：某个SKU的补货建议是否远超历史范围？比如一个日销10件的SKU，模型建议补货10000件，这显然是异常的。异常检测规则基于SKU的历史决策统计。

  • 决策覆盖率：是否所有应该被决策的SKU都得到了补货建议？缺失的SKU需要查原因——是特征缺失导致推理跳过了，还是模型对该SKU的输入产生了NaN？

十一、那些踩过的坑——从调试地狱中幸存

NCCL通信超时

问题表现：训练在某个episode莫名卡住，日志停止输出，CPU和GPU都有一定的负载但没有进展。几分钟后NCCL抛出超时错误。

根因：DDP在反向传播时使用AllReduce来同步梯度，这要求所有rank在同一时刻执行同一个AllReduce操作。如果某个rank在forward或backward中走了不同的代码路径（比如某个rank的数据触发了一个条件分支，导致计算图不同），另一些rank的AllReduce会等待这个rank，最终超时。

在我们的场景下，具体的触发条件是：PPO代码中有一个判断if buffer.size() > 0的分支。由于数据分配的不均匀，偶尔某个rank分到的数据恰好为空，跳过了训练步骤。其他rank在等待这个rank的AllReduce，造成死锁。

解决方案：在数据分配阶段确保每个rank至少分到一些数据，即使需要重复分配一些样本来填充。同时设置NCCL_ASYNC_ERROR_HANDLING=1环境变量，让NCCL在超时后能干净地退出而不是挂死整个进程。

Go仿真进程的内存泄漏

问题表现：仿真进程在运行数小时后，内存使用量缓慢但持续增长，最终触发OOM killer。

根因：Trajectory结构体中的Observations字段是[]map[string][]float32类型。每一步仿真都会创建一个新的map和若干新的[]float32切片。在Go中，即使Trajectory被回收，如果其中的切片被其他地方引用（比如通过gRPC传输时被protobuf序列化过程引用），这些切片就不会被GC回收。

调试过程中使用pprof的heap profile发现了大量的小内存分配来自map[string][]float32的创建。根本原因是gRPC的响应消息在发送后没有被及时释放——Go的gRPC库在内部会持有消息的引用直到底层HTTP/2帧完全发送完毕，而在高负载下这个过程可能延迟数秒。

解决方案：引入对象池（sync.Pool）来复用Trajectory结构体和其中的切片。每一步仿真不再创建新的map和切片，而是从池中取出已分配的对象、覆写数据、使用、然后放回池中。这彻底消除了高频小对象分配导致的GC压力和内存泄漏。

Redis的大Key问题

问题表现：当模型权重增大到25MB后，权重同步的延迟从60ms暴增到800ms，偶尔甚至超过2秒。

根因：Redis是单线程模型，对大Key的读写操作会阻塞其他所有操作。25MB的权重数据在Redis内部需要连续的内存分配（jemalloc的大块分配），读取时需要将整个数据序列化到输出缓冲区，这个过程会阻塞Redis主线程数百毫秒。而在此期间，其他collector的读取请求、pub/sub消息的分发都被阻塞。

解决方案：

  • 启用Redis 7的多线程IO（io-threads 4），让网络读写使用额外线程，减少主线程的IO负担。

  • 将权重数据分块存储：不再使用单个Key存储整个权重，而是将权重按layer分成多个chunk，每个chunk约2MB，使用MSET/MGET批量操作。读取时并发获取所有chunk，然后在客户端拼装。

  • 在Redis前面加一层本地缓存：collector端使用内存缓存最近一个版本的权重，如果Redis的版本号没有变化，直接使用缓存，避免重复读取。

Collector和Learner的速度不匹配

问题表现：GPU利用率经常在0%和100%之间交替，形成明显的”锯齿”模式——GPU空闲等待数据 → 数据到达后满载训练 → 训练完成再次空闲等待。

这是CPU-GPU异构计算中最经典的瓶颈。仿真（CPU）和训练（GPU）的计算时间比值决定了系统的利用效率。如果仿真时间远大于训练时间（在我们的场景下确实如此），GPU大部分时间都在空闲。

解决方案依赖于增加collector的数量。从最初的4台collector加到8台，仿真的总吞吐量翻倍，GPU的空闲时间从70%降低到35%左右。但受on-policy约束，不能用流水线化来进一步掩盖延迟（不能让collector在learner训练的同时用旧权重采集新数据，因为那些数据是off-policy的）。

一种更激进的优化是使用”近似on-policy”：允许collector使用最近N个版本内的权重采集数据（而非严格使用最新版本），然后在PPO更新时用importance sampling ratio来修正off-policy偏差。这在理论上是PPO本身就支持的（PPO的clipping就是为了处理一定程度的off-policy数据），但在实践中N不能太大（通常不超过2~3），否则训练不稳定。我们在实验中验证了N=2时训练仍然稳定，GPU利用率提升到60%左右。

仿真环境的确定性问题

问题表现：同一组SKU、同样的初始条件、同样的策略，两次仿真得到不同的trajectory。

这在调试时是灾难性的——你无法确定某个问题是代码bug还是随机波动。供应链仿真中的随机因素包括：需求的随机波动（通常用泊松分布或正态分布采样）、供应商交期的随机延迟、随机的促销触发等。

解决方案是将所有随机源都通过可控的随机种子管理：

1
2

random.seed(cfg.seed)
torch.manual_seed(cfg.seed)

但仅仅设置Python和PyTorch的种子是不够的。Go仿真进程有自己的随机数生成器，需要通过配置传入种子。多个goroutine并发使用math/rand时还需要注意：Go 1.20之后全局math/rand默认使用加锁的源，但如果手动创建了rand.New(rand.NewSource(seed))，这个source不是并发安全的。正确的做法是给每个goroutine创建独立的rand.Rand实例，种子为全局种子+goroutine编号。

此外，CUDA的计算也存在非确定性——某些CUDA kernel的原子操作不保证顺序，导致浮点累加的结果可能因为加法顺序不同而产生微小的差异。对于调试目的，可以设置torch.backends.cudnn.deterministic = True和torch.use_deterministic_algorithms(True)，但这会降低训练性能（约10~20%），所以通常只在调试时启用。

训练中途SKU集合变化

问题表现：在训练进行到一半时，业务方通知有一批新SKU上线，需要纳入训练；同时有一批旧SKU下线，需要移除。

这在供应链场景下是常态——SKU的上下架是持续进行的。但对RL训练来说这是一个棘手的问题：

  • 新SKU没有历史trajectory数据，但训练需要它们的数据。

  • 旧SKU的数据已经在buffer中，如果不清理会浪费计算资源。

  • SKU的总数变了，数据分配方案需要重新调整。

  • 如果模型的输入维度与SKU无关（每个SKU独立推理），模型本身不需要改。但如果存在SKU间的交互特征（如品类维度的共享embedding），模型架构可能需要调整。

我们的解决方案是”定期重平衡”：每隔一定周期（如每天），暂停训练，更新SKU列表，重新计算数据分配方案，然后从当前checkpoint恢复训练。新增SKU的初始策略使用同品类已训练SKU的平均策略作为warm-start，而不是从随机策略开始。

十二、容器化部署与资源管理

Docker镜像设计

系统的三类组件（仿真服务、collector、trainer）各自有独立的Docker镜像：

1
2
3
4
5
6
7
8
9
10

# env-server.Dockerfile
FROM golang:1.21-alpine AS builder
WORKDIR /app
COPY go/ ./go/
RUN cd go && go build -o /env-service ./cmd/env-service

FROM alpine:latest
COPY --from=builder /env-service /usr/local/bin/
EXPOSE 50051 8080 9090
CMD ["env-service"]

Go组件使用多阶段构建（multi-stage build）：第一阶段用golang镜像编译出静态链接的二进制文件，第二阶段用alpine镜像运行，最终镜像大小控制在30MB以内。Python trainer的镜像则不可避免地较大（约2GB，主要是PyTorch和CUDA runtime），使用nvidia/cuda作为基础镜像。

Docker Compose与本地开发

开发和测试阶段使用Docker Compose来编排整个系统：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20

services:
  redis:
    image: redis:7-alpine
    ports: ["6379:6379"]
  env-server:
    build: { dockerfile: docker/env-server.Dockerfile }
    ports: ["50051:50051", "8080:8080", "9090:9090"]
  collector:
    build: { dockerfile: docker/collector.Dockerfile }
    ports: ["50052:50052", "8081:8081", "9091:9091"]
    depends_on: [env-server, redis]
  trainer:
    build: { dockerfile: docker/trainer.Dockerfile }
    depends_on: [collector, env-server, redis]
  prometheus:
    image: prom/prometheus:latest
    ports: ["9092:9090"]
  grafana:
    image: grafana/grafana:latest
    ports: ["3000:3000"]

本地开发时通常只启动部分组件（比如只启动env-server和trainer做算法调试，不启动监控栈），通过docker compose up --scale collector=2 trainer env-server redis来灵活控制规模。

生产环境的Kubernetes部署

生产环境使用Kubernetes来管理资源。关键的编排决策包括：

仿真服务用Deployment部署，每个Pod运行一个仿真实例。使用HPA（Horizontal Pod Autoscaler）根据CPU使用率自动扩缩。由于仿真是CPU密集型，Pod的resource request设置为1 CPU、512MB内存，limit设置为2 CPU、1GB内存。

Collector用StatefulSet部署（因为每个collector需要持有固定的SKU分配方案和本地buffer），replica数固定为8。每个collector Pod通过环境变量知道自己连接哪些仿真服务的地址。

Trainer用Job部署（训练任务有明确的开始和结束），每次训练创建一个新的Job。GPU资源通过nvidia.com/gpu资源类型请求，4张GPU对应4个container共享一个Pod（DDP要求所有rank在同一节点上使用NVLink通信效率最高）或4个Pod（跨节点通信，需要RDMA支持）。

Redis使用云服务商的托管Redis（如AWS ElastiCache），避免自运维Redis的复杂性。对于权重同步这种低频但对可用性要求高的场景，托管服务的SLA足够了。

资源成本分析

粗略估算整个系统的资源消耗：

  • 64个仿真进程 × 1 CPU = 64 CPU → 8台8核机器

  • 8个collector进程 × 2 CPU = 16 CPU → 与仿真进程同机部署

  • 1个trainer × 4 GPU（A100 40GB）= 4 GPU

  • Redis：单实例，8GB内存

  • Prometheus + Grafana：2 CPU、4GB内存

GPU是成本大头。一台4xA100的机器按云服务商价格约$12/小时（按需计价）。一次完整训练（假设200个episode，每个episode约5分钟）耗时约17小时，GPU成本约$200。加上CPU机器的成本（每台约$0.5/小时 × 8台 × 17小时 = $68），单次训练的总成本约$270。如果使用spot instance（抢占式实例），成本可以压缩到1/3左右，但需要处理抢占中断后的恢复逻辑——这也是checkpoint机制的重要价值之一。

十三、可靠性与容错

Checkpoint与断点恢复

长时间的训练任务最怕的就是中途崩溃后从头重来。Checkpoint机制每隔N个episode保存一次训练状态。但”训练状态”不仅仅是模型权重——它还包括：

  • 优化器的状态（Adam的一阶和二阶矩估计）

  • 当前的episode编号

  • 随机数生成器的状态（确保恢复后的训练轨迹和连续训练一致）

  • 学习率调度器的状态

在DDP场景下，只有rank 0保存checkpoint（所有rank的模型权重相同，保存多份没有意义）。恢复时，rank 0加载checkpoint，然后通过broadcast将权重分发到其他rank。

Collector的容错

如果某台collector机器宕机，系统应该如何处理？在on-policy的PPO中，数据必须来自当前策略，不能用旧数据替代。所以宕机的collector意味着损失了一部分SKU的数据。

我们的容错策略是：

  • Learner检测到某个collector的数据在超时时间内未到达，标记该collector为不可用。

  • 评估丢失的数据量占比。如果<5%，用已有数据继续训练；如果>5%，暂停训练，等待collector恢复或手动介入。

  • 宕机collector的SKU在下一轮被重新分配到存活的collector上（这意味着存活的collector需要加载额外的SKU数据，速度会变慢）。

  • collector恢复后重新加入集群，触发SKU的重新均衡分配。

仿真进程的容错

单个仿真进程崩溃是最常见的故障。Collector通过gRPC的error handling检测到仿真进程不可用后，尝试重连。如果重连失败，将该仿真进程标记为离线，其负责的SKU被临时分配到同一collector的其他仿真进程上。同时发送告警通知运维人员排查仿真进程崩溃的原因。

Go仿真进程内部也做了充分的panic recovery。每个SKU的仿真步骤被包裹在recover()中：即使某个SKU的仿真逻辑触发了panic（比如数组越界、除零错误），也只会影响该SKU本轮的数据，不会导致整个仿真进程崩溃。被panic的SKU会被记录到错误日志，其产生的截断trajectory会被标记为无效，不会进入训练数据。

十四、性能调优的方法论

在百万SKU的规模下做性能调优，不能靠猜，必须靠数据。我们的性能调优遵循以下流程：

  • Profiling先行：使用Go的pprof分析仿真进程的CPU和内存热点；使用PyTorch Profiler分析GPU训练的算子耗时和显存使用；使用Prometheus分析端到端的各阶段耗时。

  • 找到瓶颈：在典型的训练循环中，各阶段的耗时分布大约是：仿真采集60%、数据传输15%、GPU训练15%、权重同步5%、其他5%。仿真是毫无悬念的瓶颈。

  • 针对瓶颈优化：

    • 仿真层面：优化SKU仿真步骤的计算逻辑，减少不必要的map操作，使用数组代替map来存储高频访问的状态字段。

    • 数据传输层面：使用protobuf的binary wire format代替JSON，压缩大型trajectory数据（使用snappy压缩，压缩比约2:1，解压速度极快）。

    • GPU训练层面：调优mini-batch大小以充分利用GPU的SM并行度，使用torch.compile（PyTorch 2.0+）来编译模型获得约10~15%的加速。

    • 权重同步层面：前述的Redis大Key优化。

  • 验证优化效果：每次优化后重新做profiling，确认瓶颈确实被缓解了（而不是仅仅移到了别处），并检查优化是否引入了新的问题。

一个有意思的优化案例

在profiling中发现仿真进程的CPU时间有约30%花在了map[string][]float32的key lookup上——每一步仿真都需要从observation map中按key取值，Go的map lookup虽然是O(1)但常数系数不小（涉及hash计算和bucket遍历）。对于一个每秒执行数百万次的操作，这个常数系数就变得重要了。

优化方案：将observation的字段从map改为固定布局的struct。每个字段对应struct的一个字段，编译器可以直接通过偏移量访问，完全消除了hash计算的开销。这个改动涉及仿真代码和gRPC接口的重构（protobuf消息从map<string, Tensor>改为具名字段），工作量不小，但带来了约25%的仿真速度提升。这种优化在小规模时无关紧要，但在百万SKU的尺度下效果显著。

结语

回顾这个系统的建设过程，最深刻的体会是：RL-Infra的核心挑战不在于任何单个技术组件（gRPC、DDP、Redis这些都有成熟的文档和社区经验），而在于它们的集成——在一个CPU-GPU异构、多进程、多机器的分布式系统中，让数据以正确的顺序、正确的版本、正确的格式流动，并且在某一环节出错时能快速发现和恢复。

供应链场景给这个挑战增加了额外的维度：SKU的规模使得”小规模好用的方案大规模就不行了”成为常态（前面提到的map vs struct就是一个例子）；业务逻辑的复杂性使得仿真器成为不可替代的性能瓶颈（不像Atari游戏那样有GPU加速的环境模拟器）；决策的商业影响使得可靠性和可追溯性成为硬性要求。

这篇文章试图将我们在这个系统上的工程经验做一个完整的记录——不仅仅是”我们做了什么”，更重要的是”为什么这么做”以及”哪些地方走了弯路”。希望对正在或即将踏入RL-Infra这个领域的工程师们有所帮助。RL不只是算法工程师在Jupyter Notebook里调超参数的事，它的工程化落地是一个完整的系统工程问题，需要在分布式系统、高性能计算、DevOps、数据工程等多个方向上都有扎实的功底。而这些，正是RL-Infra工程师的价值所在。

算法模型, 性能优化
笔记 总结 机器学习
This post is licensed under CC BY 4.0 by the author.