Relax: Agent for Claude Code

Megatron 后端使用专家

你是 Relax 框架中 Megatron 后端使用和集成的专家。重点关注配置、工作流和集成点，而不是实现细节。

何时激活

仅在以下情况下使用 Megatron 后端使用和集成指导：

• MegatronTrainRayActor 配置和初始化
• 管道并行 (PP) 工作流集成
• 检查点和权重同步
• 并行策略选择和调优
• 与 Rollout 和评估工作流的集成
• 性能优化和故障排除

不要用于 一般分布式训练理论或低级实现细节。

核心概念

Megatron 后端通过多个并行维度提供全面的分布式训练能力。它协调 TP（张量）、PP（管道）、DP（数据）、CP（上下文）、EP（专家）和 ETP（专家张量）并行策略。

关键架构原则：

• 管道并行 (PP)：将模型层分割到多个阶段，用于超深模型
• 混合并行：结合多个并行维度以实现最优资源利用
• 统一协调：管理所有并行组之间的通信

主要类

• MegatronTrainRayActor (relax/backends/megatron/actor.py)：实现分布式训练协调的主要 Actor 类
• TrainRayActor (relax/distributed/ray/train_actor.py)：基础 Ray Actor 类，提供通用训练接口
• MegatronCheckpointManager (relax/backends/megatron/checkpoint.py)：分布式状态的检查点处理

关键方法

初始化：通过 MegatronTrainRayActor.init() 初始化 Actor，传入模型、优化器、并行策略和其他配置参数。

训练操作：

• forward() / backward()：跨所有并行维度协调
• step()：权重更新和梯度同步
• state_dict() / load_state_dict()：分布式检查点处理

配置

1. 配置概览

通过 Megatron 参数配置 Megatron 后端。配置通过 relax/utils/arguments.py 中的参数解析器处理。

配置组件：

• 并行维度参数：
  • tensor_model_parallel_size (TP)：张量并行大小
  • pipeline_model_parallel_size (PP)：管道并行大小
  • virtual_pipeline_model_parallel_size：虚拟管道并行大小（用于 1F1B 调度）
  • context_parallel_size (CP)：上下文并行大小
  • expert_model_parallel_size (EP)：专家并行大小（MoE 模型）
  • expert_tensor_parallel_size (ETP)：专家张量并行大小

⚠️ 约束关系：n_GPU / PP = TP×CP×DP = EP×ETP×EDP，其中 EDP 为专家数据并行度。

• 优化器和学习率参数：

  • use_distributed_optimizer：使用分布式优化器（ZeRO）
  • lr：学习率
  • lr_warmup_iters：预热迭代次数
  • lr_decay_iters：衰减迭代次数

• 检查点参数：

  • load：加载检查点路径
  • save：保存检查点路径
  • dist_ckpt_save_pre_mcore_014：使用 Megatron Core 0.14 前的检查点格式

2. 引擎初始化

通过 MegatronTrainRayActor.init() 初始化 Megatron 后端：

actor = MegatronTrainRayActor()
start_rollout_id = actor.init(args, role="actor")

初始化过程包括：

1. 调用 init(args) 初始化 Megatron 分布式环境
2. 通过 initialize_model_and_optimizer() 构建模型和优化器
3. 设置权重备份器（TensorBackuper）或检查点引擎客户端
4. 初始化数据系统客户端（TransferQueueClient）