Run DiffusionGemma and LLaDA2 pipeline tests through the shared PipelineTesterMixin - #14661
Open
kashif wants to merge 10 commits into
Open
Run DiffusionGemma and LLaDA2 pipeline tests through the shared PipelineTesterMixin#14661kashif wants to merge 10 commits into
kashif wants to merge 10 commits into
Conversation
…dows a pipeline dir
…ma's canvas init torch.randint/multinomial/rand need the generator and the sampled tensor on the same device. A CPU generator (the portable default recommended for reproducible pipeline calls) broke on any accelerator.
it already defaults to None and works without one; just wasn't declared, so save/load silently dropped it.
follows the audioldm2 pattern for a non-image output; skips the handful of tests that assume guidance/latents/generator-lists.
same treatment as DiffusionGemma; keeps the existing regression tests as-is alongside it.
missed this one earlier; two more torch.rand calls with the same CPU-generator-on-CUDA mismatch.
was duplicating the same one-liner inline.
was skipped; adapt it to use input_ids instead of prompt so it actually runs.
Open
6 tasks
github-actions
Bot
added
size/L
sayakpaul
reviewed
Aug 31, 2026
|
The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update. |
per review: it's a one-liner, not worth a shared helper.
# Conflicts: # tests/pipelines/llada2/test_llada2.py
sayakpaul
approved these changes
Sep 1, 2026
Member
|
@yiyixuxu could you also review the changes introduced in the core files? |
sayakpaul
reviewed
Sep 1, 2026
| token = flat_logits.argmax(dim=-1, keepdim=True) | ||
| else: | ||
| scaled_probs = torch.softmax(flat_logits.float() / temperature, dim=-1) | ||
| token = torch.multinomial(scaled_probs, num_samples=1, generator=generator) |
Member
There was a problem hiding this comment.
Do we use this scheduler in any of the pipelines for dLLM? If not, let's remove these changes.
Contributor
Author
There was a problem hiding this comment.
Yes, DiffusionGemma supports it directly and our tests exercise it, so it's in scope.
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
4 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.
What does this PR do?
Follow-up to #14652. Runs
DiffusionGemmaandLLaDA2pipeline tests through the sharedPipelineTesterMixin, the same wayaudioldm2does it for a non-image output, skipping the tests that assume image/video semantics (CFG,latents, per-row generator lists). Existing hand-written generation/regression tests stay as-is alongside.Running the mixin's save/load and to-device tests on a real accelerator caught two real bugs, fixed here too:
DiffusionGemma's canvas init and all three block-diffusion schedulers calledtorch.randint/multinomial/randwith a CPU generator on a CUDA tensor, whichtorchrejects. Fixed likerandn_tensordoes: sample on the generator's device, move after.LLaDA2Pipeline.tokenizerdefaults toNoneand works fine without one, but wasn't declared in_optional_components, so save/load silently dropped it.Includes akshan-main's #14652 commit as a base — needed for DiffusionGemma's own save/load tests to pass.
Self-review notes (
/code-review medium): caught two moretorch.randcalls inadd_noisewith the same generator/device bug (fixed), and DiffusionGemma's pipeline duplicating the device-selection logic instead of reusing the new shared helper (fixed). Also flagged that the_optional_componentsfix wasn't actually exercised since that mixin test was skipped — adapted it to run withinput_idsinstead.Before submitting
self-reviewskill on the diff?Who can review?
@sayakpaul