Library use of ltx_pipelines without inference_mode retains the text encoder through autograd

Aperta
  • Qualcosa non va?
#310 0 commenti 0 reazioni 0 assegnatari Vedi su GitHub

Nessuno ha ancora preso questa issue.

Valutazione

Difficoltà
2/5
Tempo stimato
1-3 ore
Idoneità per principianti
35/100
Tipo di issue
Bug
Chiarezza
Abbastanza chiara
Stato di attività
Tranquilla
Stack tecnologico
python, pytorch
Ambito
ai, machine-learning

Direzione di ricerca

Start with PromptEncoder.__call__ in utils/blocks.py (around line 778) and the pipeline __call__ methods in distilled.py and a2vid_two_stage.py, using ti2vid_two_stages_hq.py as the reference that already uses inference mode. Reproducing the retention needs a large GPU and the full split checkpoints, and the fix choice (decorator or docs) is a maintainer decision. Done means grad mode is disabled during encoder invocation, with a regression check on the library path.

Scritto dal modello di indicizzazione a partire dal testo della issue.

Descrizione

Problem

Calling the native pipeline as a library can leave autograd enabled. Prompt encoding then retains the large text encoder and intermediate activations through its raw outputs, even after the model context exits. On our RTX 5090 under WSL, the earlier reference-release probe reported roughly 48 GiB allocated on a nominal 32 GiB card. Later decode budgeting reported usable_bytes=0, or the run failed with OOM.

This appears to be reference retention through autograd, rather than proof of an unbounded allocator leak. The WSL run exceeded physical VRAM without an immediate error; host-memory paging is the observed operational interpretation, not something established by torch.cuda.memory_allocated() alone.

Environment and scope

  • Windows 11 host, RTX 5090, native Python inference under WSL.
  • Python 3.14.7; PyTorch 2.13.0+cu132; editable ltx-pipelines 1.3.0 installation, verified from installed package metadata.
  • Local LTX-2 revision: b0b59a759853a69f67e292106db53db58326a930 (includes a local kernels build change).
  • Split LTX-2.5 assets: distilled NVFP4 transformer, BF16 Gemma4 12B encoder with projection, video VAE and audio VAE.
  • The reference-release result was recorded by an earlier local run on 06/09/2026. This draft does not claim a new reproduction of the unsafe 48 GiB case. Successful follow-up A2V drivers explicitly use inference mode.

Minimal reference-release reproduction

Adapted from our saved te_probe_reference_release.py. Run only with sufficient memory and no competing GPU job. Set ASSETS to the local split-checkpoint directory. No transformer generation is needed.

import gc
from pathlib import Path
import torch
from ltx_pipelines.utils.blocks import PromptEncoder
from ltx_pipelines.utils.model_paths import ModelPaths

ASSETS = Path("/path/to/ltx-2.5-split-assets")
paths = ModelPaths.from_split(
    transformer_path=str(ASSETS / "ltx-2.5-22b-distilled-transformer-nvfp4.safetensors"),
    text_encoder_path=str(ASSETS / "gemma4-12b-with-proj-ltx-2.5-bf16.safetensors"),
    video_vae_path=str(ASSETS / "ltx-2.5-video-vae-bf16.safetensors"),
    audio_vae_path=str(ASSETS / "ltx-2.5-audio-vae-bf16.safetensors"),
)

def memory(label):
    torch.cuda.synchronize()
    print(label, "allocated GiB", torch.cuda.memory_allocated() / 2**30,
          "reserved GiB", torch.cuda.memory_reserved() / 2**30)

enc = PromptEncoder(paths, torch.bfloat16, torch.device("cuda"))
memory("start")
with enc._text_encoder_ctx() as te:
    memory("encoder built")
    raw = te.encode(["Medium shot in a concrete stairwell. A lone man walks toward the camera."])
    memory("after encode")
memory("after encoder context exit")
gc.collect()
torch.cuda.empty_cache()
memory("after collection and empty_cache")
del raw
gc.collect()
torch.cuda.empty_cache()
memory("after dropping raw outputs")
del enc
gc.collect()
torch.cuda.empty_cache()
memory("after dropping encoder wrapper")

The private context is used only to separate the phases of PromptEncoder.__call__. The public library entry points expose the same grad-mode trap.

Observed versus expected

The earlier probe retained roughly 48 GiB after encoding. Leaving the context and calling gc.collect() / empty_cache() did not free the retained tensors. Dropping raw released the retained allocation. A driver wrapped in @torch.inference_mode() completed A2V generation without this retention failure.

Expected: inference-only library calls should either enter inference mode themselves or prominently document the caller requirement. .eval() alone does not disable autograd.

Source findings

At the revision above:

Entry point Inference-mode decorator
utils/blocks.py:PromptEncoder.__call__ (line 778) Absent
distilled.py:DistilledPipeline.__call__ (line 187) Absent
a2vid_two_stage.py:A2VidPipelineTwoStage.__call__ (line 151) Absent
distilled.py:main (line 320) Present
a2vid_two_stage.py:main (line 417) Present
ti2vid_two_stages_hq.py:TI2VidTwoStagesHQPipeline.__call__ (line 173) Already present

This is not a claim that every pipeline class lacks the decorator. The HQ path already demonstrates the safer library convention. PromptEncoder builds the encoder in eval mode, encodes into raw_outputs, then passes those outputs to the embeddings processor. Outside inference mode, those outputs can keep the computation graph alive.

Proposed fixes

  1. Wrap PromptEncoder.__call__ and the uncovered inference pipeline __call__ methods in @torch.inference_mode(), following the HQ class. Review whether any supported caller intentionally requires gradients before making this API choice.
  2. Alternatively, document the requirement prominently in packages/ltx-pipelines/README.md, beside the first library example:
with torch.inference_mode():
    result = pipeline(...)  # include prompt encoding inside this scope

A focused regression check could confirm grad mode is disabled during encoder invocation and that the public library path does not retain the encoder via output autograd graphs. The CLI-only path will not catch this difference.

No issue has been posted. No upstream or local pipeline source was changed for this draft.

Lingua principale
Python
Stelle
10k
Fork
1.6k
Nessuna PR unita negli ultimi 30g

Preparare l'ambiente

Questo progetto non fornisce container di sviluppo, Dockerfile né guida per i contributori, quindi l'ambiente è a tuo carico: parti dal suo README e consulta la nostra guida al primo contributo per i passaggi generali.

Come iniziare

  1. Leggi tutta la issue e poi la guida ai contributi del progetto.
  2. Commenta sulla issue per dire che te ne occupi tu — evita che due persone facciano lo stesso lavoro.
  3. Fai un fork del repository e lavora su un branch.
  4. Apri una pull request che faccia riferimento al numero della issue.

Altre issue di Lightricks/LTX-2

Tutte le issue di Lightricks/LTX-2

Issue simili

Altre issue su Python

Ricevi le nuove issue nella tua casella

Un breve riepilogo di issue GitHub adatte ai principianti.