Fix LTX Video skip_layer_mask device mismatch under mmgp offloading#2015
Open
Jnalley123 wants to merge 1 commit into
Open
Fix LTX Video skip_layer_mask device mismatch under mmgp offloading#2015Jnalley123 wants to merge 1 commit into
Jnalley123 wants to merge 1 commit into
Conversation
create_skip_layer_mask builds the mask on the transformer's device, which reports cpu while mmgp keeps the weights offloaded to RAM. Any STG generation that actually skips blocks then crashes with "Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu!" when the cpu mask meets cuda activations (seen with LTX Video 0.9.8 13B). Move the mask to the activations' device at the two consumption points, covering all four SkipLayerStrategy paths: AttentionSkip, AttentionValues and Residual via the reshape in AttnProcessor2_0, and TransformerBlock in BasicTransformerBlock. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What
Moves
skip_layer_maskto the activations' device at its consumption points inmodels/ltx_video/models/transformers/attention.py:AttnProcessor2_0.__call__, at thereshape(batch_size, 1, 1)step — covers theAttentionSkip,AttentionValues, andResidualstrategies downstream;BasicTransformerBlock.forward, at theview(-1, 1, 1)step — covers theTransformerBlockstrategy.Why
Transformer3DModel.create_skip_layer_maskbuilds the mask withdevice=self.device. Under mmgp offloading the transformer's parameters live in CPU RAM between calls, soself.devicereportscpuand the mask is created there, while activations run on cuda. Any STG generation that actually skips blocks (an all-ones mask is discarded early) crashes with:Reproduced with LTX Video 0.9.8 13B (
ltxv_0.9.8_13B_dev_quanto_bf16_int8.safetensors) on Windows, torch 2.10.0+cu130, mmgp 3.7.6. Fixing at the consumption points (rather than at mask creation) keeps the fix correct wherever mmgp happens to be keeping the weights at the time.Notes for review
AttentionSkipandAttentionValuesbranches, bundled with unrelated AMD/ROCm changes. This PR is the standalone, complete version: theResidualandTransformerBlockbranches have the identical latent bug and are covered here..to()is a no-op when the mask is already on the right device, so non-offloaded setups are unaffected.🤖 Generated with Claude Code