Fix DSVAE's forward path and the arguments it drops - #8413
Conversation
DSVAE.forward tests self.cuda_graph_created, which DSVAE never sets. __init__ sets decoder_cuda_graph_created, encoder_cuda_graph_created and all_cuda_graph_created, and _create_cuda_graph sets all_cuda_graph_created, so with the default enable_cuda_graph=True the first call raises AttributeError: 'DSVAE' object has no attribute 'cuda_graph_created'. all_cuda_graph_created is written twice and read nowhere. _forward is a copy of DSUNet._forward: it requires timestamp and encoder_hidden_states, which a VAE does not take, and passes them positionally into AutoencoderKL.forward(sample, sample_posterior, return_dict, generator). Take the parameters that method actually has and pass them by keyword. _decode accepts generator and drops it. Every diffusers pipeline calls vae.decode(latents, return_dict=False, generator=generator), and a stochastic decoder such as ConsistencyDecoderVAE samples with it. Forward it when one is given; not every VAE's decode takes the argument. Signed-off-by: Vineeth Sai <vineethsai4444@gmail.com>
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 4849081cf5
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| # VAE's `decode` accepts the argument. | ||
| if generator is None: | ||
| return self.vae.decode(x, return_dict=return_dict) | ||
| return self.vae.decode(x, return_dict=return_dict, generator=generator) |
There was a problem hiding this comment.
With CUDA graphs enabled, the first decode call stores this torch.Generator in static_decoder_kwargs, while _graph_replay_decoder updates only tensor kwargs before replaying the already captured graph. A later call with a different generator therefore still samples using the first call's generator, so stochastic decoders produce results determined by an earlier request rather than the caller's seed; the same non-tensor-capture issue is introduced for forward's generator argument. Exclude generator-bearing calls from graph replay or rebuild/refresh the graph when the generator changes.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
The container half is real and is now #8488: _graph_replay_decoder matches on torch.is_tensor, so anything nested is never refreshed. That helper walks dicts, lists and tuples at all six replay sites.
The generator itself is a different matter and I am leaving it alone here. A torch.Generator is not a tensor, so there is nothing to copy into, and a CUDA graph cannot capture RNG state that way at all. Forwarding it, which is what this PR does, is what the pipelines need for the non-graph path. Making seeded decodes correct under capture needs the graph pool to take a generator, which is a bigger change than this PR.
| ds_vae = DSVAE(vae, enable_cuda_graph=False) | ||
| generator = torch.Generator() | ||
|
|
||
| ds_vae.decode(torch.zeros(1), return_dict=False, generator=generator) |
There was a problem hiding this comment.
This test disables CUDA graphs and uses a CPU recording stub, so it never exercises the inference path changed by this commit under CUDA capture/replay. In particular, it cannot detect generator or argument behavior once the VAE has been graph-captured; add and run an actual-device integration test using a real VAE/pipeline for this inference-level contract.
AGENTS.md reference: AGENTS.md:L35-L36
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Not adding one. A real VAE under CUDA capture needs a GPU, and I do not have one, so I would be adding a test I cannot run.
What this PR changes is which arguments DSVAE hands to the wrapped VAE. The recording stub pins exactly that, and it fails on master. The capture path is unchanged by this diff.
1 participant
What breaks
DSVAEis the wrapperdeepspeed.init_inferenceswaps in for a diffusers VAE, so every Stable Diffusion encode/decode goes through it. Three things in it do not match what they wrap.1.
forwardreads a flagDSVAEnever sets.__init__setsdecoder_cuda_graph_created,encoder_cuda_graph_createdandall_cuda_graph_created;_create_cuda_graphsetsall_cuda_graph_created.forwardtestsself.cuda_graph_created, which onlyDSUNethas.enable_cuda_graphdefaults toTrue, so that is the default path, andall_cuda_graph_createdis written twice and read nowhere.2.
_forwardis a copy ofDSUNet._forward.A VAE has no timestep and no encoder hidden states.
AutoencoderKL.forwardis(sample, sample_posterior=False, return_dict=True, generator=None), soDSVAE(sample)raisesTypeError: missing 2 required positional arguments, and satisfying it landstimestamponsample_posterior,encoder_hidden_statesonreturn_dict, andreturn_dictongenerator.3.
_decodeacceptsgeneratorand drops it.StableDiffusionPipeline.__call__callsself.vae.decode(latents / scaling_factor, return_dict=False, generator=generator).AutoencoderKLignores it, butConsistencyDecoderVAEandAutoencoderTinysample with it, so their decode is silently non-reproducible behindinit_inference.What changed
forwardreadsall_cuda_graph_created, the flag the class sets. That makes the dead flag live and the default path work._forwardtakesAutoencoderKL.forward's parameters and passes them by keyword._decodeforwardsgeneratorwhen one was given. Only when given, sinceAutoencoderKLTemporalDecoder.decodehas no such argument.Test
New
tests/unit/inference/test_ds_vae.py, three plain functions drivingDSVAEwith a recording stub whose signatures are copied fromAutoencoderKL. No accelerator, no model download.3 failed before, 3 passed after (
AttributeError,TypeError, andgeneratorarriving asNone), run on torch 2.11.0 against deepspeed 0.19.6, whosevae.pyis byte-identical to master.yapf --style .style.yapf --diffclean on both files.