| """Lazy prompt-sharded access for large Predictor training datasets.""" |
|
|
| from __future__ import annotations |
|
|
| from pathlib import Path |
| from typing import Any |
|
|
| import torch |
| from safetensors import safe_open |
| from torch.utils.data import Dataset |
|
|
|
|
| class LazyLayer17Dataset(Dataset): |
| """Read only the tensors needed by one Layer-17 sample on demand.""" |
|
|
| def __init__( |
| self, |
| root: str | Path, |
| layer_id: int = 17, |
| include_previous_hidden: bool = True, |
| ) -> None: |
| self.root = Path(root).resolve() |
| self.layer_id = int(layer_id) |
| self.include_previous_hidden = bool(include_previous_hidden) |
|
|
| def __len__(self) -> int: |
| return 1000 |
|
|
| def __getitem__(self, index: tuple[int, int, int]) -> dict[str, Any]: |
| prompt_id, chunk, target_step = (int(value) for value in index) |
| if not 0 <= prompt_id < 1000: |
| raise ValueError(index) |
| if chunk < 1 or chunk > 6 or target_step < 1 or target_step > 3: |
| raise ValueError(index) |
|
|
| anchor_step = target_step - 1 |
| prompt_dir = self.root / f"prompt_{prompt_id:04d}" |
| trajectory_path = prompt_dir / "trajectory.safetensors" |
| context_path = prompt_dir / "chunk0_context" / "trajectory.safetensors" |
|
|
| with safe_open(trajectory_path, framework="pt", device="cpu") as handle: |
| prefix = f"chunk_{chunk:02d}_step_{target_step:02d}" |
| anchor_prefix = f"chunk_{chunk:02d}_step_{anchor_step:02d}" |
| noisy = handle.get_tensor(f"{prefix}_noisy_latent").squeeze(0) |
| timestep = handle.get_tensor(f"{prefix}_timestep").squeeze(0) |
| anchor_timestep = handle.get_tensor( |
| f"{anchor_prefix}_timestep" |
| ).squeeze(0) |
| anchor_hidden = handle.get_tensor( |
| f"{anchor_prefix}_final_hidden" |
| ).squeeze(0) |
| target_hidden = handle.get_tensor( |
| f"{prefix}_final_hidden" |
| ).squeeze(0) |
| target_flow = handle.get_tensor(f"{prefix}_flow").squeeze(0) |
| if self.include_previous_hidden and chunk > 1: |
| previous_hidden = handle.get_tensor( |
| f"chunk_{chunk - 1:02d}_step_{target_step:02d}_final_hidden" |
| ).squeeze(0) |
| else: |
| previous_hidden = None |
| if self.include_previous_hidden and previous_hidden is None: |
| with safe_open(context_path, framework="pt", device="cpu") as handle: |
| previous_hidden = handle.get_tensor( |
| f"chunk_00_step_{target_step:02d}_final_hidden" |
| ).squeeze(0) |
| elif not self.include_previous_hidden: |
| |
| |
| previous_hidden = torch.zeros_like(anchor_hidden) |
|
|
| prefeature_path = ( |
| prompt_dir |
| / "clean_prefeatures" |
| / f"block_{self.layer_id:02d}.safetensors" |
| ) |
| context_prefeature_path = ( |
| prompt_dir |
| / "chunk0_context" |
| / "clean_prefeatures" |
| / f"block_{self.layer_id:02d}.safetensors" |
| ) |
| history_parts = [] |
| with safe_open( |
| context_prefeature_path, framework="pt", device="cpu" |
| ) as context_handle, safe_open( |
| prefeature_path, framework="pt", device="cpu" |
| ) as handle: |
| for history_chunk in range(chunk): |
| source = context_handle if history_chunk == 0 else handle |
| history_parts.append( |
| source.get_tensor(f"chunk_{history_chunk:02d}").squeeze(0) |
| ) |
| history_prefeature = torch.cat(history_parts, dim=0).contiguous() |
|
|
| cross_path = prompt_dir / "cross_attention.safetensors" |
| with safe_open(cross_path, framework="pt", device="cpu") as handle: |
| cross_k = handle.get_tensor( |
| f"block_{self.layer_id:02d}_k" |
| ).squeeze(0) |
| cross_v = handle.get_tensor( |
| f"block_{self.layer_id:02d}_v" |
| ).squeeze(0) |
|
|
| return { |
| "prompt_id": prompt_id, |
| "chunk": chunk, |
| "target_step": target_step, |
| "noisy_latent": noisy, |
| "timestep": timestep, |
| "anchor_timestep": anchor_timestep, |
| "anchor_distance": ( |
| timestep.float() - anchor_timestep.float() |
| ).abs().mean(), |
| "anchor_hidden": anchor_hidden, |
| "previous_hidden": previous_hidden, |
| "target_hidden": target_hidden, |
| "target_flow": target_flow, |
| "history_prefeature": history_prefeature, |
| "cross_k": cross_k, |
| "cross_v": cross_v, |
| } |
|
|
|
|
| def collate_lazy_samples(samples: list[dict[str, Any]]) -> dict[str, Any]: |
| if not samples: |
| raise ValueError("Empty batch") |
| chunks = {sample["chunk"] for sample in samples} |
| target_steps = {sample["target_step"] for sample in samples} |
| if len(chunks) != 1 or len(target_steps) != 1: |
| raise ValueError("A batch must share chunk and target_step") |
| return { |
| "prompt_ids": [sample["prompt_id"] for sample in samples], |
| "chunk": samples[0]["chunk"], |
| "anchor_step": samples[0]["target_step"] - 1, |
| "target_step": samples[0]["target_step"], |
| **{ |
| key: torch.stack([sample[key] for sample in samples]) |
| for key in ( |
| "noisy_latent", |
| "timestep", |
| "anchor_timestep", |
| "anchor_distance", |
| "anchor_hidden", |
| "previous_hidden", |
| "target_hidden", |
| "target_flow", |
| "history_prefeature", |
| "cross_k", |
| "cross_v", |
| ) |
| }, |
| } |
|
|
|
|
| class ScheduledBatchSampler: |
| def __init__(self, batches: list[list[tuple[int, int, int]]]) -> None: |
| self.batches = batches |
|
|
| def __iter__(self): |
| return iter(self.batches) |
|
|
| def __len__(self) -> int: |
| return len(self.batches) |
|
|
