Repository: pytorch/torchtune
Pull request #2719: New ``OptimizerInBackward``
Author: login: joecummings
Labels: name: CLA Signed
description: This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed.
Created: 2025-05-12T01:24:53Z | Merged: 2025-05-19T21:37:46Z

# ``OptimizerInBackward``

### Motivation

Previously, we utilized a process very similar to the one outlined in [this blog post](https://github.com/pytorch/torchtune/pull/2719/files). Although it worked, it involved lots of if/else switching in the recipe based on what composed well with optimizer fused in backwards and what did not.

### Goal

Simplify our recipes (729 LOC -> 677 LOC)! This PR provides a canonical ``OptimizerInBackward`` class that can be used as a drop-in replacement for any PyTorch optimizer. It integrates into the `full_finetune_single_device.py` recipe and adds tests + updates documentation.

### Testing

![Screenshot 2025-05-15 at 7 16 23 PM](https://github.com/user-attachments/assets/e33a4d81-ab28-4fcb-9ffe-5e070c414394)



### To-do

- [ ] Integrate into `full_finetune_distributed.py`
- [ ] Integrate into rest of recipes
- [ ] Deprecate the utility functions used to create the optim in backward hooks previously


Diff:
diff --git a/docs/source/api_ref_modules.rst b/docs/source/api_ref_modules.rst
index 41afe0fbe5..81437283c3 100644
--- a/docs/source/api_ref_modules.rst
+++ b/docs/source/api_ref_modules.rst
@@ -40,6 +40,15 @@ Losses
     loss.ReverseKLWithChunkedOutputLoss
     loss.SymmetricKLWithChunkedOutputLoss
 
+Optimizers
+----------
+
+.. autosummary::
+    :toctree: generated/
+    :nosignatures:
+
+    optim.OptimizerInBackward
+
 Base Tokenizers
 ---------------
 Base tokenizers are tokenizer models that perform the direct encoding of text
diff --git a/docs/source/tutorials/memory_optimizations.rst b/docs/source/tutorials/memory_optimizations.rst
index 2206fd4c0d..5966d8a628 100644
--- a/docs/source/tutorials/memory_optimizations.rst
+++ b/docs/source/tutorials/memory_optimizations.rst
@@ -218,10 +218,28 @@ To understand how this works, we encourage you to read through the relevant PyTo
 
 *Sounds great! How do I use it?*
 
-.. todo ref full finetune recipe doc
+In torchtune, you can enable this feature using :class:`OptimizerInBackward`, which wraps around a normal PyTorch optimizer like so and can be used
+as a drop-in replacement for your normal Optimizer:
 
-In torchtune, you can enable this feature using the ``optimizer_in_bwd`` flag. This feature works best when using a stateful optimizer
-with a model with a lot of parameters, and when you don't need to use :ref:`gradient accumulation <glossary_grad_accm>`.
+.. code-block:: python
+
+  OptimizerInBackward(
+    model.parameters(),
+    torch.optim.AdamW,
+    lr=0.01,
+  )
+
+In order to use this OOTB with a torchtune recipe and config, simply specify your chosen optimizer and then add the key ``optimizer_in_bwd: True``, which will
+construct a :class:`OptimizerInBackward` from the base optimizer class.
+
+.. code-block:: yaml
+
+  optimizer:
+    _component_: torch.optim.AdamW
+    lr: 0.01
+  optimizer_in_bwd: True
+
+This feature works best when using a stateful optimizer with a model with a lot of parameters, and when you don't need to use :ref:`gradient accumulation <glossary_grad_accm>`.
 You won't see meaningful impact when finetuning LoRA recipes, since in this case the number of parameters being updated are small.
 
 .. _glossary_cpu_offload:
diff --git a/recipes/full_finetune_single_device.py b/recipes/full_finetune_single_device.py
index bae96a1942..16aa0dbb0e 100644
--- a/recipes/full_finetune_single_device.py
+++ b/recipes/full_finetune_single_device.py
@@ -24,6 +24,7 @@
 from torchtune.datasets import ConcatDataset
 from torchtune.modules.embedding_utils import resize_token_embeddings
 from torchtune.modules.loss import SFTLoss
+from torchtune.modules.optim import OptimizerInBackward
 from torchtune.recipe_interfaces import FTRecipeInterface
 from torchtune.training import DummyProfiler, PROFILER_KEY
 from torchtune.training.checkpointing._checkpoint_client import (
@@ -31,7 +32,6 @@
     TrainingProgress,
 )
 from torchtune.training.lr_schedulers import get_lr
-from torchtune.training.memory import OptimizerInBackwardWrapper
 
 from tqdm import tqdm
 
@@ -146,21 +146,17 @@ def __init__(self, cfg: DictConfig) -> None:
         # Training cfg
         self._resume_from_checkpoint = cfg.resume_from_checkpoint
         self._gradient_accumulation_steps = cfg.gradient_accumulation_steps
-        self._optimizer_in_bwd = cfg.optimizer_in_bwd
         self._clip_grad_norm = cfg.get("clip_grad_norm", None)
+        self.optimizer_in_bwd = cfg.optimizer_in_bwd
 
-        # Optimizer in backward is not compatible with gradient accumulation or gradient clipping
-        if self._optimizer_in_bwd:
-            if self._clip_grad_norm is not None:
-                raise RuntimeError(
-                    "Gradient clipping is not supported with optimizer in bwd."
-                    "Please set clip_grad_norm=None, or optimizer_in_bwd=False."
-                )
-            if self._gradient_accumulation_steps > 1:
-                raise RuntimeError(
-                    "Gradient accumulation is not supported with optimizer in bwd."
-                    "Please set gradient_accumulation_steps=1, or optimizer_in_bwd=False."
-                )
+        if self.optimizer_in_bwd and self._clip_grad_norm is not None:
+            raise RuntimeError(
+                "Gradient clipping is not supported with optimizer_in_bwd: True"
+            )
+        if self.optimizer_in_bwd and self._gradient_accumulation_steps > 1:
+            raise RuntimeError(
+                "Gradient accumulation is not supported with optimizer_in_bwd: True"
+            )
 
         self._checkpoint_client = CheckpointClient(cfg)
         self._enable_async_checkpointing = cfg.get("enable_async_checkpointing", False)
@@ -276,9 +272,8 @@ def setup(self, cfg: DictConfig) -> None:
 
         # _setup_optimizer should take in ckpt_dict only if training is resumed from
         # checkpoint. Transforming the opt state dict is handled by this method
-        self._optimizer = self._setup_optimizer(
+        self.optimizer = self._setup_optimizer(
             cfg_optimizer=cfg.optimizer,
-            optimizer_in_bwd=cfg.optimizer_in_bwd,
             opt_state_dict=(
                 ckpt_dict[training.OPT_KEY] if training.OPT_KEY in ckpt_dict else None
             ),
@@ -293,10 +288,10 @@ def setup(self, cfg: DictConfig) -> None:
                 try:
                     ckpt_dict = self._checkpoint_client.load_distributed_checkpoint(
                         self._model,
-                        self._optimizer,
+                        self.optimizer,
                     )
                 except Exception as e:
-                    log.warning(
+                    self._logger.warning(
                         f"Failed to load distributed checkpoint: {e}. Training will start from the base checkpoint."
                     )
 
@@ -343,13 +338,13 @@ def setup(self, cfg: DictConfig) -> None:
         # Setup lr scheduler if a component is included in the config
         lr_scheduler_cfg = cfg.get("lr_scheduler", None)
         if lr_scheduler_cfg is not None:
-            self._lr_scheduler = self._setup_lr_scheduler(
+            self.lr_scheduler = self._setup_lr_scheduler(
                 cfg_lr_scheduler=lr_scheduler_cfg,
                 num_training_steps=self.total_epochs * self._steps_per_epoch,
                 last_epoch=self.global_step - 1,
             )
         else:
-            self._lr_scheduler = None
+            self.lr_scheduler = None
 
         # Set up profiler, returns DummyProfiler (nullcontext object with no-op `step` method)
         # if cfg is missing profiler key or if `cfg.profiler.enabled = False`
@@ -432,32 +427,21 @@ def _setup_model(
     def _setup_optimizer(
         self,
         cfg_optimizer: DictConfig,
-        optimizer_in_bwd: bool = False,
         opt_state_dict: Optional[dict[str, Any]] = None,
-    ) -> Union[Optimizer, OptimizerInBackwardWrapper]:
-        """
-        Set up the optimizer. This method also handles loading the optimizer state_dict, if specified.
-        """
-        if optimizer_in_bwd:
-            # Maintain a dict of optims for every parameter.
-            optim_dict = {
-                p: config.instantiate(cfg_optimizer, [p])
-                for p in self._model.parameters()
-            }
-            # Register optimizer step hooks on the model to run optimizer in backward.
-            training.register_optim_in_bwd_hooks(
-                model=self._model, optim_dict=optim_dict
-            )
-            # Create a wrapper for checkpoint save/load of optimizer states when running in backward.
-            optimizer = training.create_optim_in_bwd_wrapper(
-                model=self._model, optim_dict=optim_dict
+    ) -> Optimizer:
+        if self.optimizer_in_bwd:
+            optimizer_cls = _get_component_from_path(cfg_optimizer.pop("_component_"))
+            optimizer = OptimizerInBackward(
+                params=self._model.parameters(),
+                optimizer_cls=optimizer_cls,
+                **cfg_optimizer,
             )
         else:
-            optimizer = config.instantiate(cfg_optimizer, self._model.parameters())
-
+            optimizer = config.instantiate(
+                cfg_optimizer, params=self._model.parameters()
+            )
         if opt_state_dict:
             optimizer.load_state_dict(opt_state_dict)
-
         self._logger.info("Optimizer is initialized.")
         return optimizer
 
@@ -467,27 +451,12 @@ def _setup_lr_scheduler(
         num_training_steps: int,
         last_epoch: int,
     ) -> LambdaLR:
-        """
-        Set up the learning rate scheduler based on the provided configuration.
-        It handles both standard optimization and optimizer-in-backward cases, and supports
-        schedulers from both torchtune.modules and torch.optim.
-        """
-        if isinstance(self._optimizer, OptimizerInBackwardWrapper):
-            # Use the first optimizer from the wrapper to represent the learning rate
-            optimizer = next(iter(self._optimizer.optim_map.values()))
-        else:
-            optimizer = self._optimizer
-
         lr_scheduler = config.instantiate(
             cfg_lr_scheduler,
-            optimizer,
+            self.optimizer,
             num_training_steps=num_training_steps,
             last_epoch=last_epoch,
         )
-
-        if isinstance(self._optimizer, OptimizerInBackwardWrapper):
-            self._optimizer.set_lr_scheduler(lr_scheduler)
-
         self._logger.info("Learning rate scheduler is initialized.")
         return lr_scheduler
 
@@ -552,7 +521,7 @@ def save_checkpoint(self, epoch: int) -> None:
         """
         self._checkpoint_client.save_checkpoint(
             model=self._model,
-            optimizer=self._optimizer,
+            optimizer=self.optimizer,
             training_progress=TrainingProgress(
                 seed=self.seed,
                 epochs_run=self.epochs_run,
@@ -585,25 +554,17 @@ def _loss_step(self, batch: dict[str, torch.Tensor]) -> torch.Tensor:
         return loss
 
     def train(self) -> None:
-        """
-        The core training loop. Supports training on subsets of the dataset using the
-        ``max_steps_per_epoch``.
-        """
-        if not isinstance(self._optimizer, OptimizerInBackwardWrapper):
-            self._optimizer.zero_grad()
-
-        # Initialize tokens count and running loss (for grad accumulation)
+        self.optimizer.zero_grad()
         t0 = time.perf_counter()
-        running_loss = 0
-        num_tokens = 0
-
+        running_loss, num_tokens = 0.0, 0
         self._profiler.start()
-        # self.epochs_run should be non-zero when we're resuming from a checkpoint
+
         for curr_epoch in range(self.epochs_run, self.total_epochs):
             pbar = tqdm(total=self._steps_per_epoch)
             self._dataloader.sampler.set_epoch(curr_epoch)
+
             for idx, batch in enumerate(self._dataloader):
-                # Start tracking CUDA memory for active steps for just the first epoch
+                # Optionally start memory profiling
                 if (
                     curr_epoch == 0
                     and self.profiler_profile_memory
@@ -611,73 +572,64 @@ def train(self) -> None:
                     and self._device.type == "cuda"
                 ):
                     torch.cuda.memory._record_memory_history()
-                utils.batch_to_device(batch, self._device)
 
-                # Calculate the number of unmasked tokens in the current batch
-                # and increment the total number of tokens seen in the step
+                utils.batch_to_device(batch, self._device)
                 current_num_tokens = (
                     batch["labels"] != self._loss_fn.ignore_index
                 ).sum()
                 num_tokens += current_num_tokens
 
-                # Loss is normalized by default so we multiply by the number of tokens
-                # This way we can normalize by the total number of tokens if we're accumulating gradients
-                current_loss = self._loss_step(batch) * current_num_tokens
-                running_loss += current_loss
-                # If using opt in backwards, we need to unnormalize before stepping
-                if isinstance(self._optimizer, OptimizerInBackwardWrapper):
-                    current_loss = current_loss * (1 / current_num_tokens)
+                # Mean loss for in-bwd optimizers, else multiply by token count
+                loss_factor = current_num_tokens if not self.optimizer_in_bwd else 1.0
+                current_loss = self._loss_step(batch) * loss_factor
+
+                running_loss += current_loss.detach()
                 current_loss.backward()
 
+                # Take a normal optimizer step
                 if (idx + 1) % self._gradient_accumulation_steps == 0:
-                    # If we're not using optimizer in backwards, we step the optimizer like normal
-                    if not isinstance(self._optimizer, OptimizerInBackwardWrapper):
-                        training.scale_grads(self._model, 1 / num_tokens)
-                        if self._clip_grad_norm is not None:
+                    grad_norm = None
+                    if not self.optimizer_in_bwd:
+                        training.scale_grads(self._model, 1.0 / num_tokens)
+                        if self._clip_grad_norm:
                             grad_norm = torch.nn.utils.clip_grad_norm_(
-                                self._model.parameters(),
-                                max_norm=float(self._clip_grad_norm),
+                                self._model.parameters(), float(self._clip_grad_norm)
                             )
-                        self._optimizer.step()
-                        self._optimizer.zero_grad(set_to_none=True)
 
-                    # Need to fix `lr_scheduler.step()` before `optimizer.step()` warning
-                    if self._lr_scheduler is not None:
-                        self._lr_scheduler.step()
-                    self.global_step += 1
+                    # This will be a no-op for optim in bwd, but prevents a warning w/ LR Scheduler
+                    self.optimizer.step()
+                    self.optimizer.zero_grad(set_to_none=True)
 
-                    loss_to_log = running_loss.detach().item() / num_tokens
+                    if self.lr_scheduler is not None:
+                        self.lr_scheduler.step()
+
+                    self.global_step += 1
+                    loss_value = (
+                        running_loss
+                        / (num_tokens if not self.optimizer_in_bwd else 1.0)
+                    ).item()
                     pbar.update(1)
                     pbar.set_description(
-                        f"{curr_epoch + 1}|{self.global_step}|Loss: {loss_to_log}"
+                        f"{curr_epoch + 1}|{self.global_step}|Loss: {loss_value:.4f}"
                     )
 
                     if self.global_step % self._log_every_n_steps == 0:
                         time_per_step = time.perf_counter() - t0
                         log_dict = {
-                            "loss": loss_to_log,
-                            # For optim in backward, this assumes all optimizers have the same LR. This is currently
-                            # true since we don't expose the ability to configure this yet.
-                            "lr": get_lr(self._optimizer),
-                            "tokens_per_second_per_gpu": num_tokens / time_per_step,
+                            "loss": loss_value,
+                            "lr": get_lr(self.optimizer),
+                            "tokens_per_second_per_gpu": (num_tokens / time_per_step),
                         }
                         if self._device.type != "cpu" and self._log_peak_memory_stats:
-                            log_dict.update(
-                                training.get_memory_stats(device=self._device)
-                            )
-                        if self._clip_grad_norm is not None:
-                            log_dict.update({"grad_norm": grad_norm})
-                        self._metric_logger.log_dict(
-                            log_dict,
-                            step=self.global_step,
-                        )
-
-                    # Reset running stats for the next step
-                    running_loss = 0
-                    num_tokens = 0
+                            log_dict.update(training.get_memory_stats(self._device))
+                        if grad_norm is not None:
+                            log_dict["grad_norm"] = grad_norm
+                        self._metric_logger.log_dict(log_dict, step=self.global_step)
+
+                    running_loss, num_tokens = 0.0, 0
                     t0 = time.perf_counter()
 
-                # Stop tracking CUDA memory now that active steps are complete
+                # Optionally stop memory profiling
                 if (
                     curr_epoch == 0
                     and self.profiler_profile_memory
@@ -689,12 +641,8 @@ def train(self) -> None:
                 ):
                     torch.cuda.memory._record_memory_history(enabled=None)
 
-                # Step the profiler
-                # Note we are stepping each batch, which might not include optimizer step in the trace
-                # if the schedule cycle doesn't align with gradient accumulation.
                 self._profiler.step()
 
-                # Check if we should stop training for this epoch
                 if (
                     (idx + 1) // self._gradient_accumulation_steps
                 ) == self.max_steps_per_epoch:
diff --git a/tests/torchtune/modules/test_optim.py b/tests/torchtune/modules/test_optim.py
new file mode 100644
index 0000000000..64b3608dea
--- /dev/null
+++ b/tests/torchtune/modules/test_optim.py
@@ -0,0 +1,143 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+# test_optim.py
+import pytest
+import torch
+
+import torch.nn as nn
+from torch.optim import AdamW
+from torch.optim.lr_scheduler import StepLR
+from torchtune.modules.optim import OptimizerInBackward
+
+
+class TestOptimizerInBackward:
+    @pytest.fixture
+    def dummy_model(self):
+        return nn.Sequential(
+            nn.Linear(10, 10),
+            nn.ReLU(),
+            nn.Linear(10, 10),
+            nn.ReLU(),
+        )
+
+    @pytest.fixture
+    def dummy_input_and_target(self) -> tuple[torch.Tensor, torch.Tensor]:
+        return torch.randn(10, 10), torch.randn(10, 10)
+
+    def test_adamw(self, dummy_model, dummy_input_and_target):
+        """Basic test with AdamW optimizer just to confirm it runs"""
+        optimizer = OptimizerInBackward(
+            dummy_model.parameters(),
+            AdamW,
+            lr=0.1,
+        )
+        initial_optimizer_sd = optimizer.state_dict()
+        optimizer.zero_grad()
+        output = dummy_model(dummy_input_and_target[0])
+        loss = nn.functional.mse_loss(output, dummy_input_and_target[1])
+        loss.backward()
+        post_loss_optimizer_sd = optimizer.state_dict()
+
+        # Check that the optimizer state has been updated
+        assert initial_optimizer_sd != post_loss_optimizer_sd
+
+        # Check that optimizer.step() doesn't update anything
+        optimizer.step()
+        assert optimizer.state_dict() == post_loss_optimizer_sd
+
+    def test_adamw_with_scheduler(self, dummy_model, dummy_input_and_target):
+        """Test that it works when paired with a learning rate scheduler."""
+        initial_lr = 0.1
+        optimizer = OptimizerInBackward(
+            dummy_model.parameters(),
+            AdamW,
+            lr=initial_lr,
+        )
+        # Assert that initial_lr has NOT been set pre-LR scheduler init
+        for param_group in optimizer.param_groups:
+            assert "initial_lr" not in param_group
+
+        scheduler = StepLR(optimizer, step_size=1, gamma=0.1)
+        # Assert that initial_lr has ben assigned to optimizer param groups
+        for param_group in optimizer.param_groups:
+            assert param_group["initial_lr"] == initial_lr
+
+        optimizer.zero_grad()
+        output = dummy_model(dummy_input_and_target[0])
+        loss = nn.functional.mse_loss(output, dummy_input_and_target[1])
+        loss.backward()
+
+        optimizer.step()  # Do this so lr scheduler doesn't complain
+        scheduler.step()
+        post_step_lr = scheduler.get_last_lr()[0]
+        assert post_step_lr == initial_lr * 0.1  # lr should be multiplied by gamma
+
+    def test_throws_error_when_step_called_with_closure(self, dummy_model):
+        """We don't implement with closure yet."""
+        optimizer = OptimizerInBackward(
+            dummy_model.parameters(),
+            AdamW,
+            lr=0.1,
+        )
+        with pytest.raises(RuntimeError):
+            optimizer.step(lambda: None)
+
+    def test_state_dict_save_load(self, dummy_model, dummy_input_and_target):
+        optimizer = OptimizerInBackward(
+            dummy_model.parameters(),
+            AdamW,
+            lr=0.1,
+        )
+        output = dummy_model(dummy_input_and_target[0])
+        loss = nn.functional.mse_loss(output, dummy_input_and_target[1])
+        loss.backward()
+        state_dict = optimizer.state_dict()
+        # Initialize a new OptimizerInBackward but with the same parameters
+        new_optimizer = OptimizerInBackward(
+            dummy_model.parameters(),
+            AdamW,
+            lr=0.1,
+        )
+        # It shouldn't match the old optimizer state dict to start
+        assert new_optimizer.state_dict() != state_dict
+        new_optimizer.load_state_dict(state_dict)
+        # But after loading, it should match
+        assert new_optimizer.state_dict() == state_dict
+
+    def test_state_dict_save_load_with_scheduler(
+        self, dummy_model, dummy_input_and_target
+    ):
+        optimizer = OptimizerInBackward(
+            dummy_model.parameters(),
+            AdamW,
+            lr=0.1,
+        )
+        scheduler = StepLR(optimizer, step_size=1, gamma=0.1)
+        output = dummy_model(dummy_input_and_target[0])
+        loss = nn.functional.mse_loss(output, dummy_input_and_target[1])
+        loss.backward()
+        optimizer.step()  # Do this so lr scheduler doesn't complain
+        scheduler.step()
+
+        for param_group in optimizer.param_groups:
+            assert param_group["initial_lr"] == 0.1
+
+        opt_state_dict = optimizer.state_dict()
+
+        # Initialize a new OptimizerInBackward but with the same parameters
+        new_optimizer = OptimizerInBackward(
+            dummy_model.parameters(),
+            AdamW,
+            lr=0.1,
+        )
+        new_optimizer.load_state_dict(opt_state_dict)
+        # Create a new LR scheduler with the new optimizer
+        new_scheduler = StepLR(new_optimizer, step_size=1, gamma=0.1, last_epoch=0)
+
+        # Now check that the last learning rates match
+        scheduler.step()  # <- THIS IS BAD, BUT NEEDED UNTIL https://github.com/pytorch/pytorch/pull/149312 LANDS
+        assert scheduler.get_last_lr() == new_scheduler.get_last_lr()
diff --git a/torchtune/modules/optim.py b/torchtune/modules/optim.py
new file mode 100644
index 0000000000..4e7d53d45c
--- /dev/null
+++ b/torchtune/modules/optim.py
@@ -0,0 +1,84 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+from typing import Iterable
+
+import torch
+from torch.optim import Optimizer
+
+__all__ = ["OptimizerInBackward"]
+
+
+class OptimizerInBackward(Optimizer):
+    """Wraps a standard PyTorch optimizer class to perform parameter updates inside the backward pass.
+
+    For each parameter, a separate optimizer instance is created, allowing in-place updates
+    during gradient computation. This reduces peak memory usage by freeing gradients immediately
+    after they are applied.
+
+    Args:
+        params (Iterable[torch.nn.Parameter]): Model parameters to optimize.
+        optimizer_cls (type[Optimizer]): The optimizer class to use in backward pass (e.g., AdamW).
+        **optimizer_kwargs: Additional arguments passed to the optimizer constructor.
+
+    Example:
+        >>> model = MyModel()
+        >>> optimizer = OptimizerInBackward(model.parameters(), torch.optim.AdamW, lr=1e-3)
+        >>> for input, target in dataloader:
+        >>>     optimizer.zero_grad()
+        >>>     output = model(input)
+        >>>     loss = loss_fn(output, target)
+        >>>     loss.backward()  # optimizer updates happen inside backward
+    """
+
+    def __init__(
+        self,
+        params: Iterable[torch.nn.Parameter],
+        optimizer_cls: type[Optimizer],
+        **optimizer_kwargs,
+    ):
+        params = list(params)  # Make a copy of params
+        self._optimizers = {}
+        self._param_to_index = {}
+
+        for idx, p in enumerate(params):
+            self._param_to_index[p] = idx
+            if p.requires_grad:
+                self._optimizers[idx] = optimizer_cls([p], **optimizer_kwargs)
+                p.register_post_accumulate_grad_hook(
+                    lambda param=p: self._step_and_clear(param)
+                )
+
+        super().__init__(params, optimizer_kwargs)
+
+    def _step_and_clear(self, param):
+        idx = self._param_to_index[param]
+        self._optimizers[idx].step()
+        self._optimizers[idx].zero_grad()
+
+    def step(self, closure=None):
+        if closure is not None:
+            raise RuntimeError("OptimizerInBackward does not support closures.")
+
+    def zero_grad(self, set_to_none=True):
+        for opt in self._optimizers.values():
+            opt.zero_grad(set_to_none=set_to_none)
+
+    def state_dict(self):
+        return {
+            "param_groups": super().state_dict()["param_groups"],
+            "optimizers": {
+                str(idx): opt.state_dict() for idx, opt in self._optimizers.items()
+            },
+            "state": {},  # State is handled by individual optimizers
+        }
+
+    def load_state_dict(self, state_dict):
+        super().load_state_dict(
+            {"param_groups": state_dict["param_groups"], "state": {}}
+        )
+        for idx, opt in self._optimizers.items():
+            opt.load_state_dict(state_dict["optimizers"][str(idx)])
diff --git a/torchtune/training/checkpointing/_checkpoint_client.py b/torchtune/training/checkpointing/_checkpoint_client.py
index f5f90f4170..64957eec85 100644
--- a/torchtune/training/checkpointing/_checkpoint_client.py
+++ b/torchtune/training/checkpointing/_checkpoint_client.py
@@ -19,6 +19,7 @@
     StateDictOptions,
 )
 from torchtune import config, training, utils
+from torchtune.modules.optim import OptimizerInBackward
 from torchtune.modules.peft import (
     get_adapter_state_dict,
     get_merged_lora_ckpt,
@@ -221,7 +222,9 @@ def _save_checkpoint_sync(
         """
         intermediate_checkpoint = epoch + 1 < training_progress.total_epochs
         checkpointer = self._get_checkpointer()
-        no_dist = not isinstance(checkpointer, DistributedCheckpointer)
+        is_not_distributed_checkpointer = not isinstance(
+            checkpointer, DistributedCheckpointer
+        )
 
         # final dict passed onto the checkpointer
         checkpoint_dict = {}
@@ -235,7 +238,7 @@ def _save_checkpoint_sync(
         model_state_dict = {}
         optim_state_dict = {}
 
-        if no_dist and not single_device:
+        if is_not_distributed_checkpointer and not single_device:
             # To prevent GPU memory from spiking during checkpoint save,
             # we consolidate the full model and optim state dicts on CPU for rank 0
             model_state_dict = training.gather_cpu_state_dict(
@@ -248,7 +251,7 @@ def _save_checkpoint_sync(
                 log.info(
                     f"Getting full model state dict took {time.perf_counter() - cp_start:.2f} secs"
                 )
-        elif no_dist:
+        elif is_not_distributed_checkpointer:
             model_state_dict = {k: v.cpu() for k, v in model.state_dict().items()}
         else:
             model_state_dict = model.state_dict()
@@ -258,21 +261,21 @@ def _save_checkpoint_sync(
                 log.info("Getting optimizer state dict...")
                 optim_start = time.perf_counter()
 
-            if no_dist:
-                if not self._optimizer_in_bwd:
-                    optim_state_dict = training.get_full_optimizer_state_dict(
-                        model,
-                        optimizer,
-                        self._is_rank_zero,
-                        device=self._device,
-                    )
-                else:
+            if is_not_distributed_checkpointer:
+                # This check can be removed once we fully migrate over to ``OptimizerInBackward``
+                if isinstance(optimizer, OptimizerInBackwardWrapper):
                     for param, opt in optimizer.optim_map.items():
                         optim_state_dict[
                             param
                         ] = training.get_full_optimizer_state_dict(
                             model, opt, self._is_rank_zero, device=self._device
                         )
+                elif isinstance(optimizer, OptimizerInBackward):
+                    optim_state_dict = optimizer.state_dict()
+                else:
+                    optim_state_dict = training.get_full_optimizer_state_dict(
+                        model, optimizer, self._is_rank_zero, device=self._device
+                    )
             else:
                 optim_state_dict = optimizer.state_dict()
 
@@ -322,7 +325,7 @@ def _save_checkpoint_helper():
 
         # Now that we have the model and optim state dict, create the actual checkpoint dict
         # to be sent to the checkpointer and ultimately written to file
-        if no_dist and not single_device:
+        if is_not_distributed_checkpointer and not single_device:
             if self._is_rank_zero:
                 _save_checkpoint_helper()
 


Reviews:
user: login: salmanmohammadi
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-15T23:58:10Z
user: login: joecummings
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-05-16T13:23:17Z
user: login: salmanmohammadi
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-19T09:31:00Z
user: login: salmanmohammadi
body: Really clean work
state: APPROVED
author_association: CONTRIBUTOR
submitted_at: 2025-05-19T09:31:18Z
user: login: joecummings
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-05-19T15:15:37Z

Review comments:
diff_hunk: @@ -218,10 +218,25 @@ To understand how this works, we encourage you to read through the relevant PyTo
 
 *Sounds great! How do I use it?*
 
-.. todo ref full finetune recipe doc
+In torchtune, you can enable this feature using :class:`OptimizerInBackward`, which wraps around a normal PyTorch optimizer like so and can be used
+as a drop-in replacement for your normal Optimizer:
 
-In torchtune, you can enable this feature using the ``optimizer_in_bwd`` flag. This feature works best when using a stateful optimizer
-with a model with a lot of parameters, and when you don't need to use :ref:`gradient accumulation <glossary_grad_accm>`.
+.. code-block:: python
+
+  OptimizerInBackward(
path: docs/source/tutorials/memory_optimizations.rst
user: login: salmanmohammadi
body: I know this was already discussed but looking at it here it's confusing that we have a class that's meant to be a drop in optimizer, but if you actually want to use it in the recipe we aren't using torchtune's regular instantiation-based configuation for this component, but rely on logic inside the recipe itself to construct the optimizer. 

I don't have a neat solution that doesn't involve breaking all the configs unfortunately. Maybe we don't advertise this as a separate modular component but as a core recipe feature? Power users can bear the brunt of implementing it in their own recipes
created_at: 2025-05-15T23:58:10Z
author_association: CONTRIBUTOR
diff_hunk: @@ -218,10 +218,25 @@ To understand how this works, we encourage you to read through the relevant PyTo
 
 *Sounds great! How do I use it?*
 
-.. todo ref full finetune recipe doc
+In torchtune, you can enable this feature using :class:`OptimizerInBackward`, which wraps around a normal PyTorch optimizer like so and can be used
+as a drop-in replacement for your normal Optimizer:
 
-In torchtune, you can enable this feature using the ``optimizer_in_bwd`` flag. This feature works best when using a stateful optimizer
-with a model with a lot of parameters, and when you don't need to use :ref:`gradient accumulation <glossary_grad_accm>`.
+.. code-block:: python
+
+  OptimizerInBackward(
path: docs/source/tutorials/memory_optimizations.rst
user: login: joecummings
body: Fair， I'll work on the wording here.
created_at: 2025-05-16T13:23:17Z
author_association: MEMBER
diff_hunk: @@ -258,21 +261,21 @@ def _save_checkpoint_sync(
                 log.info("Getting optimizer state dict...")
                 optim_start = time.perf_counter()
 
-            if no_dist:
-                if not self._optimizer_in_bwd:
-                    optim_state_dict = training.get_full_optimizer_state_dict(
-                        model,
-                        optimizer,
-                        self._is_rank_zero,
-                        device=self._device,
-                    )
-                else:
+            if is_not_distributed_checkpointer:
+                # This check can be removed once we fully migrate over to ``OptimizerInBackward``
path: torchtune/training/checkpointing/_checkpoint_client.py
user: login: salmanmohammadi
body: Worth raising an issue to follow up on this?
created_at: 2025-05-19T09:30:59Z
author_association: CONTRIBUTOR
diff_hunk: @@ -258,21 +261,21 @@ def _save_checkpoint_sync(
                 log.info("Getting optimizer state dict...")
                 optim_start = time.perf_counter()
 
-            if no_dist:
-                if not self._optimizer_in_bwd:
-                    optim_state_dict = training.get_full_optimizer_state_dict(
-                        model,
-                        optimizer,
-                        self._is_rank_zero,
-                        device=self._device,
-                    )
-                else:
+            if is_not_distributed_checkpointer:
+                # This check can be removed once we fully migrate over to ``OptimizerInBackward``
path: torchtune/training/checkpointing/_checkpoint_client.py
user: login: joecummings
body: Yeah fair, I might just opt to raise a meta-Issue tracking how we integrate this into the rest of our recipes and any cleanup that needs to happen to deprecate the old material.
created_at: 2025-05-19T15:15:37Z
author_association: MEMBER

PR comments:
user: login: pytorch-bot[bot]
created_at: 2025-05-12T01:24:56Z
body: <!-- drci-comment-start -->

## :link: Helpful Links
### :test_tube: See artifacts and rendered test results at [hud.pytorch.org/pr/pytorch/torchtune/2719](https://hud.pytorch.org/pr/pytorch/torchtune/2719)
* :page_facing_up: Preview [Python docs built from this PR](https://docs-preview.pytorch.org/pytorch/torchtune/2719/index.html)

Note: Links to docs will display an error until the docs builds have been completed.


## :white_check_mark: No Failures
As of commit ba8a97523b8fc171e51fc0f74591e644bb9fc84c with merge base e5ee1b2fcd25a411a4d0889849c1528189d56616 (<sub><sub><img alt="image" width=70 src="https://img.shields.io/date/1747409477?label=&color=FFFFFF&style=flat-square"></sub></sub>):
:green_heart: Looks good so far! There are no failures yet. :green_heart:


This comment was automatically generated by Dr. CI and updates every 15 minutes.
<!-- drci-comment-end -->
author_association: NONE
user: login: codecov-commenter
created_at: 2025-05-15T00:43:58Z
body: ## [Codecov](https://app.codecov.io/gh/pytorch/torchtune/pull/2719?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pytorch) Report
Attention: Patch coverage is `74.10072%` with `36 lines` in your changes missing coverage. Please review.
> Project coverage is 60.12%. Comparing base [(`c8e670b`)](https://app.codecov.io/gh/pytorch/torchtune/commit/c8e670b960092ffb2401845276c33ad50706378b?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pytorch) to head [(`5674f9a`)](https://app.codecov.io/gh/pytorch/torchtune/commit/5674f9ae0a7412601b7b4377802e860a5d8397f8?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pytorch).
> Report is 4 commits behind head on main.

| [Files with missing lines](https://app.codecov.io/gh/pytorch/torchtune/pull/2719?dropdown=coverage&src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pytorch) | Patch % | Lines |
|---|---|---|
| [recipes/full\_finetune\_single\_device.py](https://app.codecov.io/gh/pytorch/torchtune/pull/2719?src=pr&el=tree&filepath=recipes%2Ffull_finetune_single_device.py&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pytorch#diff-cmVjaXBlcy9mdWxsX2ZpbmV0dW5lX3NpbmdsZV9kZXZpY2UucHk=) | 0.00% | [35 Missing :warning: ](https://app.codecov.io/gh/pytorch/torchtune/pull/2719?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pytorch) |
| [torchtune/modules/optim.py](https://app.codecov.io/gh/pytorch/torchtune/pull/2719?src=pr&el=tree&filepath=torchtune%2Fmodules%2Foptim.py&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pytorch#diff-dG9yY2h0dW5lL21vZHVsZXMvb3B0aW0ucHk=) | 97.05% | [1 Missing :warning: ](https://app.codecov.io/gh/pytorch/torchtune/pull/2719?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pytorch) |

<details><summary>Additional details and impacted files</summary>


```diff
@@            Coverage Diff             @@
##             main    #2719      +/-   ##
==========================================
- Coverage   60.64%   60.12%   -0.52%     
==========================================
  Files         428      432       +4     
  Lines       26091    26520     +429     
==========================================
+ Hits        15823    15946     +123     
- Misses      10268    10574     +306     
```

</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/pytorch/torchtune/pull/2719?dropdown=coverage&src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pytorch).   
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pytorch).

<details><summary> :rocket: New features to boost your workflow: </summary>

- :snowflake: [Test Analytics](https://docs.codecov.com/docs/test-analytics): Detect flaky tests, report on failures, and find test suite problems.
</details>
author_association: NONE
user: login: krammnic
created_at: 2025-05-19T18:08:10Z
body: Amazing work! (I've attempted this, but my RFC was rejected, so I'm really glad to see this merged.)
author_association: COLLABORATOR