Skip to content

[models] optimize: call text_encoder.model instead of text_encoder to… - #214

Open
xucqX wants to merge 1 commit into
X-GenGroup:mainfrom
xucqX:optimize/ltx2-skip-lm-head-in-text-encoding
Open

[models] optimize: call text_encoder.model instead of text_encoder to…#214
xucqX wants to merge 1 commit into
X-GenGroup:mainfrom
xucqX:optimize/ltx2-skip-lm-head-in-text-encoding

Conversation

@xucqX

@xucqX xucqX commented Aug 24, 2026

Copy link
Copy Markdown

In _encode_text of LTX2_T2AV_Adapter, calling self.pipeline.text_encoder(...) runs the full Gemma3ForConditionalGeneration.forward(), which unnecessarily computes logits through the LM head. Since we only need hidden_states, switching to self.pipeline.text_encoder.model(...) (i.e., Gemma3Model) skips the LM head projection, saving compute and memory while producing identical hidden states.

@Jayce-Ping

Copy link
Copy Markdown
Collaborator

I don't think this optimization is necessary. My reasons:

  1. Maintaining encoder.__call__ is intentional to support future features that require a distributed backend wrapper. For example, when the text-encoder and the transformer are trained simultaneously.

  2. The encode_prompt function is usually used in the preprocessing stage, where the computational overhead is quite small compared to training, calling encoder.model() directly will bypass the backend forward/backward hooks.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants