[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:
Will Lin
2026-08-21 10:49:47 +00:00
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,