diff --git a/fastvideo/dataset/parquet_dataset_streaming_style.py b/fastvideo/dataset/parquet_dataset_streaming_style.py index 3fe820852..339bdff2d 100644 --- a/fastvideo/dataset/parquet_dataset_streaming_style.py +++ b/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. diff --git a/fastvideo/tests/dataset/test_parquet_dataset_streaming_style.py b/fastvideo/tests/dataset/test_parquet_dataset_streaming_style.py index 5faf764ec..797445504 100644 --- a/fastvideo/tests/dataset/test_parquet_dataset_streaming_style.py +++ b/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: diff --git a/fastvideo/tests/train/utils/test_checkpoint.py b/fastvideo/tests/train/utils/test_checkpoint.py index 021b3b634..457d068cd 100644 --- a/fastvideo/tests/train/utils/test_checkpoint.py +++ b/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 # --------------------------------------------------------------------------- diff --git a/fastvideo/train/utils/checkpoint.py b/fastvideo/train/utils/checkpoint.py index b8af0ee8d..9fcdd9c2a 100644 --- a/fastvideo/train/utils/checkpoint.py +++ b/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, diff --git a/scripts/train/migrate_streaming_dataloader_checkpoint.py b/scripts/train/migrate_streaming_dataloader_checkpoint.py new file mode 100644 index 000000000..4a3c44229 --- /dev/null +++ b/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() diff --git a/scripts/train/validate_streaming_dataloader_sidecars.py b/scripts/train/validate_streaming_dataloader_sidecars.py new file mode 100644 index 000000000..8457c25ed --- /dev/null +++ b/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()