Class ErrorBuffer

utils/error_buffer.py:8–293

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

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)