# Onescience Parallel

> >

- **Type:** Skill
- **Install:** `agentstack add skill-onescience-ai-oneskills-onescience-parallel`
- **Verified:** Yes — security-reviewed for prompt injection and unsafe behavior
- **Seller:** [onescience-ai](https://agentstack.voostack.com/s/onescience-ai)
- **Installs:** 0
- **Category:** [Agent Skills](https://agentstack.voostack.com/c/agent-skills)
- **Latest version:** 0.1.0
- **License:** MIT
- **Upstream author:** [onescience-ai](https://github.com/onescience-ai)
- **Source:** https://github.com/onescience-ai/OneSkills/tree/master/skills/onescience-parallel

## Install

```sh
agentstack add skill-onescience-ai-oneskills-onescience-parallel
```

Requires the [AgentStack CLI](https://agentstack.voostack.com/docs/cli). Works with Claude Code, Cursor, and any MCP-compatible agent.

## About

# 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 = 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`
```python
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
```python
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）
```python
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` 中注册

```python
# 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 模式（两层组合） **：
```python
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)
```

**单层使用**：
```python
# 需要完整输出时
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 还原为全维度）
           • 原始输入数据
```

**最容易踩的坑**：

```python
# ❌  错误： 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 切分
)
```

**检查命令**：
```bash
# 找出所有  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：训练接口对接

必须使用框架导入：
```python
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

```python
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

```python
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

```python
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 入口

```python
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 对象！ **

**错误做法**：
```python
class MyModel_stage0(Module):
    def __init__(self, yaml_cfg, img_size, ..., config=None):  # ❌  yaml_cfg 不应传入
```

**正确做法**：
```python
class MyModel_stage0(Module):
    def __init__(self, img_size, patch_size, embed_dim, num_heads, ..., config=None):  # ✅
```

**训练脚本 `build_model` 函数**：从 yaml 中提取标量后逐一传入，不传整个 cfg 对象：
```python
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`
 等多个参数。

```python
# 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！ **

典型错误：
```python
TypeError: dropout(): argument 'input' (position 1) must be Tensor, not tuple
```

**错误写法**：
```python
x = self.proj_drop(self.proj(x))  # ❌  proj 返回  tuple，直接传  dropout
```

**正确写法**：
```python
x = self.proj(x)[0]               # ✅  取  [0] 提取  tensor
x = self.proj_drop(x)
```

**强制性检测命令**：
```bash
# 检查所有  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]` 索引传入。

**错误做法**：
```python
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 一致）：
```python
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 一致）：
```python
# 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**：
```bash
# 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 没有真正张量并行！ **

**错误做法**（只改一半）：
```python
self.attn = nn.MultiheadAttention(...)  # ❌  内部没有张量并行！
self.proj = RowParallelLinear(...)     # 只有这层并行
```

**正确做法**（参考 Pangu EarthDistributedAttention3D）：
```python
# 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.

## Pricing

- **Free** — Free

## Security capabilities

Automated source analysis of v0.1.0 — what this tool can access:

- **Network access:** no
- **Filesystem access:** no
- **Shell / process execution:** no
- **Environment & secrets:** no
- **Dynamic code execution:** no

*"Yes" means the capability is present in the source — more access means more to trust, not that it is unsafe.*


## Versions

- **0.1.0** — security scan: passed — Imported from the upstream source.

## Links

- Listing page: https://agentstack.voostack.com/l/skill-onescience-ai-oneskills-onescience-parallel
- Seller: https://agentstack.voostack.com/s/onescience-ai
- Browse the marketplace: https://agentstack.voostack.com/browse

---
Listed on AgentStack — the marketplace for AI agent skills and MCP servers. Every listing is security-reviewed. Creators keep 70%.