Skip to content

Fix broken group offloading with block_level for models with standalone layers #12692

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
DN6 merged 14 commits into from
Dec 5, 2025
Merged

Fix broken group offloading with block_level for models with standalone layers #12692

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
14 commits
Select commit Hold shift + click to select a range
ad1fc37
fix: group offloading to support standalone computational layers in b…
rycerzes Nov 21, 2025
59b6b67
test: for models with standalone and deeply nested layers in block-le…
rycerzes Nov 21, 2025
fa94f37
feat: support for block-level offloading in group offloading config
rycerzes Nov 24, 2025
fb8a741
fix: group offload block modules to AutoencoderKL and AutoencoderKLWan
rycerzes Nov 24, 2025
e71d91e
fix: update group offloading tests to use AutoencoderKL and adjust in…
rycerzes Nov 24, 2025
09dd19b
Merge branch 'main' into fix/broken-group-offloading-using-block_level
rycerzes Nov 24, 2025
e771143
refactor: streamline block offloading logic
rycerzes Dec 3, 2025
ab9b249
Apply style fixes
github-actions[bot] Dec 3, 2025
26bccde
update tests
DN6 Dec 4, 2025
f305934
update
DN6 Dec 4, 2025
cf65ae3
fix for failing tests
DN6 Dec 5, 2025
09a7b0a
clean up
DN6 Dec 5, 2025
c888aac
revert to use skip_keys
DN6 Dec 5, 2025
4bd3384
clean up
DN6 Dec 5, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Prev Previous commit
Next Next commit
fix: group offload block modules to AutoencoderKL and AutoencoderKLWan
  • Loading branch information
@rycerzes
rycerzes committed Nov 24, 2025
commit fb8a74195dbb9fb2b0677f906450cb67f601e8bf
1 change: 1 addition & 0 deletions src/diffusers/models/autoencoders/autoencoder_kl.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -72,6 +72,7 @@ class AutoencoderKL(ModelMixin, AutoencoderMixin, ConfigMixin, FromOriginalModel

_supports_gradient_checkpointing = True
_no_split_modules = ["BasicTransformerBlock", "ResnetBlock2D"]
_group_offload_block_modules = ["quant_conv", "post_quant_conv", "encoder", "decoder"]

@register_to_config
def __init__(
Expand Down
1 change: 1 addition & 0 deletions src/diffusers/models/autoencoders/autoencoder_kl_wan.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -964,6 +964,7 @@ class AutoencoderKLWan(ModelMixin, AutoencoderMixin, ConfigMixin, FromOriginalMo
# keys toignore when AlignDeviceHook moves inputs/outputs between devices
# these are shared mutable state modified in-place
_skip_keys = ["feat_cache", "feat_idx"]
_group_offload_block_modules = ["quant_conv", "post_quant_conv", "encoder", "decoder"]

@register_to_config
def __init__(
Expand Down