机群元数据架构 — 概览
Megascale 机群元数据架构是为多切片 TPU 机群命名并描述它的数据模型。它位于每切片 ICI 互连(由 tpunetd 配置的环面芯片网络)之上,位于 XLA collectives 层之下。它为运行时回答四个问题:
- 机群中有什么? — 有多少个切片、每个切片有多少台主机、每个切片是什么形状。
- 我是谁? — 此进程的
(slice_id, host_id)及其incarnation_id。 - 如何到达对等方? — 每主机 DCN 端点表。
- 如何在机群范围内为芯片命名? — 从扁平设备 id 到
(slice, logical-device, core-type)地址的转换。
两个坐标系统
该架构有意将两个坐标系统分开,并且只在一个点将它们连接起来。
| 层 | 包 | 所有者 | 芯片命名方式 |
|---|---|---|---|
| DCN / 跨切片 | xla.megascale.runtime.* | Megascale coordinator | (slice_id, host_id) + 每切片逻辑设备 id |
| ICI / 切片内 | superpod.routing.* + tpu.* | tpunetd | ToroidalTopology 中的 ChipCoordinate(int64 向量) |
从 ICI 层跨到 DCN 层的唯一数据项是每切片形状,它编码为 TpuTopologyArgsProto,并承载在每个 SliceInfo 中。DCN 层把该形状视为一个经过验证的不透明 blob,它从不重新推导芯片坐标。
权威对象
实时的内存中机群模型是一个 C++ 类:
text
xla::megascale::runtime::MultiSliceTopologyAndLocation
```text
请注意,名称**没有** `Proto` 后缀。它的序列化形式是 `MultiSliceTopologyAndLocationProto`。所有需要机群视图的对象,包括 `Communicator`、`MegaScaleMultiSliceConfig`、`PjrtExecutableContext`、编译器调度器以及跨主机 barrier,都会通过 `const&` 或 `shared_ptr<const ...>` 持有此类。
它有两种构建方式:
- `FromProto(MultiSliceTopologyAndLocationProto)` — 反序列化 coordinator 在 bootstrap 时广播的内容。
- ctor `(local_slice_id, local_host_id, …, incarnation_id)`,随后调用 `AddSlice(slice_id, shared_ptr<const tpu::TpuTopology>)` — 增量组装。
并通过 `ToProto()` 序列化回去。
## 各部分所在位置
| 关注点 | 页面 |
|---------|------|
| fleet → slice → host → chip → core 层级 | [拓扑模型](topology-model.md) |
| `(slice_id, host_id, incarnation_id)` 身份 | [主机身份](host-identity.md) |
| 扁平设备 id → `(slice, device, core)` | [全局寻址](global-addressing.md) |
| 两个坐标系统的比较 | [ICI 与 DCN](ici-vs-dcn.md) |
| 3D 环面切片描述符 | [切片形状](slice-shape.md) |
| 每台主机在启动时发布的内容 | [Bootstrap 交换](bootstrap-exchange.md) |
| barrier / 错误聚合器如何使用它 | [Barrier 与错误用法](barrier-error-usage.md) |
| 逐消息 wire 解码 | [字段解码](field-decode.md) |
## 来源
该架构是从 `libtpu.so`(libtpu-0.0.40)内部的 `protodesc_cold` FileDescriptorProto 池中恢复的,并与链接的 C++ 类符号、访问器方法名称以及 rodata 类型字符串进行了交叉核对。机群元数据分布在四个跨切片 proto 文件(`transport.proto`、`common/runtime.proto`、`megascale/addresses.proto`、`megascale/dcn_topology.proto`)以及每切片的 `tpu_topology_args.proto` 中。