onescience-parallel

# Pipeline Parallel 改造 Skill ## 重要原则（必读） 1. **复用 OneScience 模块**：所有基础模块必须使用 `onescience` 已有实现，禁止重复实现 2. **先读后写**：开始前必须阅读 `context.md` 和 `architecture.md` 3. **参考已有实现**：必须参考 `examples/earth/pangu_weather_distributed/` 的 pangu 实现 4. **保持参数一致性**：进行并行改造时，各 Stage 类的 `__init__` 参数签名和内部初始化逻辑应尽可能与原模型保持完全一致（如 `config`、 `mask` 等参数的处 理），避免随意更改参数名或逻辑。 5. **严禁循环导入**：在创建 Distributed 模块（如 `{StyleName}DistributedFuser`）时， **禁止**在模块内部导入 `OneFuser`、 `OneTransformer` 等顶层包装类 ，因为这些包装类通常已经导入了你的 Distributed 模块，会导致 `ImportError`。应直接导入具体的子模块类（如 `from .{stylename}distributedlocalsiefuser im port {StyleName}DistributedLocalSIEFuser`）。 --- ## 改造三步流程 ``` 步骤 1：模型拆分 (PP) → 步骤 2： TP 并行模块 → 步骤 3：训练接口对接 ``` --- ## 步骤 1：模型拆分（Pipeline Parallel） ### 1.1 切分策略 按前向执行顺序找**计算串行边界**，切点满足： - 数据依赖最少（只有一个 tensor 出口） - 跨 stage 传输的中间 tensor 尽量小 - 各 stage 计算量尽量均衡 典型 4-stage 切分（U-Net/Encoder-Decoder 结构）： | Stage | 内容 | 说明 | |-------|------|------| | 0 | Embedding + Encoder 前半 | 首阶段，产生 skip connection | | 1 | Downsample + 中间层 | 下采样后的计算密集区 | | 2 | 解码层 + Upsample | 解码器前半段 | | 3 | Decoder 后半 + Recovery | 消费 skip，产出最终结果 | ### 1.2 每个 Stage 的必要属性 ```python class MyModel_stageN(Module): def __init__(self, original_arg1, original_arg2, ..., megatron_config=None): """ 参数签名应尽可能与原模型保持一致。 如果原模型第一个参数是 config (yaml配置 )，则保持不变； 额外传入的 Megatron 核心配置建议命名为 megatron_config 避免冲突。 """ super().__init__(meta=MetaData()) # ① 必须有这三个属性，均设为 None self.pre_process = ...