How could I implement gradient accumulation?

Daraan · June 17, 2025, 1:04pm

1. Severity of the issue: (select one)
High: Completely blocks me.

For my own research as well as recreation I would like to work with gradient accumulation. From my research it does not appear that ray/rllib supports this in the respective framework learners.

Among others I have compute_gradients and prevent optim.zero_grad for my accumulation steps and likely return an empty dict in these steps. However I am not sure if that is sufficient and what side-effects might appear.

I’ll update this post with my findings - however, I would appreciate any further input that could help me reach a implementation supporting gradient accumulation and avoid pitfalls with RLlib’s framework.

github.com/ray-project/ray

rllib/core/learner/torch/torch_learner.py

eed7e024d


      
          def compute_gradients(
              self, loss_per_module: Dict[ModuleID, TensorType], **kwargs
          ) -> ParamDict:
              for optim in self._optimizer_parameters:
                  # `set_to_none=True` is a faster way to zero out the gradients.
                  optim.zero_grad(set_to_none=True)

Daraan · September 3, 2025, 9:31am

I was able to write and test a small Learner in my project: GitHub - Daraan/ray_utilities: ray & RLlib tools for unified code across different repositories. Experiments with dynamic hyperparameters that updates the gradients after multiple steps and accumulates it in between.

github.com/Daraan/ray_utilities

ray_utilities/learners/ppo_torch_learner_with_gradient_accumulation.py

main

"""NOTE: Experimental, edge cases not tested yet."""

from __future__ import annotations

import logging
from typing import TYPE_CHECKING, Any, Collection

from ray.rllib.algorithms.ppo.torch.ppo_torch_learner import PPOTorchLearner
from ray.rllib.utils.annotations import override
from ray.rllib.utils.typing import ParamDict

if TYPE_CHECKING:
    import torch
    from ray.rllib.utils.typing import (
        ModuleID,
        ParamDict,
        ParamRef,
        TensorType,
    )

This file has been truncated. show original

The file and class is a standalone and usable in any rllib+torch algorithm.

Lars_Simon_Zehnder · September 3, 2025, 9:58am

Hi @Daraan , thanks for the question and for sharing your code here. You made everything as intended: create your own Learner and then override the compute_gradients - great work!

I just flew over it and something was suspicious to me: it looks as if the zero_grad is called in the last accumulation step at which updating is happening. As a result only the last gradient is used while the history of accumulated gradients is removed.

Let me know, if this is the case or not.

Daraan · September 3, 2025, 4:28pm

Thank you so much for your review. You are totally correct here. Tricky to test it correctly.
I’ve updated it.

Thank you!

Lars_Simon_Zehnder · September 3, 2025, 4:48pm

Great, I could help! Very nice example!

Keep up the good work!

Topic		Replies	Views
Independent gradient update for each loss RLlib	2	318	March 13, 2021
How does rllib parallelise gradient computation and updating? Configure Algorithm, Training, Evaluation, Scaling	0	277	September 27, 2023
Creating saliency maps / activation maximization with trained policy RLlib	3	364	October 6, 2022
[RLlib] Pytorch multiple optimizers support RLlib	1	592	January 4, 2023
Very slow gradient descent on remote workers RLlib	14	2462	June 8, 2021

How could I implement gradient accumulation?

Related topics