wehub-resource-sync
322 lines
14 KiB
Python
322 lines
14 KiB
Python
# Adopted from https://github.com/vita-epfl/Stable-Video-Infinity
|
||
# SPDX-License-Identifier: Apache-2.0
|
||
|
||
import random
|
||
import torch
|
||
|
||
|
||
class ErrorBuffer:
|
||
"""Bucketed ring buffer for storing prediction errors on CPU.
|
||
|
||
Two layouts are supported:
|
||
|
||
* **1D (timestep-only)** — when ``num_blocks <= 0``. Buckets are keyed by
|
||
the diffusion timestep. This is the original SVI behavior.
|
||
|
||
* **2D (position × timestep)** — when ``num_blocks > 0``. Each entry is
|
||
keyed by both the global block position along the sequence and the
|
||
timestep. Inject paths can then choose:
|
||
- ``sample(pos, t)``: match BOTH position and timestep
|
||
(E_vid / E_noise — noise-level dependent errors)
|
||
- ``sample_pos_any_t(pos)``: match position, sample uniformly across
|
||
timesteps (E_img — position-dependent context corruption that is
|
||
agnostic to the current denoising step)
|
||
- ``sample_global()``: legacy fallback, samples uniformly everywhere
|
||
|
||
The 2D layout encodes the teacher-forcing insight that ``noisy_suffix[i]``
|
||
looks at clean_prefix[0..i] during training but at model rollouts during
|
||
inference; storing prediction errors per-position therefore lets later
|
||
blocks self-feed larger errors without any manual position ramp.
|
||
|
||
**Sharded timestep buckets** (``shard_size > 1``):
|
||
Each rank only allocates the timestep buckets it owns
|
||
(``t_bucket % shard_size == shard_rank``), reducing per-rank CPU memory
|
||
by ~``shard_size`` times. Typically ``shard_rank/shard_size`` are set to
|
||
``sp_rank/sp_size`` so that sharding is per-SP-rank and saving follows
|
||
the same per-SP-rank pattern as the 2D position split. On ``add()``,
|
||
non-owned buckets are silently skipped; on ``sample()``, non-owned buckets
|
||
are remapped to the nearest owned one.
|
||
"""
|
||
|
||
def __init__(
|
||
self,
|
||
num_buckets=40,
|
||
max_size_per_bucket=50,
|
||
num_train_timesteps=1000,
|
||
modulate_factor=0.3,
|
||
replacement_strategy="random",
|
||
num_blocks=0,
|
||
global_block_offset=0,
|
||
shard_rank=0,
|
||
shard_size=1,
|
||
):
|
||
self.num_buckets = num_buckets
|
||
self.max_size = max_size_per_bucket
|
||
self.num_train_timesteps = num_train_timesteps
|
||
self.modulate_factor = modulate_factor
|
||
self.replacement_strategy = replacement_strategy
|
||
self.bucket_width = num_train_timesteps / num_buckets
|
||
self.num_blocks = int(num_blocks) if num_blocks else 0
|
||
# ``global_block_offset`` is only used for stats / debug display so
|
||
# users can tell which absolute positions of the full sequence this
|
||
# buffer covers (the LAST SP rank carries the highest accumulated
|
||
# error positions). It does NOT participate in bucket keying.
|
||
self.global_block_offset = int(global_block_offset)
|
||
|
||
self.shard_rank = int(shard_rank)
|
||
self.shard_size = max(int(shard_size), 1)
|
||
self._owned_t_buckets = sorted(
|
||
t for t in range(num_buckets) if t % self.shard_size == self.shard_rank
|
||
)
|
||
|
||
if self.num_blocks > 0:
|
||
self.buckets = {
|
||
(p, t): []
|
||
for p in range(self.num_blocks)
|
||
for t in self._owned_t_buckets
|
||
}
|
||
else:
|
||
self.buckets = {t: [] for t in self._owned_t_buckets}
|
||
self.total_added = 0
|
||
|
||
# ------------------------------------------------------------------ keys
|
||
def _t_bucket(self, timestep_index):
|
||
b = int(timestep_index / self.bucket_width)
|
||
return max(0, min(b, self.num_buckets - 1))
|
||
|
||
def _is_owned_t(self, t_bucket):
|
||
return self.shard_size <= 1 or (t_bucket % self.shard_size == self.shard_rank)
|
||
|
||
def _nearest_owned_t(self, t_bucket):
|
||
"""Remap ``t_bucket`` to the closest owned timestep bucket."""
|
||
if self.shard_size <= 1 or self._is_owned_t(t_bucket):
|
||
return t_bucket
|
||
fwd = (self.shard_rank - t_bucket % self.shard_size) % self.shard_size
|
||
bwd = self.shard_size - fwd
|
||
t_up, t_down = t_bucket + fwd, t_bucket - bwd
|
||
up_ok = 0 <= t_up < self.num_buckets
|
||
down_ok = 0 <= t_down < self.num_buckets
|
||
if up_ok and down_ok:
|
||
return t_up if fwd <= bwd else t_down
|
||
return t_up if up_ok else t_down
|
||
|
||
def _make_key(self, t_bucket, block_pos):
|
||
if self.num_blocks > 0:
|
||
assert block_pos is not None, "block_pos required when num_blocks>0"
|
||
p = max(0, min(int(block_pos), self.num_blocks - 1))
|
||
return (p, t_bucket)
|
||
return t_bucket
|
||
|
||
# ------------------------------------------------------------------ add
|
||
def add(self, error_block, timestep_index, block_pos=None):
|
||
"""Store a single block error into the matching bucket.
|
||
|
||
Args:
|
||
error_block: (block_size, C, H, W) tensor
|
||
timestep_index: int, raw index in [0, num_train_timesteps)
|
||
block_pos: int, global block position; required iff num_blocks>0
|
||
"""
|
||
t = self._t_bucket(timestep_index)
|
||
if not self._is_owned_t(t):
|
||
return
|
||
key = self._make_key(t, block_pos)
|
||
# Store in the source dtype on CPU to match SVI (which keeps bf16),
|
||
# cutting buffer memory in half vs. casting to fp32.
|
||
entry = error_block.detach().to("cpu", copy=True)
|
||
|
||
buf = self.buckets[key]
|
||
if len(buf) < self.max_size:
|
||
buf.append(entry)
|
||
else:
|
||
if self.replacement_strategy == "fifo":
|
||
buf.pop(0)
|
||
buf.append(entry)
|
||
elif self.replacement_strategy == "l2":
|
||
stacked = torch.stack(buf)
|
||
dists = (stacked - entry.unsqueeze(0)).flatten(1).norm(dim=1)
|
||
most_similar = torch.argmin(dists).item()
|
||
buf[most_similar] = entry
|
||
else: # "random" (default)
|
||
idx = random.randint(0, self.max_size - 1)
|
||
buf[idx] = entry
|
||
self.total_added += 1
|
||
|
||
# ------------------------------------------------------------------ sample
|
||
def sample(self, timestep_index, device, dtype, block_pos=None):
|
||
"""Sample one entry matching (block_pos, timestep_index) when 2D, or
|
||
just timestep_index when 1D. Non-owned timestep buckets are
|
||
transparently remapped to the nearest owned one. Returns None if the
|
||
(remapped) bucket is empty."""
|
||
t = self._nearest_owned_t(self._t_bucket(timestep_index))
|
||
key = self._make_key(t, block_pos)
|
||
buf = self.buckets[key]
|
||
if not buf:
|
||
return None
|
||
err = random.choice(buf)
|
||
return self._modulate(err).to(device=device, dtype=dtype)
|
||
|
||
def sample_pos_any_t(self, block_pos, device, dtype):
|
||
"""For 2D buffers: sample at the given position, with random timestep.
|
||
|
||
This is the natural choice for context (E_img) injection — the clean
|
||
prefix is the result of a full ODE rollout so its accumulated error
|
||
could have originated at any timestep, but its magnitude scales with
|
||
position along the sequence.
|
||
|
||
Falls back to ``sample_global`` when the buffer is 1D.
|
||
Only owned timestep buckets are scanned.
|
||
"""
|
||
if self.num_blocks <= 0:
|
||
return self.sample_global(device, dtype)
|
||
p = max(0, min(int(block_pos), self.num_blocks - 1))
|
||
all_entries = []
|
||
for t in self._owned_t_buckets:
|
||
all_entries.extend(self.buckets[(p, t)])
|
||
if not all_entries:
|
||
return None
|
||
err = random.choice(all_entries)
|
||
return self._modulate(err).to(device=device, dtype=dtype)
|
||
|
||
def sample_global(self, device, dtype):
|
||
"""Sample one entry uniformly from ALL buckets (legacy SVI E_img)."""
|
||
all_entries = []
|
||
for buf in self.buckets.values():
|
||
all_entries.extend(buf)
|
||
if not all_entries:
|
||
return None
|
||
err = random.choice(all_entries)
|
||
return self._modulate(err).to(device=device, dtype=dtype)
|
||
|
||
# ------------------------------------------------------------------ misc
|
||
def _modulate(self, err):
|
||
if self.modulate_factor > 0:
|
||
lo = 1.0 - self.modulate_factor
|
||
hi = 1.0 + self.modulate_factor
|
||
err = err * random.uniform(lo, hi)
|
||
return err
|
||
|
||
def is_empty(self):
|
||
return self.total_added == 0
|
||
|
||
def has_pos(self, block_pos):
|
||
"""Whether ANY owned timestep bucket at ``block_pos`` has samples (2D only)."""
|
||
if self.num_blocks <= 0:
|
||
return not self.is_empty()
|
||
p = max(0, min(int(block_pos), self.num_blocks - 1))
|
||
return any(len(self.buckets[(p, t)]) > 0 for t in self._owned_t_buckets)
|
||
|
||
def stats(self):
|
||
filled = sum(1 for b in self.buckets.values() if len(b) > 0)
|
||
total = sum(len(b) for b in self.buckets.values())
|
||
num_owned_t = len(self._owned_t_buckets)
|
||
denom = self.num_blocks * num_owned_t if self.num_blocks > 0 else num_owned_t
|
||
out = {
|
||
"total_added": self.total_added,
|
||
"filled_buckets": f"{filled}/{denom}",
|
||
"total_entries": total,
|
||
}
|
||
if self.shard_size > 1:
|
||
out["shard"] = f"shard_rank={self.shard_rank}/{self.shard_size} ({num_owned_t}/{self.num_buckets} t-buckets)"
|
||
if self.num_blocks > 0:
|
||
lo = self.global_block_offset
|
||
hi = self.global_block_offset + self.num_blocks
|
||
out["global_block_range"] = f"[{lo},{hi})"
|
||
return out
|
||
|
||
def state_dict(self):
|
||
# Keys are tuples (pos, t) when 2D — torch.save handles them fine
|
||
# via pickle. We serialize the bucket layout so loaders can validate.
|
||
return {
|
||
"buckets": {k: list(v) for k, v in self.buckets.items()},
|
||
"total_added": self.total_added,
|
||
"num_blocks": self.num_blocks,
|
||
"num_buckets": self.num_buckets,
|
||
"global_block_offset": self.global_block_offset,
|
||
"shard_rank": self.shard_rank,
|
||
"shard_size": self.shard_size,
|
||
}
|
||
|
||
def load_state_dict(self, state, strict_offset=True):
|
||
"""Restore buckets from a serialized state.
|
||
|
||
Args:
|
||
state: dict produced by ``state_dict``.
|
||
strict_offset: when True (default) and the buffer is 2D,
|
||
refuse to load if the saved ``global_block_offset`` does
|
||
not match the current one. This prevents the silent
|
||
position-misalignment bug under SP, where a checkpoint
|
||
saved by SP rank 0 (covering global blocks ``[0, B)``)
|
||
would otherwise be loaded into SP rank 1 (which expects
|
||
``[B, 2B)``) and corrupt position-bucketed sampling.
|
||
Pass ``strict_offset=False`` only for backward-compat
|
||
with checkpoints saved before this field existed.
|
||
"""
|
||
if self.num_blocks > 0 and strict_offset:
|
||
saved_off = state.get("global_block_offset", None)
|
||
if saved_off is None:
|
||
raise RuntimeError(
|
||
"Refusing to load: this is a 2D position-bucketed buffer "
|
||
"but the checkpoint has no `global_block_offset` field. "
|
||
"Pass strict_offset=False if you accept the misalignment risk."
|
||
)
|
||
if int(saved_off) != self.global_block_offset:
|
||
raise RuntimeError(
|
||
f"Refusing to load: checkpoint covers global blocks "
|
||
f"starting at {saved_off}, but this rank covers blocks "
|
||
f"starting at {self.global_block_offset}. Make sure each "
|
||
f"SP rank loads its own per-rank checkpoint file."
|
||
)
|
||
# Shard check: warn but don't crash if shard layout changed (e.g.
|
||
# resuming a non-sharded checkpoint into a sharded buffer is fine —
|
||
# we just load whichever buckets overlap).
|
||
saved_shard_size = int(state.get("shard_size", state.get("dp_size", 1)))
|
||
saved_shard_rank = int(state.get("shard_rank", state.get("dp_rank", 0)))
|
||
if saved_shard_size != self.shard_size or saved_shard_rank != self.shard_rank:
|
||
import logging
|
||
logging.warning(
|
||
f"[ErrorBuffer] Shard layout changed: checkpoint was "
|
||
f"shard_rank={saved_shard_rank}/{saved_shard_size}, current is "
|
||
f"shard_rank={self.shard_rank}/{self.shard_size}. "
|
||
f"Loading overlapping buckets only."
|
||
)
|
||
|
||
saved = state["buckets"]
|
||
# Lenient match: ignore keys that don't exist in the current layout.
|
||
for k in self.buckets:
|
||
if k in saved:
|
||
self.buckets[k] = saved[k]
|
||
elif isinstance(k, tuple):
|
||
# Try string-form key from older serializations
|
||
continue
|
||
elif str(k) in saved:
|
||
self.buckets[k] = saved[str(k)]
|
||
self.total_added = int(state.get("total_added", 0))
|
||
|
||
|
||
def build_error_buffer(config, num_blocks=0, global_block_offset=0,
|
||
shard_rank=0, shard_size=1):
|
||
"""Build an ErrorBuffer from an OmegaConf/dict config node.
|
||
|
||
When ``num_blocks > 0`` the buffer becomes 2D (position × timestep),
|
||
enabling teacher-forcing-aware position-dependent error injection.
|
||
Pass ``global_block_offset`` so logs can identify which absolute slice
|
||
of the full sequence this rank's buffer covers (e.g. the last SP rank
|
||
is responsible for the most error-accumulated tail blocks).
|
||
|
||
``shard_rank`` / ``shard_size`` shard timestep buckets: each rank only
|
||
allocates the buckets it owns, reducing per-rank CPU memory by
|
||
~``shard_size`` times. Typically set to ``(sp_rank, sp_size)``.
|
||
"""
|
||
cfg = config if isinstance(config, dict) else dict(config)
|
||
return ErrorBuffer(
|
||
num_buckets=cfg.get("num_buckets", 40),
|
||
max_size_per_bucket=cfg.get("buffer_size_per_bucket", 50),
|
||
num_train_timesteps=cfg.get("num_train_timesteps", 1000),
|
||
modulate_factor=cfg.get("modulate_factor", 0.3),
|
||
replacement_strategy=cfg.get("replacement_strategy", "random"),
|
||
num_blocks=num_blocks,
|
||
global_block_offset=global_block_offset,
|
||
shard_rank=shard_rank,
|
||
shard_size=shard_size,
|
||
)
|