AICB开启和关闭AIOB功能生成负载文件逻辑不一致

[jayjay-cc]

SimAI/aicb/workload_generator/SimAI_training_workload_generator.py中，同样的不启用SP，启用AIOB时backward_comm = "NONE"，而不开AIOB时backward_comm = "ALLREDUCE"，哪一种情况才是对的？

[tianhao909]

【人工回复】

应该要参考你具体想仿真的megatron框架的实现情况再对齐一下， 目前是0.8core版本，可能有些陈旧， 你们可以参考megatron框架更新一下，以真实框架的情况为准

[tianhao909]

【Agent A回复】【仅供参考】

🇨🇳 中文回复版本

1. 需求背景总结
在此 issue 中，用户发现在 SimAI/aicb/workload_generator/SimAI_training_workload_generator.py 内部生成层级负载时，对于未开启 Sequence Parallelism (SP) 的场景，开启 AIOB 和不开启 AIOB 时反向传播的张量并行通信（backward_comm）被赋予了不同的值。不开启 AIOB 时为 "ALLREDUCE"，而开启 AIOB 时却为 "NONE"。用户希望确认对于非 SP 场景，哪一种设定才是正确的。

2. 涉及文件/模块

• 文件：aicb/workload_generator/SimAI_training_workload_generator.py
• 模块：SIMAI_workload 类的 workload_generate_aiob() 和 workload_generate() 方法。
• 具体位置：处理非序列并行 (enable_sequence_parallel == False) 情况的 else 分支逻辑（约第 501 行）。

3. 根因分析或实现思路

• 结论：不开 AIOB 时的 ALLREDUCE 才是正确的。
• 根本原因：在经典 1D Tensor Parallelism （张量并行且未开启 SP）的 Transformer 块（如 Attention 或 MLP）设计中：
  • 前向传播 (Forward)：Column Linear 在各个 GPU 上独立计算（不通信），然后 Row Linear 计算完毕后，将切片结果累加，因此会执行一次 ALLREDUCE。
  • 反向传播 (Backward)：Row Linear 对权重求导不产生通信，但 Column Linear 对上一层输入的梯度求导后，其结果是对特征维度的分片（partial sums），必须跨卡执行一次 ALLREDUCE 才能进行完整的前部分网络反向计算。
• 不论是否开启用于优化 DP/ZeRO 通信开销的 AIOB，单个模型块内部针对 TP 相关的输入张量反向通信步骤都是不可省略的。开启 AIOB 时将 backward_comm 设为 "NONE" 属于代码中的遗漏 bug，这会导致生成的记录少算了一次反向传播时的通信开销。

4. 最小修改方案
修改 aicb/workload_generator/SimAI_training_workload_generator.py 函数 workload_generate_aiob() 中非 SP 场景相关的代码分支逻辑：

修改前：

                    else:
                        forward_comm = "ALLREDUCE"
                        backward_comm = "NONE"

修改后：

                    else:
                        forward_comm = "ALLREDUCE"
                        backward_comm = "ALLREDUCE"

5. 测试前置条件

• 此测试并不依赖任何额外的模型或通信数据。
• 仅需使用原有的工作流脚本，通过控制命令变量进行比对生成：
  • 设置 --tensor-model-parallel-size > 1 (如 2 或 4)。
  • 务必不添加 --enable-sequence-parallel 标志。
  • 加上和去掉 --aiob-enable 进行两次脚本调用（生成两个 .txt 负载文件）。
  • 对比检查两个生成的负载日志文件中对应诸如 attention, mlp 等 Work_Item 的相关行，确保 backward_comm 字段是否对齐并均正确显示为 ALLREDUCE。

6. README 是否需要修改及建议内容

• 不需要修改 README。 该问题系通信开销建模逻辑缺陷的内部代码修正，无需对外增加功能开关或更改配置文档。

---

🇬🇧 English Reply Version

1. Requirements Background Summary
The user identified an inconsistency within SimAI_training_workload_generator.py when generating workload traces for the non-Sequence Parallelism (SP) scenario. Specifically, when AIOB is enabled, the backward pass tensor parallel communication (backward_comm) is set to "NONE", whereas it is correctly set to "ALLREDUCE" when AIOB is disabled. The goal is to clarify which logic mathematically constructs the accurate communication overhead for non-SP 1D Tensor Parallel routines.

2. Involved Files/Modules

• File: aicb/workload_generator/SimAI_training_workload_generator.py
• Modules/Classes: The workload_generate_aiob() and workload_generate() methods inside the SIMAI_workload class.
• Specific location: The else logic block handling the non-SP scenario (enable_sequence_parallel == False, around line 501).

3. Root Cause Analysis & Implementation Logic

• Conclusion: Setting backward_comm = "ALLREDUCE" (the logic without AIOB) is the logically correct behavior.
• Root Cause: Under underneath a traditional 1D Tensor Parallelism configuration inside a Transformer block (Attention or MLP) with SP disabled:
  • Forward Pass: The Column Linear operation applies independently, but the output from the subsequent Row Linear layer produces partial sums which necessitate one ALLREDUCE.
  • Backward Pass: While Row Linear backpropagation requires no synchronization, computing the gradients corresponding to the inputs entering the Column Linear generates partial derivative splits across GPUs. Another synchronous ALLREDUCE is essential to merge these gradients.
• Therefore, the requirement for a Tensor Parallel backward input synchronization remains unconditionally true, regardless of AIOB configuration (which primarily handles Data Parallel optimizations). Bypassing this by passing "NONE" inside workload_generate_aiob() is a bug resulting in underestimated communication simulations.

4. Minimal Modification Plan
Modify the else branch representing the lack of sequence parallelism logic in the workload_generate_aiob() method located in aicb/workload_generator/SimAI_training_workload_generator.py:

Before Modification:

                    else:
                        forward_comm = "ALLREDUCE"
                        backward_comm = "NONE"

After Modification:

                    else:
                        forward_comm = "ALLREDUCE"
                        backward_comm = "ALLREDUCE"

5. Testing Prerequisites

• No external model weights or real runtime collective data are necessary.
• Generating and verifying outcomes involves merely executing the generator script:
  • Configure runtime setup enforcing Tensor Parallelism (e.g., --tensor-model-parallel-size 2).
  • Do strictly NOT pass the --enable-sequence-parallel flag.
  • Evaluate tracing workloads sequentially with and without --aiob-enable.
  • Compare the resultant output .txt trace logs to assure the backward_comm property corresponds identically to ALLREDUCE for essential layers irrespective of the AIOB status.

6. README Modification Check

• No README updates are necessary. This minor patch repairs an internal communication omission calculation and does not inherently change any application behavior or introduce new flags.

[tianhao909]

【Agent B回复】【仅供参考】

结构化分析（中文）

1. 需求背景总结

用户在 不启用序列并行（SP）、其它条件相同的前提下，对比 开启 AIOB 与 关闭 AIOB 两种路径生成的 SimAI 负载里 backward_comm 字段：前者某些行为 NONE，后者为 ALLREDUCE，希望弄清 哪一种语义正确，以及这是否属于生成逻辑不一致/缺陷。

SimAI 仓库中 aicb 为 git submodule，实际源码在 aliyun/aicb；分析以 aicb master 上该文件为准（与 issue 中路径一致）。

2. 涉及文件/模块

位置	说明
aicb/workload_generator/SimAI_training_workload_generator.py	SIMAI_workload.workload_generate()（aiob_enable=False）与 workload_generate_aiob()（aiob_enable=True）；入口 if __name__ == "__main__" 根据 args.aiob_enable 二选一
同文件 Work_Item dataclass	forward_comm / backward_comm / *_comm_size 等字段定义
（若做仿真验证）SimAI/astra-sim 侧 workload 解析与统计逻辑	需确认是否对 forward_comm 与 backward_comm 分别累加，以判断是否重复计数

3. 根因分析或实现思路

需要 先区分两类「行」，否则容易误判为矛盾：

（A）layernorm 这一行（无 SP 时插入的 layernorm grad sync）

在 当前 aicb master 中，两条路径对 name=="layernorm" 且 not enable_sequence_parallel 时 一致 使用：

• backward_comm="ALLREDUCE"
• backward_comm_size=2 * total_params

即：若 issue 截图里 AIOB 侧 layernorm 仍是 NONE，更可能是旧版本、或对比时行没对齐（看成了别的行）。

（B）Transformer 层循环里、无 SP 的分支（else，与 SP 分支相对）

这里存在 编码约定差异：

• workload_generate_aiob（约 488–495 行）：tensor_model_parallel_size > 1 时
  forward_comm="ALLREDUCE"，backward_comm="NONE"（TP>1 时把 TP 相关 allreduce 主要写在 forward 侧）。
• workload_generate（约 760–763 行）：无 SP 的 else 里
  forward_comm="ALLREDUCE" 且 backward_comm="ALLREDUCE"。

因此，在 同一层、无 SP、TP>1 时，会出现你观察到的现象：AIOB 路径 backward_comm=NONE，非 AIOB 路径 backward_comm=ALLREDUCE。这不是「少建模了一次通信」，而是 同一类 TP 通信放在 forward_comm 还是同时放在两侧 的 表示法不一致。

语义上哪种「对」，取决于 SimAI 仿真器是否把 forward_comm 与 backward_comm 各算一遍网络负载：

• 若 分别累加：非 AIOB 路径对同一层写两次 ALLREDUCE 可能 重复计数；AIOB 路径（仅 forward 为 ALLREDUCE）更接近 单次 collective。
• 若 只对某一侧计数或有特殊规则：需对照 astra-sim/SimAI 解析代码才能下结论。

（C）阶段前置：grad / grad_norm 结构不同

• AIOB：拆成 grad_gather、grad_param_comm、grad_param_compute 等，大量 backward_comm="NONE"，数据并行相关 collective 在 dp_comm 等字段。
• 非 AIOB：以 grad_norm 等单行表达，通信拆在不同字段。

因此 前几行 也会出现「一边全是 NONE、另一边不是」——这往往是 阶段拆分方式不同，不是同一行的直接对比。

4. 最小修改方案（规划，不写具体 patch）

1. 先定规范：在 aicb 或 SimAI 文档中明确：无 SP 时，TP 的 allreduce 应 只出现在 forward_comm、还是 backward_comm、或两者各表示前后向之一。
2. 再二选一统一实现（改动尽量局部、只动 SimAI_training_workload_generator.py 中无 SP 的 else 分支之一）：
   • 方案 A：把 workload_generate 的无 SP 分支改成与 workload_generate_aiob 一致（例如 TP>1 时 backward_comm="NONE"），避免潜在双重 ALLREDUCE。
   • 方案 B：把 workload_generate_aiob 改成与 workload_generate 一致（两侧皆 ALLREDUCE），若仿真器只认这种写法。
3. 建议：在改代码前 在 SimAI 仓库中 grep/阅读 workload 解析处对 forward_comm/backward_comm 的使用，用 单次最小用例（固定 TP、无 SP）跑仿真对比总通信量，再定方案 A 或 B。

5. 测试前置条件

• 不依赖 大规模真实集群；需要 可复现的 CLI 参数（同一模型、enable_sequence_parallel=False、固定 tensor_model_parallel_size、有无 --aiob_enable）。
• 若开 AIOB：需要 comp_filepath / 算子耗时文件（extract_averages 等），否则 AIOB 路径可能退化为默认时间，但不影响 对比 backward_comm 字符串 本身。
• 建议：对同一参数生成两个 .txt workload，diff 或通过脚本只比较 layernorm 行与「无 SP 下某一 transformer 层」行。

6. README 是否需要修改及建议内容

• 需要（建议）：若在 issue 或代码中确认「无 SP + TP>1」应用 统一编码，应在 aicb README 或 SimAI 文档 中用一小节说明：Work_Item 中 TP allreduce 与 forward_comm/backward_comm 的对应关系，以及 AIOB 与 non-AIOB 是否必须逐字段一致。
• 非必须：若仅内部约定、且仿真器已兼容两种历史格式，可只在 issue/注释中说明。

7. 缺失信息与假设

• 缺失：SimAI/astra-sim 解析 workload 文本时 对 forward_comm 与 backward_comm 的 精确语义（是否求和、是否与 dp_comm 独立）。
• 假设：issue 中截图基于的 aicb commit 与当前 master 可能不同；以仓库 HEAD 为准时，layernorm 行在两条路径下已对齐为 ALLREDUCE。
• 假设：用户问的「不启用 SP」同时包含 TP=1 与 TP>1；TP=1 时 AIOB 无 SP 分支里 forward_comm/backward_comm 常为 NONE，与「ALLREDUCE」对比的是 不同配置。

---

Structured Analysis (English)

1. Background

The author compares SimAI workload files generated with AIOB enabled vs disabled under sequence parallel (SP) off, and notices backward_comm is sometimes NONE vs ALLREDUCE. They ask which behavior is correct and whether this is an inconsistency. The canonical source for SimAI/aicb/... is the aliyun/aicb submodule; this analysis uses SimAI_training_workload_generator.py on aicb master.

Note: Issue #251 has no comments; the “latest question” is the issue body itself.

2. Files / modules

• aicb/workload_generator/SimAI_training_workload_generator.py: workload_generate() vs workload_generate_aiob(), selected by args.aiob_enable in __main__.
• Work_Item fields: forward_comm, backward_comm, sizes, dp_comm, etc.
• SimAI/astra-sim workload consumer (if validating end-to-end): how it aggregates forward vs backward comm fields.

3. Root cause / reasoning

Compare the right rows:

(A) layernorm row (inserted when not enable_sequence_parallel)
On current aicb master, both paths use backward_comm="ALLREDUCE" and backward_comm_size=2*total_params for that row. So if screenshots show NONE for AIOB on layernorm, that may be an older revision or row misalignment.

(B) Non-SP transformer layer loop (else vs SP branch)
Here lies the real mismatch when TP > 1:

• workload_generate_aiob: forward_comm="ALLREDUCE", backward_comm="NONE".
• workload_generate: forward_comm="ALLREDUCE" and backward_comm="ALLREDUCE".

So for the same layer, SP off, TP>1, you see exactly NONE vs ALLREDUCE on backward_comm. That is mainly a representation split (TP collective attributed to forward vs duplicated on both sides), not necessarily “missing comm” in AIOB—check whether the simulator double-counts when both are ALLREDUCE.

(C) Preamble (grad_* vs grad_norm)
AIOB splits the grad phase into multiple items with dp_comm; the non-AIOB path uses a different shape. Early lines will differ even when both are “correct” structurally.

4. Minimal change plan (no code)

1. Define the intended convention for non-SP TP layers in Work_Item (single field vs dual ALLREDUCE).
2. Align one of: workload_generate’s non-SP else block or workload_generate_aiob’s non-SP else block with the other, after confirming simulator semantics (avoid double-counting).
3. Validate with one fixed config (SP off, fixed TP) and two generated files + optional sim totals.

5. Test prerequisites

• Same CLI except aiob_enable; SP off; fixed TP.
• AIOB may need compute trace file for realistic timings; not required for string-level backward_comm comparison.
• Diff or script-based comparison of named rows (layernorm, one transformer layer).

6. README

• Recommended: document the forward_comm / backward_comm convention for TP without SP and state whether AIOB/non-AIOB outputs must match field-for-field.

7. Gaps and assumptions

• Gap: Exact SimAI parser behavior for summing forward_comm and backward_comm.
• Assumption: Issue screenshots may not match current master for layernorm.
• Assumption: TP=1 yields NONE on both comm fields in the AIOB non-SP branch—comparing that to ALLREDUCE mixes different hyperparameters.

[tianhao909]

【Agent C回复】【仅供参考】

分析结论（仅分析，暂不改代码）

背景理解

从 issue 描述看，当前问题是：

• 在 不开启 SP 的情况下，
• 开启 AIOB 生成的 workload 中 backward_comm = "NONE"
• 关闭 AIOB 生成的 workload 中 backward_comm = "ALLREDUCE"

如果除了 AIOB 开关外，其余并行配置完全一致，那么这两种结果在语义上不应该同时成立。因为 AIOB 更像是 compute/profile 侧增强功能，正常情况下应影响的是算子时间/profile 数据，而不应改变同一并行配置下的通信类型定义。

初步判断

基于当前信息，更可能是 workload generator 在 AIOB / 非 AIOB 两条路径中的通信字段生成逻辑不一致，而不是 AIOB 本身应该改变 backward communication 语义。

也就是说，更合理的设计应当是：

• backward_comm 由并行策略决定（如 TP / SP / EP / DP 等）
• AIOB 只影响 compute 时间来源
• 在相同并行配置下，AIOB 开关不应单独导致 NONE 与 ALLREDUCE 的切换

当前推测的根因

更像是以下两种情况之一：

1. AIOB 分支绕过了公共通信推导逻辑，导致 backward_comm 落回默认值 NONE
2. AIOB 分支和非 AIOB 分支各自维护了一套 if/else，在某个条件分支上发生了漂移

最小修改建议

建议优先检查 aicb/workload_generator/SimAI_training_workload_generator.py 中：

• 非 AIOB 路径如何设置 backward_comm
• AIOB 路径如何设置 backward_comm

最小修复方式建议为：

• 抽出统一的 backward_comm 推导逻辑
• AIOB / 非 AIOB 共用该逻辑
• AIOB 仅保留 compute/profile 相关差异

这样改动最小，也最符合职责边界。

还缺少的信息

为了进一步确认哪一种输出“理论上正确”，还需要补充以下信息：

1. issue 截图对应的是哪一类层（attention / mlp / moe / column / row 等）？
2. backward_comm 在生成器中最终映射的是 input_grad_comm 还是 weight_grad_comm？
3. 对应的完整生成命令是什么？
4. AIOB 模式下是否读取了额外 profile 文件？

当前结论

在现有信息下，更倾向于认为这是一个生成逻辑不一致的问题。
也就是说，AIOB 开关本身不应决定 backward_comm 是否为 ALLREDUCE；如果同配置下出现这种差异，更像是实现上的分支不一致，需要统一逻辑。

Analysis Summary (plan only, no code changes yet)

Background understanding

From the issue description, the current problem is:

• with SP disabled,
• when AIOB is enabled, the generated workload uses backward_comm = "NONE"
• when AIOB is disabled, the generated workload uses backward_comm = "ALLREDUCE"

If all parallelism settings are otherwise identical, these two results should not both be correct semantically. This is because AIOB appears to be a compute/profile-side enhancement, so under normal design it should affect operator timing/profile data, rather than change the communication type definition under the same parallel configuration.

Initial judgment

Based on the currently available information, it is more likely that the workload generator has inconsistent communication-field generation logic between the AIOB and non-AIOB paths, rather than AIOB being intended to change backward communication semantics.

A more reasonable design would be:

• backward_comm is determined by the parallelism strategy (TP / SP / EP / DP, etc.)
• AIOB only affects the source of compute timing
• under the same parallel configuration, toggling AIOB should not independently switch backward_comm between NONE and ALLREDUCE

Suspected root cause

This looks more like one of the following:

1. the AIOB branch bypasses the shared communication-derivation logic, causing backward_comm to fall back to the default NONE
2. the AIOB and non-AIOB branches maintain separate if/else logic, and one branch has drifted

Minimal fix suggestion

The first place to inspect should be aicb/workload_generator/SimAI_training_workload_generator.py, specifically:

• how backward_comm is set in the non-AIOB path
• how backward_comm is set in the AIOB path

The recommended minimal fix is:

• extract one shared backward_comm derivation logic
• make both AIOB and non-AIOB paths reuse it
• keep AIOB-specific differences limited to compute/profile handling

This gives the smallest change and the cleanest responsibility boundary.

Missing information

To further confirm which output is “theoretically correct”, the following information is still needed:

1. what exact layer type is shown in the screenshots (attention / mlp / moe / column / row, etc.)?
2. does backward_comm eventually map to input_grad_comm or weight_grad_comm in the generator?
3. what is the exact full generation command?
4. does AIOB mode read any additional profile file?

Current conclusion

With the currently available information, this looks more like a workload-generation inconsistency bug.
In other words, the AIOB switch itself should not decide whether backward_comm is ALLREDUCE or not; if this difference appears under the same configuration, it is more likely caused by diverged implementation logic that should be unified.