[GritHat]

Problem

diffusion_generate in constrained_diffusion/eval/dllm/models/dream/generate_constrained.py is decorated with @torch.no_grad() (line 172 on main, 2682359), but it is not a generator function: it creates the _sample(...) generator (line 254) and returns it unconsumed (line 274). Applied to a plain function, torch.no_grad() only wraps that call (decorate_context in torch/utils/_contextlib.py); it re-enters the context on every resumption only when the decorated function is itself a generator function (_wrap_generator). So the model forward passes inside _sample (line 380) run later, when the caller iterates the generator, in whatever grad mode the caller has.

The callers do not disable gradients: DreamModel.generate_constrained and generate_unconstrained (dream/model.py, lines 120 and 181) iterate the generator directly, and so does the evaluation harness above them. The model is in eval() mode, which does not affect autograd, and its parameters have requires_grad=True as loaded. As a result, every Dream forward pass in the evaluation builds an autograd graph that is never used, and the saved activations of one step stay alive until the next step's forward pass replaces them.

This is specific to the Dream path. In upstream Dream, _sample returns the final tensor, so the decorator on diffusion_generate covered it; that stopped being true when _sample became a generator here. The LLaDA path is fine: generate in llada/generate_constrained.py is itself a generator function carrying the decorator. dream/generate.py and the MRI _sample are not generators.

What it costs

Measured on an RTX 3060 (12 GB) with Dream-Coder-v0-Instruct-7B loaded in 8-bit (bitsandbytes), one forward pass:

• sequence length 553 (a jsonschema prompt plus 256 new tokens): 8.28 GiB peak with autograd off, CUDA out of memory (peak above 11.3 GiB) with autograd on;
• sequence length 200: 9.50 GiB peak with autograd on; the logits are identical bit for bit in the two modes.

In bf16 the extra memory at the longest jsonschema sequence (763 tokens) is roughly 11–15 GiB on top of the 14 GiB of weights (estimated by running the harness with the model truncated to 2–8 layers and extrapolating to 28: 11.5 GiB with fused SDPA, 15.4 GiB with the math fallback). On the 48 GB GPUs assumed by run_dllm.py this goes unnoticed; on a 24 GB card the bf16 harness does not fit, and on 12 GB in 8-bit the first jsonschema instance already fails. In the constrained branch the OOM is caught per instance (generic_inference.py, line 172), so the instance is dropped from the results without stopping the run.

Outputs are not affected: autograd does not change the computed values, and with the same seed the generated ids are identical with and without the change (see below).

Time is not where it hurts. Per forward pass on the same GPU (median of 15 passes timed with CUDA events, the previous step's graph kept alive as in _sample): in bf16, with the model truncated to 2–8 layers and extrapolated linearly to 28, 470.3 ms with autograd off vs 470.8 ms with it on at 763 tokens (+0.1 %), and 392.1 vs 392.2 ms at 553 tokens; in 8-bit with the full model at 200 tokens, +1.3 %, which is within the run-to-run noise. So the runtime figures in the paper should not be affected in any measurable way, at least on a GPU where the kernel time dominates; the cost of the bug is memory.

Fix

Add @torch.no_grad() to _sample (one line). The decorator on diffusion_generate stays, the generator interface and the sampling logic are unchanged, and this is the same pattern as the LLaDA path. yield from result in diffusion_generate would also work, but it would turn diffusion_generate itself into a generator and defer the argument preparation to the first iteration.

Test

tests/test_dream_generate_no_grad.py drives the real diffusion_generate with a tiny stand-in model (CPU, well under a second) and records torch.is_grad_enabled() and logits.requires_grad inside every forward pass. On main it fails ([True, True]); with this change it passes ([False, False]), all masked positions are filled and the caller's grad mode is restored afterwards. It needs only torch, which the package already pulls in through accelerate, so it runs in the existing CI matrix (3.11–3.13).

Validation

• pytest tests/test_dream_generate_no_grad.py (Python 3.13, torch 2.8.0): fails on main with assert [True, True] == [False, False], passes on this branch.

• Reproduction with a small script that applies the decorator at runtime (not part of the PR): on main the forward passes report grad enabled [True, True]; with torch.no_grad() on _sample, [False, False]; the generated ids are identical in the two runs.

• pytest tests on this branch: 407 passed, 8 skipped (the 406 of main plus the new test).

• ruff check and ruff format --check (ruff 0.7.0, as in .pre-commit-config.yaml): clean on the two files.

• The evaluation harness itself (generic_inference.main), 20 jsonschema instances (--limit 20, seed 0, 32 steps, 256 new tokens, 60 s timeout), Dream-Coder-7B in 8-bit on the RTX 3060 (12 GB), same code except for this one line:

  • without the change: the first instance (567 tokens) fails with CUDA out of memory at 11.2 GiB allocated; the memory is not released afterwards and the remaining 19 instances fail immediately as well. The run ends after 31 s with 0 of 20 instances and no results file; the only trace is the Error generating code for instance … lines in the log. Reproduced twice.
  • with the change: 20 of 20 instances complete, none times out, peak allocated 9.5 GiB (10.3 GiB reserved), 173 s in total.
  • the same 20 instances with 4-bit (nf4) weights instead of 8-bit: still 0 of 20 without the change (out of memory at the first instance, 11.3 GiB peak with about 5 GiB of weights), 20 of 20 with it (peak 6.8 GiB, 216 s).

  A paired comparison of the outputs on this GPU is therefore not possible; the evidence that the values do not change is the bit-identical logits above and the identical ids in the reproduction.

• Forward-pass timing (CUDA events, median of 15 passes, RTX 3060): bf16 with the model truncated to 2, 4, 6 and 8 layers and extrapolated to 28 (the per-layer cost is 15.65 ms without autograd and 15.67 ms with it at 763 tokens): +0.1 % at 763 tokens, 0.0 % at 553; full model in 8-bit at 200 tokens: +1.3 %, with about 1 % of noise between two repeated no-grad series; at 300 and 400 tokens the autograd-on pass runs out of memory.

[nielstron]

The test case is overkill IMO for this one line change. Thanks for catching this though!

[GritHat]

Thanks for the review, I've removed the test case and kept the fix minimal.

I appreciate you taking the time to look at this!