环境信息

- 昇腾硬件：8 × 910B2
- CANN：25.3.RC1
- torch_npu：2.7.1
- 代码版本：master 分支，HEAD 852ddffd（2026-08-07）
- 模型：Qwen3.5-9B（LoRA r=8）

问题描述

当 loss_cfg.loss_type: per_token_loss 时，训练侧会用 PrefetchGradAccDataLoader 包裹基础 dataloader。其 load_state_dict 签名为 keyword-only（**kwargs），而 train_engine.py load() 在 resume 时以位置参数调用它，二者不匹配，断点续训时全 rank 崩溃 TypeError，退出码 1，无法从 checkpoint 恢复训练。

根因

mindspeed_mm/fsdp/data/dataloader/dataloader.py，PrefetchGradAccDataLoader.load_state_dict 为 keyword-only：

https://gitcode.com/Ascend/MindSpeed-MM/blob/master/mindspeed_mm/fsdp/data/dataloader/dataloader.py#L252-L253

def load_state_dict(self, **kwargs):                 # ← keyword-only
    self.base_dataloader.load_state_dict(**kwargs)

mindspeed_mm/fsdp/train/train_engine.py load() 在 resume 时以位置参数调用：

https://gitcode.com/Ascend/MindSpeed-MM/blob/master/mindspeed_mm/fsdp/train/train_engine.py#L377-L378

if self.train_dataloader is not None:
    self.train_dataloader.load_state_dict(state["extra_state"]["train_dataloader"])  # ← 位置传参

激活条件（关键）

PrefetchGradAccDataLoader 仅在 loss_type == "per_token_loss" 时启用，否则训练用的是来自 torchdata 库的基础 StatefulDataLoader（其 load_state_dict(state_dict) 接受位置参数，调用不冲突，故不崩）：

https://gitcode.com/Ascend/MindSpeed-MM/blob/master/mindspeed_mm/fsdp/train/trainer.py#L381-L385

train_dataloader = build_dataloader(train_dataset)
if args.features.loss_cfg.loss_type == "per_token_loss":      # ← 仅此 loss_type 才包裹
    train_dataloader = PrefetchGradAccDataLoader(
        train_dataloader, grad_acc_step=args.training.gradient_accumulation_steps
    )

因此本 bug 仅在 per_token_loss 配置下浮现（合法且常见的配置，含 Intern-S2 等官方示例）。default loss_type 下用基础 StatefulDataLoader，位置调用 OK，不崩——这也是为何该 bug 非普遍触发但确属真实缺陷。

复现步骤

1. Qwen3.5-9B LoRA 配置，关键字段：

features:
  loss_cfg:
    loss_type: per_token_loss     # ← 激活 PrefetchGradAccDataLoader，触发条件
training:
  lora:
    enable: true
    r: 8
  save: /path/to/save
  no_save_optim: false            # run1 存全量（含 optim/rng），便于 resume
  no_save_rng: false
  no_load_optim: true             # run2 从 run1 的 dcp resume
  no_load_rng: true
  save_interval: 2
  train_iters: 3
  load: /path/to/save             # run2 resume 用

2. run1 跑 3 iter 存全量 DCP（含 optim/rng）；run2 设 train_iters: 5 并指向 run1 的 save 目录续训：

torchrun --nproc_per_node 8 mindspeed_mm/fsdp/train/trainer.py <config_run2>.yaml

预期行为

run2 从 run1 的 iter_0000003 干净 resume，iteration 与 consumed_train_samples 连续（如 iter4/5 → iter5/5，consumed 64 → 80），rc=0。

实际行为

run2 模型权重成功 Loaded checkpoint 后，到 dataloader 恢复处即崩：

[rank0]: Loaded checkpoint from /path/to/save/iter_0000003
[rank0]:   File ".../train_engine.py", line 378, in load
[rank0]:     self.train_dataloader.load_state_dict(state["extra_state"]["train_dataloader"])
[rank0]: TypeError: PrefetchGradAccDataLoader.load_state_dict() takes 1 positional argument but 2 were given

退出码 1，全 rank 崩。任何 per_token_loss 配置的断点续训都会触发。

建议修复

让 PrefetchGradAccDataLoader.load_state_dict 接受位置参数，与调用方及被包裹的基础 StatefulDataLoader（均按位置传参）签名对齐：

     def state_dict(self):
         return self.base_dataloader.state_dict()

-    def load_state_dict(self, **kwargs):
-        self.base_dataloader.load_state_dict(**kwargs)
+    def load_state_dict(self, state_dict=None):
+        # Accept a positional state_dict (the train_engine caller passes it
+        # positionally). Forward to the wrapped StatefulDataLoader which also
+        # takes a single positional state_dict argument.
+        if state_dict is not None:
+            self.base_dataloader.load_state_dict(state_dict)

该签名同时兼容位置调用（train_engine.py:378 现状）与原 keyword 调用（若存在），不引入行为变化。