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 等参数的处

理），避免随意更改参数名或逻辑。

1. 严禁循环导入：在创建 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 的必要属性

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 = None
        self.share_embeddings_and_output_weights = None
        self.input_tensor = None          # Stage 0 也要有

        # ② 必须初始化  config（Megatron get_model_config() 需要）
        # 使用传入的  megatron_config，若无则从  args 获取
        if megatron_config is None:
            args = get_args()
            megatron_config = core_transformer_config_from_args(args)
        self.config = megatron_config

        # ③ 保持原模型的初始化逻辑
        self.arg1 = original_arg1
        # ... 原模型的参数处理  ...

    def set_input_tensor(self, input_tensor):
        """Megatron pipeline 调度钩子，所有  Stage 都必须实现 """
        self.input_tensor = input_tensor

1.3 forward 方法规范

Stage 0（首阶段）：直接处理输入，不读 input_tensor

def forward(self, x):
    x = self.embed(x)
    x = self.block1(x)
    skip = x
    meta = torch.tensor([B, Pl, Lat, Lon], device=x.device, dtype=x.dtype)
    return (x, skip, meta)   # 必须返回  tuple

Stage 1/2（中间阶段）：解包 input_tensor，原样透传 skip/meta

def forward(self, x):
    if self.input_tensor is not None:
        # 必须处理  list 类型（recv_forward 返回  list）
        if isinstance(self.input_tensor, list):
            x, skip, meta = tuple(self.input_tensor)
        else:
            x, skip, meta = self.input_tensor
    x = self.process(x)
    return (x, skip, meta)   # skip/meta 原样透传

Stage 3（末阶段）：消费 skip，返回最终结果（不是 tuple）

def forward(self, x):
    if self.input_tensor is not None:
        x, skip, meta = ...
    B = int(meta[0].item())   # 从  meta 恢复  shape
    x = torch.concat([x, skip], dim=-1)
    return self.recovery(x)   # 返回最终结果，不打包  tuple

1.4 关键细节

• shape meta 必须是 tensor，不能是 Python int（pipeline 通信只支持 tensor）
• Pipeline 并行中 micro-batch size = 1，不能使用动态 batch size
• drop_path 必须基于总深度计算，各 stage 取对应切片，保证权重可从原模型加载：

``python drop_path = np.linspace(0, 0.2, l1d + l2d + l3d + l4d).tolist() stage0_drop_path = drop_path[:l1d] # 前 l1d 个 stage1_drop_path = drop_path[l1d:l1d+l2d] ``

• 重计算层用 checkpoint() 包裹节省显存

> 完整 Stage 代码模板 → 见 references/stage_templates.md

---

步骤 2： TP 并行模块改造

2.1 检测哪些层需要 TP

从 Stage 类开始，逐层检查调用链，遇到以下模式需创建 Distributed 版本：

| 原始代码 | 替换方案 | |---------|---------| | nn.Linear(in, out) | ColumnParallelLinear 或 RowParallelLinear | | Mlp(in, hidden) | DistributedMlp(config=config) | | nn.MultiheadAttention(...) | 手动实现： qkv→Column， proj→Row | | 调用含上述模式的模块 | 为该模块创建 Distributed 版本 |

2.2 创建 Distributed 模块的步骤

通用命名规则： {ModuleStyle}Distributed{ModuleType}

• {ModuleStyle}：原模块风格名（如 Pangu、 Earth、 Xihe 等）
• {ModuleType}：模块类型（Fuser、 Transformer、 Attention 等）

创建步骤（以 {StyleName 风格为例）：

1. 创建 {StyleName}DistributedFuser → src/onescience/modules/fuser/{stylename}distributedfuser.py
2. 创建 {StyleName}DistributedTransformer → src/onescience/modules/transformer/
3. 创建 {StyleName}DistributedAttention → src/onescience/modules/attention/
4. 在对应的 onefuser.py / onetransformer.py / oneattention.py 中注册

# onefuser.py 注册示例
_FUSER_REGISTRY = {
    "{StyleName}Fuser": {StyleName}Fuser,                   # 原版本
    "{StyleName}DistributedFuser": {StyleName}DistributedFuser,  # Distributed 版本
}

2.3 config 传递链路（关键）

config 必须从 Stage 一路传递到每个并行层：

Stage(config)
  → OneFuser(style="{StyleName}DistributedFuser", config=config)
    → {StyleName}DistributedFuser(config=config)
      → OneTransformer(style="{StyleName}DistributedTransformer", config=config)
        → {StyleName}DistributedAttention(config=config)
          → ColumnParallelLinear(config=config)  ← 最终使用
          → RowParallelLinear(config=config)
        → DistributedMlp(config=config)          ← 最终使用

> 说明：将 {StyleName} 替换为实际模块风格名，如： Pangu、 Earth、 Xihe 等。

2.4 并行线性层使用规范

MLP 模式（两层组合） ：

self.fc1 = ColumnParallelLinear(
    input_size=in_features, output_size=hidden_features,
    config=config, bias=True,
    gather_output=False,       # 不聚合，传给  fc2
)
self.fc2 = RowParallelLinear(
    input_size=hidden_features, output_size=in_features,
    config=config, bias=True,
    input_is_parallel=True,    # 输入已按  TP 切分
)
# forward: x, _ = self.fc1(x); x, _ = self.fc2(x)

单层使用：

# 需要完整输出时
self.proj = ColumnParallelLinear(..., gather_output=True)
# 或
self.proj = RowParallelLinear(..., input_is_parallel=False)

input_is_parallel 判定决策树（必须逐层检查！） ：

拿到一个输入张量  → 问：它的最后一维是否已被  TP 切分？
  │
  ├─ ✅  是（每个  rank 只有  dim // tp）
  │    → 使用  input_is_parallel=True
  │    → Weight 形状 : (output_size // tp, input_size)
  │    → 适用场景：
  │        • 来自  ColumnParallelLinear 的输出
  │        • 来自  Embedding 经  ColumnParallelLinear 处理
  │
  └─ ❌  否（每个  rank 持有完整  dim）
       → 使用  input_is_parallel=False
       → Weight 形状 : (output_size, input_size // tp)
       → 适用场景：
           • 来自  nn.Linear 的输出
           • torch.cat(...) 的拼接结果
           • 来自  RowParallelLinear 的输出（已  all-reduce 还原为全维度）
           • 原始输入数据

最容易踩的坑：

# ❌  错误： out 来自  RowParallelLinear（已  all-reduce 全维度）
#          x 是原始全维度输入
#          拼接后  input 为全维度，但错误地设了  input_is_parallel=True
out, _ = self.attn_proj(out)            # RowParallelLinear → 输出已是全维度！
x_concat = torch.cat([out, x], dim=-1)  # 拼接后  (N, 2*dim)，全维度
x_out, _ = self.concat_proj(x_concat)   # ❌  若  concat_proj.input_is_parallel=True → 崩溃！

# ✅  正确：
self.concat_proj = RowParallelLinear(
    input_size=2 * dim, output_size=dim,
    config=config, bias=True,
    input_is_parallel=False,  # ← 输入是拼接结果，未被  TP 切分
)

检查命令：

# 找出所有  RowParallelLinear，检查  input_is_parallel 是否正确
grep -n "RowParallelLinear" src/onescience/modules/**/*.py -A 3 | grep -E "RowParallelLinear|input_is_parallel"

| 输入来源 | input_is_parallel | 说明 | |---------|---------------------|------| | ColumnParallelLinear 输出 | True | Column 输出已按 TP 切分 | | 手动 Q/K/V attention 输出 | True | 各 head 计算在 TP 切分后的维度上进行 | | RowParallelLinear 输出（all-reduce 后） | False | 已还原为全维度！ | | torch.cat() 拼接结果 | False | 全维度拼接 | | nn.Linear / 原始输入 | False | 全维度 |

> 完整并行层参数说明 → 见 references/parallel_linear.md

---

步骤 3：训练接口对接

必须使用框架导入：

from onescience.distributed.megatron.training import pretrain, get_args
from onescience.distributed.megatron.core import mpu
from onescience.distributed.megatron.training.arguments import core_transformer_config_from_args
from onescience.distributed.pipelinetensorshapeconfig import PipelineTensorShapeConfig

3.1 model_provider

def model_provider(pre_process=False, post_process=True):
    config = para_init()
    pp_rank = mpu.get_pipeline_model_parallel_rank()
    stages = {0: MyModel_stage0, 1: MyModel_stage1, 2: MyModel_stage2, 3: MyModel_stage3}
    model = stages[pp_rank](config=config, ...)

    # 配置跨  stage 通信  tensor 的  shape（必须与  forward 返回的  tuple 一一对应）
    pp_config = PipelineTensorShapeConfig(
        num_stages=4,
        stage_shapes=[
            # stage0→1: forward 返回  (x, skip) 的  shape
            [[1, seq_len, dim], [1, seq_len, dim]],
            # stage1→2
            [[1, seq_len//2, dim*2], [1, seq_len, dim]],
            # stage2→3
            [[1, seq_len, dim], [1, seq_len, dim]],
        ]
    )
    get_args().pipeline_tensor_shape_config = pp_config
    return model

stage_shapes 填写规则：

• 长度 = num_stages - 1（最后一个 stage 不发送）
• 每项 shape 与该 stage forward 返回 tuple 中各 tensor 严格对应
• batch_size 填实际值（通常为 1）

3.2 forwardstepfunc

def forward_step_func(data_iterator, model):
    data = next(data_iterator)                  # 直接用，不要  try-except
    invar = data[0].cuda().float()             # 必须转  float32
    outvar = data[1].cuda().float()

    output = model(invar)

    def loss_func(output):                     # 通过闭包捕获  outvar
        loss = LpLoss()(outvar, output)
        num_tokens = torch.tensor(1, device="cuda")
        reporting_loss = torch.cat([loss.clone().detach().view(1), num_tokens.view(1)])
        return loss, num_tokens, {'lm loss': reporting_loss}

    return output, partial(loss_func)

注意：只有 pipeline 的第一个和最后一个 stage 实际需要数据，但 forward_step_func 中可以统一调用 next(data_iterator)，框架会处理数据分发。

3.3 dataset_provider

def train_valid_test_dataset_provider(train_val_test_num_samples):
    # 必须返回  Dataset，不是  DataLoader
    train_dataset = MyDataset(mode='train')
    val_dataset   = MyDataset(mode='val')
    test_dataset  = MyDataset(mode='test')
    return train_dataset, val_dataset, test_dataset

3.4 pretrain 入口

if __name__ == "__main__":
    train_valid_test_dataset_provider.is_distributed = True  # 必须设置

    pretrain(
        train_valid_test_dataset_provider=train_valid_test_dataset_provider,
        model_provider=model_provider,
        model_type=None,                          # 必须
        forward_step_func=forward_step_func,      # 注意参数名
        args_defaults={'dataloader_type': 'cyclic'}  # 必须
    )

---

重要补充约束

文件结构对应原则

> 与原模型文件结构一一对应 > > 原模型有 N 个模块文件，分布式就对应创建 N 个分布式模块文件： > > 例 1： Pangu 模式（单文件） > `` > 原文件 : pangufuser.py (1个文件 ) > 分布式 : pangudistributedfuser.py (对应 1个文件 ) > ` > > **例 2： Xihe 模式（多文件） ** > ` > 原文件 : xihefuse.py > xihelocalsiefuser.py (共 3个文件 ) > xiheglobalsiefuser.py > ------------------------------------------------ > 分布式 : xihedistributedfuser.py > xihedistributedlocalsiefuser.py (对应 3个文件 ) > xihedistributedglobalsiefuser.py > ` > > **规则**：每个 {StyleName}Xxx.py 对应创建 {StyleName}DistributedXxx.py` > > 所有分布式模块都要在对应的 registry 中注册。

训练脚本位置原则

> 参考原模型的训练脚本所在目录，创建对应 _distributed / _distributed_4stage 子目录： > > `` > examples/{domain}/ > {model_name}/ # 原训练脚本目录 > {model_name}_distributed/ # 分布式训练脚本目录 > {model_name}_distributed_4stage/ # 4阶段流水线训练脚本目录 > ` > > **例**： examples/earth/panguweather/ → examples/earth/panguweatherdistributed4stage/`

原则 8： Stage 类参数接口规范

⚠️ Stage 的 __init__ 只接收标量参数 + config（megatron config），禁止传入 yaml config 对象！

错误做法：

class MyModel_stage0(Module):
    def __init__(self, yaml_cfg, img_size, ..., config=None):  # ❌  yaml_cfg 不应传入

正确做法：

class MyModel_stage0(Module):
    def __init__(self, img_size, patch_size, embed_dim, num_heads, ..., config=None):  # ✅

训练脚本 build_model 函数：从 yaml 中提取标量后逐一传入，不传整个 cfg 对象：

model = stage_cls(
    img_size=cfg_data.dataset.img_size,
    patch_size=cfg.patch_size,
    embed_dim=cfg.embed_dim,
    num_heads=cfg.num_heads,
    # ... 其他标量参数
    config=config,   # ✅  只传  megatron config
)

{StyleName}DistributedFuser 的 num_heads 接口规范：

使用单个 num_heads 参数，所有子模块（local、 global 等）共用，由 Stage 按 num_heads[i] 索引传入。禁止拆分为 num_heads_local / num_heads_global 等多个参数。

# Stage 调用
OneFuser(style="{StyleName}DistributedFuser", num_heads=num_heads[i], ...)

# DistributedFuser 内部
class {StyleName}DistributedFuser(nn.Module):
    def __init__(self, ..., num_heads=8, ...):  # ✅  单一参数
        self.local_blocks = ...Fuser(num_heads=num_heads, ...)
        self.global_blocks = ...Fuser(num_heads=num_heads, ...)

强制约束： num_heads 所有值必须能被 tp_size 整除。

---

原则 9：并行线性层返回值处理

⚠️ ColumnParallelLinear / RowParallelLinear 返回的是 tuple，不是 tensor！

典型错误：

TypeError: dropout(): argument 'input' (position 1) must be Tensor, not tuple

错误写法：

x = self.proj_drop(self.proj(x))  # ❌  proj 返回  tuple，直接传  dropout

正确写法：

x = self.proj(x)[0]               # ✅  取  [0] 提取  tensor
x = self.proj_drop(x)

强制性检测命令：

# 检查所有  forward 中的调用
grep -n "\.proj(" src/onescience/modules/**/*.py
grep -n "RowParallelLinear\|ColumnParallelLinear" src/onescience/modules/**/*.py -A 1

| 检查点 | 要求 | |--------|------| | self.proj(x)[0] | ✅ 必须加 [0] | | self.fc1(x)[0] | ✅ 必须加 [0] | | 所有并行层调用后立即接 [0] | ✅ |

原则 9： DistributedFuser 参数接口统一原则

⚠️ {StyleName}DistributedFuser 必须使用单个 num_heads 参数，禁止拆分为 num_heads_local / num_heads_global！

参考 PanguDistributedFuser 的接口规范：所有子模块（local、 global）共用同一个 num_heads，由外部 Stage 按 num_heads[i] 索引传入。

错误做法：

class XiheDistributedFuser(nn.Module):
    def __init__(self, ..., num_heads_local=6, num_heads_global=12, ...):  # ❌
        self.local_blocks = XiheLocalFuser(num_heads=num_heads_local, ...)
        self.global_blocks = XiheGlobalFuser(num_heads=num_heads_global, ...)

正确做法（与 PanguDistributedFuser 一致）：

class XiheDistributedFuser(nn.Module):
    def __init__(self, ..., num_heads=6, ...):  # ✅  单一参数
        self.local_blocks = XiheLocalFuser(num_heads=num_heads, ...)
        self.global_blocks = XiheGlobalFuser(num_heads=num_heads, ...)

Stage 调用方式（与 pangudistributed4stage 一致）：

# stage0
OneFuser(style="XiheDistributedFuser", num_heads=num_heads[0], ...)
# stage1
OneFuser(style="XiheDistributedFuser", num_heads=num_heads[1], ...)

强制约束： num_heads 必须能被 tp_size 整除， config.yaml 中的 num_heads 所有值都必须满足此条件。

---

原则 9： Pipeline Parallel Stage 负载均衡原则

⚠️ OOM 90% 来自 stage 计算量不均衡！

典型错误：

stage0: 1 个  Fuser (num_local=1)
stage1: 1 个  Fuser (num_local=2)
stage2: 2 个  Fuser (num_local=2) ❌  是  stage0 的  4 倍！
stage3: 1 个  Fuser (num_local=1)

均衡要求：每个 stage 的计算量差异控制在 ±20% 以内

强制性 Checklist：

# 1. 统计每个  stage 的  Fuser 数量
grep -n "OneFuser" src/onescience/models/*_distributed*/*.py

# 2. 检查每个  Fuser 的  num_local / num_global
grep -n "num_local\|num_global" src/onescience/models/*_distributed*/*.py

4 stage 通用负载均衡配置： | Stage | 计算量 | numlocal / numglobal | |-------|--------|------------------------| | stage0 | 输入 embedding + 第一个 Fuser | 统一 = 1 | | stage1 | Downsample + 中间 Fuser | 统一 = 1 | | stage2 | 中间 Fuser | 统一 = 1 | | stage3 | Upsample + Skip connection + Recovery | 统一 = 1 |

内存优化额外技巧：

• ❌ 不要创建不必要的 tensor 副本： x1 = x.contiguous(); x = x.contiguous()
• ✅ 共享 tensor： x = x.contiguous(); x1 = x
• ✅ 所有 Fuser 都必须用 checkpoint

原则 10： MHA 内部完整张量并行原则

⚠️ 仅把 proj 改成并行 = MHA 没有真正张量并行！

错误做法（只改一半）：

self.attn = nn.MultiheadAttention(...)  # ❌  内部没有张量并行！
self.proj = RowParallelLinear(...)     # 只有这层并行

正确做法（参考 Pangu EarthDistributedAttention3D）：

# Q, K, V 投影全部用  ColumnParallelLinear
self.q_proj = ColumnParallelLinear(dim, dim, config=config, ...)
self.k_proj = ColumnParallelLinear(dim, dim, config=config, ...)
self.v_proj = ColumnParallelLinear(dim, dim, config=config, ...)

# 输出用  RowParallelLinear
self.proj = RowParallelLinear(dim, dim, config=config, input_is_parallel=True, ...)

# 每个  rank 只算自己分到的  heads
self.num_heads_per_r

…

## Source & license

This open-source skill is cataloged on AgentStack and links to its original source — we do not rehost the code.

- **Author:** [onescience-ai](https://github.com/onescience-ai)
- **Source:** [onescience-ai/OneSkills](https://github.com/onescience-ai/OneSkills)
- **License:** MIT

Install and usage instructions live in the source repository linked above.