


修改分布式机器学习训练与推理代码而不把它改坏
在修改需要多个进程或加速器来训练、运行机器学习模型的软件时使用。它会在有风险的改动前先确认,让权重版本保持一致,让每个 rank 以相同顺序调用集合通信,保护配置兼容性,并报告跳过了哪些测试。
要使用这个AI规则
## 编辑前先读代码,修改共享配置或公共 API 前先征得同意
编辑前,阅读相关实现、相邻模块的既有模式、配置契约和测试。使用项目的环境管理器、格式化工具、代码检查工具和 pre-commit 工具链;不得绕过钩子。
绝不硬编码凭据、机器路径、服务端点或集群拓扑。不得臆测调度器设置,也不得把重建加速器驱动程序和运行时技术栈当作排查问题的捷径。
执行以下操作前须先取得明确决定:
- 更改共享配置结构;
- 添加依赖项;
- 改变启动器或调度器的行为;
- 删除、重命名公共 API,或对其进行不兼容的更改。
如果无法使用所需硬件或访问集群,应说明跳过了哪些集成检查及其原因。不得把未经验证的分布式假设直接转化为代码。
按子系统分派工作:训练引擎、推理引擎、算法、启动器、调度器、数据集、奖励、工作流或分布式调试。如果项目已有专门的实现指南或领域评审流程,应遵循该指南或流程。
## 让各子系统职责分开,并给数据标上权重版本
将编排、rollout 生成、训练、推理、基础设施、模型适配、奖励、数据集和跨领域工具分别作为不同职责。
训练器负责协调循环、数据加载和工作流。Rollout 工作流定义异步的回合生成过程。推理引擎生成输出并接收权重更新;训练引擎使用 rollout 张量、计算更新并发布新的权重版本。
在异步系统中,为生成的数据和更新附加明确的权重版本元数据。绝不使用模型版本暗中不匹配的数据进行训练,也不得以不匹配的模型版本作为评估对象。
通过标准跟踪器发出指标,将产物持久化到配置的输出根目录下,并使用规范的检查点保存与恢复路径。可观测性和恢复能力是执行契约的一部分。
优先采用组合与委托,而不是深层继承。扩展边界应足够小,以便独立测试引擎、工作流和模型适配器。
## 编写避免设备到主机同步的高性能加速器代码
遵循所属模块现有的命名、类型和日志约定。使用明确的类型提示和导入;绝不使用通配符导入。仅在确实需要重量级可选依赖项的函数或代码路径中导入这些依赖项。
使用项目日志记录器,而不是直接打印。日志记录中应包含组件和分布式 rank 上下文,以便并发输出仍可追溯到来源。
在加速器热点路径中:
- 避免设备到主机的同步,例如为记录日志而提取标量值或列表值;
- 批量执行张量操作,不要在 Python 中逐个遍历元素;
- 明确设置设备和数据类型;
- 在边界处断言张量形状符合预期。
保持 rollout 工作流非阻塞。等待异步引擎调用和文件操作完成;不得在异步回合路径中引入同步 I/O。
开发新功能时,遵循最相近的可用模块模式,添加有针对性的测试,并避免无关重构。
## 在保持兼容的前提下修改配置字段和公共 API
配置字段应依次排列为必填字段、常用可选字段、少用可选字段,最后是内部字段。在初始化时进行验证,并抛出清晰的异常,其中应指出无效值或违反的约束。
新增公共字段时提供安全的默认值。移除前先标记为弃用,并避免更改既有字段的类型。公共配置类必须记录约束和可接受的值。
对于取值范围有限的命令行选项,应使用受约束的类型,而不是自由格式的字符串。复用规范的配置数据类,不要创建并行的配置结构。
将共享配置、启动器、调度器和公共 API 名称视为影响重大的表面。在变更说明中明确兼容性影响,并在进行破坏性或不兼容的编辑前取得批准。
如果无法根据当前契约确定正确行为,应留下范围明确的后续处理标记,并报告尚未解决的约束,而不是自行杜撰默认值。
## 让进程组和集合通信在所有 rank 上保持一致
绝不在导入时创建全局进程组。当进程组成员身份有影响时,应显式传入预期的进程组,并针对该进程组查询 rank。
设备网格的维度必须与声明的并行模型保持一致,在适用时包括数据分片、张量、上下文、专家和专家张量轴。
集合通信操作是全局契约:
- 所有参与的 rank 必须以相同顺序调用同一个 all-reduce;
- 应有意识地选择归约操作,例如求和或求平均值;
- 广播必须指定明确的源 rank;
- 屏障是诊断工具,不是常规同步策略。
排查卡死时,检查集合通信序列是否不一致。排查结果错误时,检查归约语义和权重版本。排查内存不足故障时,在削减模型功能前,先检查分片、张量放置和设备所有权。
仅在有针对性的调试中启用详细的分布式、通信库和同步启动诊断;理解问题后,应将其移除或禁用。
## 测试分布式代码,并在拉取请求中写明跳过的测试
测试名称应体现行为、条件和预期结果,并按准备、执行、断言的结构组织测试。为慢速、异步、加速器和分布式测试添加标记,使默认 CI 能选择正确的测试套件。
如果无法使用加速器或多节点环境,应明确说明原因并妥善跳过相应测试。使用框架的伪进程组测试单元级分布式行为;不要模拟深层分片或分布式张量内部机制,因为这些机制的语义正是测试要验证的内容。
比较张量时,应明确指定相对容差和绝对容差。优先使用受管理的临时路径和限定作用域的 monkeypatch。
先运行范围较窄的单元测试套件;硬件条件允许时,再运行直接受影响的加速器或分布式测试套件。提交前执行完整的格式化和 pre-commit 工具链。
通过规范的生成器生成命令行参考文档。通过项目封装程序构建发布文档,由它执行所需的静态设置和打包;不要使用会遗漏项目步骤的底层文档命令。
在拉取请求中说明问题、风险区域、确切的测试命令、跳过的测试套件及原因,以及任何硬件限制。提交主题应遵循约定并使用祈使语气,同时在评审前移除开发中的临时提交。
