Per-component offload: streaming text encoder while transformers stay GPU-resident (warm-serving OOM at 96 GB)

Đang mở
  • Có gì không ổn?
#232 1 bình luận 0 reaction 0 người được giao Xem trên GitHub

Chưa có ai nhận issue này.

Đánh giá

Độ khó
4/5
Thời gian dự kiến
3-5 ngày
Mức phù hợp với người mới
52/100
Loại issue
Tính năng
Độ rõ ràng
Khá rõ ràng
Mức độ hoạt động
Ít trao đổi
Công nghệ
python, pytorch
Lĩnh vực
machine-learning

Hướng nghiên cứu

Start with packages/ltx-pipelines/src/ltx_pipelines/utils/blocks.py, especially PromptEncoder.__call__ and _text_encoder_ctx, then trace how pipeline constructors pass offload_mode to the prompt encoder and diffusion stages. Check the related issues (#164, #143, #140, and #152) for context. Done means callers can stream the text encoder while keeping transformers GPU-resident, with the default behavior preserved.

Do mô hình lập chỉ mục viết ra từ nội dung của issue.

Mô tả

Setup

Self-hosted serving of LTX-2.3 (22B) on a single RTX PRO 6000 (96 GB), bf16. To keep warm latency low we keep both stage transformers GPU-resident between requests (~70.8 GB) and cache text embeddings per prompt.

Symptom

The first request with a new prompt on a warm container OOMs:

70.8 GB resident transformers + ~23 GB Gemma full-GPU build → 93.75 GB → CUDA OOM

PromptEncoder.__call__ builds the full Gemma text encoder on GPU per call (packages/ltx-pipelines/src/ltx_pipelines/utils/blocks.py — _text_encoder_ctx returns gpu_model(self._build_text_encoder()) when offload_mode == NONE). With embedding caching, every repeated-prompt request skips this entirely, so benchmarks stay green and the OOM only appears in production on the first unseen prompt.

Why --offload cpu doesn't solve it

The streaming text-encoder builder already exists on every PromptEncoder instance and works well — but offload_mode is a single global knob: every pipeline factory passes the same value to the PromptEncoder and all DiffusionStages. Setting --offload cpu to fix the text encoder also layer-streams the 22B transformers, which defeats persistent-GPU serving (and per the README, offload disables FP8 quantization as a further coupling).

So today there is no supported way to get: streaming text encoder + GPU-resident transformers.

Request

Per-component offload granularity — for example a text_encoder_offload_mode parameter on the pipeline constructors (defaulting to the global offload_mode), or an explicit offload_mode override accepted by PromptEncoder consumers. Even exposing it only at the Python API level (not CLI) would cover the serving use case.

Workaround we run (in case it helps others)

We wrap the prompt-encode call and flip _offload_mode to OffloadMode.CPU only for cache-miss calls when free VRAM is below a threshold, restoring it after. Same weights/dtype → identical embeddings; the request that used to die completes at 78.6 GB peak (~+10 s once per prompt, then cached). Happy to share details or a patch if useful.

Related: #164 (Gemma VRAM pain, smaller-encoder angle), #143 (Gemma co-residency OOM), #140 (offload question), #152 (adjacent external-serving OOM).

Ngôn ngữ chính
Python
Star
10k
Fork
1.6k
Không có pull request nào được merge trong 30 ngày

Chuẩn bị môi trường

Dự án này không cung cấp dev container, Dockerfile hay hướng dẫn đóng góp, nên bạn cần tự thiết lập môi trường: hãy bắt đầu từ README và xem hướng dẫn đóng góp lần đầu của chúng tôi để biết các bước chung.

Bắt đầu từ đâu

  1. Đọc hết issue, rồi đọc hướng dẫn đóng góp của dự án.
  2. Bình luận trên issue rằng bạn sẽ nhận — tránh hai người làm cùng một việc.
  3. Fork repository và làm thay đổi trên một nhánh.
  4. Mở pull request có tham chiếu số hiệu của issue.

Issue khác của Lightricks/LTX-2

Tất cả issue của Lightricks/LTX-2

Issue tương tự

Thêm issue về Python

Nhận issue mới trong hộp thư của bạn

Bản tóm tắt ngắn những issue GitHub phù hợp với người mới.