onescience-parallel

name	onescience-parallel
description	将 PyTorch 模型改造为支持 Pipeline Parallelism (PP) + Tensor Parallelism (TP) 的分布式训练模型。 适用场景：单机单卡模型改多卡分布式、模型拆分为多个 pipeline stage、替换 nn.Linear 为并行线性层、 编写 forward_step_func / model_provider / dataset_provider、配置 PipelineTensorShapeConfig、 创建 Distributed 版本模块（DistributedFuser / DistributedAttention / DistributedMlp）。 当用户提到以下任何关键词时，务必使用此 skill： "流水线并行"、"pipeline parallel"、"模型并行改造"、"分布式训练"、"stage 拆分"、 "ColumnParallelLinear"、"RowParallelLinear"、"pretrain 接口"、"forward_step_func"、 "model_provider"、"TP 并行"、"张量并行"、"多卡训练"、"Megatron"。
type	executor

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 的必要属性

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 取对应切片，保证权重可从原模型加载：

  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 forward_step_func

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/pangu_weather/ → examples/earth/pangu_weather_distributed_4stage/

原则 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 调用方式（与 pangu_distributed_4stage 一致）：

# 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	计算量	num_local / num_global
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_rank = self.num_heads // self.tp_size

强制性 Checklist：

# 检查是否还在用原生  MHA
grep -n "nn.MultiheadAttention" src/onescience/modules/**/*distributed*.py

# 检查  Q/K/V 是否都并行化了
grep -n "q_proj\|k_proj\|v_proj" src/onescience/modules/**/*distributed*.py

检查点	要求
nn.MultiheadAttention	❌ 必须删除，不能在分布式模块中使用
q_proj = ColumnParallelLinear	✅
k_proj = ColumnParallelLinear	✅
v_proj = ColumnParallelLinear	✅
proj = RowParallelLinear(input_is_parallel=True)	✅
num_heads_per_rank	✅ 必须存在

原则 11：流水线张量传递原则

⚠️ **流水线之间只传特征张量，不传非张量，不传可计算的信息！ **

核心思想：跨 GPU 通信带宽是最宝贵的资源

错误做法：

return (x, x1, mask1, mask2)  # ❌  传  4 个  tensor
# mask1, mask2 都是确定性可计算的！不需要传！

正确做法：

# stage0 -> stage1 -> stage2 -> stage3
# 所有  stage 之间只传  (x, x1) 两个特征张量
# mask 在每个  stage 内部根据分辨率计算

优化前后对比：

	传递张量数量	通信量比例
优化前	3-4 个	100% + 额外开销
优化后	2 个	~50%

强制性 Checklist：

# 检查每个  stage 的  return 语句张量数量
grep -n "return " src/onescience/models/*_4stage/*.py

# 检查是否在传递  mask
grep -n "mask" src/onescience/models/*_4stage/*.py | grep "return\|input_tensor"

检查点	要求
流水线传 mask	❌ 绝对禁止
每个 stage 传特征数量	控制在 2-3 个以内
非张量跨 stage 传递	❌ 绝对禁止
可计算的信息（如分辨率、常数模式）	本地计算，不传

⚠️ **致命陷阱 1：模型和训练脚本不同步！ **

模型  stage return (x, x1)                  ✅  2 个张量
训练脚本  PipelineTensorShapeConfig:        ❌  还写着  3-4 个张量！

强制性同步检查：

# 模型文件
grep "return (" src/onescience/models/*_4stage/*.py

# 训练脚本
grep -A5 "stage_shapes" examples/*/*_4stage/train*.py

# 两者张量数量必须一一对应！

⚠️ **致命陷阱 2：不传 mask ≠ 改变 mask 逻辑！ **

错误做法（本地逻辑变了）：

# 为了 "简化 "，直接变成全 1 mask
mask_coarse = torch.ones((H_out, W_out))  # ❌  原模型逻辑被改变了！

正确做法（本地计算完全复现原逻辑）：

# mask 不跨  stage 传递，但是每个  stage 内部完全复现原计算逻辑
self.mask_full = mask_full  # __init__ 时保存常量

def forward(self, x):
    mask2 = None
    if self.mask_full is not None:
        # ✅  完全复现原模型的  change_mask 逻辑
        mask2 = self.change_mask(self.mask_full, x, h_out=H_out, w_out=W_out)

项目	要求
跨 stage 传递 mask	❌ 禁止
mask 计算逻辑	✅ 必须与原模型 100% 一致
change_mask 函数	✅ 每个需要的 stage 都要有
mask_full is None 判断	✅ 必须保留分支

⚠️ **致命陷阱 3： mask_full 可能是字符串路径！ **

错误：

patch = mask_full[i*2:(i+1)*2, j*2:(j+1)*2]
# TypeError: string indices must be integers
# 因为  mask_full = "dataset/mask.npy" 是路径！

正确：所有用到 mask_full 的地方都必须先处理类型：

if not torch.is_tensor(mask_full):
    if isinstance(mask_full, str):  # 路径字符串
        mask_full = np.load(mask_full)
    mask_full = torch.tensor(mask_full, dtype=torch.float32)

强制检查：

grep -n "isinstance.*str" src/onescience/models/*_4stage/*.py
# 必须每个  stage 都有字符串路径处理！

⚠️ **致命陷阱 4： Module 签名严格校验！ **

OneScience 的 Module 基类在 __new__ 时会做严格的参数绑定检查：

# Module.__new__ 中会执行：
bound_args = sig.bind_partial(*args, **kwargs)
# ❌  只要有一个多余参数，就直接  TypeError!

解决方案：所有 Stage 的 init 必须加kwargs**：

def __init__(
    self,
    config,
    img_size=(...),
    ...,
    megatron_config=None,
    **kwargs,  # ✅  必须加！吸收所有多余参数
):

没有**kwargs	有**kwargs
❌ 直接 TypeError	✅ 参数静默吸收
❌ 训练脚本传新参数就崩	✅ 前向兼容

强制检查：

grep -A10 "class.*_stage" src/onescience/models/*/*.py | grep "**kwargs"
# 每个  Stage 都要有！

⚠️ **致命陷阱 5：只有 stage0 需要读数据！其他 stage 会死锁！ **

**90% 的人都会踩的坑！ **

# ❌  错误写法：所有  stage 都在调用  next(data_iterator)
def forward_step_func(data_iterator, model):
    invar, outvar = next(data_iterator)  # ❌  stage1/2/3 卡在这里！
    output = model(invar)

**每个 stage 运行在独立进程中！它们都在试图读数据，但数据 iterator 是有状态的！ ** stage0 读完数据后， stage1 就拿不到了，直接死锁！

**✅ 正确写法： **

def forward_step_func(data_iterator, model):
    from onescience.distributed.megatron.core import parallel_state

---

原则 12： Fuser 模块 forward 方法统一接口规范

⚠️ **Fuser 模块必须同时支持 dict 和 object 属性访问！ **

错误做法（会导致 KeyError/AttributeError）：

def forward(self, obj):
    if isinstance(obj, dict):
        x = obj["x"]
        mask = obj["mask"]  # ❌  KeyError if mask not present!
    else:
        x = obj.x
        mask = obj.mask     # ❌  AttributeError if mask not present!
    ...
    x = blk(x, mask)        # ❌  子模块期望接收  dict！

**✅ 正确写法： **

def forward(self, obj):
    if isinstance(obj, dict):
        x = obj["x"]
        mask = obj.get("mask")           # ✅  安全获取，不存在返回  None
    else:
        x = obj.x
        mask = getattr(obj, 'mask', None) # ✅  安全获取，不存在返回  None
        obj = {"x": x, "mask": mask}     # ✅  统一转成  dict

    for blk in self.blocks:
        x = blk(obj)                      # ✅  始终传递  dict 给子模块
        obj["x"] = x

    return x

强制检查：

grep -A15 "def forward.*obj" src/onescience/modules/fuser/*distributed*.py
# 检查是否使用了  .get() 和  getattr(..., None)

检查点	要求
obj["mask"]	❌ 必须改成 obj.get("mask")
obj.mask	❌ 必须改成 getattr(obj, 'mask', None)
blk(x, mask)	❌ 必须改成 blk(obj)
子模块调用	✅ 始终传递 dict 格式
is_first_stage = parallel_state.is_pipeline_first_stage()

if is_first_stage:

只有 stage0 读真实数据

invar, outvar = next(data_iterator) output = model(invar) else:

stage1/2/3 不传输入，由 set_input_tensor 负责

output = model(None)

**配合每个  Stage 的  forward： **
| Stage | 写法  |
|-------|------|
| stage0 | `def forward(self, x):` 使用真实的  x |
| stage1/2/3 | `def forward(self, x=None):`**忽略  x，只用  input_tensor！ ** |

```python
# stage1/2/3 forward 第一行：
assert self.input_tensor is not None, "必须通过 set_input_tensor 设置 !"
x, x1 = tuple(self.input_tensor) # 从流水线前一阶段来的 tensor

模块调用链深度分析方法

从 Stage 类开始， 逐层追踪完整调用链，画出模块依赖树，所有叶子节点都要检查：

MyModel_stageN
 ↳ OneFuser(style="MyDistributedFuser")
 ↳ MyDistributedFuser
 ↳ MyDistributedSubModuleA
 ↳ OneTransformer(style="...Distributed...")
 ↳ OneAttention(style="...Distributed...") ✅ 都要并行
 ↳ OneMlp(style="...Distributed...")

关键检查点：

1. 所有含 nn.Linear 的子模块都要有分布式版本
2. 从 Stage 到最内层并行层， config 参数必须完整传递
3. 优先复用已有的 Earth/Pangu 分布式模块（如 TransformerBlock）

---

严格检测流程（必须执行，避免漏改）

第一步：文件级枚举检测

必须运行此命令列出该模型的所有相关文件：

# 1. 找出所有带模型名的源文件
find src -name "*{model_name}*.py" | sort
find examples -name "*{model_name}*" | sort

# 2. 通用检测命令示例（替换 {model_prefix} 为实际前缀）
ls src/onescience/modules/**/{model_prefix}*.py
ls src/onescience/modules/**/*{ModelPrefix}*.py

检测结果必须做成对照表（将 {Model} 替换为实际模型名）：

类型	原文件	分布式文件命名规范	是否已创建
Transformer	{model}transformer.py	{model}distributedtransformer*.py	☐
Fuser	{model}fuse*.py	{model}distributedfuser*.py	☐
Attention	{model}*attention.py	{model}distributed*attention.py	☐
MLP	{model}mlp.py	{model}distributedmlp.py	☐

**命名规范： **原文件名 + distributed + 可选后缀

• 例： xihefeaturegroupattention.py → xihedistributedfeaturegroupattention.py
• 例： pangutransformer3d.py → pangudistributedtransformer3d.py

违反此步骤后果：漏改文件

第二步：调用链递归检测

必须从 Stage 入口递归追踪所有调用：

Stage 层入口
 ↳ 检查每个 OneFuser / OneAttention / OneMlp / OneTransformer
 ↳ 进入对应分布式模块的 __init__
 ↳ 检查内部所有子模块调用
 ↳ 是否都使用了 Distributed 版本？
 ↳ config 是否层层传递了？

**递归检测 checklist（必须逐项打勾） **：

• 所有 OneTransformer(style="...") 都使用分布式 style
• 所有 OneAttention(style="...") 都使用分布式 style
• 所有 OneMlp(style="...") 都使用分布式 style
• 所有 Mlp() 都替换为 DistributedMlp()
• 所有 nn.Linear 都替换为并行 Linear
• 每个分布式模块都接收 config=None 参数
• config 参数都传递到所有子模块调用

第三步： Registry 完整性检测

必须检查所有 registry 文件：

grep -l "_REGISTRY" src/onescience/modules/**/one*.py

Registry 文件	分布式模块是否已注册
onetransformer.py	✅
oneattention.py	✅
onemlp.py	✅
onefuser.py	✅

---

检查清单：

• ✅ 第一步：文件枚举检测，生成对照表
• ✅ 第二步：调用链递归追踪，逐层检查
• 为每个叶子节点创建 Distributed 版本
• 验证 config 在整个链路中正确传递
• ✅ 第三步： Registry 完整性检查

---

改造检查清单

Stage 拆分

• 每个 Stage 包含 pre_process=None、 share_embeddings_and_output_weights=None、 config、 input_tensor
• 每个 Stage 实现 set_input_tensor
• 首 Stage 不读 input_tensor（直接处理原始输入）
• 中间 Stage 解包后原样透传skip/meta
• 末 Stage 返回最终结果（非 tuple）
• shape meta 以 torch.tensor 传递，不用 Python int
• drop_path 切片与原模型保持一致
• 重计算层用 checkpoint() 包裹

TP 并行

• 所有 Distributed 模块已创建并注册到对应 One* 文件
• config 在整个调用链路中正确传递
• 双层组合（MLP 模式）： gather_output=False + input_is_parallel=True
• 单层独立：明确设置 gather_output=True 或 input_is_parallel=False
• concat 后的 tensor：使用 input_is_parallel=False（拼接结果未被 TP 切分）
• RowParallelLinear 输出后再次拼接：输出已 all-reduce 还原为全维度，再次拼接 /使用前应视为全维度
• 所有 RowParallelLinear 的 input_is_parallel 已逐层核实：对照 2.4 节的判定决策树逐一检查

训练接口

• stage_shapes 与各 Stage forward 返回的 tensor shape 完全一致
• dataset_provider 返回 Dataset，不是 DataLoader
• loss_func 通过闭包 /partial 捕获 label
• 数据转换为 .float()
• is_distributed = True 已设置
• args_defaults={'dataloader_type': 'cyclic'} 已配置

---

常见错误速查

错误信息	原因	解决方案
_get_attr_wrapped_model couldn't find attribute config	Stage 缺少 self.config	所有 Stage 添加 config 初始化
mat1 and mat2 shapes cannot be multiplied	input_is_parallel 设置错误	见下方详细排查步骤
Input type (double) and bias type (c10::Half)	数据类型不匹配	加 .float() 转换
RuntimeError: list 解包错误	未处理 list 类型的 input_tensor	添加 isinstance(self.input_tensor, list) 判断

mat1 and mat2 shapes cannot be multiplied 排查步骤

此错误 100% 是 RowParallelLinear 的 input_is_parallel 设置与实际输入张量的 TP 状态不匹配：

1. 确认错误发生位置：找到 traceback 中报错的 RowParallelLinear 层名称
2. 检查输入来源：该层的输入是从哪里来的？
3. 判定 TP 状态：输入张量的最后一维是全维度还是已切分？
4. 修正参数：
   • 若输入已切分 → input_is_parallel=True（权重形状 : output//tp × input_size）
   • 若输入全维度 → input_is_parallel=False（权重形状 : output × input//tp）

典型错误模式：

错误 traceback:
 .../xihedistributedfeatureungroupattention.py:95 in forward
 x_out, _ = self.concat_proj(x_concat)
 RuntimeError: mat1 and mat2 shapes cannot be multiplied (N × 2*dim and dim × dim)

原因 : x_concat = cat(RowParallelLinear_output, original_input)
 RowParallelLinear 的输出已经 all-reduce 还原为全维度
 cat 拼接后也是全维度 → 应使用 input_is_parallel=False

---

参考资料

• references/stage_templates.md — 完整 Stage 0/1/2/3 代码模板
• references/parallel_linear.md — 并行线性层详细参数和通信模式
• references/train_interface.md — pretrain 接口完整模板（model_provider / forward_step_func / dataset_provider）

参考实现（代码库中）：

• 模型： src/onescience/models/pangu_distributed_4stage/
• 训练： examples/earth/pangu_weather_distributed/train_distributed_4stage.py
• Distributed Fuser： src/onescience/modules/fuser/pangudistributedfuser.py
• Distributed Attention： src/onescience/modules/attention/earthdistributedattention3d.py