wehub/nvlabs--longlive
1
0
Fork 0
Code Issues Pull Requests Actions Deployments Models Agent Notes Packages Projects Releases Wiki Activity

chore: import upstream snapshot with attribution

Browse Source
This commit is contained in:
wehub-resource-sync
2026-07-13 12:31:40 +08:00
commit 4df75a30cd
234 changed files with 41407 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
utils/error_buffer.py
+321
View File
@@ -0,0 +1,321 @@
# 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,
)