fix: stop REINS/q-steering/av-decouple forcing host syncs in the H3 block loop
comfy_aimdo's dynamic-VRAM block prefetch pipelines weight cast/unpin against CUDA stream timing; a CPU sync (bool()/item() of a CUDA tensor) or a surprise allocation mid-loop can race it into reclaiming a weight buffer a still-queued kernel needs -- a silent segfault, no Python exception possible. All three H3 rating-driven dit hooks did this on every block, every step: - h3_repr_steering's injection guard and h3_q_steering's injection guard both did `if torch.isfinite(x) and x > 0:` (forces a host sync) before writing the steered rows, and h3_repr_steering cloned the whole block output first. Replaced with a GPU-only torch.where guard; dropped the now-unnecessary clone in h3_repr_steering (out is this hook's own private tensor already). - h3_repr_steering.capture() now .cpu()s its descriptor immediately instead of leaving it on GPU competing for VRAM headroom for the rest of the run. - h3_av_decouple's attention override re-checked `.any()` on its video/audio/ other masks every single attention call; now decided once per unique seq_len alongside the mask itself (typically once per generation). Unvalidated against a real GPU (none available here) -- worth confirming on the rental that REINS no longer crashes silently under dynamic VRAM offload. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
This commit is contained in:
digital-garbage
co-authored by
Claude Sonnet 5
Claude Sonnet 5
parent
c70e16cde8
commit
f5729a96df
+7
-1
@@ -120,7 +120,13 @@ def capture(hidden_state, video_mask):
|
||||
size needs adaptive_avg_pool1d for."""
|
||||
if video_mask is None or not bool(video_mask.any()):
|
||||
return None
|
||||
return hidden_state[video_mask].detach().float().mean(dim=0)
|
||||
# .cpu() immediately -- these accumulate across all 50 CANDIDATE_BLOCKS, every
|
||||
# denoising step, and otherwise sit on the GPU for the rest of the run competing
|
||||
# for VRAM headroom with the block-weight prefetch pipeline (comfy_aimdo). See
|
||||
# project_reward_model_rework memory: that pressure was implicated in a silent
|
||||
# segfault under dynamic VRAM offload. The descriptor only ever gets saved to
|
||||
# disk (torch.save) or diffed against another CPU tensor -- never used on-GPU.
|
||||
return hidden_state[video_mask].detach().float().mean(dim=0).cpu()
|
||||
|
||||
|
||||
# --- persistence -------------------------------------------------------------------------
|
||||
|
||||
+36
-15
@@ -6932,10 +6932,18 @@ class FunPackLTXAVSceneChainSampler:
|
||||
if direction is not None and _strength > 0.0:
|
||||
rows = out[mask]
|
||||
row_norm = rows.detach().float().norm(dim=-1).mean()
|
||||
if torch.isfinite(row_norm) and row_norm > 0:
|
||||
out = out.clone()
|
||||
out[mask] = rows + (direction.to(out.dtype).to(out.device)
|
||||
* _strength * row_norm).to(rows.dtype)
|
||||
# GPU-only guard (no .item()/bool() host sync): a host sync here
|
||||
# forces the CPU thread to wait mid-block-loop, which throws off
|
||||
# the async weight-prefetch pipeline's own stream-wait timing
|
||||
# (comfy_aimdo/model_prefetch.py) enough to segfault on a
|
||||
# VRAM-constrained dynamic-offload run -- see project_reward_model
|
||||
# -rework memory. torch.where no-ops the update when row_norm is
|
||||
# 0/NaN/inf instead of skipping it, at the cost of the multiply
|
||||
# always running -- cheaper than a stall that can crash the run.
|
||||
bump = (direction.to(out.dtype).to(out.device)
|
||||
* _strength * row_norm.to(out.dtype))
|
||||
safe = torch.isfinite(row_norm) & (row_norm > 0)
|
||||
out[mask] = torch.where(safe, rows + bump.to(rows.dtype), rows)
|
||||
return {"img": out}
|
||||
return _tag_dit_hook(_hook, inner)
|
||||
|
||||
@@ -7062,10 +7070,20 @@ class FunPackLTXAVSceneChainSampler:
|
||||
seq_len = q.shape[-2]
|
||||
cached = _mask_cache.get(seq_len)
|
||||
if cached is None:
|
||||
cached = (_rs.video_mask_from_mod_segments(seg, seq_len, q.device),
|
||||
_rs.audio_mask_from_mod_segments(seg, seq_len, q.device))
|
||||
_v = _rs.video_mask_from_mod_segments(seg, seq_len, q.device)
|
||||
_a = _rs.audio_mask_from_mod_segments(seg, seq_len, q.device)
|
||||
# .any() decided ONCE per unique seq_len (one host sync, typically
|
||||
# once for the whole generation) instead of on every attention call,
|
||||
# every block, every step -- see h3_repr_steering's capture() fix and
|
||||
# project_reward_model_rework memory for why a host sync in this loop
|
||||
# risks a silent segfault under comfy_aimdo's dynamic-VRAM prefetch.
|
||||
_other = ~(_v | _a) if (_v is not None and _a is not None) else None
|
||||
cached = (_v, _a, _other,
|
||||
bool(_v.any()) if _v is not None else False,
|
||||
bool(_a.any()) if _a is not None else False,
|
||||
bool(_other.any()) if _other is not None else False)
|
||||
_mask_cache[seq_len] = cached
|
||||
vmask, amask = cached
|
||||
vmask, amask, other, has_v, has_a, has_other = cached
|
||||
if vmask is None or amask is None or k.shape[-2] != seq_len:
|
||||
# Captured mod_segments but couldn't turn it into rows to bias -- the
|
||||
# penalty is silently NOT applying this generation. Loud once, not once
|
||||
@@ -7079,7 +7097,6 @@ class FunPackLTXAVSceneChainSampler:
|
||||
f"generation is running the ORDINARY joint attention. The "
|
||||
f"penalty is not being applied; this is not saturation.")
|
||||
return _next(func, q, k, v, heads, mask=mask, skip_reshape=skip_reshape, **kwargs)
|
||||
other = ~(vmask | amask)
|
||||
# `strength` IS the raw softmax logit penalty, not a normalized 0-1 fraction --
|
||||
# there's no local way to know H3's actual trained attention-logit scale (the
|
||||
# checkpoint isn't on this machine, see the method docstring), so a 0-1 dial
|
||||
@@ -7094,19 +7111,19 @@ class FunPackLTXAVSceneChainSampler:
|
||||
f"seq_len={seq_len}. (Printed once per generation, not per block.)")
|
||||
out = torch.empty(q.shape[0], seq_len, heads * q.shape[-1],
|
||||
dtype=q.dtype, device=q.device)
|
||||
if amask.any():
|
||||
if has_a:
|
||||
bias = torch.zeros(seq_len, dtype=q.dtype, device=q.device)
|
||||
bias[vmask] = bias_val
|
||||
out[:, amask, :] = _next(func, q[:, :, amask, :], k, v, heads,
|
||||
mask=bias.view(1, 1, 1, seq_len),
|
||||
skip_reshape=skip_reshape, **kwargs)
|
||||
if vmask.any():
|
||||
if has_v:
|
||||
bias = torch.zeros(seq_len, dtype=q.dtype, device=q.device)
|
||||
bias[amask] = bias_val
|
||||
out[:, vmask, :] = _next(func, q[:, :, vmask, :], k, v, heads,
|
||||
mask=bias.view(1, 1, 1, seq_len),
|
||||
skip_reshape=skip_reshape, **kwargs)
|
||||
if other.any():
|
||||
if has_other:
|
||||
out[:, other, :] = _next(func, q[:, :, other, :], k, v, heads,
|
||||
mask=None, skip_reshape=skip_reshape, **kwargs)
|
||||
return out
|
||||
@@ -7359,10 +7376,14 @@ class FunPackLTXAVSceneChainSampler:
|
||||
**kwargs)
|
||||
rows = q[:, :, vmask, :]
|
||||
row_norm = rows.detach().float().norm(dim=-1).mean()
|
||||
if torch.isfinite(row_norm) and row_norm > 0:
|
||||
dir_per_head = direction.to(q.dtype).to(q.device).view(heads, 1, head_dim)
|
||||
q = q.clone()
|
||||
q[:, :, vmask, :] = rows + dir_per_head * _strength * row_norm
|
||||
# GPU-only guard (no .item()/bool() host sync) -- same fix as
|
||||
# h3_repr_steering's hook, see that one for why a host sync here can
|
||||
# segfault under comfy_aimdo's dynamic-VRAM block prefetch.
|
||||
dir_per_head = direction.to(q.dtype).to(q.device).view(heads, 1, head_dim)
|
||||
bump = dir_per_head * _strength * row_norm.to(q.dtype)
|
||||
safe = torch.isfinite(row_norm) & (row_norm > 0)
|
||||
q = q.clone()
|
||||
q[:, :, vmask, :] = torch.where(safe, rows + bump, rows)
|
||||
return _next(func, q, k, v, heads, mask=mask, skip_reshape=skip_reshape,
|
||||
**kwargs)
|
||||
|
||||
|
||||
Reference in New Issue
Block a user