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

Open
  • Is something wrong?
#310 0 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

Assessment

Difficulty
2/5
Estimated time
1-3 hours
Newbie friendliness
35/100
Issue type
Bug
Clarity
Mostly clear
Activity status
Quiet
Tech stack
python, pytorch
Domain
ai, machine-learning

Research direction

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.

Written by the indexing model from the issue text.

Description

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.

Dominant language
Python
Stars
10k
Forks
1.6k
No merged PRs in 30d

Getting set up

This project ships no dev container, Dockerfile or contributing guide, so setting up is up to you: start from its README, and see our first-contribution guide for the general steps.

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

More from Lightricks/LTX-2

All issues in Lightricks/LTX-2

Similar issues

More Python issues

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.