ZO技术报告

分割训练稳定性与性能优化技术报告

1. 背景与目标

本项目为 ZO Split Learning 多客户端训练系统：

server 部署在本机（热点提供方）
client 部署在 4 台 Jetson Xavier NX
客户端通过无线热点连接 server，执行分割训练与资源心跳上报

在长期运行中出现以下问题：

训练在中后期明显变慢，早期可推进，后期趋于停滞
客户端在线/离线状态抖动（session 重建频繁）
某些阶段出现“单客户端推进、其他客户端卡住”
前端曲线显示不完整、历史可视化不稳定

本轮工作的目标是：

提升训练稳定性（避免后期失稳）
保留监控能力（尤其 Jetson/jtop 数据）
保障前端曲线可用性与可解释性
在不改变核心算法严谨性的前提下，尽可能提高可持续训练时长

2. 算法与训练任务说明

2.1 算法框架（ZO Split Learning）

本项目采用零阶优化（Zero-Order, ZO）与 Split Learning 结合的训练方式。

核心思想是在不直接反向传播到服务端参数的情况下，通过扰动估计实现参数更新。

客户端侧（Jetson）
- 执行 ViT.embeddings 得到 token 级中间特征
- 将中间特征发送至服务端
- 接收服务端返回的正常/扰动特征后，计算任务头损失并更新本地头部参数
服务端侧（本机）
- 执行 encoder + layernorm
- 对参数进行随机扰动，计算扰动前后特征
- 基于客户端回传的损失差完成零阶更新

2.2 当前 split 切分点与通信特征

当前切分点位于 ViT 的 embeddings 之后，意味着每个 step 都要传输较大的 token 张量。

该设计在有线或低丢包环境中可运行，但在 4 客户端共享无线热点场景下，对稳定性更敏感。

2.3 训练任务定义

数据集：CIFAR-10（客户端本地加载）
训练集子集：约 5000 样本（用于本轮实验）
客户端数量：4
目标：在分割训练架构下稳定推进训练并提升分类精度，同时保持可观测性（监控与曲线）

2.4 为什么本任务容易出现“累积型退化”

该任务不是单纯计算密集，而是“计算 + 通信 + 监控”耦合：

每 step 的训练请求都包含中间特征传输
多客户端并发会放大无线链路抖动
客户端还需并行进行资源采集与心跳上报

因此，随着运行时间增长，系统更容易从“可用”进入“高抖动/高延迟/高CPU”的状态。

3. 关键现象与证据

3.1 训练层面

早期可训练，后期退化明显（非“启动即故障”）
不同 batch size 下失稳时刻不同（例如 32 时更早，16 时更晚），说明是累积型问题

3.2 网络与系统层面

心跳 session recreated 在部分阶段高频出现（后经调度降低）
某些时段热点功能异常（无法连上/开启），疑似链路高压后驱动或设备状态退化

3.3 Jetson 资源证据

排查输出显示：

CUDA 可用（torch.cuda.is_available() == True）
python 进程 CPU 占用异常高，线程数一度非常高
同时期 ksoftirqd 与 TCP 重传并不总是高，说明并非每次都由网络重传风暴主导

结论：CPU 压力来源是复合型，包含应用层开销（训练+序列化+采集）与部分阶段链路抖动放大。

4. 根因分析（分层）

4.1 通信-计算耦合过强

当前 split 点在 ViT embeddings 之后，客户端每 step 都要发送较大的中间张量，长期累计通信量大。

即便单次可承受，持续并发时会放大链路抖动的影响。

4.2 协议重试与并发放大效应

训练请求曾采用长超时 + 循环重试，遇到链路抖动时会产生“放大器”效应
多客户端并发会争抢有限无线链路，导致时延和失败率上升

4.3 监控链路开销（jtop）

已确认历史上存在 jtop 稳定性问题。

“不崩溃”不代表“低开销”，若采集方式不当，可能导致运行越久 CPU 越高。

4.4 前后端数据窗口策略导致观测偏差

历史曲线曾被窗口裁剪（仅显示尾部数据），造成“看起来丢点/从后段开始画”的错觉。

5. 已实施的优化项

以下均已在代码中落地。

4.1 协议幂等化（防重复更新）

为 forward_zo/update_zo 引入 requestId
服务端增加去重缓存与 in-flight 合并，避免同请求重试导致重复计算/重复更新

4.2 调度器（稳定优先）

服务端增加全局训练调度（限流 + 公平放行）
后续优化为“仅在有等待请求的客户端之间轮询”，避免固定轮询导致单客户端首步卡死
调度间隔可通过环境变量配置（例如 10s、6s、4s）

4.3 会话化历史与可视化修复

训练历史按 run/session 归档，支持会话选择查看
修复前端增量更新时 Jetson 扩展字段未同步的问题
曲线保存支持按会话标识命名

4.4 jtop 采集链路重构

由“每次采样创建上下文”改为“持久实例复用 + 冷却重连”
减少频繁创建/销毁带来的 CPU 与线程压力
退出时显式释放采集资源

4.5 评估与训练参数调优

当前评估策略调整为：每 50 step 评估一次，每次 64 batch（避免全量评估造成假卡顿）
跳过 step=0 评估，解决“启动后长时间无输出”的体验问题

6. 当前配置摘要（关键）

5.1 客户端

BATCH_SIZE = 32
EVAL_INTERVAL = 50
EVAL_MAX_BATCHES = 64
心跳与资源采集已做降压与容错

5.2 服务端

启用协议去重与 in-flight 合并
启用调度器（建议通过环境变量动态调参）
- SCHEDULER_MIN_FORWARD_INTERVAL（当前可根据稳定性设为 4~6s）
- SCHEDULER_ACTIVE_STEP_TTL

7. 阶段性效果

8. 仍需关注的问题

无线链路天然不稳定：高并发大包长期运行仍有风险
jtop 命令行工具本身（TUI）仍可能出现上游兼容问题（不影响当前前端采集链路）
若追求更高上限与更高精度收敛，仍建议考虑通信侧进一步减压（需评估算法严谨性约束）

9. 后续建议（按优先级）

P0（持续稳定）

继续采用可配置调度间隔，建议 4~6s 作为平衡区间
增加调度器状态可视化（等待队列长度、放行速率、超时回收次数）

P1（可观测性）

增加周期性协议指标：
- forward cache hit
- update dedup hit
- per-client 等待时延

P2（长期优化）

在不改算法数学定义前提下，评估通信实现优化（编码/打包方式）
若允许算法侧变更，再评估 split 点或通信压缩方案

10. 结论

本次优化证明：

问题并非单点故障，而是“训练机制、协议策略、监控采集、无线链路”共同作用下的累积型退化。

通过协议幂等、调度限流、jtop 采集重构、评估策略调整与前端修复，系统已从“中后期易失稳”进入“可持续推进”的状态。

后续以“可观测调优 + 渐进提速”为主线，可在稳定前提下继续逼近更优训练效率与收敛效果。