在 SageMaker HyperPod 上使用 NVIDIA Cosmos 3 构建物理 AI 模型工厂

构建物理人工智能系统需要连续的管道,而不是单一的训练工作。本文展示了如何在 Amazon EKS 上的持久、弹性 Amazon SageMaker HyperPod 集群上运行该模型工厂(合成数据生成、训练后和使用 NVIDIA Cosmos 3 进行闭环评估),并以 GPU 吞吐量作为重要指标。

来源:亚马逊云科技 _机器学习

物理 AI 系统,例如将现实世界数据转化为物理动作的机器人或自动驾驶汽车 (AV),无法在单个训练作业中构建。相反,它需要一个连续的管道:生成合成数据、训练后感知和策略模型的循环,以便系统了解其周围环境并可以采取行动,并在闭环模拟中评估两者。连续运行该管道是物理人工智能模型工厂的工作,将新的现实世界数据流转变成更好的模型。

本文展示了如何在 Amazon SageMaker HyperPod 上使用 NVIDIA Cosmos 3 构建物理 AI 模型工厂,内容包括:

  • Cosmos 3 的独特之处在于:混合变压器 (MoT) 设计,具有每层联合注意力和刻意的训练与推理不对称性。
  • 为什么设计选择可以通过 Amazon Elastic Kubernetes Service (Amazon EKS) 清晰地映射到 Amazon SageMaker HyperPod。
  • 集群和共享多 TB 存储层设置。
  • 针对三个代表性工作负载进行分布式后训练,并在公共 DROID 数据集上对机器人策略阶段进行完整的端到端演练。
  • 随附的存储库包含每个阶段的清单和配置文件。您可以在 awsome-distributed-ai GitHub 存储库中找到可运行的代码,包括基础设施模板和作业清单,可将此设计转变为工作集群。

    运行循环是一种容量承诺。逐步获取 GPU 会增加这种规模的可变性:可用性和交付时间可能会有所不同,并且您获得的容量可能会落在远离您数据的可用区或 AWS 区域。无论是通过针对有限活动的灵活培训计划还是针对开放式活动的容量预留,将容量投入到整个循环中都可以避免这种流失。由于无论管道是否取得进展,您都需要为该容量付费,因此控制成本的指标并不是任何一项作业的峰值吞吐量。它是 GPU 有效吞吐量:整个循环中每个预留 GPU 小时的有用管道进度。

    物理 AI 管道通常为每个阶段提供单独的 GPU 容量:一组节点用于生成合成数据,另一组用于后期训练,另一组用于评估,每个节点都有自己的生命周期来启动和拆除。 NVIDIA Cosmos 3 使这一点变得不必要。作为一个开放的全模式世界基础模型,Cosmos 3 将视频、图像、动作和声音视为单个令牌流。它以三种模式运行相同的 Transformer 主干:用于合成视频生成的前向动态世界模型、逆向动态动作标记器和可部署的动作策略。由于一个模型系列涵盖生成、训练后和评估,因此这些阶段变成了在单个集群控制平面下调度到一个持久、弹性 GPU 节点池上的三个工作负载。它是分时容量,而不是每个阶段都有一个单独的池。 NVIDIA 在 Linux 基金会的 OpenMDW-1.1 许可下发布了它,并在 Cosmos 3 技术报告中描述了该架构。

    1. Cosmos 3 的工作原理

    世界模型的常见模式是将扩散变换器视频生成器与提供文本调节的单独视觉语言模型配对。 Cosmos 3 采用了不同的方法:一个主干可以处理这两种情况,并在每一层进行集成。这种集成使其成为端到端物理 AI 模型工厂的引擎。三种架构选择定义了它:

  • 图 1:Cosmos 3 共享令牌流及其通过每层注意力加入的两位专家
  • (来源:Cosmos 3:物理人工智能的全模态世界模型)
  • 4.1 先决条件
  • 一个令牌流。每种模态都馈入单个共享序列,因此一个模型可以跨模态读取和生成。模型为理解而读取的图像、它生成的像素以及姿势增量和抓取状态的紧凑的每个实施例向量都有自己的编码器。视觉变换器 (ViT) 处理图像理解,而冻结的 Wan2.2 视频变分自动编码器 (VAE) 处理像素生成。这个动作向量可以让同一个模型同时驱动自动驾驶汽车和机器人手臂。该序列将自回归 (AR) 区域(它读取的文本和视觉)置于扩散区域(它产生的视频、音频和动作)之前。

    两名专家,在每一层加入 (MoT)。每一层都运行一个预测下一个标记的推理器和一个对视频、音频和动作进行去噪的生成器。双流注意力将它们结合在一起,因此生成过程在每一层都基于推理器的输出,而不仅仅是最后一次。常见的替代方案将扩散变压器 (DiT) 连接到视觉语言模型 (VLM) 上,并交叉处理其最终输出一次。 Cosmos 3 将推理机中的一代一路向下。

    推理时不对称。训练和部署运行的工作量不同。训练运行完整的去噪计划并将视频解码回像素,因为预测的视频是损失的一部分。在机器人上,同一模型运行一些降噪步骤并完全跳过视频解码。视频潜伏仍然在内部产生以支撑动作,但只有动作标记被解码为机器人执行的关节位置。

    下图显示了一个视图中的前两个选择:共享令牌流和通过注意力连接的两个专家。自回归 (AR) 子序列(模型读取以理解的文本和视觉标记)和扩散模型 (DM) 子序列(它生成的视频、音频和动作标记)穿过共享的 Reasoner 和 Generator 塔。右侧的注意力掩码显示了两位专家的不同之处:DM 查询同时关注 AR 和 DM 密钥(完全注意力),而 AR 查询保持因果关系并且从不看到扩散标记。

    三种操作模式,一种架构

  • 中间训练的基础检查点通过更改哪些标记作为噪声开始来运行三个作业。然后,训练后将检查点专门化为单一模式和控制频率。
  • 正向动态(世界模型)。动作干净,视频嘈杂。 “鉴于这个框架和这个动作,接下来会发生什么?”这是合成数据引擎,呈扇形展开以生成长尾驾驶场景或真实收集无法负担的罕见操作交互。
  • 逆动态(动作标签)。视频干净,动作嘈杂。 “鉴于这两个帧,什么动作导致了变化?”将未标记的视频(原始远程操作录音、第三人称机器人视频、YouTube 驾驶镜头)转换为动作标记的训练数据。
  • 策略(部署的机器人)。两者都有噪声,以 3 视图图像和本体感觉为条件。它输出 32 个未来关节位置,预测视频帧作为副产品,为动作预测奠定基础。

    图 2:一个检查点的三种操作模式,设置令牌以噪音开始

    2. 从单一模型到永久模型工厂

    生产机器人或 AV 的团队不会运行任何微调工作负载。它运行一个循环:摄取真实数据,管理它,用合成数据增强它,训练后,在闭环模拟中评估,部署策略,收集更多数据,然后重复。

    图 3:作为四级飞轮的物理 AI 模型工厂

    3. 我们正在构建什么

    4. SageMaker HyperPod EKS 上的集群设置

    该模型系列有两层:Cosmos3-Nano(16B 参数,位于密集 8B 参数 Qwen3-VL 主干上)和 Cosmos3-Super(64B 参数,位于密集 32B 参数 Qwen3-VL 主干上)。 Cosmos3-Nano-Policy-DROID 等任务变体建立在这些层之上。 NVIDIA 还发布了 Cosmos3-Edge,这是一个用于设备上部署的紧凑型 4B 层(以 Jetson Thor 和 Orin 为基准)。 Edge 与 Nano 和 Super 共享相同的物理世界预训练数据,但建立在从头开始训练的密集约 2B 主干上,而不是从 Qwen3-VL 初始化,因此它是一个单独的权重谱系:您可以直接针对目标硬件对 Edge 进行后训练,而不是将 Nano 检查点缩小到其中。

    下图并排显示了三种模式,实线框表示干净(已知)的标记,虚线框表示模型去噪的噪声标记。前向动态保持动作和当前帧干净,并对未来视频进行降噪。逆动态保持视频干净并对动作进行降噪。该策略仅看到第一帧是干净的,并对机器人将执行的动作进行降噪。相同的架构运行所有三个,只有干净与嘈杂令牌的模式发生变化。在基础检查点中,所有三种模式均可用;训练后的变体(例如 Cosmos3-Nano-Policy-DROID)专门用于 15 Hz、32 步范围的策略模式。

  • 循环有四个阶段。 (1) 将真实世界的物理 AI 数据(DROID、BridgeData2、AV 传感器日志)提取并整理到 Amazon Simple Storage Service (Amazon S3) 和 Amazon FSx for Lustre 上的共享语料库中。 (2) Cosmos3-Super 教师生成合成数据来扩充该语料库。 (3) 综合的合成语料库和真实语料库对可部署的 Cosmos3-Nano 策略进行后训练,并在 Nano 层和 Super 层都应用了视觉微调。 (4)在闭环模拟中对策略进行评估,其失败成为新的生成目标,重新进入下一轮语料库。
  • 理想情况下,该循环不会停止,每个阶段都会在新数据到达时再次运行。这种节奏使得成本驱动因素是 GPU 有效吞吐量(每保留 GPU 小时的有用管道进度),而不是任何一项作业的峰值吞吐量。当各阶段共享一个池时,良好吞吐量最高,因此在不同集群之间重新配置或移动数据时会损失很少的 GPU 时间。 Cosmos 3 使这成为可能:它将三个模型类(世界模拟生成器、策略、感知模型)统一为一个以不同模式运行的模型。为了支持该飞轮,下面的集群必须匹配该形状:一个控制平面上有一个持久池,而不是每个作业都有不同的计算环境。
  • Amazon EKS 上的 Amazon SageMaker HyperPod 恰好提供了这种形状。 Cosmos 3 背后的每个架构选择都会产生具体的集群需求。单令牌流和 64B MoT 结构使训练成为需要低延迟互连的长序列、多节点作业。训练与推理的不对称性将生成、训练后和评估保留在一个模型和一个存储层上,因此它们可以分时共享一个已提交的容量池,而不是逐步对其进行碎片化。连续运行飞轮需要预留并持续监控的容量。四个 Amazon SageMaker HyperPod 属性依次满足这些需求:
  • 该解决方案对三个代表性工作负载进行后训练,每个工作负载都是飞轮的一个阶段,并端到端地演练生成和评估阶段。所有三个都在具有真实检查点的 p5en.48xlarge(8 个 NVIDIA H200 GPU)节点上端到端运行。这三个是机器人操作策略和两个视觉感知微调工作负载。
  • 在第一个作业运行之前,您需要做好以下准备:

    11. 结论

    一个集群适用于所有阶段。由于 Amazon SageMaker HyperPod 使用 EKS 编排集群,因此循环的三个引擎作为普通 Kubernetes 工作负载在单个共享 GPU 池上运行。生成在 vLLM-Omni 服务器上运行,在 torchrun 下的 cosmos-framework 上进行后训练(完全分片数据并行 (FSDP2) 加上 Ulysses 上下文并行),并在单 GPU 策略服务器上进行评估。它们还共享一个存储层。 Amazon FSx for Lustre 文件系统可通过 Elastic Fabric Adapter (EFA) 访问,并通过数据存储库关联 (DRA) 由 Amazon S3 存储桶提供支持,安装一次后即可从同一路径为所有三个阶段提供服务。生成写入合成剪辑,训练后读取它们,策略服务器从同一卷加载其检查点。阶段之间无需重新配置步骤或 TB 级数据迁移,并且区域锁定的 AV 数据保留在一个区域内集群中。

    运行状况检查、自动恢复容量。连续循环需要已配置并主动监控的容量:生成是突发性的并主导 GPU 时间,并且训练后在许多节点上运行数天。 Amazon SageMaker HyperPod 持续检测故障节点并自动重新启动或替换它们,您可以通过灵活的培训计划提前投入该容量。然后,其托管作业自动恢复将工作程序故障转变为有限恢复。 Kubeflow PyTorchJob 重新创建了 pod gang、NCCL 重组以及来自最新 PyTorch 分布式检查点 (DCP) 的 cosmos 框架简历。因此,节点故障最多花费一个检查点间隔的重做工作加上节点替换和重新安排延迟,而不是丢失运行。

    EFA 已连接多节点 NCCL。Cosmos 3 打包在一起进行训练的长序列(视频潜伏加文本加动作,每个令牌数万个)将 64B 层推向 FSDP2 之上的上下文并行。每一层都发布跨节点集体。手动支持的是通常的多节点时间接收器:匹配 EFA 堆栈、NCCL 插件以及精确的 torch 和 NCCL 版本以及 cosmos-framework 引脚。 Amazon SageMaker HyperPod 出厂时已预先配置,当与 AWS Deep Learning Containers (DLC) 映像(其 torch 和 aws-ofi-nccl 版本与框架引脚匹配)配对时,NCCL over EFA 配置为开箱即用。

    可选:许多实施例的任务治理。如果工厂同时服务多种机器人类型或 AV 变体,Amazon SageMaker HyperPod 任务治理(基于 Kueue 构建)会将池划分为具有配额、优先级和抢占功能的命名空间范围的队列。然后,数十个异构作业共享一项容量预留,而不是临时争夺该容量,这样可以通过让原本空闲的 GPU 跨项目保持忙碌来提高吞吐量。单一实施例程序可以跳过它,但一旦许多作业竞争同一个池,任务治理就会得到回报。

    Amazon SageMaker HyperPod 和更轻的选项之间的选择取决于工作单元。一次性微调工作负载不一定需要 Amazon SageMaker HyperPod 集群的弹性和持久性。临时托管训练作业(例如 Amazon SageMaker AI 训练作业)就足够了,因为短期运行很少会发生节点故障。 Amazon SageMaker HyperPod 非常适合持续的 Cosmos 3 飞轮,其中生成连续运行,后期训练是多节点且长期运行的,评估位于同一存储层上,并且从统计上来看,故障很频繁。

    作业提交.kubectl 针对集群进行配置并安装了 Kubeflow Training Operator,以便识别 PyTorchJob 自定义资源。

    关于作者

    虽然这里没有具体介绍 AV 后期训练,但 Cosmos3 的基础模型是在公共合成驾驶语料库(SDG-DriveSim、Hugging Face 上的 nvidia/PhysicalAI-WorldModel-Synthetic-Autonomous-Driving-Scenarios 数据集)上进行训练的。其每个实施例的动作投影旨在扩展到 AV 自我姿势动作空间,因此本文中介绍的相同配方和集群设置也适用于 AV 后期训练。

    训练堆栈是 NVIDIA 的 cosmos-framework,运行时无需分叉或对框架包进行源代码编辑,因此上游更新可以干净利落地进行。它使用 FSDP2 进行训练,并随着序列长度和节点数量的增长扩展到混合分片数据并行性 (HSDP) 和上下文并行性。 6.2 节介绍了如何在每层设置这些并行性选择。

  • 本指南自始至终使用 p5en.48xlarge 作为参考实例类型。我们没有发布跨实例排名,而是为您提供了一种可以在您自己的硬件上运行的良好吞吐量方法(第 7 节)。它建立在每步时间、GPU 饱和度和可配置模型 FLOP 利用率 (MFU) 的基础上。借助它,您可以根据自己的测量来调整所选平台的大小,包括 NVIDIA Blackwell 平台(B200、B300 和机架级 GB200/GB300 NVL72)。
  • 设置集群分为满足先决条件、通过 EFA 启用 NCCL、打开深度运行状况检查和自动恢复、选择基本训练映像以及暂存和验证结果。
  • 您可以使用 Amazon SageMaker HyperPod-EKS Terraform 模块满足大部分先决条件。这些模块配置 EKS 协调的 Amazon SageMaker HyperPod 集群、虚拟私有云 (VPC) 和支持 EFA 的安全组、Amazon FSx for Lustre 文件系统和 CSI 驱动程序、Kubeflow 训练运算符以及可观测性附加组件。如果您愿意,还可以使用 Amazon SageMaker AI 控制台通过 AWS CloudFormation 创建这些资源。或者,您也可以携带自己的同等物品。
  • Cluster.由 Amazon EKS 编排的 Amazon SageMaker HyperPod 集群,在单个可用区中具有由 p5en.48xlarge 节点组成的 GPU 实例组(请参阅使用 Amazon EKS 编排创建 Amazon SageMaker HyperPod 集群)。在这种规模下,GPU 容量是约束条件:p5en 很少能按需提供,因此需要制定灵活的训练计划或容量预留来保护节点。
  • 服务配额。由于 GPU 容量限制了可用性,因此在扩展集群之前,请通过 AWS Service Quotas 请求为目标区域中所选实例类型提供足够的服务配额。
  • 存储。已安装 FSx for Lustre CSI 驱动程序,并在与 GPU 节点相同的 VPC 和子网中附加 Amazon FSx for Lustre 文件系统。 Terraform 模块在您打开 FSx 模块时进行配置,或者您可以附加现有文件系统。

    Credentials.A Hugging Face 读取令牌,该令牌帐户接受了 nvidia/Cosmos-Guardrail1 许可证,存储为名为 hf-token 的 Kubernetes 机密,因为生成路径和策略服务路径在启动时会拉取此门控 Guardrail 存储库。

    4.2 通过 EFA 启用 NCCL

    4.3 深度健康检查和自动恢复

    4.4 选择基础训练图像

    DLC 上的 DROID 视频解码路径中还出现了两个构建时问题:对于 torchcodec 来说 FFmpeg 版本太旧,以及缺少共享 libpython。两者都是打包问题,并已将干净的修复程序烘焙到随附存储库中的 Dockerfile 中。

    4.5 暂存和验证

    kubectl 获取节点

    kubectl 获取 pods -n kubeflow

    5.存储层接线内森·阿诺德

    EFA 为 NCCL 提供了用于多节点集合的内核旁路、远程直接内存访问 (RDMA) 传输能力,并且每个 p5en.48xlarge 节点通告 16 个 EFA 网络接口卡 (NIC)。在 Amazon SageMaker HyperPod EKS 上,EFA 驱动程序(来自 Deep Learning AMI)和 EFA 设备插件(由 HyperPod 服务预安装)已就位。 Pod 规范仅请求 vpc.amazonaws.com/efa 资源以及 GPU(请参阅示例清单)。训练映像必须携带针对框架使用的同一 NCCL 版本构建的 aws-ofi-nccl 插件,这正是此示例基于 AWS DLC 构建的原因(第 4.4 节)。

    硬件上存在的 EFA 与 NCCL 实际使用它的不同,因此请验证传输而不是假设它。训练清单在训练开始前运行一个简短的诊断前导码,并且在 NCCL 调试登录后,健康的多节点运行会报告使用 GPUDirect RDMA 作为所选传输的 EFA。 TCP 的回退在日志中显示为 NET/Socket,这意味着集合体正在错误的传输上运行。要在实际运行之前验证端到端互连,请运行标准多节点 NCCL 测试(例如 all_reduce_perf)并确认实现的总线带宽(请参阅 NCCL 测试指南)。有关确切的诊断命令和完整日志签名,请参阅存储库自述文件。

    每个节点的运行状况监控代理持续运行基本的被动检查(DCGM 策略违规、nvidia-smi 错误、GPU 计数验证),而深度运行状况检查(DCGM 4 级诊断和 NCCL/EFA 基准测试)在节点加入或集群更新时运行。当您打开自动节点恢复时,来自任何这些来源的故障都会触发 Amazon SageMaker HyperPod 重新启动或替换有故障的实例,并在替换准备就绪后自动恢复从最后一个检查点重新启动作业。

    让分布式训练跨节点运行可能会消耗大量的设置时间,因此值得将基本图像的选择视为深思熟虑的决定而不是假设。这一选择取决于一个问题:NCCL 集体是否跨节点使用 EFA,以及 Cosmos 框架固定的 CUDA 轮子是否加载?该框架的虚拟环境 (venv) 引脚 torch==2.10.0+cu130 (CUDA 13) 及其 CUDA 轮(flash-attn、transformer-engine、natten)仅针对 CPython 3.13 发布。这些引脚以两种方式驱动基础映像的选择:映像的 NCCL 必须与火炬轮捆绑的 NCCL 相匹配,EFA 才能正常工作,并且映像必须提供 CPython 3.13 环境供火炬轮安装。

    不匹配的基础映像可能会阻止多节点。即使 EFA 本身功能齐全,通用 GPU PyTorch 基础映像也可以在初始化时验证单节点,但无法通过 EFA 进行跨节点 NCCL(例如 fi_getinfo() 无可用数据)。根本原因是版本矩阵不匹配。 cosmos-framework venv 的 torch 捆绑了特定的 NCCL(此处为 2.28.9),但如果基础映像捆绑的 aws-ofi-nccl 插件是针对不同的 NCCL 构建的,则插件和运行时不会对齐。设置 NCCL_NET_PLUGIN=none 可以避免错误,但只能将跨节点流量丢弃到 TCP 而不是 EFA,这对于多节点性能来说是不可能的。

    简而言之,通过 cosmos-framework 版本兼容性来选择基础镜像,并通过 EFA 验证 NCCL,而不是通过品牌或熟悉程度。对于此框架版本,版本匹配的 AWS DLC 可能是一个省力的路径。

    选择基础映像并准备好必备集群后,您可以在运行工作负载之前暂存映像和存储并验证结果。每个步骤都由随附存储库中的代码支持,因此您可以运行模板而不是手动组装资源。

    版本匹配的 AWS Deep Learning Containers 映像删除了这项工作。PyTorch 的 AWS Deep Learning Containers (DLC) 附带 torch 2.10.0+cu130,与 cosmos-framework pin 完全匹配。它还捆绑了一个经过 AWS 调整、版本匹配的 EFA 堆栈(EFA 1.47.0、libfabric 2.4、aws-ofi-nccl 1.18.0、GDRCopy 2.5.1)。由于 DLC 附带 venv 安装的相同火炬轮,因此 DLC 中的 NCCL 和 aws-ofi-nccl 构建一致。然后,多节点 EFA 配置为无需重建插件或版本不匹配即可工作。

    构建训练映像并将其推送到 Amazon Elastic Container Registry (Amazon ECR),并应用存储类和可选的 Amazon S3 数据存储库关联,以便数据集和基本检查点在首次访问时合并到 /fsx 中:

    ./build-push.sh

    envsubst <存储/存储-fsx-efa-sc.yaml | kubectl 应用-f-

    envsubst < storage/storage-fsx-dra.yaml | kubectl apply -f -

    在提交作业之前,您需要验证集群是否已准备就绪:确认每个 GPU 节点均显示“就绪”并且 Kubeflow 训练操作员 Pod 显示“正在运行”。

    完成配置和验证后,第 6.1 节将介绍准备数据、启动机器人策略作业、监控它并验证其输出,第 10 节将介绍如何在完成后拆除工作负载和集群。

    飞轮在各个阶段之间移动数 TB 的数据集,因此存储层是一流的设计决策,而不是事后的想法。

    暂存:Hugging Face 到 S3 到 FSx for Lustre。数据集和基本检查点从 Hugging Face 暂存到区域内 Amazon S3 存储桶,然后通过 DRA 将其附加到 FSx for Lustre 文件系统。 FSx for Lustre 向每个 pod 提供一个位于 /fsx 的 POSIX 命名空间,并且 DRA 在首次访问时从 S3 延迟加载对象或按需预加载它们。

    两种 I/O 机制。工作负载不会以相同的方式对存储施加压力,因此有助于将它们视为两种不同的机制。机器人策略数据(LeRobot/DROID 数据集)是元数据和小文件体系:许多小型 Parquet 分片和短视频剪辑,其中请求率和延迟比原始带宽更重要。视频 SFT 数据是一种带宽和大文件体系,其中持续吞吐量占主导地位。由于这两种制度的方向不同,因此正确的后端取决于访问模式,而不是单一的最佳选择判决。这两种机制均由此处的一个 FSx for Lustre 文件系统提供服务。每个目录的 Lustre 调整(条带数量和大小、渐进式文件布局和客户端预读)是一种进一步的优化,您可以对每个访问模式进行分层,而不是此示例预设的内容。

    6. 启动分布式后训练

    本节端到端运行机器人策略阶段,然后解释作业如何映射到 Kubernetes、如何在 H200 上配置并行性以及弹性检查点如何工作。

    6.1 端到端运行机器人策略阶段

    机器人策略工作负载的训练后可归结为一个简短的序列:准备数据和基本检查点、启动分布式作业、监控它并验证输出。本演练使用 SageMaker HyperPod 清单,在 kubectl 应用环境变量之前使用 envsubst 渲染它们。接下来的小节解释了每个步骤背后的机制。

    python -m cosmos_framework.scripts.convert_model_to_dcp \--checkpoint-path Cosmos3-Nano \-o $BASE_CHECKPOINT_PATH然后,您可以通过渲染和应用其清单来启动分布式作业,该清单将图像、/fsx 卷和 torchrun 启动连接到 Kubeflow PyTorchJob(第 6.2 节):6.4 弹性检查点

    FSx for Lustre 加 EFA,进行冷热权衡。为 Lustre 文件系统配置支持 EFA 的 PERSISTENT_2 FSx,其大小适合吞吐量目标 (1000 MBps/TiB),因为容量决定了聚合吞吐量上限:示例的 9.6 TiB 文件系统的聚合最高约为 9.4 GB/s。 EFA 可以提高每个客户的上限。非 EFA 文件系统的每个客户端实例的上限为 100 Gbps,而支持 EFA 的文件系统通过 EFA 可以达到每个客户端 700 Gbps。通过在支持 EFA 的 NVIDIA GPU 实例(例如 p5en)上使用 GPUDirect Storage,其速度高达 1200 Gbps。因此,AWS 建议对任何高于 10 GBps 的文件系统启用 EFA。将这些视为记录的上限而不是测量的吞吐量:在您自己的集群上使用 fio(安装在训练映像中)进行基准测试,并注意单个对象存储服务器 (OST) 的流量上限为 5 Gbps,因此每个客户端的高速率需要跨多个 OST 进行条带化。后端在首次接触时的原始吞吐量方面也存在很大差异:本地 NVMe 速度最快,FSx for Lustre 次之,而来自 Amazon S3 的冷读取速度仍然较慢。这个差距就是冷启动和首次接触成本。工作集页面缓存预热后,将从 RAM 提供训练读取,因此后端不再是瓶颈。因此,后端选择对于冷启动和大于 RAM 的工作集很重要,而不是对于缓存驻留数据集的热重用。 cosmos-framework 数据加载器强化了这一点:后台工作人员在 GPU 计算当前步骤时预取和解码即将到来的批次,因此一旦工作集预热,步骤将保持计算限制而不是 I/O 限制。

    首先准备数据和基本检查点。将 Amazon FSx for Lustre DRA 指向保存公共 DROID 数据集的 Amazon S3 存储桶,以便数据集在首次读取时合并到 /fsx 中。 Hugging Face 发布的检查点作为 Diffusers/safetensors 提供,因此您可以将其转换为 cosmos-framework 加载的 DCP 格式(CPU 工作没问题,但 64B Super 需要数百 GB RAM):

    envsubst < hyperpod-eks/train-multi-node-dlc.yaml | kubectl apply -f -

    您可以通过两种方式监控进度。 kubectl 显示调度和 Pod 运行状况,第 7 节中的 Goodput 仪表板显示运行过程中的损失、每步时间和 GPU 饱和度:

    kubectl 获取 pytorchjob cosmos3-droid-policy-hp

    kubectl 日志 -f cosmos3-droid-policy-hp-worker-0

    您可以通过确认运行按配置的时间间隔将 DCP 检查点写入 /fsx 上的 $IMAGINAIRE_OUTPUT_ROOT 来验证输出,并且随着步骤时间保持稳定,仪表板上的训练损失呈下降趋势。由于检查点落在共享卷上,因此第 9 节中的评估服务器可以直接加载它,从而闭合循环。

    6.2 作业如何映射到 Kubernetes

    第 6.1 节中启动的清单将 torchrun 包装在 Kubeflow PyTorchJob 中,一个主副本和用于 N 节点运行的 N−1 个工作副本。 PyTorchJob 控制器将标准 PyTorch 交会变量(协调器地址和端口、每个 pod 排名和世界大小)注入到每个 pod 中,然后 torchrun 为每个 GPU 生成一个进程,以通过 EFA 形成全局进程组。在 2 节点 p5en 运行中,这是跨 2 个节点的 16 个等级。

    6.3 H200 并行度配置

  • 该示例还在导入时应用了两个小的运行时补丁,而不是编辑框架:防止高排名计数的梯度范数监控中的空分片边缘情况,以及第 7 节中介绍的 OpenTelemetry 指标桥。两者都是从示例中应用的,因此上游框架更新仍然会干净地下降。
  • 7. 衡量重要因素:可以运行的良好投入方法7.1 故障恢复的工作原理
  • 三个训练后工作负载按层划分。 16B Nano 完全符合可部署策略,因为它是运送到机器人的模型,并且足够小,可以负担得起全参数训练。 64B Super 是合成数据生成器,它适应 LoRA,而不是完全重新训练:您很少需要重新学习 64B 教师,只需将其转移到您的领域(您的相机、照明、对象类或场景组合)。冻结主干和训练 16 级适配器会崩溃优化器和指数移动平均 (EMA) 内存,将检查点从 64B 快照缩小到兆字节的适配器张量,并让一个冻结的超级基础通过交换适配器来服务多个域。所有这些都提高了大型层面的产出。
  • 这些层还推动了并行策略,将 H200 的 141 GB 高带宽内存 (HBM) 作为最大限度减少通信的杠杆。每个 GPU 更多的 HBM 意味着更少的激进分片和更少的每步集合。 Nano 工作负载运行纯 FSDP2,分片程度设置为全局大小。 Super 在 FSDP2 之上添加了上下文并行度 2,因为长压缩序列使注意激活记忆成为限制器。 Ulysses 上下文并行性将序列分割到 GPU 上,每个注意力层只有几个 all-to-all。在节点数较多时,将复制度设置为 1 以上,以便在跨集群 all-gather 流量成为瓶颈时切换到 HSDP。

    cosmos-framework 将检查点编写为 PyTorch 分布式检查点 (DCP),并且操作策略检查点行为预设在 action_policy_public_lerobot 实验的检查点块中。对于弹性热启动来说,三个设置很重要:

    dcp_async_mode_enabled=False,因此默认情况下保存是同步的。对于异步 DCP,将其设置为 True,这样可以忽略稳态检查点停顿,同时限制故障时丢失的工作。

    strict_resume=False,因此新初始化的动作头(或 LoRA 适配器)可以在模型的其余部分从转换后的基础检查点热启动时进行初始化(第 6.1 节)。

    keys_to_skip_loading 列出了加载期间不需要的张量(动作头和基本模型的 EMA 权重),因此它们会重新初始化而不是出错。

    第四个设置 ckpt_type 在同一实验中默认为 dcp,但每次运行都会被覆盖(例如,冒烟测试的虚拟)。更改前面三个设置中的任何一个都意味着编辑实验配置,而不是传递运行时标志。

    异步 DCP 将引脚大致模型大小保存在主机共享内存中,因此请慷慨地设置 pod 的 /dev/shm 卷 sizeLimit。例如,256Gi 运行良好,因为 P5en 具有大约 2 TiB 的主机内存。在 64Gi 下,保存可能会引发内存不足 (OOM) 错误。

    与 HyperPod 自动恢复相结合,Pod 重启或节点替换会自动从上一个检查点继续。训练清单还会在启动时检查现有检查点并从中恢复,而不是从头开始。

    此部分为您提供了一个可以在您自己的集群上重现的指标设置。它基于 Amazon SageMaker HyperPod EKS 可观测性插件构建,并使用随附存储库附带的仪表板和桥接器将 GPU 遥测和 cosmos-framework 训练指标统一在一个 Amazon Managed Grafana 窗格中。仪表板是一个可导入的 Grafana 模型,cosmos3-goodput-dashboard.json。要加载它,请打开 Grafana 工作区,选择“仪表板”、“新建”、“导入”,然后上传文件,如可观察性自述文件中所述。

    图 4:cosmos-framework 训练器窗格,从框架回调桥接

    8. 调查结果和最佳实践

    基础设施和 GPU 指标,默认情况下。流式多处理器 (SM) 和 HBM 利用率、NCCL 和 EFA 流量以及节点运行状况均通过 DCGM 导出器和节点导出器流入 Amazon Managed Service for Prometheus,无需自定义检测。 Amazon SageMaker HyperPod 可观测性插件无需额外配置即可提供集群、节点和作业视图,并在 Grafana 中与 cosmos-framework 指标一起呈现。

  • Cosmos 框架指标,桥接到本机堆栈中。cosmos 框架通过其回调发出损失、每步计时器、MFU、梯度范数和序列打包统计数据,默认为权重和偏差 (W&B)。为了将它们放置在与 GPU 指标相同的 Prometheus 工作区中且不依赖 W&B,存储库附带了一个小型 OpenTelemetry Protocol (OTLP) 桥。该桥将这些框架标量镜像到 Amazon SageMaker HyperPod 可观测性插件的集群内 OTLP 收集器(通过 gRPC 访问的 hyperpod-otel-collector 服务)。它还添加了 Cosmos 特定的 MFU 回调,该回调根据您配置的峰值 FLOPS 常量公开 MFU。要打开它,请在训练 Pod 上设置一个环境变量:将 OTEL_EXPORTER_OTLP_ENDPOINT 设置为可观测性插件的集群内 OTLP 端点 (http://hyperpod-otel-collector.hyperpod-observability.svc:4317)。如果不设置它,默认路径将保持不变。该存储库还附带 Grafana 仪表板和 OTLP 桥,可将两个指标源一起呈现。
  • 以下两个窗格是单次运行的说明性捕获,而不是从中读取数字的性能结果。第一个窗格是 cosmos-framework trainer 视图。它显示了训练损失、步长、每个 GPU 实现的 TFLOPS、模型 FLOP 利用率(此处根据框架的默认每个 GPU 峰值计算)、迭代吞吐量、梯度范数和序列打包令牌长度,所有这些都从框架回调桥接。
  • 第二个窗格是来自 Amazon SageMaker HyperPod 可观测性插件(使用 DCGM 导出器)的 GPU 和基础设施视图:图形引擎活动、GPU 利用率、使用的帧缓冲区 (HBM) 和 GPU 功率(按 GPU 报告),以便立即看到空闲或驱动不足的设备。

    图 5:使用 DCGM 导出器的可观察性插件的 GPU 和基础设施窗格

    此设置显示两层信号,作为方法而不是发布的分数:

    微层:每 GPU 饱和度和步进效率。DCGM 图形引擎活动和 HBM 利用率直接来自可观测性附加组件,作为直接测量的硬件信号。 cosmos-framework 回调添加了每个 GPU 实现的 TFLOPS、迭代吞吐量、梯度范数和序列打包令牌长度。 MFU 也会显示出来,根据您为加速器和精度设置的可配置的每 GPU 峰值 FLOPS 常数进行计算。

    宏观层:Goodput 框架。将微层饱和度与 Goodput 分数相结合,以推断整个飞轮的有效利用率,而不是任何一步的峰值吞吐量。良好吞吐量分数是在初始化和调度、检查点停顿以及重新启动和恢复之后,花费在有用的前进进度上的 GPU 小时的份额。

    7.2 选择良好输出最优检查点间隔

    这些发现故意是相对的和基于制度的(哪个工作负载最重,哪个后端选择很重要),因此无论您在那里测量的绝对数字如何,它们都保留在您自己的集群上。

    9. 飞轮的其他阶段:生成和评估

    相同的集群、图像和存储层运行生成和评估阶段,从而形成闭环。

  • 埃里克·萨利赫
  • 在此堆栈上,长期运行期间的单个工作单元或节点故障无需人工干预即可解决,并且路径值得跟踪,因为它决定了故障的成本。节点发生故障后,Amazon SageMaker HyperPod 运行状况监控代理会检测到损坏的节点,并且其节点恢复系统会重新启动或替换该节点。然后,PyTorchJob 的托管自动恢复(使用 sagemaker.amazonaws.com/enable-job-auto-resume 注释)将 pod 组重新创建到正常容量。 torchrun 重新会合,cosmos-framework 会从最新的检查点自动恢复。训练从最后一个检查点继续,而不是从头开始。
  • 恢复分解为两个部分:(b1) 节点替换和 Pod 重新调度,以及 (b2) 检查点重新加载以及追赶故障点。尽管 b2 是框架级的并且跨平台的行为相同,但 Amazon SageMaker HyperPod 的节点自动替换和作业自动恢复解决了 b1 部分,否则该部分需要人工检测、替换和重新启动。通过 b1 自动化和 b2 以检查点间隔为界,故障最多花费一个检查点间隔的重做工作加上 b1 重新安排延迟,而不是丢失运行,并且它是优化检查点间隔大小的基础。
  • 检查点间隔是训练保存其状态的频率。保存得太少,失败就会导致大量工作白费。保存过于频繁,保存本身就会消耗有用的计算。异步模式下的 cosmos-framework 检查点:GPU 快速将模型状态复制到主机内存,训练继续进行,同时对 FSx 的写入在后台完成。因此,保存的稳态成本很小,并且单个工作人员的故障大约会导致一个检查点间隔的工作丢失加上重新安排和重新加载的时间。
  • 有一个已知的最佳间隔可以平衡这两种成本,由 Young/Daly 公式给出:间隔 ≈ √(2 × C × MTBF),其中 C 是检查点节省成本,MTBF 是平均故障间隔时间。这两个输入都是您在自己的集群上测量而不是猜测的东西。对于已发布的主播,Meta 的 Llama 3 405B 预训练在 16,384 个 H100 GPU 上大约每三个小时就会出现一次中断。将其视为数量级插图,而不是要重复使用的数字,因为它是 H100,其规模比该样本的 p5en 簇大得多。 C 是您可以从 cosmos-framework 步骤指标中读取的保存持续时间。 MTBF 是在整个队列中观察到的节点故障之间的平均运行时间,该估计值会随着运行历史记录的累积而变得更加清晰。
  • 例如,如果节省成本为 30 秒,MTBF 为 24 小时(86,400 秒),则最佳间隔为 √(2 × 30 × 86,400) ≈ 2,300 秒,或大约每 40 分钟一次。在该时间间隔内,检查点开销计算为每 2,300 秒大约节省 30 秒,约为挂钟的 1%,这是调整时间间隔以保护的有效吞吐量。插入您自己的 C 和 MTBF,以获得最大化集群上的良好吞吐量的时间间隔,并在其中任何一个发生变化时重新推导它。注意 MTBF 是队列级别的量:对于 N 个节点,它大致等于每个节点的 MTBF 除以 N,因此 1024-GPU(128-节点)集群比 16-GPU(2-节点)集群更频繁地发生故障,并且需要相应更短的间隔。在求解之前,按节点数缩放每个节点的数字。

    按访问模式存储。元数据机制有利于 S3 的冷批量暂存和本地 NVMe 的热重用。带宽制度有利于 FSx 加 EFA 作为共享多级平面。后端在冷第一次接触时很重要,而不是对于缓存驻留数据集的热重用而言。

    curl -F input_reference=@clip.mp4 http://localhost:8000/v1/videos/sync --output generated.mp4

    测量的扩展和饱和度。在我们在 p5en (H200) 节点上的运行中(使用第 7 节的 goodput 仪表板记录),超级 (64B) LoRA 工作负载在从 1 到 4 个节点(8 到 32 个 GPU)扩展时保持接近持平的强扩展效率。它的线性度保持在大约 0.97-0.99,梯子上每步时间的误差在 3% 左右。根据 H200 BF16 峰值计算,模型 FLOP 利用率在计算密集型超级工作负载中下降到接近 0.50,在较轻的 Nano 视觉工作负载中下降到接近 0.24,符合预期的顺序(较大、计算密集型模型使 GPU 更充分地饱和)。这些是来自此设置的相对数字,应该在可比较的硬件上重现形状,而不是排行榜数字。

    每个工作负载的节点数量和并行度。根据每个工作负载的挂钟目标调整节点数量,并让工作负载的每步成本驱动选择:64B Super LoRA 层主导每步成本,而机器人策略工作负载是最轻的。由于 cosmos-framework 打包数据加载器拥有固定的每排名令牌预算,因此每步时间(而不是每小时迭代次数)是横向扩展时要跟踪的吞吐量信号。随着跨节点全收集流量的增长,并行策略也从纯 FSDP2 转向 HSDP(第 6.2 节)。

    管理弹性购买什么。将回报表示为恢复时间和良好投资分数改善,而不是绝对美元。自动替换加上自动恢复消除了每次节点故障时的手动检测-替换-重新启动循环,因此故障成本最多降至重做工作的一个检查点间隔加上节点替换和重新安排延迟。

    Amazon SageMaker HyperPod 何时获得回报?对于一次性小微调来说,临时托管训练作业(例如每次运行时配置并在完成后拆除的 Amazon SageMaker AI 训练作业)就足够了,因为自动恢复很少触发。 Amazon SageMaker HyperPod 非常适合持续、大规模的飞轮:许多并发作业、连续合成生成、预留容量以及统计上故障频繁的数十个节点尾部。

    生成(写入绑定)。生成在与训练不同的图像上运行:官方 vllm/vllm-omni:cosmos3 引擎(cp312 / vLLM 0.23 堆栈),而不是 cosmos-framework DLC 训练图像。它作为 OpenAI 兼容服务器(vllm 服务 nvidia/Cosmos3-Super --omni)出现,在 POST /v1/videos/sync 处接受 Cosmos3-Super 视频到视频 (V2V) 请求(请参阅生成清单)。服务器就绪后,端口转发并发布调节剪辑以获取生成的延续:

    返回的剪辑是一个合成示例,可为下一轮训练后提供数据。在这个阶段,自动驾驶团队会展开长尾驾驶场景(与 SDG-DriveSim 等语料库针对的同一类安全关键极端案例)以增强真实车队数据。生成使用自己的并行轴,无分类器指导(CFG)并行 x Ulysses x HSDP(--cfg-parallel-size,--ulysses- Degree,--use-hsdp --hsdp-shard-size),这与训练端 FSDP2 + 上下文并行旋钮不同,并且必须乘以每个节点的 GPU 数量。服务器是单节点的,因此生成不需要跨节点 EFA,并且可以作为独立服务器进行横向扩展。 Guardrails (nvidia/Cosmos-Guardrail1) 根据请求 (extra_params.guardrails) 进行切换,而不是通过服务器标志,并且必须在 Hugging Face 令牌的帐户上接受门控 Guardrail 许可证,否则启动会失败。批处理完成后删除生成作业。

    10. 成本考虑和清理

    kubectl 删除 pytorchjob cosmos3-droid-policy-hpkubectl删除-f hyperpod-eks/serve-policy.yaml

    评估(延迟限制)。轻量级单 GPU 部署通过 HTTP 提供 Cosmos 3 操作策略以进行闭环评估,使用相同的 FSx 卷,因此集群上生成的任何检查点都可以直接使用(请参阅策略服务清单)。检查点必须是本地目录(裸露的 org/repo Hugging Face id 不会被服务器解析),因此首先预下载或导出您自己的训练后检查点到 FSx。模拟器的控制循环通过 GET /info 确认服务器已启动,然后将观察结果发布到 POST /预测每个步骤并接收下一个操作块。请求携带当前相机帧和任务提示:

    # 准备情况/模型元数据

    卷曲 http://localhost:8000/info

    # 一个控制步骤:观察输入,动作块输出

    卷曲 -X POST http://localhost:8000/predict \

    -H“内容类型:application/json”\-d '{"image": "", "prompt": "拿起杯子", "domain_name": "droid", "image_size": 256}'服务器返回预测的动作块(第 1 节中描述的 32 个未来关节位置),模拟器在发送下一个观察之前应用该块。该请求/响应周期是闭环:训练后生成的检查点与评估服务器从 FSx 加载的检查点相同,因此可以评估新训练的策略,而无需在集群之间移动数据。Amazon SageMaker HyperPod 是一个持久集群:实例在属于集群时进行计费,FSx for Lustre 按预配置容量按小时计费,任何服务或可视化 Pod 只要运行就保留一个 GPU 节点。请参阅 Amazon SageMaker AI 和 Amazon FSx for Lustre 定价页面了解当前费率。要暂停会话之间的实例成本,同时保持集群配置,请将 GPU 实例组缩放为零,或完全删除集群。请参阅管理 SageMaker HyperPod 集群。完成后删除训练工作负载和策略服务部署:删除 FSx 文件系统会删除其检查点和日志的本地副本。在 S3 DRA 下写入的数据将导出回链接的存储桶并保留下来,但该路径之外(或尚未导出)的任何内容都会丢失。在删除之前确认您的 DRA 已完成导出(或下载您想要保留的内容),并注意底层 S3 存储桶将独立保留并单独计费,直到您将其也删除为止。aws fsx 删除文件系统 --文件系统-id随着物理 AI 投入生产,挑战从单一的训练运行转变为持续且经济地运行整个生命周期。本文介绍了如何在 Amazon SageMaker HyperPod (EKS) 上使用 NVIDIA Cosmos 3 作为该飞轮的基底。该底层是一个持久的、大规模的集群,具有一个共享存储层,用于合成生成、策略和感知模型的后期训练以及闭环评估,以良好吞吐量而不是单作业吞吐量作为重要指标。我们在公共 DROID 数据集上运行了机器人策略后训练阶段,在 1-4 个节点上进行了 Nano 和超级视觉微调,在版本匹配的 AWS 深度学习容器映像上验证了多节点 EFA,并提供了一个关于本机可观察性的可重现的 Goodput 仪表板,该仪表板将 Cosmos 框架和 GPU 指标统一在一个窗格中。在这些运行中,超级工作负载在该范围内保持了近乎持平的强扩展效率(大约为线性的 0.97-0.99)。由于每个阶段共享相同的集群和存储,因此训练后写入 Amazon FSx for Lustre 的检查点与评估服务器加载的检查点相同。生成的合成剪辑将落在下一轮训练后读取的同一卷上。

    这篇文章提供了一个参考架构和一个可重现的方法,您可以指向您自己的实施例和数据,而不是您所相信的排行榜。首先,探索 awsome-distributed-ai GitHub 存储库,使用 LeRobotV3ActionDataset 配方作为起点来适应您自己的数据集,并将飞轮扩展到您自己的机器人或车辆。要了解更多信息,请参阅 NVIDIA Cosmos 站点、cosmos-framework 存储库和 Amazon SageMaker HyperPod 文档。

    Nathan 是位于德克萨斯州奥斯汀的 AWS 的高级 AI/ML 专家解决方案架构师。他帮助 AWS 客户(从小型初创公司到大型企业)在 AWS 上高效地训练和部署基础模型。当他不与客户合作或修补最新的开源工具时,他喜欢与他的四只狗一起跑步和玩耍。

    Eric 是 AWS 的高级 GenAI 专家,专注于物理 AI 的基础模型训练和推理。他正在与顶级物理 AI 模型构建者和 AWS 服务团队合作,在 AWS 上实现大规模分布式训练和推理,并领导与战略客户的联合 GTM 动议。在加入 AWS 之前,Eric 领导产品团队构建企业 AI/ML 解决方案,其中包括用于微调、RAG 和托管推理的前沿生成式 AI 服务。他拥有加州大学洛杉矶分校安德森分校的商业分析硕士学位。