ASTRA-sim 通信仿真框架

ASTRA-sim 是 Georgia Tech、Meta、Intel 联合开发的分布式 AI 训练仿真器，目前最新版本为 2.2，论文发表于 ISPASS 2023（arXiv:2303.14006），是业界引用最广泛的开源 AI 通信仿真框架。

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

---

整体架构

ASTRA-sim 采用严格三层分离架构，依赖方向单向（Workload → System → Backend，无反向依赖）：

┌────────────────────────────────────────────────────────┐
│  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 后端精确验证，无需修改上层逻辑。

---

Chakra 执行图

工作负载用 Chakra Execution Trace（ET） 表示——一种 Protobuf 序列化的 DAG（有向无环图），每个 NPU 一份 .et 文件，描述该 NPU 上的算子执行顺序和数据依赖关系。

Chakra 由 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	远端内存访问

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 组合

配置示例（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 大小与 Stream 调度

chunk_size = data_size / (preferred_dataset_splits × npus_count)

每个 chunk 一个 StreamBaseline，每个 stream 包含有序的 CollectivePhase 队列（不同维度的通信相位）。inter-dimension-scheduling 控制维度间顺序：Ascending（按维度编号）/ RoundRobin / OfflineGreedy（贪心优化）。

Ring AllReduce 延迟公式

N 节点，数据量 M 时的 Ring AllReduce 延迟：

T_AllReduce = 2(N-1) × (α + M / (N × β))
            = 2(N-1)α + 2(N-1)/N × M/β

带宽利用率 = 2(N-1)/N，N 增大时趋近 100%。

多维拓扑的分相位计算（ASTRA-sim 2.0 核心增强）

对于跨越多个层级的通信组（如节点内 NVLink + 跨节点 InfiniBand），将集合通信分解为多个相位串行执行：

Phase 0（dim 0，Ring，N₀=8，β₀=NVLink 带宽）：

T₀ = 2(N₀-1)α₀ + 2(N₀-1)/N₀ × M/β₀

Phase 1（dim 1，Switch，N₁=16，β₁=IB 带宽）：

T₁ = 2(N₁-1)α₁ + 2(N₁-1)/N₁ × M/β₁

T_total = T₀ + T₁  （串行执行两个相位）

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

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

---

网络后端

Analytical 后端（分析式后端）

核心公式（α-β 模型）：

communication_delay = α + message_size / β

其中：
  α = link.latency（链路基础延迟，单位 ns）
  β = link.bandwidth（链路带宽，单位 GB/s）

两种模式：

模式	描述	精度
congestion_unaware	纯 α-β，不感知链路争用	低（5%–30%+）
congestion_aware	显式链路建模 + 排队延迟	较高

拓扑配置格式（YAML）：

# 2D 分层：NVLink（节点内）+ InfiniBand（节点间），128 NPU
topology:   [ Ring,   Switch ]
npus_count: [ 8,      16    ]   # 8 GPU/节点 × 16 节点 = 128 GPU
bandwidth:  [ 120.0,  20.0  ]   # GB/s，第 0 维 NVLink，第 1 维 IB
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）

初始化时预计算所有节点对之间的通信参数，后续仿真中直接查表，是其秒级执行速度的关键。

在 V100 上的验证：与真实 NCCL 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 对原版 NS-3 后端的改进分析）。

---

精度特性与局限

精度验证条件（原始论文）：

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

重要局限：5% 误差仅在小规模（≤16 GPU）、单操作（AllReduce）、大消息（640 MB+）条件下测量。扩展到更大规模（512 GPU）时，原版 NS-3 后端误差飙升至 530.2%，主要来源是拥塞控制建模不足和 QP 行为简化。

其他已知局限：

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

---

参考资料

• ASTRA-sim 2.0 论文 (arXiv:2303.14006)
• ASTRA-sim Analytical Backend 文档
• ASTRA-sim NS-3 Backend 文档
• ASTRA-sim GitHub
• Analytical Network Source
• TACOS 论文 (MICRO 2024)