ASTRA-sim

三层架构如何支持可插拔仿真后端

核心要点：

三层分离：Workload (Chakra ET) → System （集合通信调度） → Backend （网络后端），单向依赖

Chakra ET: MLCommons 标准 DAG 格式，每 NPU 一份 .et, COMM_COLL_NODE 携带 involved_dimensions 实现维度感知

可插拔后端：Analytical (秒级 $\alpha$-$\beta$) / NS-3 （分钟级包仿真） / Garnet (NoC)

多维拓扑：层级互联映射为维度数组，每维独立 $\alpha$/$\beta$，跨维通信分相位串行

精度边界：V100 4-16 GPU 大消息 AllReduce 5%，扩展到 512 GPU 时 NS-3 后端误差 530%+

ASTRA-sim 是 Georgia Tech / Meta / Intel 联合开发的分布式 AI 训练仿真器（最新 2.2, ISPASS 2023[1]），是业界引用最广泛的开源 AI 通信仿真框架。

GitHub: https://github.com/astra-sim/astra-sim

三层架构是什么 / 为什么有效

核心问题：三层分离的具体职责怎么划？解决了什么工程问题？

架构图：

┌────────────────────────────────────────────────────────┐
│  Workload Layer                                        │
│  Chakra ET (.et, protobuf DAG) → ETFeeder → Workload   │
├────────────────────────────────────────────────────────┤
│  System Layer                                          │
│  Sys 类: 集合通信分解 / Stream 调度 / 事件管理         │
├────────────────────────────────────────────────────────┤
│  Backend Layer (网络 / 计算 / 内存)                    │
│  AstraNetworkAPI / AstraComputeAPI / AstraRemoteMemAPI │
│  三种实现: Analytical (快) / NS-3 (慢精确) / Garnet    │
└────────────────────────────────────────────────────────┘

核心价值：同一工作负载可用秒级 Analytical 后端快速探索，也可切换到分钟级 NS-3 精确验证，无需修改上层逻辑。依赖方向单向（Workload → System → Backend，无反向），后端替换不影响工作负载。

Chakra ET 怎么描述工作负载

核心问题：怎么把 LLM 训练 / 推理工作负载转成 DAG？维度感知通信怎么实现？

Chakra Execution Trace 是 Protobuf 序列化的 DAG，每 NPU 一份 .et 文件，描述该 NPU 上的算子执行顺序和数据依赖。由 Meta 发起，现为 MLCommons 标准，可跨仿真器共享。

关键节点类型见 :

节点类型	描述
`COMP_NODE`	计算核，含 FLOP 数和 tensor 大小
`COMM_COLL_NODE`	集合通信 (AllReduce / AllGather / ReduceScatter / AllToAll)
`COMM_SEND_NODE` / `COMM_RECV_NODE`	点对点通信
`MEM_LOAD_NODE` / `MEM_STORE_NODE`	远端内存访问

@tbl-sim-astra-chakra-nodes Chakra ET 节点类型

COMM_COLL_NODE 的关键字段：

comm_type:              ALL_REDUCE | ALL_GATHER | REDUCE_SCATTER | ALL_TO_ALL
comm_size:              数据量 (字节)
involved_dimensions:    参与的网络维度列表, 如 [0, 1]
communicator_group_id:  参与的 NPU 子集

维度感知是关键：involved_dimensions 让 TP 的 AllReduce 只在芯片级维度 (dim 0) 运行，DP 的 AllReduce 只在跨机架维度 (dim 2) 运行，避免"所有 AllReduce 用同一带宽"的误差。

DAG 生成工具链：

PyTorch 执行迹 → chakra_converter → Chakra ET
MSCCLang DSL 可描述集合通信算法并生成 Chakra 表示
RAPID-LLM / DeepFlow：从模型规格自动生成 ET

集合通信怎么分解

核心问题：System Layer 怎么把一个 COMM_COLL_NODE 变成具体的 P2P send/recv 序列？

算法选择优先级：节点级自定义 > 全局自定义（TACOS / MSCCLang 合成） > 内置算法。

内置算法：

算法	适用拓扑	适用场景
`ring`	Ring	大数据量，带宽最优，$2(N-1)$ 步
`doubleBinaryTree`	任意	小数据量，$O(\log N)$ 步，低延迟
`halvingDoubling`	FullyConnected	2 的幂次节点数
`direct`	Switch	交换机拓扑，单跳直达，in-network aggregate
`hierarchicalRing`	多维拓扑	分层 Ring 组合

@tbl-sim-astra-algo-builtin 内置集合通信算法

配置示例 (system.json，每种集合通信独立指定算法):

{
  "all-reduce-implementation": "ring",
  "all-gather-implementation": "direct",
  "reduce-scatter-implementation": "direct",
  "all-to-all-implementation": "ring",
  "scheduling-policy": "FIFO",
  "preferred-dataset-splits": 64,
  "active-chunks-per-dimension": 1,
  "inter-dimension-scheduling": "baseline"
}

Chunk 大小：chunk_size = data_size / (preferred_dataset_splits × npus_count)。每个 chunk 一个 StreamBaseline，每个 stream 含有序的 CollectivePhase 队列（不同维度的通信相位）。inter-dimension-scheduling 控制维度间顺序：Ascending / RoundRobin / OfflineGreedy。

Ring AllReduce 延迟公式：

$$\begin{equation} T_{\text{AllReduce}} = 2(N-1)\alpha + \frac{2(N-1)}{N} \cdot \frac{M}{\beta} \label{eq:sim-astra-ring-ar} \end{equation}$$

带宽利用率 $2(N-1)/N \to 1$ ($N$ 增大时)。

多维拓扑怎么建模

核心问题：跨越多个层级 (NVLink + IB) 的通信怎么分相位串行？

ASTRA-sim 2.0 核心增强：将集合通信分解为多个相位串行执行。

Phase 0 (dim 0, Ring, $N_0=8$, $\beta_0$=NVLink 带宽):

$$\begin{equation} T_0 = 2(N_0-1)\alpha_0 + \frac{2(N_0-1)}{N_0} \cdot \frac{M}{\beta_0} \label{eq:sim-astra-phase-nvlink} \end{equation}$$

Phase 1 (dim 1, Switch, $N_1=16$, $\beta_1$=IB 带宽):

$$\begin{equation} T_1 = 2(N_1-1)\alpha_1 + \frac{2(N_1-1)}{N_1} \cdot \frac{M}{\beta_1} \label{eq:sim-astra-phase-ib} \end{equation}$$ $$\begin{equation} T_{\text{total}} = T_0 + T_1 \label{eq:sim-astra-phase-total} \end{equation}$$

并行策略与维度的典型映射：

TP （张量并行） → dim 0：高带宽 NVLink 维度（8 GPU / 节点）
PP （流水线并行） → dim 1: InfiniBand 跨节点
DP （数据并行） → dim 2：集群间 DCN

网络后端是怎么实现的

核心问题：Analytical 和 NS-3 后端各自的输入格式和精度边界？

Analytical 后端

核心公式 ($\alpha$-$\beta$ 模型):

$$\begin{equation} T_{\text{comm}} = \alpha + \frac{M}{\beta}, \quad \alpha = \text{link.latency (ns)}, \quad \beta = \text{link.bandwidth (GB/s)} \label{eq:sim-astra-analytical} \end{equation}$$

两种模式：

模式	描述	精度
`congestion_unaware`	纯 $\alpha$-$\beta$，不感知链路争用	低 (5%–30%+)
`congestion_aware`	显式链路建模 + 排队延迟	较高

@tbl-sim-astra-analytical-modes Analytical 后端两种模式

拓扑配置格式 (YAML):

# 2D 分层: NVLink (节点内) + InfiniBand (节点间), 128 NPU
topology:   [ Ring,   Switch ]
npus_count: [ 8,      16    ]   # 8 GPU/节点 × 16 节点
bandwidth:  [ 120.0,  20.0  ]   # GB/s
latency:    [ 200.0,  1600.0]   # ns

# 3D Torus, 64 NPU
topology:   [ Ring,   Ring,  Ring  ]
npus_count: [ 4,      4,     4     ]
bandwidth:  [ 100.0,  50.0,  25.0  ]
latency:    [ 100.0,  500.0, 2000.0]

三种网络构建块：

构建块	拓扑形状	推荐集合算法
`Ring`	每节点 2 个邻居，环形	`ring` （带宽最优）
`FullyConnected`	全连接，$N-1$ 条链路	`halvingDoubling`
`Switch`	中心交换机，单跳到所有节点	`direct` (in-network aggregate)

@tbl-sim-astra-analytical-blocks Analytical 后端三种网络构建块

初始化时预计算所有节点对的通信参数，仿真中直接查表 — 这是秒级速度的关键。V100 验证 AllReduce 平均误差 5%，比 NS-3 后端快 756× （1.70 秒 vs 21 分钟）。

NS-3 后端

物理拓扑格式（文本，每行一条链路）：

# <src> <dst> <bw_Gbps> <latency_μs> <err_rate> <queue_size>
0 1 100 1.0 0.0 1024000

逻辑拓扑 (JSON): {"logical-dims": ["8", "16"]}

集成接口 (AstraNetworkAPI):

sim_send(tag, src, dst, size, callback) — 触发 NS-3 包级仿真
sim_recv(tag, src, dst, size, callback) — 注册接收回调

支持 TCP/IP 协议栈、拥塞控制 (DCQCN / ECN)、精细包级行为。代码基于阿里巴巴 HPCC 项目，拥塞控制建模精度有限 (SimAI 已在此基础上改进，见 7.3 SimAI)。

精度在什么范围可信

核心问题：5% 误差是在什么条件下测的？扩展规模后精度怎么变？

原始论文验证条件：

参数	值
测试硬件	NVIDIA V100, NVLink, 150 GB/s 双向
集群规模	仅 4 GPU 和 16 GPU
集合通信操作	仅 AllReduce （Ring 算法）
消息大小范围	640 MB – 1.5 GB
平均误差	5% （Analytical 后端）

@tbl-sim-astra-validation 原始论文精度验证条件

重要边界：5% 误差仅在小规模 (≤16 GPU) / 单操作 (AllReduce) / 大消息 (640 MB+) 下测量。扩展到 512 GPU，原版 NS-3 后端误差飙升至 530.2%，主要来自拥塞控制建模不足和 QP 行为简化。

其他已知局限：

无异构硬件支持（同构集群假设）
计算通信重叠建模不完整
无 NVLS (NVLink-Sharp) 等节点内 in-network compute 支持
无 MoE 专家并行的完整 EP AllToAll 建模
工作负载来源依赖手写或脚本生成，无真实 GPU profiling

这些缺口正是 SimAI 在 ASTRA-sim 基础上重点补强的方向。

Takeaway

知识点	核心结论
三层架构	Workload (Chakra) → System （集合通信） → Backend （网络）单向依赖，后端可插拔
Chakra ET	MLCommons 标准 DAG, `involved_dimensions` 实现维度感知
集合通信分解	5 种内置算法 + Chunk + Stream 调度；`inter-dimension-scheduling` 控相位顺序
多维拓扑	跨层通信分相位串行，$T = \sum_i T_i$
Analytical 后端	预计算节点对参数，秒级速度，V100 大消息 AllReduce 5% 误差
NS-3 后端	包级 RoCEv2 / DCQCN，但拥塞控制建模简化，512 GPU 时 530%+ 误差
精度边界	≤16 GPU / 大消息 / 单 AllReduce 才达到 5%；大规模 / 多操作精度急剧下降

参考资料

Won et al., ASTRA-sim2.0: Modeling Hierarchical Networks and Disaggregated Systems for Large Model Training at Scale, ISPASS 2023. https://arxiv.org/abs/2303.14006

三层架构是什么 / 为什么有效​

Chakra ET 怎么描述工作负载​

集合通信怎么分解​

多维拓扑怎么建模​

网络后端是怎么实现的​

Analytical 后端​

NS-3 后端​

精度在什么范围可信​

Takeaway​

参考资料​