Fix rank-local dataloader checkpoint resume

2026-07-22 21:29:48 -07:00
parent c0cb2228ac
commit 41dcd065ae
6 changed files with 862 additions and 4 deletions
@@ -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()
@@ -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()