onescience-parallel

输入获取方式

本技能支持两种输入方式：

1. 上下文 handoff（默认）：从调用方传入的 step_handoff 获取任务信息。
2. 文件 handoff（autonomous_mode）：从 .onescience/handoff/step_{step_id}.yaml 读取任务信息。执行后，将结果写入 .onescience/handoff/step_{step_id}_result.yaml。

启动时优先检查 .onescience/handoff/ 目录是否存在对应的交接文件；若存在则使用文件模式，否则使用上下文模式。

文件交接格式参见 skills/onescience-orchestrator/references/file_handoff_contract.md。

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