feat: package issue 918 Numba roadmap independently - #929
Conversation
Numba 效率复测：统计口径、前后对比与推断测试环境
测试提交：PR #929 统计口径
前后对比G1 热切片：新自动 assembled Numba vs 旧手写 Numba
串行路径反而变快：2048 env 为 G1 新 Numba vs 新 NumPy fallback
Motion Tracking 完整
|
| env 数 | 旧手写 Numba | 新 plan Numba | 变化 |
|---|---|---|---|
| 2048 | 0.145 ms（约 14.12M env/s） | 0.162 ms（约 12.64M env/s） | −10.5% |
| 8192 | 0.425 ms（约 19.28M env/s） | 0.389 ms（约 21.05M env/s） | +9.3% |
G1 真实 collector（新 NumPy vs 新 Numba）
| env 数 | collector NumPy | collector Numba | collector 变化 | update_state 变化 |
|---|---|---|---|---|
| 2048 | 0.058M steps/s | 0.059M steps/s | +1.5% | 3.814 → 2.576 ms（1.48×） |
| 8192 | 0.118M steps/s | 0.126M steps/s | +6.7% | 9.227 → 4.506 ms（2.05×） |
性能下降的推断
- 通用 fused kernel 的并行友好性低于旧手写 kernel。 新实现通过
inputs/parameters/workspace/observations的 tuple 和通用 output routing 执行，每个 term 都有运行时 scale 分支；旧实现是固定参数、固定布局、直接写 observation 的任务专用 kernel。前者给 Numba/LLVM 的常量传播、内联、别名分析和寄存器/缓存优化空间更小。 - 每步热切片前后增加了 host-side 工作和中间缓冲。 例如 G1 的
dof_pos_diff会在 binder 中单独生成，随后再被 observation kernel 读取；新 runtime 还要做 input rebinding、scale 同步、workspace/日志管理。单线程时这些开销被 fused 计算收益覆盖，高线程时变成不可并行的固定成本。 - 数据搬运/写回更重，导致内存带宽和缓存竞争提前饱和。 2048/8192 行在 32/64 线程下已出现调度、缓存和写回开销主导的迹象；新实现的并行效率只有 G1 32T 的约 6.6%/9.0%、64T 的约 2.9%/4.3%，而旧实现约为 25.0%/36.9% 和 13.7%/27.0%。
- 线程数不是越大越好。 新实现的最佳点已经从旧实现的 64T 降到 32T；继续使用 64T 会增加 Numba
prange调度和共享缓存压力。 - 端到端收益受非 Numba 部分限制。 8192 env 时
update_state已提升 2.05×，但 physics 约 29 ms、collector 其余部分约 31 ms，因此总 collector 只提升 6.7%。2048 env 的 8 个测量样本中，other_ms还从 22.50 增至 23.10 ms，不能据此断言稳定的端到端回退。
优化建议
- 先补分段计时：分别记录
_sync_scales、_bind_inputs/中间数组、runtime.execute、日志归约；同时记录 kernel-only 的 1/2/4/8/16/32/64T 曲线，先确认退化是在 host binder、Numba kernel 还是日志路径。 - 按已解析 plan 做专门化：materialization 时把 active scales、term 顺序、observation offset 固化，生成无运行时
if scales[i]的 kernel；尽量传递直接的 typed array 参数，减少嵌套 tuple/动态索引。 - 减少中间写回：消除
dof_pos_diff等可以在 observation 写出时直接计算的临时数组；对同一输入的多个输出做 fused/common-subexpression 计算，避免重复遍历。 - 优化并行布局：评估
prangechunk/scheduling 和 threading layer；对log_scratch做 cache-line padding 或在 benchmark 中关闭日志归约，确认是否存在 false sharing；按 N 自适应选择线程上限，不能固定 64/128T。 - 提高 e2e 统计可信度：每个配置使用至少 100 个测量 step、多个独立进程/重复 run，报告 median/p95 和 CPU utilization，再决定 collector gate。
- 先修复可用性阻断：为 SAC/APPO/PPO Motrix/MuJoCo owner 配置补齐
term_plan，否则 Motion Tracking 的 Numba e2e 仍会直接构造失败。
当前结论：PR #929 相对 NumPy fallback 有明显热切片收益，但相对旧手写 Numba 的高线程扩展明显回退，且端到端收益有限；建议完成上述分解和优化后再评估合并。
补充复测：165.245.137.171 双路 160 核服务器补充上一条性能评论，在
测试通过 detached worktree 执行，使用服务器已有 venv 并设置 服务器有其他训练任务，测试开始时约占 5–8 个 CPU 核。由于这是双路 NUMA 机器，额外做了：
1. G1 热切片：默认全机调度，3 次独立进程的最佳吞吐中位数统计仍为 reward + termination + actor/critic observation assembly，单位
新方案相对同提交 NumPy fallback 仍有约 10.4×–12.2× 的热切片加速，但不能据此认定相对旧手写 Numba 全面更快：2048 env 有收益，8192 env 有回退。 2. NUMA 敏感性：G1 8192 env
这组数据不能简单解读为 node 0 上“新方案快 2.5×”；它首先证明当前 benchmark/runtime 强烈受 CPU/memory NUMA placement 影响。观察上，新方案在两个 socket 间较稳定，而旧手写路径差异很大；默认调度仍会因为 worker 和 first-touch memory placement 产生变化。 线程扩展也再次显示：新 G1 在这台机器上的最佳点是 8–16T，32T 后没有收益，跨过单 node 的 80T 后继续下降；160T 时 2048/8192 env 只有约 3.20M/6.89M env/s。固定使用 64/128/160T 不合理。 3. Motion Tracking：完整
|
| env 数 | 旧手写 Numba | 新 plan Numba | 新方案吞吐变化 |
|---|---|---|---|
| 2048 | 0.368 ms @16T | 0.365 ms @16T | +0.8% |
| 8192 | 1.054 ms @32t | 0.988 ms @32t | +6.6% |
本机上 Motion 完整数组路径没有复现 10.13.1.125 的 2048-env 小幅回退，但收益也很有限；说明该差异具有硬件/NUMA依赖。
4. G1 collector：固定 node 0、100-step active window
新实现相对自身 NumPy fallback：
| env 数 | NumPy collector | 新 Numba collector | collector 变化 | update_state 变化 |
|---|---|---|---|---|
| 2048 | 80.59K steps/s | 84.54K steps/s | +4.9% | 3.094 → 1.520 ms（2.03×） |
| 8192 | 99.31K steps/s | 106.24K steps/s | +7.0% | 8.648 → 2.793 ms（3.10×） |
直接比较新旧 Numba collector：
| env 数 | 旧手写 Numba collector | 新 Numba collector | 新/旧变化 |
|---|---|---|---|
| 2048 | 88.06K steps/s | 84.54K steps/s | −4.0% |
| 8192 | 106.44K steps/s | 106.24K steps/s | −0.2% |
因此在控制 NUMA 后，新旧 Numba 的真实 collector 基本同档；新方案相对 NumPy 的 update_state 加速明确，但大部分收益仍被 physics、actor/replay 和 bookkeeping 稀释。
5. Motion e2e 阻断仍然存在
sac/g1_motion_tracking/motrixsim 在 PR #929 上仍构造失败：
ValueError: MotionTracking fused Numba execution requires env.term_plan
即 owner config 覆盖问题在第二台服务器上稳定复现，不是 10.13.1.125 的环境偶发问题。
更新后的结论与建议
- 确认的收益：相对 NumPy fallback，G1 热切片约 10×–12×，100-step collector 约 +5%–7%；Motion 完整数组路径在该服务器约 +1%–7%。
- 不能确认的收益：相对旧手写 Numba，结果随 env 规模、socket 和 NUMA placement 改变；G1 8192 默认调度仍回退约 13%，固定 node 后 collector 基本持平。
- 新增判断：此前的高线程回退不仅是 kernel 结构问题，也叠加了未控制的 NUMA/first-touch/worker placement；但即使固定单 node，新方案最佳点仍只有 16–32T，说明并行扩展瓶颈真实存在。
建议在继续优化前先把 benchmark gate 固定为：显式 CPU affinity + memory policy、记录 threading layer/CPU utilization、每配置多进程重复并报告 median；runtime 为每个 num_envs 选择线程上限。代码侧继续拆分 binder/kernel/log 计时，减少中间缓冲和通用 tuple routing，并修复 Motion owner config 的 term_plan 覆盖。
合并建议不变：当前有明确的 NumPy fallback 收益，但没有稳定优于旧手写 Numba，且 Motion e2e 仍不可用，暂不建议合并。
原始结果保存在服务器：/tmp/unilab-pr929-bench.AU2tpk/results/。
补充：
|
| env_num | 10.13.1.125 旧 | 新 | 变化 | 165.245.137.171 旧 | 新 | 变化 |
|---|---|---|---|---|---|---|
| 16,384 | 73.14 | 31.78 | -56.5% | 24.36 | 15.50 | -36.4% |
| 32,768 | 85.20 | 34.36 | -59.7% | 30.06 | 18.90 | -37.1% |
| 65,536 | 80.30 | 35.43 | -55.9% | 28.15 | 23.63 | -16.1% |
| 131,072 | 52.85 | 35.30 | -33.2% | 29.21 | 24.15 | -17.3% |
| 262,144 | 36.96 | 31.24 | -15.5% | 37.05 | 26.06 | -29.7% |
G1 在两台机器、全部 2^{14}–2^{18} 规模均未超过旧手写 kernel。Threadripper 上规模增大后差距收窄，但到 2^{18} 仍慢 15.5%；因此不能把当前 assembled plan 宣称为 G1 的性能改进。
Motion（完整 update_state，ms；新吞吐相对旧）
| env_num | 10.13.1.125 旧 ms | 新 ms | 吞吐变化 | 165.245.137.171 旧 ms | 新 ms | 吞吐变化 |
|---|---|---|---|---|---|---|
| 16,384 | 0.896 | 0.841 | +6.6% | 1.654 | 1.688 | -2.0% |
| 32,768 | 3.005 | 1.673 | +79.6% | 3.803 | 2.537 | +49.9% |
| 65,536 | 7.035 | 3.883 | +81.2% | 6.794 | 4.232 | +60.6% |
| 131,072 | 13.047 | 8.911 | +46.4% | 12.852 | 9.924 | +29.5% |
| 262,144 | 29.572 | 18.771 | +57.5% | 37.275 | 23.259 | +60.3% |
Motion 从 2^{15} 开始在两台服务器均有明确大规模收益（约 +29% 至 +81%）；2^{14} 基本持平。
结论与后续建议
- 性能表现是任务相关的，不能用统一的“assembled Numba 更快”概括：Motion 大规模收益明确，G1 则在所有规模回退。
- G1 应保留/恢复任务专门化优化，优先排查 observation routing、中间 buffer、tuple indirection，以及未能常量化/融合的 kernel 分支。
- Motion 的 preamble/workspace/observation 融合方案在大规模下有效，可作为 G1 优化的参考。
- 合并建议不变：Motion e2e 中
env.term_plan阻断仍未修复，且 G1 相对旧实现存在稳定回退；在这两项解决并重新验证前，不建议合并 PR feat: package issue 918 Numba roadmap independently #929。
NumPy / 手写 Numba / assembled Numba 统一时延对比已将三条路径统一放入同一张表，单位均为 ms/次批量 提交与路径：
物理含义与计时范围G1 walk 每次调用包括：
Motion tracking 每次调用包括：
三条路径都使用同一批已经物化的 deterministic synthetic arrays；因此时间包含上述数组计算、临时/工作区读写和输出数组写入。路径实现差异也包含在内：NumPy 包含 Python reward/observation dispatch 与 NumPy 向量化操作；手写 Numba 包含旧的任务专用手写 kernel；assembled Numba 包含新的 plan runtime dispatch、生成的融合 kernel 以及其运行时 workspace/output 操作。G1 的 NumPy 路径还包含 synthetic backend 的内存数组访问，但不代表真实后端 getter 成本。 以下内容不在计时范围内：MuJoCo/Motrix 物理积分、真实后端状态读取/数据搬运、环境构造、reset、domain randomization/RNG/noise、motion sampling、state replacement、policy inference、collector/replay/learner，以及 Numba 首次 JIT 编译（编译在正式计时前完成）。本轮也没有运行 collector e2e。 统计口径
结论
|
1 participant
Relates to #918
Purpose
This PR repackages the complete #918 Numba roadmap as one independently reviewable change. It is intentionally not merged automatically; maintainers should review and manually merge it into
refactor/issue-902-mjwarp-squashed-delivery.Branch / history separation
feat/issue-918-numba-independentb8ecb79e, the refactor branch state before Roadmap: 在 G1 walk 与 motion tracking 跑通结构化、自动融合的 NumPy/Numba term #918 startedrefactor/issue-902-mjwarp-squashed-deliveryb8ecb79e09702431After manual merge, the refactor branch will contain the mjwarp baseline plus this independently reviewable Numba roadmap.
Included roadmap work
The PR does not add backend, runner,
NpEnvlifecycle, training ABI, checkpoint, sim2sim, or regular CI changes, and does not expand beyond the two pilot task families.Validation
make test-allScope