Fix CompatibleDiT len() crash with cache_model + torch.compile (#502)#588
Open
Deaquay wants to merge 1 commit into
Open
Fix CompatibleDiT len() crash with cache_model + torch.compile (#502)#588Deaquay wants to merge 1 commit into
Deaquay wants to merge 1 commit into
Conversation
) When cache_model=True is combined with torch.compile, cached DiT/VAE modules are wrapped in torch._dynamo.OptimizedModule. Several call sites do implicit truthiness checks on these wrapped modules which invoke __bool__, fall back to __len__, and raise TypeError because the underlying nn.Module subclass does not implement __len__. This patch replaces those checks with explicit is None / is not None identity checks on potentially-compiled module objects. generation_utils.py already uses this convention (lines 799, 807); this brings generation_phases.py and model_configuration.py in line with it. Sites patched: - src/core/model_configuration.py:596, 616 (if cached_model:) - src/core/generation_phases.py:298, 620, 896 (if runner.X and ...) - src/core/generation_phases.py:310, 632 (and not ctx[...][cached_X]) - src/core/generation_phases.py:319, 641 (ctx[...][cached_X] or ...) Tested on RTX 5090 with 7B fp8 model: without the patch, gen 2 of any session with cache_model=True and torch.compile enabled crashes with TypeError. With the patch, two sequential generations succeed and the second hits the cached runner.
Author
|
Note: Issue #502 was closed by its reporter after they worked around the crash by disabling torch.compile (removing the compile-settings node from their workflow). However, the underlying bug remains — |
|
This should be merged. |
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
2 participants
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.
When
cache_model=Trueis combined withtorch.compile, cached DiT/VAE modules are wrapped intorch._dynamo.OptimizedModule. Several call sites do implicit truthiness checks on these wrapped modules, which invoke__bool__, fall back to__len__, and raiseTypeError: CompatibleDiT does not support len()because the underlying nn.Module subclass doesn't implement__len__.This PR replaces those checks with explicit
is None/is not Noneidentity checks on potentially-compiled module objects.src/core/generation_utils.pyalready uses this convention (lines 799, 807); this bringsgeneration_phases.pyandmodel_configuration.pyin line with it.Sites patched
src/core/model_configuration.py:596, 616—if cached_model:src/core/generation_phases.py:298, 620, 896—if runner.X and ...src/core/generation_phases.py:310, 632—and not ctx[...]['cached_X']:src/core/generation_phases.py:319, 641—ctx[...]['cached_X'] or ctx[...]['X_newly_cached']Testing
RTX 5090, 7B fp8 model (
seedvr2_ema_7b_fp8_e4m3fn_mixed_block35_fp16.safetensors),cache_model=True, torch.compile enabled.TypeError: CompatibleDiT does not support len()at the firstif cached_model:(line 596 ofmodel_configuration.py), then progressively at the other sites if individual sites are patched in isolation.♻️ Reusing cached runner template) and skips re-materialization.Refs #502.