Skip to content

[tests] implement base model output caching in model-level tests#14059

Merged
sayakpaul merged 5 commits into
mainfrom
cache-base-model-output
Jun 25, 2026
Merged

[tests] implement base model output caching in model-level tests#14059
sayakpaul merged 5 commits into
mainfrom
cache-base-model-output

Conversation

@sayakpaul

@sayakpaul sayakpaul commented Jun 24, 2026

Copy link
Copy Markdown
Member

Implements fixture-level caching of base model output to avoid recomputation:

  • common.py
  • memory.py

Other mixins don't have the right contract to implement it yet.

@github-actions github-actions Bot added tests size/M PR with diff < 200 LOC labels Jun 24, 2026
@sayakpaul

Copy link
Copy Markdown
Member Author

/diffusers-bot pytest tests/models/transformers/test_models_transformer_flux.py -k "offload"

@github-actions

github-actions Bot commented Jun 24, 2026

Copy link
Copy Markdown
Contributor

pytest tests/models/transformers/test_models_transformer_flux.py -k "offload" passed on GPU — view logs.

@sayakpaul sayakpaul requested review from DN6 and dg845 and removed request for DN6 June 24, 2026 20:41
@sayakpaul sayakpaul marked this pull request as ready for review June 24, 2026 20:41
@sayakpaul sayakpaul requested a review from DN6 June 24, 2026 20:43
Tests that still need a live model (e.g. to save or offload it) build their own with the same
seed, so the reloaded model's weights match this cached output.
"""
torch.manual_seed(0)

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Should BaseModelOutputMixin also expose the seed it uses to create the base_model_output fixture? I think this would allow tests to more easily match its behavior.

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I thought about it but then this change should also be propagated to the generator property, e.g.

@property
def generator(self):
return torch.Generator("cpu").manual_seed(0)

So, I think this should be done in a separate PR.

@dg845 dg845 left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for the PR! Left one comment.

@sayakpaul sayakpaul merged commit 9d52996 into main Jun 25, 2026
18 checks passed
@sayakpaul sayakpaul deleted the cache-base-model-output branch June 25, 2026 05:01
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

size/M PR with diff < 200 LOC tests

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants