File size: 3,569 Bytes
d5e0d8f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 | """Prompt trajectories and fixed FFFF targets for Predictor-v4 Stage 2."""
from __future__ import annotations
import json
from pathlib import Path
from typing import Any
import torch
from safetensors import safe_open
from torch.utils.data import Dataset
class PredictorV4TrajectoryDataset(Dataset):
"""One deterministic seven-chunk trajectory per training prompt."""
def __init__(
self,
cases_path: str | Path,
*,
data_root: str | Path,
max_cases: int | None = None,
) -> None:
self.cases_path = Path(cases_path).resolve()
self.data_root = Path(data_root).resolve()
with self.cases_path.open("r", encoding="utf-8") as handle:
cases = [json.loads(line) for line in handle if line.strip()]
if max_cases is not None:
cases = cases[: int(max_cases)]
if not cases:
raise ValueError(f"No trajectory cases in {self.cases_path}")
for position, case in enumerate(cases):
if int(case["case_id"]) != position:
raise ValueError("Stage-2 case_id values must be dense and ordered")
if int(case.get("seed", -1)) != 0:
raise ValueError(f"case {position} does not use latent seed 0")
if not str(case.get("prompt", "")).strip():
raise ValueError(f"case {position} has an empty prompt")
self.cases = cases
def __len__(self) -> int:
return len(self.cases)
def __getitem__(self, index: int) -> dict[str, Any]:
case = self.cases[index]
return {
"case_id": int(case["case_id"]),
"prompt": str(case["prompt"]),
"seed": int(case["seed"]),
}
def offline_step_path(self, case_id: int, chunk_id: int) -> Path:
path = (
self.data_root
/ "steps"
/ f"case_{int(case_id):06d}"
/ f"chunk_{int(chunk_id):02d}.safetensors"
)
if not path.is_file():
raise FileNotFoundError(path)
return path
def trajectory_collate(items: list[dict[str, Any]]) -> dict[str, Any]:
if len(items) != 1:
raise ValueError(
"Stage-2 rollout currently requires one trajectory per rank"
)
return items[0]
def load_offline_ffff_target(
path: str | Path,
*,
step_id: int,
device: torch.device,
) -> dict[str, torch.Tensor]:
"""Load one immutable Full-trajectory hidden/flow target."""
path = Path(path)
names = {
"hidden": f"step_{int(step_id)}_final_hidden",
"flow": f"step_{int(step_id)}_flow",
"timestep": f"step_{int(step_id)}_timestep",
}
with safe_open(str(path), framework="pt", device="cpu") as handle:
missing = set(names.values()).difference(handle.keys())
if missing:
raise KeyError(f"{path} lacks fixed FFFF tensors {sorted(missing)}")
result = {key: handle.get_tensor(name) for key, name in names.items()}
if result["hidden"].dtype != torch.bfloat16:
raise TypeError(f"{path}: target hidden must be BF16")
if result["flow"].dtype != torch.bfloat16:
raise TypeError(f"{path}: target flow must be BF16")
if result["timestep"].dtype != torch.int64:
raise TypeError(f"{path}: legacy offline timestep must be INT64")
return {
key: value.to(device=device, non_blocking=True)
for key, value in result.items()
}
__all__ = [
"PredictorV4TrajectoryDataset",
"load_offline_ffff_target",
"trajectory_collate",
]
|