状态（Status）： Draft
作者（Authors）： @Your_community
创建日期（Created）： YYYY-MM-DD
更新日期（Updated）： YYYY-MM-DD
相关 Issue/PR： #123（关联 Issue/PR 以便追踪背景）

1. 概述

1.1 简介

当前 PyNative 模式下，Trainer 主要面向训练流程设计，缺乏统一、便捷的模型纯前向推理入口。本提案为 PyNative 模式下的 Trainer 新增模型前向推理能力，并天然支持数据并行、模型并行、流水线并行等多种分布式并行策略。

推理由 TrainConfig 顶层字段 run_mode: predict 路由：run_mindformer.py --mode 1 下，run_mode: predict 时走 Trainer.inference()，否则走 Trainer.train()。字段名与图模式顶层 run_mode 对齐，非法值在 Trainer.__init__ 构造期 fail-fast。训练 / 推理是两个独立 msrun 进程、两份 yaml，中间靠 checkpoint 文件衔接。

1.2 动机

用户在执行评估、预测或部署前的精度验证时，往往需要自行拼接分布式推理逻辑，不仅开发门槛高，而且容易因并行策略处理不当引入精度或性能问题。本提案降低 PyNative 模式下推理任务开发成本，提供与训练一致的分布式并行体验，使 Trainer 在动态图场景下真正覆盖"训练-推理"全流程。

1.3 目标

功能目标（对齐需求交付物）：

• Trainer.inference() 推理接口：支持传入单条数据或数据集（五种输入形态，见 3.5.2.1），输入输出行为正常。
• 推理过程自动继承训练阶段的并行策略（TP / CP / DP / EP / PP），用户无需感知分布式细节即可在多卡环境完成推理。多卡 DP 分流仅在 dp_replicate > 1 且非 PP 时生效；PP 开启时各 PP 线解码全量记录、靠 sample_idx 去重（见 3.1 "DP 样本分流"），DP 并发度在该组合下不生效——这是 POC 阶段的实现选择。
• 支持加载训练保存的分布式 checkpoint（完整分布式权重 / HF 权重），正确切分到对应卡上。
• 与现有训练流程无缝衔接：训练完成后可直接调用推理接口进行评测；兼容动态图特性，不破坏已有训练功能（旧 yaml 不含 inference: 段、run_mode 缺省为 train 时行为不变）。

精度目标：使用训练配置文件正常拉起 DeepSeek-V4-Flash 模型推理任务，加载 HF 权重可推理出符合逻辑的 token；同 prompt 多入口输出逐位一致，多卡解码不失步。

性能目标：不涉及（POC 阶段聚焦功能与精度）。

非目标：

• 不做 KV cache（增量推理）——解码循环每生成一个 token 跑一轮全量前向，推理加速留待后续演进。
• 不提供采样类解码参数（temperature / top_k / top_p / repetition_penalty 等）——解码固定贪心（argmax）。
• 不做 chat template 渲染——inference() 仅做 base 续写，instruct 模型的对话模板渲染由使用方离线预处理。对 instruct 模型直接调用会得到 base 续写行为，后续演进需补 prompt 渲染器。

2. 用例分析

• UC1 训练后精度验证：训练 yaml 追加 inference: 段并置 run_mode: predict，加载训练产出的分布式权重（或 HF 权重），贪心生成验证输出。要求：推理前向与训练数值路径一致；旧 yaml 不含 inference 段且 run_mode 缺省时训练行为不变。
• UC2 批量评估：jsonl 数据集直接作为 input_data，多卡 DP 下样本自动分流并行生成（见 3.1 样本分流），批量解码、逐条落盘，输出回显原始字段供评测脚本对账。要求：多卡解码不失步；tokenize 失败的样本记 error 行不中断批量；超长样本（padded 长度超过 model.seq_length）记 error 行跳过，不截断。前向失败当前不在容错范围内：generate_batch 抛出的异常会传播出整批推理，与 UC2 "单条失败不中断批量" 的意图存在偏差。
• UC3 并行策略继承：并行策略读训练侧 parallelism.*，支持跨策略加载（ReshardLoader 自动切分）。要求：跨策略输出与原策略基线逐位一致；PP 仅支持基础调度（interleave / overlap / swap / FSDP+PP 组合 fail-fast 拒绝）。

3. 方案设计

3.1 总体方案

核心思路：推理 = 训练的前向 + 一个解码循环。模型调用、并行切分全部复用训练代码，只新增"输入归一、token 挑选、落盘"三块推理专属逻辑。

部件	来源	说明
GPTModel.construct(input_ids, position_ids)	复用训练	训练多传 labels 计 loss，推理不传即返回 logits
_apply_parallelism / parallelize_model	复用训练	TP/CP/DP/EP/PP 切分、FSDP 分片同一条路径；forward_only 参数透传至 build_stages
PipelineScheduleRuntime.run()	复用训练	PP 下前向走训练同一调度：first stage 消费输入，last stage 产出 logits
ScaledLossPipelineStage	复用训练	训练 / 推理同一 stage 类；推理传 has_backward=False + dyn_shape=True，forward-only 兜底以子类覆写 bwd_send_specs 落在 mindformers 侧，hyper-parallel 依赖库零改动
_load_inference_checkpoint	新增	forward-only 加载，不恢复 optimizer/step；checkpoint.load_path 缺失直接 ValueError（避免随机权重产出垃圾输出）
generate_batch 批量解码循环	新增	每生成一轮 token 跑一次 [bs, L] 全量前向，无 KV cache；EOS / 超长样本冻结陪跑，无需节拍器
_sample_shard_info	新增	DP 样本分流判定（见下"DP 样本分流"）
gather_vocab_parallel_logits / select_last_logits_row	新增	TP all-gather 补全 vocab；CP 按序列分片取最后一行
normalize_input / 落盘分发	新增	五种输入形态归一；按扩展名写 .jsonl / .json / .txt（无扩展名按 txt）

流程图（双线框 = 与训练共用的部件，单线框 = 前向推理新增，删除线 = 训练有但推理被裁掉的部件）：

yaml(run_mode: predict + inference: 段) ──► Trainer.__init__
        │
        ▼
┌─ Trainer.__init__ ─────────────────────────────────────────────┐
│  run_mode 校验 ∈ {train, predict},非法值 fail-fast             │
│  is_inference = (run_mode == "predict")                          │
│  is_inference 且 config.inference is None → ValueError          │
│        │                                                         │
│        ▼                                                         │
│  ┏━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┓                          │
│  ┃ configure_inference(剥 MTP 层)     ┃ build 前裁 MTP +         │
│  ┃ build_inference_tokenizer          ┃ compress_ratios 尾部     │
│  ┗━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┛                          │
│        ▼                                                         │
│  ┏━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┓                          │
│  ┃ _create_model + _apply_parallelism  ┃ 与训练同代码路径         │
│  ┃ (forward_only=True 透传 build_stages)┃ TP/CP/DP/EP/PP/FSDP    │
│  ┗━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┛                          │
│        ▼                                                         │
│  ~~数据集 / 优化器 / lr / callback / LoRA / MTP-sync~~  ✂ 裁掉  │
│  is_inference 分支:return,跳过训练态 init                       │
└────────────┬────────────────────────────────────────────────────┘
             ▼
┌─ trainer.inference() ──────────────────────────────────────────┐
│  set_train(False) → construct 返回 logits                       │
│  _validate_pp_inference()(PP 路径 fail-fast)                    │
│  _load_inference_checkpoint(HF / MF 格式自动识别)                │
│        ▼                                                         │
│  normalize_input → records 流                                    │
│  _sample_shard_info:按并行形态判分流轴                           │
│        ├─ 单卡/纯 FSDP/PP ──► 不分流,全体 rank 共享同一条 batch 流│
│        └─ dp_replicate > 1 ──► 条带分流到各副本                  │
│        ▼                                                         │
│  generate_from_records → generate_batch 批量贪心解码             │
│  rank 0 落盘(.jsonl 流式 / .json 数组 / .txt 逐行)             │
└─────────────────────────────────────────────────────────────────┘

3.2 类图

classDiagram
    direction TB

    class Trainer {
        +config: TrainConfig
        +model: PreTrainedModel
        +run_mode: Optional~str~
        +train_dataset: Optional~Dataset~
        +eval_dataset: Optional~Dataset~
        +optimizer: Optional~Any~
        +lr_scheduler: Optional~Any~
        +callbacks: Optional~List~
        +compute_loss_func: Optional~Callable~
        +inference(input_data: Optional~Union~, max_new_tokens: Optional~int~, output: Optional~str~, batch_size: Optional~int~, dataset_field: str) list
        -_validate_pp_inference()
        -_load_inference_checkpoint(model: PreTrainedModel)
    }
    class TrainConfig {
        +run_mode: str = train
        +inference: Optional~InferenceConfig~
    }
    class InferenceConfig {
        +max_new_tokens: int = 50
        +batch_size: int = 8
        +output: str = inference_result.jsonl
        +input_data: Optional~Union~
    }
    class inference_module {
        +configure_inference(config: Any)
        +build_inference_tokenizer(config: Any)
        +normalize_input(input_data: Any) list
        +generate_from_records(model: Any, tokenizer: Any, records: Any, dataset_field: str, parallel_dims: Any, max_new_tokens: int, batch_size: int, rank: int, output_path: str, pp_ctx: Any) list
        +generate_batch(model: Any, tokenizer: Any, prompts: Any, pad_token_id: int, parallel_dims: Any, max_new_tokens: int, max_seq_length: int, rank: int, pp_ctx: Any) list
        +gather_vocab_parallel_logits(logits: Any, parallel_dims: Any) Any
        +select_last_logits_row(logits: Any, last_index: int, parallel_dims: Any) Any
        +_sample_shard_info(parallel_dims: Any, pp_ctx: Any) tuple
    }
    class parallelize_model {
        +forward_only bool
    }
    class StageModelBuilder {
        +build_stages(model_config: Any, pp_mesh: Any, tp_mesh: Any, has_backward: bool, dyn_shape: bool) tuple
    }
    class ScaledLossPipelineStage {
        +bwd_send_specs(micro_index: int) list
    }
    class PipelineInferenceContext {
        +schedule: Any
        +is_last_stage: bool
        +group: Any
        +src: int
        +broadcast_tokens(token_ids: Any) list
    }
    class PreTrainedModel {
        <<abstract>>
    }
    class GPTModel {
        +config: TransformerConfig
        +vocab_size: int
        +max_sequence_length: int
        +pre_process: bool
        +post_process: bool
        +share_embeddings_and_output_weights: bool
        +position_embedding_type: str
        +layer_start: int
        +layer_end: int
        +mtp_process: bool
        +return_logits: bool
        +construct(input_ids: Tensor, position_ids: Tensor, labels: Tensor) Tensor
    }
    GPTModel --|> PreTrainedModel

    Trainer *-- TrainConfig : Trainer 自建, 从 yaml/dict 加载
    Trainer o-- PreTrainedModel : 外部传入 model 引用
    TrainConfig o-- InferenceConfig : Optional 字段, 可为 None
    Trainer ..> inference_module : 构造期与推理期两次调用模块函数
    Trainer ..> parallelize_model : forward_only 透传
    parallelize_model ..> StageModelBuilder : 路由至构建器
    StageModelBuilder ..> ScaledLossPipelineStage : 构建 stage 实例
    inference_module ..> GPTModel : 非 PP 路径直接 construct 前向
    inference_module ..> PipelineInferenceContext : PP 路径下条件持有

调用链要点：

• 构造期：Trainer.__init__（入参见类图字段区）→ TrainConfig（读 run_mode/inference）→ inference.py（configure_inference + build_inference_tokenizer）→ parallelize_model → StageModelBuilder → ScaledLossPipelineStage。run_mode="predict" 时构造期旁路训练态 init。
• 推理期：Trainer.inference() → inference.py（generate_from_records → generate_batch）；PP 路径下 generate_from_records 持有 PipelineInferenceContext 驱动 schedule.run 与 broadcast_tokens。
• ScaledLossPipelineStage 训练 / 推理共用，差异仅在构建参数。

3.3 时序图

下列 UML 时序图描述关键路径的参与者协同与消息时序。参与者与类图节点对应：

• Trainer：推理入口对象（类图 Trainer）
• inference：inference.py 模块（类图 inference_module），承载推理专属函数
• Model：GPTModel 实例（类图 GPTModel），前向入口
• Schedule：PP 调度器（类图 PipelineInferenceContext.schedule），PP 路径下驱动 stage 前向
• Rank0：落盘 rank

3.3.1 端到端推理时序（构建期 + 推理期）

从 yaml 加载到结果落盘的全链路。一次性动作（权重加载、set_train）在批循环外，loop 只保留批量解码。

sequenceDiagram
    autonumber
    participant Trainer
    participant inference as inference_module
    participant Model
    participant Schedule
    participant Rank0

    Note over Trainer: build 阶段 Trainer.__init__
    Trainer->>Trainer: 校验 run_mode 属于 train 或 predict
    alt run_mode == predict
        Trainer->>Trainer: is_inference 置 True
        opt config.inference is None
            Trainer->>Trainer: raise ValueError
        end
        Trainer->>inference: configure_inference 传 config
        Note right of inference: 剥 MTP 层 + 截 compress_ratios
        Trainer->>inference: build_inference_tokenizer 传 config
        inference->>inference: 从 load_path 检测 tokenizer 文件
        inference->>inference: 解析 pad_token_id 写回 model.config
        inference-->>Trainer: tokenizer
    else run_mode == train
        Note over Trainer: 走训练路径 行为不变
    end

    Trainer->>Model: _create_model + _apply_parallelism forward_only=is_inference
    Note right of Model: build_stages 参数 has_backward=False dyn_shape=True
    opt is_inference
        Trainer->>Trainer: return 跳过训练态 init
    end

    Note over Trainer: inference 阶段 Trainer.inference
    Trainer->>Model: set_train False
    opt PP enabled
        Trainer->>Trainer: _validate_pp_inference
        Note right of Trainer: interleave/overlap/swap 触发 NotImplementedError
    end
    Trainer->>Trainer: _load_inference_checkpoint 传 model
    alt is_hf_checkpoint load_path
        Trainer->>Trainer: load_hf_checkpoint
    else MF 格式
        Trainer->>Trainer: load_checkpoint
    end

    Trainer->>inference: generate_from_records 传 model tokenizer records pp_ctx
    opt PP enabled
        Trainer->>Trainer: 构造 PipelineInferenceContext
    end

    loop 每个 batch 在 generate_from_records 内部
        inference->>inference: _sample_shard_info 得 shard_id num_shards
        alt PP enabled
            inference->>Schedule: pp_ctx.schedule.run 传相同 kwargs
            Note right of Schedule: first stage 消费输入 last stage 产 logits
        else 非 PP 路径
            inference->>Model: model construct 传 input_ids position_ids
            Note right of Model: 不传 labels 即返 logits
        end
        inference->>inference: select_last_logits_row 与 gather_vocab_parallel_logits
        inference->>inference: argmax 取 token, EOS 冻结陪跑
    end

    inference->>Rank0: all_gather_object 汇聚 按 sample_idx 去重排序
    Rank0->>Rank0: 落盘 .jsonl / .json / .txt
    Trainer->>Trainer: barrier 然后 destroy_process_group

3.3.2 关键交互要点

端到端时序图（3.3.1）覆盖了构建期与推理期的全链路。下面三点是 PR 实现的关键交互细节，对应 inference.py 的热路径：

• 批量解码（generate_batch）：每步将活跃样本右 pad 到最长，一次 [bs, L] 前向；逐活跃样本经 select_last_logits_row（CP）与 gather_vocab_parallel_logits（TP）取最后逻辑 logits 行后 argmax。某样本 EOS 或超长即冻结为 pad 行陪跑，输出丢弃；finished[] 派生自各 rank 相同的 all-gathered logits，故 if all(finished): break 判定集体一致，无需节拍器。

• PP 单步（PipelineInferenceContext）：pp>1 时各 rank 传相同 kwargs 调 schedule.run，first stage 消费输入，last stage 产出 logits 并选出 per-sample token，再经 broadcast_object_list 广播回全部 stage 供下一步 append。src 取本线 last stage 的全局 rank（get_rank_list_along_axis("pp")[-1]，与 has_last 同源）——MindSpore broadcast 的 src 语义为全局 rank，多 PP 线下按局部 rank 或世界组广播会错源。ScaledLossPipelineStage.bwd_send_specs 覆写返回空 spec list，使调度器为非首 stage 排的 BWD_SEND 步变 no-op，训练路径不受影响。

forward-only 在 hyper-parallel 侧若仍触发梯度归约，可能需要补 execute_reduce_grad 直返的覆写；POC 测试矩阵未覆盖到此路径，作为已知待验证项。

• DP 分流汇聚（_sample_shard_info + generate_from_records finally 块）：dp_replicate > 1 且非 PP 时，各副本按 shard_id 条带式认领样本各自组批解码，节奏独立、副本间零通信；解码完成后 all_gather_object 跨 world group 汇聚各副本的 (payload, text_line) 条目，由 _merge_shard_entries 按 sample_idx 去重排序（tp/cp siblings 报相同条目取首个），rank 0 按输入顺序落盘。PP 开启时不分片，各 PP 线解码全量记录靠 sample_idx 去重——功能正确但 DP 并发度浪费。

3.4 安全隐私与 DFX 设计

• 兼容性：旧 yaml 无 inference: 段、run_mode 缺省为 train 时训练行为不变；InferenceConfig 是新加字段，训练 / 推理 yaml 互不干扰。
• 可维护性：训练侧 init 的跳过集中在 Trainer.__init__ 的 is_inference 分支；_load_inference_checkpoint 独立于 resume 导向的 _load_checkpoint，互不污染。
• 可测试性：8 卡冒烟产物 md5 作为"配置改了但行为不变"的回归基准；五种输入形态分别有 baseline。
• 可靠性：边界校验在 config 加载期 fail-fast；rank 0 单点写输出，无分布式锁需求；权重加载复用既有 is_hf_checkpoint / get_checkpoint_path 工具；tokenize 失败逐 rank 记 logger.error 并写 error 行，不中断整轮推理。
• 训练隔离：全部 is_inference 判断均为"推理旁路"方向（推理跳过训练态初始化），不设 run_mode: predict 时训练路径与既有代码逐行一致。
• 输出文件权限：写盘统一走 _open_output_file（os.open + tools/utils.FILE_PERMISSION，0o640，对齐静态图 BaseTrainer.predict_process）；O_CREAT 只对新建文件生效，覆盖已有文件不改其权限。

3.5 编程与调用设计

3.5.1 编程模型基本设计

• 环境：Python 3.11，MindSpore 2.10.0，msrun 8 worker，Ascend 910B2 × 8，CANN 9.1.0-beta.3。
• 开发约束：yaml 是调用方唯一入口（Trainer.inference() 形参只做语义有意义的覆写）；checkpoint.load_path 必填（同时作为权重与 tokenizer 目录）；PP 推理需 interleave_num=1、overlap_b_f=false、swap=false（校验放 _validate_pp_inference，fail-fast）。

3.5.2 接口定义与设计

3.5.2.1 输入格式

input_data 是唯一数据入口，支持五种输入形态，由 normalize_input 统一归一为 list[dict]：单条 str（字面 prompt，或文件路径——.jsonl 按 jsonl 解析，其他后缀按纯文本每行一条）、单条 dict、list[dict] / list[str]。其中结构化输入（单条 dict / list[dict] / .jsonl 文件）超出静态图推理入口的能力，并带出配套参数 dataset_field（record 中 prompt 字段名，默认 "instruction"，由 Trainer.inference() 形参确认）；静态图纯文本输入无此概念。混合类型 list 与非 str/dict/list 类型触发 TypeError。

3.5.2.2 Trainer.inference()

接口描述：PyNative 模型文本生成入口，贪心（argmax）解码。

接口原型：

def inference(self, input_data=None, max_new_tokens=None, output=None,
              batch_size=None, dataset_field="instruction") -> list:

输入/输出参数：

参数名称	输入/输出	类型	描述	取值范围
input_data	输入	Optional[Union[str, list, dict]]	推理输入，五种形态（见 3.5.2.1）	None 时取 yaml inference.input_data
max_new_tokens	输入	Optional[int]	生成 token 数上限	> 0；None 取 yaml
output	输入	Optional[str]	输出文件路径，按扩展名分发	.jsonl / .json / .txt / 无扩展名（按 txt）；None 取 yaml
batch_size	输入	Optional[int]	每批前向解码的样本数	> 0；None 取 yaml inference.batch_size
dataset_field	输入	str	record 中 prompt 字段名	默认 "instruction"
返回值	输出	list[(sample_idx, generated_ids)]	逐 sample 生成结果	—

异常处理：checkpoint.load_path 为空 / 缺 tokenizer 文件 / output 扩展名不识别 / max_new_tokens ≤ 0 / batch_size ≤ 0 / run_mode 非法时 ValueError；run_mode: predict 但缺 inference: 段时 ValueError。tokenize 失败的样本以 logger.error 逐 rank 记录并写 error 行，不中断批量。前向失败（generate_batch 抛异常）当前会传播出整批推理，不在容错范围内——与 UC2 "不中断批量" 意图存在偏差。

约束说明：需在 is_inference=True 的 Trainer 上调用（yaml run_mode: predict）；解码固定贪心 argmax，各 rank 从相同的 all-gathered logits 取同一 token，多卡确定可复现。

batch_size 跨 rank 一致性当前只靠文档约束（"all ranks share the batch boundaries"），代码未做 all-reduce 校验。yaml 同源时天然一致，但程序化 Trainer.inference(batch_size=X) 覆写若在部分 rank 上调用会静默 desync。建议入口加一次 broadcast(batch_size) 或 all-reduce 校验。

调用参考代码：

trainer = Trainer(config)          # yaml 含 run_mode: predict
results = trainer.inference()      # 全走 yaml
results = trainer.inference(input_data="hello", max_new_tokens=10)  # 覆写

3.5.2.3 InferenceConfig（yaml 配置项）

接口描述：inference: 段配置项，挂在 TrainConfig.inference 下。权重与 tokenizer 共用 checkpoint.load_path 单一入口（HF 权重同目录自带词表；MF 格式权重需将词表文件拷入该目录）。

字段	类型	默认	描述
max_new_tokens	int	50	生成 token 数上限
batch_size	int	8	每批 [batch, seq] 前向解码的样本数；各 rank 必须共享同一取值
output	str	"inference_result.jsonl"	输出路径，按扩展名分发：.jsonl 流式逐行 / .json 数组 / .txt（或无扩展名）一行一条
input_data	str / list / dict	null	推理输入，五种形态见 3.5.2.1

run_mode 字段挂在 TrainConfig 顶层（非 InferenceConfig），取值 train（默认）/ predict，与图模式字段名对齐，非法值构造期报错。

__post_init__ 校验：max_new_tokens > 0、batch_size > 0、output 扩展名 ∈ {.jsonl, .json, .txt, 无扩展名}。超长输入不做截断：tokenize 后 padded 长度超过 model.seq_length 的样本按 error 行记录并跳过（fail-fast，与官方 dsv4-flash generate.py 行为一致）。

3.5.2.4 normalize_input 输入形态

接口描述：统一五种输入形态为 list[dict]。

输入	处理
str	字面 prompt；若为文件路径，.jsonl 按 jsonl 解析，其他后缀按纯文本每行一条
dict	单条 record
list[dict] / list[str]	批量；list[str] 每条 wrap 成 {"instruction": s}

3.5.3 使用说明

# 训练 yaml 基础上追加:
run_mode: 'predict'                # 路由字段,顶层
checkpoint:
  load_path: /path/to/weights_with_tokenizer   # HF ckpt 自带词表;MF ckpt 需将词表文件拷入该目录
inference:
  input_data: /path/to/prompts.txt   # 或 str / dict / list
  max_new_tokens: 128
  batch_size: 8
  output: inference_result.jsonl

bash mindformers/scripts/msrun_launcher.sh \
  "python run_mindformer.py --mode 1 --config test.yaml" 8

使用约束：

• checkpoint.load_path 必填，空则 ValueError（推理在随机初始化上跑会产出垃圾输出）；该目录须同时含权重与 tokenizer 文件（tokenizer.json / tokenizer_config.json / vocab.json / merges.txt / tokenizer.model / spiece.model 任一命中即可），缺文件报错指明。
• output 扩展名不识别则 ValueError（.jsonl / .json / .txt / 无扩展名可配）。
• 训练 yaml / 推理 yaml / tokenizer 三者 vocab_size 必须一致（FSDP 按 vocab 分片，不一致 shape mismatch）。
• .json 输出收集到 list 末尾一次性 dump，sample 数 > 10k 有内存压力，大规模用 .jsonl。
• 解码固定贪心（argmax）；temperature / top_k / top_p / repetition_penalty 等采样类参数不在本方案范围（见非目标）。
• 异常逐 rank 记 logger.error（tokenize 失败写 error 行）落盘，tokenize 失败不中断整轮推理；前向失败当前会中断整批。

4. 测试设计

测试环境：8 卡 Atlas 910B2 msrun，DSV4 tiny 训练 20 步另存的权重；冒烟脚本 predict/_script/run_test.sh。

4.1 一致性回归（必跑）

• 输入构造：固定 3 条 sample prompt，8 卡默认切分跑推理。
• 预期结果：输出 md5 = f00efec1d859bca98c22e77e54da2516，作为"行为不变"改动的复现门禁。

4.2 多格式输入

• 输入构造：同一 prompt 分别用五种输入形态（字面 str / txt 文件 / .jsonl 文件 / dict / list[str]）传入。
• 预期结果：五种形态各能跑通；同 prompt 不同入口的输出 ids 逐位一致。

4.3 输出格式

• 输入构造：同一推理结果分别指定 .jsonl / .json / .txt 输出路径；另用非法扩展名（如 .xml）触发校验。
• 预期结果：.jsonl 流式逐行写、.json 数组一次性 dump、.txt 逐行写；非法扩展名构造期 ValueError；新建文件权限 0o640。

4.4 训练回归

• 输入构造：yaml 不带 inference 段、run_mode 缺省（默认 train），正常训练配置跑训练流。
• 预期结果：训练流照常、训练 yaml 零改动；训练分支隔离审计——全部 is_inference 判断均为"推理旁路"方向，训练路径与 upstream 代码逐行一致。

4.5 并行继承矩阵

• 输入构造：9 种切分组合 + 非均衡 PP：DP8 / TP=2 / EP=2 / CP=2 / TP×CP / PP=4 / TP×PP / EP×PP / PP=2 非均衡（10/6 层），各组合加载训练 20 步另存的基线权重。
• 预期结果：各组合输出 md5 与基线一致（跨策略 reshard 后逐位等价）。

4.6 DP 样本分流

• 输入构造：replicate / HSDP（data_parallel_shard: 2，4 副本）3 条样本；对照纯 FSDP（data_parallel_shard: -1）。
• 预期结果：replicate 模式下 3 条样本条带分流到各副本、各自组批自由解码，md5 与基线一致；纯 FSDP 全体共享一个 batch 流。

4.7 PP 边界

• 输入构造：run_mode: predict + PP enabled，分别置 interleave_num > 1 / overlap_b_f=true / swap=true；对照 pipeline_parallel=1（单模型路径）。
• 预期结果：三种不支持组合在 _validate_pp_inference 入口 fail-fast 拒绝（NotImplementedError），报错信息指向对应开关；pipeline_parallel=1 走原单模型路径不受影响。

4.8 PP 推理

• 输入构造：pp=2（1f1b 调度）3 样本贪心解码；对照 pp=1（单模型路径）。
• 预期结果：三样本全部生成成功，输出与 pp=1 逐字节一致——PP 切分下解码循环不失步，结果与单模型路径等价。

4.9 PP 组内广播验证

• 输入构造：pp=2×副本4 与 pp=4×副本2 两种 PP 线配置，同权重同 prompt；对照 FSDP8 基线。
• 预期结果：两种 PP 线配置输出 md5 互同 e6bf3a88b9af964318d18735c30dc2cf，且与改造前实现逐字节一致；FSDP8 基线与 PP 输出 9/10 逐字节一致，唯一分叉样本为尾部近 tie 位 token 重排（源于 FSDP 分片归约顺序，与 PP 路径无关）。

4.10 Qwen3-8B 真机权重验证

• 输入构造：Qwen3-8B 官方 HF safetensors 权重 + FSDP8，10 条 prompt（算术 / 世界知识 / 古诗词 / 中英互译），max_new_tokens: 32 贪心解码。
• 预期结果：可判定任务全部正确（算术结果对、首都知识对、诗句续写对、中英翻译对等）。

4.11 单卡 UT

• 输入构造：pytest 纯逻辑单测（不依赖多卡）。
• 预期结果：本阶段不随仓交付（单测基线另行建设），多卡行为面（真实 TP/CP/PP 通信、md5 基线、跨策略加载）由 8 卡冒烟与 ST 覆盖。

5. 缺点与风险 (可选)

6. 现有技术 (可选)

7. 未解决问题 (可选)