[bugfix] H3 conditioner: no_grad instead of inference_mode - FSDP2-sharded encode crashed (#1732 default-path gap)
Carried blocker fix discovered while benchmarking this stack (GB200 job 2594,
all four legs): with text_encoder_cpu_offload=True - the FastVideoArgs DEFAULT
and the shipped H3 example configuration - TextEncoderLoader FSDP2-fully_shards
the H3 conditioner, and #1732's @torch.inference_mode() on
MiniMaxH3Qwen3VLConditioner.encode_ids then kills the very first encode:
File torch/distributed/fsdp/_fully_shard/_fsdp_param_group.py, in
wait_for_unshard: with torch.autograd._unsafe_preserve_version_counter(t):
RuntimeError: Inference tensors do not track version counter.
FSDP2's lazy unshard runs inside the inference-mode region, so its all-gather
tensors are inference tensors, and the version-counter preservation hook
cannot read t._version. The PR's own benchmarks ran with offload disabled,
which is exactly the unexercised-default-path gap called out as finding 2 of
the pr1732 review (there for FP8; the same gap bites plain bf16 via
inference_mode).
Fix: @torch.no_grad() instead. It frees the same activation memory (the
-262 MiB claim comes from the early-return slim contract, not from
inference_mode's bookkeeping), is fully FSDP2-compatible, and as a bonus
retires review finding 7: prompt_embeds are ordinary tensors again, so any
future on-the-fly-conditioning training can backprop through them without a
clone at the stage boundary.
Verified: 1x GPU FastH3 leg boots and generates after this change (leg reruns
on GB200); the 1732 unit suites (truncation + checkpoint-fp8) still pass.
Note: fastvideo/tests/stages/test_text_encoding.py has 7 pre-existing failures
that reproduce byte-identically on plain origin/main 56d4a6074 (unrelated
generic-stage tests; not introduced by the stack, verified by A/B).
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
co-authored by
Claude Fable 5
parent
68e6ffca9e
commit
ad9cd63122
@@ -622,7 +622,14 @@ class MiniMaxH3Qwen3VLConditioner(TextEncoder[torch.Tensor]):
|
||||
f"tokens={int(mask.sum())}, features={features.shape[0]}")
|
||||
return mask
|
||||
|
||||
@torch.inference_mode()
|
||||
# no_grad, NOT inference_mode: with text_encoder_cpu_offload=True (the
|
||||
# FastVideoArgs default) the loader FSDP2-shards this conditioner, and
|
||||
# FSDP2's wait_for_unshard reads tensor._version via
|
||||
# _unsafe_preserve_version_counter - inference tensors do not track
|
||||
# version counters, so inference_mode crashes the first encode. no_grad
|
||||
# frees the same activation memory and keeps prompt_embeds ordinary
|
||||
# tensors (safe for any future backward through the conditioning).
|
||||
@torch.no_grad()
|
||||
def encode_ids(
|
||||
self,
|
||||
input_ids: torch.Tensor,
|
||||
|
||||
Reference in New Issue
Block a user