[xs1997zju]

Behavior

SP ranks currently repeat the eager DataLoader pipeline for the same data-parallel replica. Add an opt-in DataloaderConfig.sp_media_dedup=False switch: every rank still completes the existing dataset/packing initialization collectives, then only the SP leader creates the runtime sampler/collator/DataLoader. Each step broadcasts the complete collated batch over the existing SP group before the original _prepare_model_input / SP split / loss / forward path.

dataloader_cfg = DataloaderConfig(
    # existing dataset/collator/packing options
    sp_media_dedup=True,
)

Production code reads this config field. The experimental SP_MEDIA_DEDUP environment variable is parsed only by the external recipe; this PR does not introduce a production environment-variable interface.

• Disabled/SP1 paths retain the original behavior; enabled mode currently requires mapped DataloaderConfig and TP1.
• Dynamically traverse actual batch/SequenceContext fields; broadcast small metadata and tensor payloads separately, preserving repeated tensor references. Followers reconstruct the complete CPU batch; runtime mesh is bound locally by the original preparation path.
• Keep one fetching/epoch loop. Mirror iterator base-seed RNG consumption on followers, including persistent-worker reuse. XTuner's generator wrapper delays the underlying iterator creation until next().
• Followers skip DataLoader checkpoint operations; leaders retain the existing DP state aggregation and original map_location=DEVICE restore. Synchronize loader length for total-epoch mode.
• No collator/model/engine changes, new production process group, lazy-media path, scatter, async pipeline, error-coordination protocol, or experiment-side allocator hook.

Why seed_draw performs a discarded random draw

BASE creates a loader iterator on every rank. PyTorch draws an int64 base_seed during underlying iterator creation; the mapped loader supplies no dedicated generator, so this advances that rank's default CPU RNG. Dedup followers skip iterator creation. Without compensation, they would also skip this RNG consumption and later CPU random operations could move relative to BASE.

# Leader: check before original fetch creates the underlying iterator.
seed_draw = will_create_underlying_iterator()
batch = original_fetch() if sp_leader else None
# Carry a boolean event alongside the complete batch, not a seed value.
batch, epoch, seed_draw = broadcast_batch(batch, group, ...)
if not sp_leader and seed_draw:
    torch.empty((), dtype=torch.int64).random_()  # draw once, discard the result

The follower only performs the otherwise-missing random draw; it creates no worker and does not reseed. First iterator creation draws once. Persistent workers reuse the iterator across epochs without another base-seed draw; non-persistent rebuilds draw again. A new process after restart follows the same creation rule. Ordinary per-step fetches do not repeatedly compensate. This preserves the corresponding RNG progression, not a full RNG-state checkpoint/restore or a requirement that different ranks have identical states. XTuner's __iter__ is a generator, so creating its outer iterator does not yet execute PyTorch iterator creation; the current predicate still runs before the base-seed draw.

Validation

On the actual PR tree rebased onto InternLM/xtuner:main at 8e6c1824:

• 27 passed, 6 deselected, 131.47s: new leader-only loader/transport CPU tests plus existing DataLoader and GLM53 VL collator regressions. GPUs explicitly hidden. The excluded existing tests require GPU execution.
• CPU checks cover DP2×SP2, workers0/4, SP1/off, actual DataLoader state save/load and continued consumption across an epoch, RNG parity, image/video fields, nested arbitrary tensor shapes/dtypes, aliases, and launcher failure propagation.
• ruff check, ruff format --check on all four changed files, and git diff --check pass.
• The optional real tokenizer/processor media test uses XTUNER_TEST_GLM53_MODEL; it was enabled in this run rather than relying on a hard-coded machine path.

CUDA_VISIBLE_DEVICES='' PYTHONPATH="$PWD" GLOO_SOCKET_IFNAME=lo \
XTUNER_TEST_GLM53_MODEL=/path/to/local/GLM-5.3-Flash-25B \
python -m pytest tests/train/test_sp_leader_loader.py \
  tests/datasets/test_dataloader.py tests/datasets/test_glm53_vl_collator.py \
  -k 'not test_dataloader_resume_multi_process' -q

GPU accuracy, real restart, performance and memory evidence

The reviewed experiment snapshot was tested on 8×H200 with PyTorch2.9.1+cu128, a five-layer GLM-5.3-Flash-25B configuration, EP4, FP8, frozen vision/projector and four loader workers. Each SP setting compares its own BASE and SP arms; different SP settings use different length/batch configurations.

SP / DP	Pack/max length	Global batch	Continuous BASE / SP step median (s), steps6–20	R1 steady CPU PSS BASE→SP
8 / 1	524288	1	31.09 / 29.09	399.7→278.2 GiB
4 / 2	262144	2	22.55 / 21.67	340.1→197.9 GiB
2 / 4	131072	4	18.65 / 18.30	189.5→164.6 GiB

Both BASE and SP have continuous40-step references and real20+20 checkpoint/restart trajectories: save complete native training state at step20, exit, start a new torchrun with explicit checkpoint, consume steps21–40. The marker below is the restart boundary. Rank0 LLM/MTP loss and gradient-norm curves closely track; observed differences were reviewed and accepted, without claiming bitwise equality. SP4's resumed loss/gradient maximum differences exceed the single observed BASE-pair range; raw differences are included rather than describing them as a proven noise bound.

Performance is window-dependent: SP8 initial steady step medians improve about6% and resumed medians are level; SP4 resumed median is about4.1% slower, SP2 R1 about3.6% slower. These short runs do not establish a uniform speedup or a stable regression. Full wall time is separate: reported SP8 BASE20+20=29.2min vs SP32.6min. GPU peaks remain near the H200 physical limit in both arms. See all three settings' step-time/TGS tables and memory time series in the linked evidence.

Full evidence: three settings, accuracy/grad/time/memory plots and per-rank/per-step CSVs

Scope of the evidence

• GPU results belong to the experiment snapshot; the feature was then ported to current upstream main, where the CPU tests above were rerun. Bilingual-comment edits do not change runtime behavior.
• SP8 had allocation failures/retries at the memory limit. Both resumed arms use an external one-time empty_cache() after first resumed backward/before clipping; that hook is not part of this patch. Failed-attempt logs were removed by the experimental launcher, so successful runs alone do not establish robustness at that margin.
• OSS and TP>1 were not validated. BASE does not save all training RNG state.
• PSS summaries are observed equal-length R1 comparisons. The inherited sampler broadly selects training processes rather than enforcing one launcher PID, so unrelated-process exclusion is not independently guaranteed. Initialization still runs on all ranks and CPU peaks are not eliminated.

Evidence lives on a separate fork branch pinned to commit ea9118ab6e0e9ea9e47ee9fb372b4ad553a11702, keeping plots/metrics out of the four-file production/test patch.

[ShilohYu]

[可诊断性] total_step 广播块建议移到 total_epoch/total_step 断言之后

问题:新加的 total_step 广播块在两行断言之前。dedup 开启且 total_epoch、total_step 都未设时,loader_len * cast(int, self._total_epoch) 会对 None 做乘法,抛 TypeError: unsupported operand type(s),把原有的友好断言信息盖掉。

例子/实测:代码阅读确认执行顺序:广播块 trainer.py:730-734,断言在 737-739;cast 在运行时不做转换,int * None 必抛 TypeError。

[ShilohYu]

[设计] 下沉 seed_draw 谓词、收拢补抽(建议性,非阻塞)

问题: _iterator_seed_draw 直接读 self._dataloader._iterator 私有属性——torch 改内部实现时会静默失效。

建议:

1. 谓词下沉为 Dataloader 的具名方法,补抽动作收进 sp_batch_broadcast,Trainer 只调具名接口:

# datasets/dataloader.py
class Dataloader(torch.utils.data.DataLoader, BaseDataloader):
    def will_consume_base_seed(self) -> bool:
        """下一次创建底层 iterator 是否会从默认 CPU RNG 抽一次 base_seed。"""
        return not self.persistent_workers or self._iterator is None

# sp_batch_broadcast.py
def mirror_base_seed_draw() -> None:
    """follower 补一次与 BASE 对齐的默认 CPU RNG 消耗,结果丢弃。"""
    torch.empty((), dtype=torch.int64).random_()

# trainer.py::_data_iter(两处替换)
seed_draw = self._sp_media_dedup and data_iter is not None and self._dataloader.will_consume_base_seed()
...
if self._dataloader is None and seed_draw:
    mirror_base_seed_draw()

同时补一组针对谓词的单测,因为谓词的正确性依赖 torch 的私有实现细节——_iterator 只在复用分支被赋值(dataloader.py:487-492),这不是公开契约:torch 升级可以在公开行为不变的前提下让谓词静默算错,而算错的后果是 follower 的 RNG 流无声偏移(不报错、现场离得远)。单测按真值表覆盖:num_workers 取 0(非复用模式)和 4(复用模式)两种配置,各测「首迭代」与「跨 epoch」两个时刻,共 4 个用例——预期:首迭代恒 True;跨 epoch 时非复用为 True(每轮重建)、复用为 False(复用不重抽)。任一格子翻转都会立刻发现问题,而不是等 RNG 错位才暴露。