SIGN IN SIGN UP

Pretrain, finetune ANY AI model of ANY size on 1 or 10,000+ GPUs with zero code changes.

0 0 104 Python

Fix torch.compile breaking toggle_optimizer / untoggle_optimizer (#21686)

* Fix toggle_optimizer breaking under torch.compile (#21513)

`LightningModule.toggle_optimizer` and `untoggle_optimizer` mutate
`requires_grad` on parameters to implement multi-optimizer gradient
masking. Dynamo/AOTAutograd does not support `setattr()` on
`Tensor.requires_grad` because it can change a tensor's leaf-ness
mid-graph, so when the `LightningModule` is wrapped with
`torch.compile` tracing either graph-breaks with
"Unsupported: setattr() on Tensor.requires_grad" or raises a
`KeyError` on the internal `param_requires_grad_state` mapping when
the traced parameter references diverge from those held by
`trainer.optimizers`.

Decorate both helpers with `@torch.compiler.disable` (the same
pattern already used for logging bookkeeping in
`logger_connector/result.py`) so they run as opaque Python when
called from a compiled `training_step`. Eager behavior is unchanged.

Adds a CPU regression test that compiles a two-optimizer
`LightningModule` calling `toggle_optimizer` / `untoggle_optimizer`
in `training_step` and exercises one training iteration, plus a
CHANGELOG entry.

* Narrow test_toggle_untoggle to check compiler.disable attribute (#21513)

The previous regression test compiled a `LightningModule` end-to-end
and called `self.optimizers()` inside the compiled `training_step`,
which unrelated to the toggle_optimizer fix trips a separate Dynamo
limitation: tracing `self.trainer.strategy._lightning_optimizers`
raises `InternalTorchDynamoError: GetAttrVariable(...) has no type`
across all CI platforms and torch versions.

The shipped fix — `@torch.compiler.disable` on `toggle_optimizer` /
`untoggle_optimizer` — does not require a full compiled trainer run
to verify; it only guarantees Dynamo skips those two methods.
Replace the integration test with a direct attribute check that both
methods carry the `_torchdynamo_disable` marker installed by
`torch.compiler.disable`, following the same `has_dynamo(fn)` pattern
already used by `tests/utilities/test_compile.py::test_compile_uncompile`.

Toggle/untoggle functional correctness remains covered by the existing
`test_toggle_untoggle_2_optimizers_no_shared_parameters` and
`test_toggle_untoggle_3_optimizers_shared_parameters` tests in this
file.

---------

Co-authored-by: Deependu <deependujha21@gmail.com>
G
Gaurav Dubey committed
1120456eef88c7d532246a0d30a359019725b283
Parent: b4b5f6d
Committed by GitHub <noreply@github.com> on 6/1/2026, 11:32:33 AM