ggaaooppeenngg

为什么计算机科学是无限的但生命是有限的

AI Infra 的一些挑战

发表于 更新于 Disqus：

从 Grok-1.5 Infra 的技术报告中可以窥见，Grok-1.5 在基础设施方面具有以下核心优势：

  1. 先进的分布式训练框架：基于 JAX、Rust 和 Kubernetes 的技术栈，不仅确保了高性能，还能快速适配和训练新的模型架构。
  2. 卓越的可靠性和可用性：通过自研的训练协调器，系统能够智能地检测并隔离故障节点，大幅降低训练任务中断的风险。
  3. 高效的存储与数据处理：在检查点存储、数据加载和训练作业重启等环节都进行了深度优化，将训练过程中的停机时间降至最低。

另一个典型案例是 Meta 的 Research Super Compute (RSC) 超算集群：

  1. 算力规模：
    • 早期配置：760台 NVIDIA DGX A100 系统，总计 6,080 个 GPU
    • 当前规模：已升级至 24,576 张 H100 GPU，算力获得质的飞跃
  2. 网络互联：采用双网络方案
    • NVIDIA Quantum InfiniBand，带宽高达 1600 Gb/s
    • RoCE（RDMA over Converged Ethernet）作为补充互联方案
  3. 存储系统：采用自研的 Tectonic 文件系统，通过 FUSE 提供标准的 Linux 文件系统接口，确保高效的数据访问

从这些实践可以看出，现代 AI 基础设施主要围绕三大核心要素展开：

  • 计算能力（以 GPU 为核心）
  • 网络互联（RoCE 或 InfiniBand）
  • 存储系统

而在上层的编排调度领域，系统的容错能力和可靠性则成为关键考量因素。

GPU 架构与互联

在当前AI训练领域，主流的GPU型号主要是NVIDIA的A100、H100和H200系列，它们按照发布时间依次提供了更强大的算力和更优化的架构设计。关于GPU的详细架构，特别是其拓扑结构，可以参考这篇深度解析文章。

GPU互联拓扑

GPU之间的互联拓扑结构主要取决于不同总线间的传输特性，GPU之间可以通过NVIDIA专有的NVLink高速互联技术直接通信。在现代GPU集群中，主要有以下几种互联方式：

  1. NVSwitch架构：通过NVIDIA的交换架构实现所有GPU之间的全互联
  2. 走网卡，如果卡之间没有NVSwitch的话，可以绕过CPU走网卡：
    1
    GPU0 -> PCIe -> IB(InfiniBand) -> PCIe -> GPU1
    这种通信模式由NCCL（NVIDIA Collective Communications Library）负责协调和优化。

GPU分配策略

NVIDIA开源的go-gpuallocator库提供了一系列基于拓扑关系的GPU分配策略。例如，其中的NewStaticDGX1Policy专门针对DGX-1标准配置优化。考虑到单机环境下GPU组合的可能性有限，这种基于静态规则的分配策略已经能够很好地满足需求。

这些分配策略的核心目标是最小化跨总线和跨NUMA节点的通信开销，确保GPU间通信尽可能利用最高带宽的数据通路，从而提供最优的训练性能。

跨节点的通信

在分布式训练场景下，跨节点通信需要经过更长的数据传输路径：

1
GPU -> NIC -> 叶层交换机 -> 核心交换机 -> NIC -> GPU

这种通信模式面临两个主要的优化方向：

  1. 本地化优化：尽可能将相关联的GPU任务分配在物理位置相近的节点上，以减少网络延迟。

  2. 负载均衡：避免将所有任务集中在同一交换机下，防止出现网络拥塞。过度集中可能导致局部带宽饱和，反而降低整体训练效率。

这种权衡本质上是一个网络流优化问题。通过图论中的网络流算法，可以在通信延迟和带宽利用率之间找到最优平衡点，从而实现更高效的跨节点通信。

一个分布式训练的带宽瓶颈来源于带宽最低的那条路径。

利用 Kubernetes Pod 亲和性优化网络拓扑

在 Kubernetes 环境下，我们可以通过 Pod 亲和性（Affinity）和规则来优化 GPU 任务的分配。主要可以从以下几个方面入手：

拓扑感知调度：使用 topologyKey 确保相关联的 Pod 被调度到网络拓扑上接近的节点：
例如同一个分布式训练任务(training-group = group1)尽让分配在一个机架上，同交换机，同核心交换机也是类似的。

1
2
3
4
5
6
7
8
9
10
11
12
affinity:
  podAffinity:
    preferredDuringSchedulingIgnoredDuringExecution:
    - weight: 50
      podAffinityTerm:
        labelSelector:
          matchExpressions:
          - key: training-group
            operator: In
            values:
            - group1
        topologyKey: topology.kubernetes.io/rack  # 同机架优先

这种方案的优势在于：

  • 配置简单，易于理解和维护
  • 充分利用 Kubernetes 原生能力，无需额外组件
  • 可以根据实际需求灵活调整权重和策略

存储系统

AI训练中的存储系统面临着两个主要挑战：

1. 海量小文件问题

AI训练数据集通常包含大量的小文件，这对传统文件系统的性能和管理造成了巨大压力。一些现代分布式文件系统提供了很好的解决方案，例如 Meta 的 Tectonic 和与其架构类似的 JuiceFS，它们采用了以下优化方案：

元数据管理优化：

  • 使用元数据库管理文件结构，将 ls 命令转化为简单的字符串前缀匹配操作
  • 避免了传统 Linux 文件系统依赖 inode 管理的方式
  • 解决了 inode 臃肿问题（在传统系统中，一个 inode 的大小可能与文件本身相当）

2. Checkpoint 存储挑战

分布式训练中的 checkpoint 文件体积巨大，这在大语言模型训练中尤为明显：

  • 以 LLaMA-2-70B 为例，单个完整的 checkpoint 就需要 140GB 存储空间（FP16格式）
  • 训练过程中需要定期保存 checkpoint，累积存储需求可能达到 TB 甚至 PB 级别
  • 需要存储系统能够提供高带宽和低延迟的读写性能，同时保证数据的可靠性

这些挑战要求存储系统具备：

  • 强大的扩展性
  • 高效的数据压缩能力
  • 智能的数据分层存储机制
  • 可靠的数据备份和恢复能力

训练的稳定性

在大规模 AI 训练中，硬件故障是一个常见问题。特别是新型号显卡往往会有较高的故障率，再加上传统的硬件错误，这些都可能导致训练中断。因此，快速识别错误并恢复训练成为了一个关键挑战。目前主流的解决方案主要有以下两种：

基于 torchrun 的弹性训练

torchrun 提供了两种容错机制：简单重试和弹性训练。

  1. 简单重试模式：
    通过 --max-restarts 参数配置重试次数：

    1
    2
    3
    4
    5
    6
    7
    8
    torchrun \
        --nnodes=$NUM_NODES \
        --nproc-per-node=$NUM_TRAINERS \
        --max-restarts=3 \
        --rdzv-id=$JOB_ID \
        --rdzv-backend=c10d \
        --rdzv-endpoint=$HOST_NODE_ADDR \
        YOUR_TRAINING_SCRIPT.py [script args...]

  2. 弹性训练模式：
    通过设置 nnodes 的范围来支持动态节点数：

    1
    2
    3
    4
    5
    6
    7
    8
    torchrun \
        --nnodes=1:4 \  # 支持1-4个节点的动态伸缩
        --nproc-per-node=$NUM_TRAINERS \
        --max-restarts=3 \
        --rdzv-id=$JOB_ID \
        --rdzv-backend=c10d \
        --rdzv-endpoint=$HOST_NODE_ADDR \
        YOUR_TRAINING_SCRIPT.py [script args...]

弹性训练模式需要配置服务发现机制，默认使用 c10d 作为内置的节点发现服务，也支持使用 etcd 等外部服务。

当节点发生变化时，系统会自动处理以下场景：

  • 节点离开（缩容）：系统通知 agent，停止现有 workers，重新组建 WorkerGroup，使用新的 RANK 和 WORLD_SIZE 启动所有 workers
  • 节点加入（扩容）：接纳新节点，按照相同流程重组 WorkerGroup

基于 DeepSpeed 的弹性训练

DeepSpeed 提供了更细粒度的弹性训练配置：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
{
  "elasticity": {
    "enabled": true,
    "max_train_batch_size": "seqlen",
    "micro_batch_sizes": 8,
    "min_gpus": 1024,
    "max_gpus": "fixed_linear",
    "min_time": "seqlen",
    "version": 8,
    "ignore_non_elastic_batch_info": 1024,
    "num_gpus_per_node": "fixed_linear",
    "model_parallel_size": MODEL_PARALLEL_SIZE
  }
}

DeepSpeed 的特点是：

  • 支持动态调整 batch size
  • 以 GPU 为粒度进行弹性伸缩（而不是节点级别）
  • 提供更丰富的训练参数配置

弹性训练控制器

要实现完整的弹性训练支持，控制器需要：

  1. 依赖服务发现机制进行节点注册和健康检查
  2. 动态调整弹性策略（如 min_nodes、max_nodes 等参数）

对于简单的降级场景，通过静态配置即可实现：

  • 将 max_nodes 设置为总资源规格
  • 将 min_nodes 设置为最小运行要求（如设置为 1:4 表示支持 1-4 张显卡的动态伸缩）

节点的问题发现

可以使用node-promblem-detector，
node-problem-detector 是一个用于在集群管理栈的上游层次中使各个节点问题可见的守护进程。它在每个节点上运行，检测节点问题并将其报告给 apiserver。

总结

To train our largest Llama 3 models, we combined three types of parallelization: data parallelization, model parallelization, and pipeline parallelization. Our most efficient implementation achieves a compute utilization of over 400 TFLOPS per GPU when trained on 16K GPUs simultaneously. We performed training runs on two custom-built 24K GPU clusters. To maximize GPU uptime, we developed an advanced new training stack that automates error detection, handling, and maintenance. We also greatly improved our hardware reliability and detection mechanisms for silent data corruption, and we developed new scalable storage systems that reduce overheads of checkpointing and rollback. Those improvements resulted in an overall effective training time of more than 95%. Combined, these improvements increased the efficiency of Llama 3 training by ~three times compared to Llama 2.

LLAMA3 的技术博客揭示了许多令人振奋的优化成果，这些优化背后蕴含着大量值得深入研究的技术细节。虽然我们可能难以直接接触如此大规模的训练集群及其面临的挑战，但这些技术进展仍然为整个 AI 基础设施领域提供了宝贵的参考和启发。