Fix rank-local dataloader checkpoint resume

fastvideo/dataset/parquet_dataset_streaming_style.py

@@ -211,6 +211,107 @@ def _contiguous_row_group_shards(row_groups: list[dict[str, Any]], num_shards: i
     return shards
 
 
+def reconstruct_streaming_dataset_state(
+    manifest: dict[str, Any],
+    *,
+    global_rank: int,
+    world_size: int,
+    sp_world_size: int,
+    num_workers: int,
+    batch_size: int,
+    read_batch_size: int,
+    seed: int,
+    shuffle_row_groups: bool,
+    yielded_samples: int,
+    worker_id: int = 0,
+) -> dict[str, Any]:
+    """Reconstruct an epoch-zero streaming cursor without reading Parquet.
+
+    This is intentionally strict and is meant for audited migration of legacy
+    checkpoints whose rank-local dataloader state was incorrectly stored in a
+    shared DCP key.  It refuses epoch rollover because the exact wrapper state
+    at that boundary depends on whether the iterator has observed StopIteration.
+    """
+    world_size = int(world_size)
+    sp_world_size = int(sp_world_size)
+    global_rank = int(global_rank)
+    num_workers = max(int(num_workers), 1)
+    batch_size = int(batch_size)
+    read_batch_size = int(read_batch_size)
+    yielded_samples = int(yielded_samples)
+    worker_id = int(worker_id)
+    if world_size < 1 or sp_world_size < 1 or world_size % sp_world_size:
+        raise ValueError("world_size must be positive and divisible by sp_world_size")
+    if not 0 <= global_rank < world_size:
+        raise ValueError("global_rank is outside the requested world")
+    if not 0 <= worker_id < num_workers:
+        raise ValueError("worker_id is outside the requested worker layout")
+    if batch_size < 1 or read_batch_size < 1:
+        raise ValueError("batch_size and read_batch_size must be positive")
+    if yielded_samples < 0 or yielded_samples % batch_size:
+        raise ValueError("yielded_samples must be a non-negative batch multiple")
+
+    row_groups = manifest.get("row_groups")
+    if not isinstance(row_groups, list) or not row_groups:
+        raise ValueError("Streaming manifest has no row groups")
+    num_sp_groups = world_size // sp_world_size
+    shards = _contiguous_row_group_shards(
+        row_groups,
+        num_sp_groups * num_workers,
+    )
+    shard_rows = [sum(int(item["rows"]) for item in shard) for shard in shards]
+    samples_per_worker = min(shard_rows) // batch_size * batch_size
+    if samples_per_worker == 0:
+        raise ValueError("Dataset is too small for the requested topology")
+    if yielded_samples >= samples_per_worker:
+        raise ValueError(
+            "Legacy cursor reconstruction only supports the first epoch and "
+            "requires yielded_samples < samples_per_worker"
+        )
+
+    sp_group_index = global_rank // sp_world_size
+    shard_index = sp_group_index * num_workers + worker_id
+    shard = list(shards[shard_index])
+    if shuffle_row_groups:
+        random.Random(int(seed)).shuffle(shard)
+
+    row_group_position = 0
+    row_offset = 0
+    remaining = yielded_samples
+    if remaining:
+        for position, item in enumerate(shard):
+            rows = int(item["rows"])
+            if remaining <= rows:
+                row_group_position = position
+                row_offset = remaining
+                break
+            remaining -= rows
+        else:
+            raise ValueError("yielded_samples exceeds the selected streaming shard")
+
+    return {
+        "version": _STATE_VERSION,
+        "manifest_fingerprint": str(manifest["fingerprint"]),
+        "topology": {
+            "global_rank": global_rank,
+            "world_size": world_size,
+            "sp_world_size": sp_world_size,
+            "sp_group_index": sp_group_index,
+            "num_workers": num_workers,
+            "batch_size": batch_size,
+            "read_batch_size": read_batch_size,
+            "seed": int(seed),
+            "shuffle_row_groups": bool(shuffle_row_groups),
+            "samples_per_worker": samples_per_worker,
+        },
+        "worker_id": worker_id,
+        "epoch": 0,
+        "row_group_position": row_group_position,
+        "row_offset": row_offset,
+        "yielded_samples": yielded_samples,
+    }
+
+
 class LatentsParquetStreamingDataset(IterableDataset):
     """Sequential row-group loader with deterministic DP/SP sharding.
 

fastvideo/tests/dataset/test_parquet_dataset_streaming_style.py

@@ -112,6 +112,57 @@ def test_state_dict_resumes_at_exact_next_batch(
     assert [item["id"] for item in actual] == [item["id"] for item in expected]
 
 
+def test_reconstructed_state_matches_continuous_cursor(
+    parquet_root: Path, tmp_path: Path, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    _patch_dist(monkeypatch)
+    continuous = _dataset(parquet_root, tmp_path / "owned-cache")
+    iterator = iter(continuous)
+    for _ in range(3):
+        next(iterator)
+    expected_state = continuous.state_dict()
+    expected_next = next(iterator)["info_list"]
+
+    reconstructed_state = streaming.reconstruct_streaming_dataset_state(
+        continuous.manifest,
+        global_rank=0,
+        world_size=1,
+        sp_world_size=1,
+        num_workers=0,
+        batch_size=2,
+        read_batch_size=2,
+        seed=42,
+        shuffle_row_groups=False,
+        yielded_samples=6,
+    )
+    assert reconstructed_state == expected_state
+
+    resumed = _dataset(parquet_root, tmp_path / "owned-cache")
+    resumed.load_state_dict(reconstructed_state)
+    actual_next = next(iter(resumed))["info_list"]
+    assert [item["id"] for item in actual_next] == [item["id"] for item in expected_next]
+
+
+def test_reconstructed_state_refuses_epoch_boundary(
+    parquet_root: Path, tmp_path: Path, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    _patch_dist(monkeypatch)
+    dataset = _dataset(parquet_root, tmp_path / "owned-cache")
+    with pytest.raises(ValueError, match="first epoch"):
+        streaming.reconstruct_streaming_dataset_state(
+            dataset.manifest,
+            global_rank=0,
+            world_size=1,
+            sp_world_size=1,
+            num_workers=0,
+            batch_size=2,
+            read_batch_size=2,
+            seed=42,
+            shuffle_row_groups=False,
+            yielded_samples=dataset.samples_per_worker,
+        )
+
+
 def test_odd_row_group_tails_are_carried_into_next_group(
     parquet_root: Path, tmp_path: Path, monkeypatch: pytest.MonkeyPatch
 ) -> None:

fastvideo/tests/train/utils/test_checkpoint.py

@@ -14,6 +14,10 @@ from pathlib import Path
 from typing import Any
 
 import pytest
+import torch
+import torch.distributed.checkpoint as dcp
+
+import fastvideo.train.utils.checkpoint as checkpoint_module
 
 from fastvideo.train.utils.checkpoint import (
     CheckpointConfig,
@@ -86,6 +90,44 @@ class _MissingLoad:
         return {}
 
 
+class _Method:
+
+    def checkpoint_state(self) -> dict[str, Any]:
+        return {"method": _Full()}
+
+
+class _RecordingStateful:
+
+    def __init__(self, state: dict[str, Any] | None = None) -> None:
+        self.state = dict(state or {})
+        self.loaded: dict[str, Any] | None = None
+
+    def state_dict(self) -> dict[str, Any]:
+        return dict(self.state)
+
+    def load_state_dict(self, state: dict[str, Any]) -> None:
+        self.loaded = dict(state)
+
+
+class _LoaderWithDataset(_RecordingStateful):
+
+    def __init__(self) -> None:
+        super().__init__()
+        self.dataset = _RecordingStateful()
+
+
+class _TensorStateful:
+
+    def __init__(self, value: int) -> None:
+        self.value = torch.tensor([value], dtype=torch.int64)
+
+    def state_dict(self) -> dict[str, Any]:
+        return {"value": self.value}
+
+    def load_state_dict(self, state: dict[str, Any]) -> None:
+        self.value = state["value"]
+
+
 def test_is_stateful_true_for_full_object() -> None:
     assert _is_stateful(_Full()) is True
 
@@ -98,6 +140,150 @@ def test_is_stateful_false_when_missing_load_state_dict() -> None:
     assert _is_stateful(_MissingLoad()) is False
 
 
+def test_dataloader_is_not_part_of_shared_dcp_state(tmp_path: Path) -> None:
+    manager = CheckpointManager(
+        method=_Method(),
+        dataloader=_RecordingStateful({"cursor": 7}),
+        output_dir=str(tmp_path),
+        config=CheckpointConfig(save_steps=1, keep_last=0),
+    )
+    assert set(manager._build_states()) == {"method"}
+
+
+def test_dcp_partial_load_ignores_legacy_dataloader_key(tmp_path: Path) -> None:
+    checkpoint_dir = tmp_path / "dcp"
+    dcp.save(
+        {
+            "method": _TensorStateful(7),
+            "dataloader": _TensorStateful(99),
+        },
+        checkpoint_id=str(checkpoint_dir),
+    )
+    restored = _TensorStateful(-1)
+    dcp.load({"method": restored}, checkpoint_id=str(checkpoint_dir))
+    assert restored.value.item() == 7
+
+
+def test_legacy_single_rank_dataloader_dcp_fallback(tmp_path: Path) -> None:
+    checkpoint_dir = _make_checkpoint_dir(tmp_path, 10)
+    dcp.save(
+        {"dataloader": _TensorStateful(42)},
+        checkpoint_id=str(checkpoint_dir / "dcp"),
+    )
+    restored = _TensorStateful(-1)
+    manager = CheckpointManager(
+        method=_Method(),
+        dataloader=restored,
+        output_dir=str(tmp_path),
+        config=CheckpointConfig(save_steps=1, keep_last=0),
+    )
+    manager._load_dataloader_snapshot(checkpoint_dir, 10)
+    assert restored.value.item() == 42
+
+
+def test_rank_local_dataloader_sidecar_roundtrip(
+    tmp_path: Path, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    monkeypatch.setattr(checkpoint_module, "_rank", lambda: 3)
+    monkeypatch.setattr(checkpoint_module, "_world_size", lambda: 8)
+    checkpoint_dir = _make_checkpoint_dir(tmp_path, 20)
+    original = _RecordingStateful({"cursor": 11})
+    manager = CheckpointManager(
+        method=_Method(),
+        dataloader=original,
+        output_dir=str(tmp_path),
+        config=CheckpointConfig(save_steps=1, keep_last=0),
+    )
+    manager._save_dataloader_snapshot(checkpoint_dir, 20)
+    path = checkpoint_dir / "dataloader_state_rank3.pt"
+    assert path.is_file()
+    assert not list(checkpoint_dir.glob(".*.tmp"))
+
+    restored = _RecordingStateful()
+    restored_manager = CheckpointManager(
+        method=_Method(),
+        dataloader=restored,
+        output_dir=str(tmp_path),
+        config=CheckpointConfig(save_steps=1, keep_last=0),
+    )
+    restored_manager._load_dataloader_snapshot(checkpoint_dir, 20)
+    assert restored.loaded == {"cursor": 11}
+
+
+def test_dataset_only_migration_sidecar_uses_loader_dataset(
+    tmp_path: Path, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    monkeypatch.setattr(checkpoint_module, "_rank", lambda: 1)
+    monkeypatch.setattr(checkpoint_module, "_world_size", lambda: 4)
+    checkpoint_dir = _make_checkpoint_dir(tmp_path, 10)
+    migration_dir = tmp_path / "migration"
+    migration_dir.mkdir()
+    torch.save(
+        {
+            "version": 1,
+            "rank": 1,
+            "world_size": 4,
+            "step": 10,
+            "state_kind": "dataset",
+            "state": {"cursor": 99},
+        },
+        migration_dir / "dataloader_state_rank1.pt",
+    )
+    monkeypatch.setenv("FASTVIDEO_DATALOADER_STATE_DIR", str(migration_dir))
+    loader = _LoaderWithDataset()
+    manager = CheckpointManager(
+        method=_Method(),
+        dataloader=loader,
+        output_dir=str(tmp_path),
+        config=CheckpointConfig(save_steps=1, keep_last=0),
+    )
+    manager._load_dataloader_snapshot(checkpoint_dir, 10)
+    assert loader.loaded is None
+    assert loader.dataset.loaded == {"cursor": 99}
+
+
+def test_dataloader_sidecar_metadata_mismatch_is_rejected(
+    tmp_path: Path, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    monkeypatch.setattr(checkpoint_module, "_rank", lambda: 0)
+    monkeypatch.setattr(checkpoint_module, "_world_size", lambda: 2)
+    checkpoint_dir = _make_checkpoint_dir(tmp_path, 10)
+    torch.save(
+        {
+            "version": 1,
+            "rank": 1,
+            "world_size": 2,
+            "step": 10,
+            "state_kind": "dataloader",
+            "state": {},
+        },
+        checkpoint_dir / "dataloader_state_rank0.pt",
+    )
+    manager = CheckpointManager(
+        method=_Method(),
+        dataloader=_RecordingStateful(),
+        output_dir=str(tmp_path),
+        config=CheckpointConfig(save_steps=1, keep_last=0),
+    )
+    with pytest.raises(ValueError, match="rank mismatch"):
+        manager._load_dataloader_snapshot(checkpoint_dir, 10)
+
+
+def test_missing_rank_local_dataloader_sidecar_fails_closed(
+    tmp_path: Path, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    monkeypatch.setattr(checkpoint_module, "_world_size", lambda: 2)
+    checkpoint_dir = _make_checkpoint_dir(tmp_path, 10)
+    manager = CheckpointManager(
+        method=_Method(),
+        dataloader=_RecordingStateful(),
+        output_dir=str(tmp_path),
+        config=CheckpointConfig(save_steps=1, keep_last=0),
+    )
+    with pytest.raises(RuntimeError, match="Legacy multi-rank DCP dataloader state is unsafe"):
+        manager._load_dataloader_snapshot(checkpoint_dir, 10)
+
+
 # ---------------------------------------------------------------------------
 # B. _parse_step_from_dir
 # ---------------------------------------------------------------------------

fastvideo/train/utils/checkpoint.py

@@ -7,6 +7,7 @@ import os
 import random
 import re
 import shutil
+import tempfile
 from dataclasses import dataclass
 from pathlib import Path
 from typing import Any
@@ -27,6 +28,8 @@ from fastvideo.logger import init_logger
 logger = init_logger(__name__)
 
 _CHECKPOINT_DIR_RE = re.compile(r"^checkpoint-(\d+)$")
+_DATALOADER_STATE_VERSION = 1
+_DATALOADER_STATE_OVERRIDE_ENV = "FASTVIDEO_DATALOADER_STATE_DIR"
 
 
 def _is_stateful(obj: Any) -> bool:
@@ -39,11 +42,36 @@ def _rank() -> int:
     return 0
 
 
+def _world_size() -> int:
+    if dist.is_available() and dist.is_initialized():
+        return int(dist.get_world_size())
+    return 1
+
+
 def _barrier() -> None:
     if dist.is_available() and dist.is_initialized():
         dist.barrier()
 
 
+def _atomic_torch_save(payload: Any, path: Path) -> None:
+    """Atomically publish a rank-local torch payload."""
+    path.parent.mkdir(parents=True, exist_ok=True)
+    fd, temporary_path = tempfile.mkstemp(
+        dir=path.parent,
+        prefix=f".{path.name}.",
+        suffix=".tmp",
+    )
+    try:
+        with os.fdopen(fd, "wb") as handle:
+            torch.save(payload, handle)
+            handle.flush()
+            os.fsync(handle.fileno())
+        os.replace(temporary_path, path)
+    finally:
+        if os.path.exists(temporary_path):
+            os.unlink(temporary_path)
+
+
 def _parse_step_from_dir(checkpoint_dir: Path) -> int:
     match = _CHECKPOINT_DIR_RE.match(checkpoint_dir.name)
     if not match:
@@ -214,10 +242,6 @@ class CheckpointManager:
     def _build_states(self) -> dict[str, Any]:
         states: dict[str, Any] = self.method.checkpoint_state()
 
-        # Dataloader (optional but recommended for exact resume).
-        if _is_stateful(self.dataloader):
-            states["dataloader"] = self.dataloader
-
         # Callback state (e.g. EMA shadow weights, validation RNG).
         if self._callbacks is not None and _is_stateful(self._callbacks):
             states["callbacks"] = _CallbackStateWrapper(self._callbacks, )
@@ -263,6 +287,12 @@ class CheckpointManager:
         dcp.save(states, checkpoint_id=str(dcp_dir))
         _barrier()
 
+        # Dataloader state is rank-local.  Saving it under one shared DCP key
+        # lets DCP deduplicate fields from different ranks into an invalid
+        # chimera.  Keep one atomic sidecar per rank instead.
+        self._save_dataloader_snapshot(checkpoint_dir, step)
+        _barrier()
+
         if self._callbacks is not None:
             self._callbacks.on_checkpoint_save(
                 self.method,
@@ -325,6 +355,115 @@ class CheckpointManager:
             checkpoint_dir / f"rng_state_rank{rank}.pt",
         )
 
+    def _save_dataloader_snapshot(
+        self,
+        checkpoint_dir: Path,
+        step: int,
+    ) -> None:
+        if not _is_stateful(self.dataloader):
+            return
+        rank = _rank()
+        payload = {
+            "version": _DATALOADER_STATE_VERSION,
+            "rank": rank,
+            "world_size": _world_size(),
+            "step": int(step),
+            "state_kind": "dataloader",
+            "state": self.dataloader.state_dict(),
+        }
+        path = checkpoint_dir / f"dataloader_state_rank{rank}.pt"
+        _atomic_torch_save(payload, path)
+        logger.info("Saved rank-local dataloader state to %s", path)
+
+    def _resolve_dataloader_snapshot(
+        self,
+        checkpoint_dir: Path,
+    ) -> Path | None:
+        rank = _rank()
+        default_path = checkpoint_dir / f"dataloader_state_rank{rank}.pt"
+        if default_path.is_file():
+            return default_path
+
+        override_dir = os.environ.get(_DATALOADER_STATE_OVERRIDE_ENV, "").strip()
+        if override_dir:
+            override_path = (Path(os.path.expanduser(override_dir)).resolve() /
+                             f"dataloader_state_rank{rank}.pt")
+            if override_path.is_file():
+                return override_path
+            raise FileNotFoundError(
+                f"{_DATALOADER_STATE_OVERRIDE_ENV} is set, but rank {rank} "
+                f"state is missing: {override_path}"
+            )
+
+        # A single-rank legacy DCP cannot mix state from different ranks, so
+        # loading its original dataloader key remains safe and preserves
+        # backwards compatibility.
+        if _world_size() == 1:
+            return None
+
+        raise RuntimeError(
+            "Checkpoint has no rank-local dataloader state sidecar. Legacy "
+            "multi-rank DCP dataloader state is unsafe because identically "
+            "named rank-local fields can be deduplicated across ranks. "
+            f"Create audited migration sidecars and set {_DATALOADER_STATE_OVERRIDE_ENV}."
+        )
+
+    def _load_dataloader_snapshot(
+        self,
+        checkpoint_dir: Path,
+        step: int,
+    ) -> None:
+        if not _is_stateful(self.dataloader):
+            return
+        path = self._resolve_dataloader_snapshot(checkpoint_dir)
+        if path is None:
+            logger.warning(
+                "Loading legacy single-rank dataloader state from DCP at %s",
+                checkpoint_dir,
+            )
+            dcp.load(
+                {"dataloader": self.dataloader},
+                checkpoint_id=str(checkpoint_dir / "dcp"),
+            )
+            return
+        payload = torch.load(path, map_location="cpu", weights_only=False)
+        if not isinstance(payload, dict):
+            raise ValueError(f"Invalid dataloader state payload: {path}")
+
+        expected = {
+            "version": _DATALOADER_STATE_VERSION,
+            "rank": _rank(),
+            "world_size": _world_size(),
+            "step": int(step),
+        }
+        for field, expected_value in expected.items():
+            actual = payload.get(field)
+            if actual != expected_value:
+                raise ValueError(
+                    f"Dataloader state {field} mismatch in {path}: "
+                    f"expected {expected_value!r}, got {actual!r}"
+                )
+
+        state_kind = payload.get("state_kind")
+        state = payload.get("state")
+        if not isinstance(state, dict):
+            raise ValueError(f"Dataloader state is not a dictionary: {path}")
+        if state_kind == "dataloader":
+            self.dataloader.load_state_dict(state)
+        elif state_kind == "dataset":
+            dataset = getattr(self.dataloader, "dataset", None)
+            if not _is_stateful(dataset):
+                raise ValueError(
+                    "Dataset-only migration state requires a stateful "
+                    f"dataloader.dataset: {path}"
+                )
+            dataset.load_state_dict(state)
+        else:
+            raise ValueError(
+                f"Unsupported dataloader state_kind {state_kind!r} in {path}"
+            )
+        logger.info("Restored rank-local dataloader state from %s", path)
+
     def load_rng_snapshot(
         self,
         checkpoint_path: str,
@@ -388,6 +527,7 @@ class CheckpointManager:
         states = self._build_states()
         logger.info("Loading Phase 2 checkpoint from %s", resolved)
         dcp.load(states, checkpoint_id=str(resolved / "dcp"))
+        self._load_dataloader_snapshot(resolved, step)
         if self._callbacks is not None:
             self._callbacks.on_checkpoint_load(
                 self.method,

scripts/train/migrate_streaming_dataloader_checkpoint.py

@@ -0,0 +1,227 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Create audited rank-local dataloader sidecars for a legacy checkpoint.
+
+The legacy checkpoint is never modified.  This utility only supports the
+deterministic, single-process-per-rank OpenVid streaming layout used by the
+GB200 run and deliberately refuses ambiguous epoch-boundary migrations.
+"""
+
+from __future__ import annotations
+
+import argparse
+import hashlib
+import json
+import os
+from datetime import datetime, timezone
+from pathlib import Path
+from typing import Any
+
+import torch
+import yaml
+
+from fastvideo.dataset.parquet_dataset_streaming_style import (
+    reconstruct_streaming_dataset_state,
+)
+from fastvideo.train.utils.checkpoint import (
+    _DATALOADER_STATE_VERSION,
+    _atomic_torch_save,
+)
+
+
+def _sha256(path: Path) -> str:
+    digest = hashlib.sha256()
+    with path.open("rb") as handle:
+        for chunk in iter(lambda: handle.read(1024 * 1024), b""):
+            digest.update(chunk)
+    return digest.hexdigest()
+
+
+def _nested(mapping: dict[str, Any], *keys: str) -> Any:
+    value: Any = mapping
+    for key in keys:
+        if not isinstance(value, dict) or key not in value:
+            raise KeyError(".".join(keys))
+        value = value[key]
+    return value
+
+
+def _write_json_atomic(path: Path, payload: dict[str, Any]) -> None:
+    temporary = path.with_name(f".{path.name}.tmp")
+    with temporary.open("w", encoding="utf-8") as handle:
+        json.dump(payload, handle, indent=2, sort_keys=True)
+        handle.flush()
+        os.fsync(handle.fileno())
+    os.replace(temporary, path)
+
+
+def _validate_checkpoint_config(
+    checkpoint_config: dict[str, Any],
+    run_config: dict[str, Any],
+) -> None:
+    critical_paths = (
+        ("training", "distributed", "num_gpus"),
+        ("training", "distributed", "sp_size"),
+        ("training", "data", "data_path"),
+        ("training", "data", "dataloader_type"),
+        ("training", "data", "streaming_manifest_path"),
+        ("training", "data", "streaming_read_batch_size"),
+        ("training", "data", "streaming_shuffle_row_groups"),
+        ("training", "data", "dataloader_num_workers"),
+        ("training", "data", "train_batch_size"),
+        ("training", "data", "seed"),
+        ("training", "loop", "gradient_accumulation_steps"),
+    )
+    mismatches = []
+    for path in critical_paths:
+        checkpoint_value = _nested(checkpoint_config, *path)
+        run_value = _nested(run_config, *path)
+        if checkpoint_value != run_value:
+            mismatches.append((".".join(path), checkpoint_value, run_value))
+    if mismatches:
+        formatted = "; ".join(
+            f"{path}: checkpoint={old!r}, config={new!r}"
+            for path, old, new in mismatches
+        )
+        raise ValueError(f"Run config does not match checkpoint metadata: {formatted}")
+
+
+def migrate(args: argparse.Namespace) -> dict[str, Any]:
+    config_path = Path(args.config).expanduser().resolve()
+    checkpoint_dir = Path(args.checkpoint).expanduser().resolve()
+    output_dir = Path(args.output_dir).expanduser().resolve()
+    if output_dir == checkpoint_dir or checkpoint_dir in output_dir.parents:
+        raise ValueError("Migration output must be outside the legacy checkpoint")
+    if not (checkpoint_dir / "dcp").is_dir():
+        raise FileNotFoundError(f"Missing DCP directory: {checkpoint_dir / 'dcp'}")
+
+    metadata_path = checkpoint_dir / "metadata.json"
+    with metadata_path.open(encoding="utf-8") as handle:
+        metadata = json.load(handle)
+    step = int(metadata["step"])
+    if step != int(args.expected_step):
+        raise ValueError(f"Checkpoint step mismatch: expected {args.expected_step}, got {step}")
+
+    with config_path.open(encoding="utf-8") as handle:
+        run_config = yaml.safe_load(handle)
+    if not isinstance(run_config, dict):
+        raise ValueError("Run config must be a YAML mapping")
+    checkpoint_config = metadata.get("config")
+    if not isinstance(checkpoint_config, dict):
+        raise ValueError("Legacy checkpoint metadata has no embedded run config")
+    _validate_checkpoint_config(checkpoint_config, run_config)
+
+    distributed = _nested(run_config, "training", "distributed")
+    data = _nested(run_config, "training", "data")
+    loop = _nested(run_config, "training", "loop")
+    if str(data["dataloader_type"]).strip().lower() != "streaming":
+        raise ValueError("This migration only supports the streaming dataloader")
+    if int(data["dataloader_num_workers"]) != 0:
+        raise ValueError("Legacy migration currently requires dataloader_num_workers=0")
+
+    world_size = int(distributed["num_gpus"])
+    if world_size != int(args.expected_world_size):
+        raise ValueError(
+            f"World-size mismatch: expected {args.expected_world_size}, config has {world_size}"
+        )
+    sp_world_size = int(distributed["sp_size"])
+    batch_size = int(data["train_batch_size"])
+    grad_accum = int(loop["gradient_accumulation_steps"])
+    consumed_batches = int(args.consumed_batches_per_rank)
+    expected_batches = step * grad_accum
+    if consumed_batches != expected_batches:
+        raise ValueError(
+            f"Consumed batch count mismatch: explicitly supplied {consumed_batches}, "
+            f"but step*gradient_accumulation_steps is {expected_batches}"
+        )
+    yielded_samples = consumed_batches * batch_size
+
+    manifest_path = Path(data["streaming_manifest_path"]).expanduser().resolve()
+    with manifest_path.open(encoding="utf-8") as handle:
+        manifest = json.load(handle)
+    dataset_root = os.path.realpath(os.path.expanduser(str(data["data_path"])))
+    if manifest.get("dataset_root") != dataset_root:
+        raise ValueError("Streaming manifest belongs to a different dataset root")
+
+    output_dir.mkdir(parents=True, exist_ok=True)
+    existing = list(output_dir.glob("dataloader_state_rank*.pt"))
+    if existing and not args.overwrite:
+        raise FileExistsError(
+            f"Migration output already contains {len(existing)} sidecars; use --overwrite"
+        )
+
+    sidecars = []
+    for rank in range(world_size):
+        dataset_state = reconstruct_streaming_dataset_state(
+            manifest,
+            global_rank=rank,
+            world_size=world_size,
+            sp_world_size=sp_world_size,
+            num_workers=0,
+            batch_size=batch_size,
+            read_batch_size=int(data["streaming_read_batch_size"]),
+            seed=int(data["seed"]),
+            shuffle_row_groups=bool(data["streaming_shuffle_row_groups"]),
+            yielded_samples=yielded_samples,
+        )
+        payload = {
+            "version": _DATALOADER_STATE_VERSION,
+            "rank": rank,
+            "world_size": world_size,
+            "step": step,
+            "state_kind": "dataset",
+            "state": dataset_state,
+        }
+        path = output_dir / f"dataloader_state_rank{rank}.pt"
+        _atomic_torch_save(payload, path)
+        loaded = torch.load(path, map_location="cpu", weights_only=False)
+        if loaded != payload:
+            raise RuntimeError(f"Sidecar round-trip mismatch: {path}")
+        sidecars.append(
+            {
+                "rank": rank,
+                "path": path.name,
+                "sha256": _sha256(path),
+                "row_group_position": dataset_state["row_group_position"],
+                "row_offset": dataset_state["row_offset"],
+                "yielded_samples": dataset_state["yielded_samples"],
+            }
+        )
+
+    receipt = {
+        "schema": "fastvideo-legacy-streaming-dataloader-migration/v1",
+        "created_at_utc": datetime.now(timezone.utc).isoformat(),
+        "source_checkpoint": str(checkpoint_dir),
+        "source_checkpoint_metadata_sha256": _sha256(metadata_path),
+        "source_config": str(config_path),
+        "source_config_sha256": _sha256(config_path),
+        "manifest": str(manifest_path),
+        "manifest_fingerprint": manifest["fingerprint"],
+        "step": step,
+        "world_size": world_size,
+        "sp_world_size": sp_world_size,
+        "batch_size": batch_size,
+        "gradient_accumulation_steps": grad_accum,
+        "consumed_batches_per_rank": consumed_batches,
+        "yielded_samples_per_rank": yielded_samples,
+        "sidecars": sidecars,
+    }
+    _write_json_atomic(output_dir / "migration_receipt.json", receipt)
+    return receipt
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--config", required=True)
+    parser.add_argument("--checkpoint", required=True)
+    parser.add_argument("--output-dir", required=True)
+    parser.add_argument("--expected-step", required=True, type=int)
+    parser.add_argument("--expected-world-size", required=True, type=int)
+    parser.add_argument("--consumed-batches-per-rank", required=True, type=int)
+    parser.add_argument("--overwrite", action="store_true")
+    args = parser.parse_args()
+    receipt = migrate(args)
+    print(json.dumps(receipt, indent=2, sort_keys=True))
+
+
+if __name__ == "__main__":
+    main()

scripts/train/validate_streaming_dataloader_sidecars.py

@@ -0,0 +1,153 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Validate migrated T2V streaming sidecars under the real process topology."""
+
+from __future__ import annotations
+
+import argparse
+import hashlib
+import json
+import os
+from datetime import datetime, timezone
+from pathlib import Path
+from typing import Any
+
+import pyarrow.parquet as pq
+import torch
+import torch.distributed as dist
+import yaml
+
+from fastvideo.dataset.dataloader.schema import pyarrow_schema_t2v
+from fastvideo.dataset.parquet_dataset_streaming_style import (
+    build_parquet_streaming_style_dataloader,
+)
+from fastvideo.distributed import (
+    cleanup_dist_env_and_memory,
+    maybe_init_distributed_environment_and_model_parallel,
+)
+from fastvideo.train.utils.checkpoint import CheckpointConfig, CheckpointManager
+
+
+def _sha256(path: Path) -> str:
+    digest = hashlib.sha256()
+    with path.open("rb") as handle:
+        for chunk in iter(lambda: handle.read(1024 * 1024), b""):
+            digest.update(chunk)
+    return digest.hexdigest()
+
+
+def _next_manifest_id(dataset: Any, state: dict[str, Any]) -> str:
+    shard = dataset._worker_shard(0)
+    position = int(state["row_group_position"])
+    offset = int(state["row_offset"])
+    item = shard[position]
+    if offset == int(item["rows"]):
+        position += 1
+        offset = 0
+        item = shard[position]
+    parquet_file = pq.ParquetFile(os.path.join(dataset.path, item["file"]))
+    ids = parquet_file.read_row_group(int(item["row_group"]), columns=["id"]).column("id")
+    return str(ids[offset].as_py())
+
+
+def _write_json_atomic(path: Path, payload: dict[str, Any]) -> None:
+    path.parent.mkdir(parents=True, exist_ok=True)
+    temporary = path.with_name(f".{path.name}.tmp")
+    with temporary.open("w", encoding="utf-8") as handle:
+        json.dump(payload, handle, indent=2, sort_keys=True)
+        handle.flush()
+        os.fsync(handle.fileno())
+    os.replace(temporary, path)
+
+
+def validate(args: argparse.Namespace) -> None:
+    maybe_init_distributed_environment_and_model_parallel(tp_size=1, sp_size=1)
+    try:
+        rank = dist.get_rank()
+        world_size = dist.get_world_size()
+        config_path = Path(args.config).expanduser().resolve()
+        checkpoint_dir = Path(args.checkpoint).expanduser().resolve()
+        sidecar_dir = Path(args.sidecar_dir).expanduser().resolve()
+        with config_path.open(encoding="utf-8") as handle:
+            config = yaml.safe_load(handle)
+        training = config["training"]
+        distributed = training["distributed"]
+        data = training["data"]
+        if world_size != int(distributed["num_gpus"]):
+            raise ValueError(
+                f"Runtime world_size={world_size} does not match config num_gpus={distributed['num_gpus']}"
+            )
+        if int(distributed["sp_size"]) != 1:
+            raise ValueError("This validator currently requires sp_size=1")
+
+        dataset, dataloader = build_parquet_streaming_style_dataloader(
+            path=str(data["data_path"]),
+            batch_size=int(data["train_batch_size"]),
+            num_data_workers=int(data["dataloader_num_workers"]),
+            parquet_schema=pyarrow_schema_t2v,
+            manifest_path=str(data["streaming_manifest_path"]),
+            cfg_rate=float(data["training_cfg_rate"]),
+            text_padding_length=512,
+            seed=int(data["seed"]),
+            read_batch_size=int(data["streaming_read_batch_size"]),
+            shuffle_row_groups=bool(data["streaming_shuffle_row_groups"]),
+        )
+        sidecar_path = sidecar_dir / f"dataloader_state_rank{rank}.pt"
+        payload = torch.load(sidecar_path, map_location="cpu", weights_only=False)
+        expected_id = _next_manifest_id(dataset, payload["state"])
+
+        os.environ["FASTVIDEO_DATALOADER_STATE_DIR"] = str(sidecar_dir)
+        manager = CheckpointManager(
+            method=object(),
+            dataloader=dataloader,
+            output_dir=str(checkpoint_dir.parent),
+            config=CheckpointConfig(save_steps=0, keep_last=0),
+        )
+        manager._load_dataloader_snapshot(checkpoint_dir, int(args.expected_step))
+        actual_id = str(next(iter(dataloader))["info_list"][0]["id"])
+        if actual_id != expected_id:
+            raise AssertionError(
+                f"Rank {rank} resumed at {actual_id!r}; expected {expected_id!r}"
+            )
+
+        result = {
+            "rank": rank,
+            "sidecar": sidecar_path.name,
+            "sidecar_sha256": _sha256(sidecar_path),
+            "next_id": actual_id,
+            "row_group_position": payload["state"]["row_group_position"],
+            "row_offset": payload["state"]["row_offset"],
+        }
+        gathered: list[dict[str, Any] | None] = [None] * world_size
+        dist.all_gather_object(gathered, result)
+        if rank == 0:
+            receipt = {
+                "schema": "fastvideo-streaming-dataloader-sidecar-validation/v1",
+                "created_at_utc": datetime.now(timezone.utc).isoformat(),
+                "checkpoint": str(checkpoint_dir),
+                "checkpoint_metadata_sha256": _sha256(checkpoint_dir / "metadata.json"),
+                "config": str(config_path),
+                "config_sha256": _sha256(config_path),
+                "sidecar_dir": str(sidecar_dir),
+                "step": int(args.expected_step),
+                "world_size": world_size,
+                "results": gathered,
+            }
+            _write_json_atomic(Path(args.output).expanduser().resolve(), receipt)
+            print(json.dumps(receipt, indent=2, sort_keys=True))
+        dist.barrier()
+    finally:
+        cleanup_dist_env_and_memory()
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--config", required=True)
+    parser.add_argument("--checkpoint", required=True)
+    parser.add_argument("--sidecar-dir", required=True)
+    parser.add_argument("--expected-step", required=True, type=int)
+    parser.add_argument("--output", required=True)
+    validate(parser.parse_args())
+
+
+if __name__ == "__main__":
+    main()