Review and potentially simplify the implementation

[IvanUkhov]

Hello,

Thank you for the package! As discussed in keras-team/tf-keras#107, I am wondering if it would be possible to make use of the following example to make the package even better:

https://blog.ivanukhov.com/2024/01/31/gradient-accumulation.html

[andreped]

As discussed in the other thread, it would be of interest to test the proposed solution as it works for the latest implementation of the Optimizer class and may work in distributed settings.

Are you able to have a look at this, @tno123?

---

EDIT: For now I suggest first testing if using the implementation above in a distributed setting works out-of-the-box using the latest TF version (TF==2.15). You can make a Jupyter Notebook and test it for free on Kaggle where you will have access to 2 GPUs.

You can likely base the notebook on this test:
https://github.com/andreped/GradientAccumulator/blob/main/tests/test_optimizer_distribute.py

[tno123]

Sure, I'll take a look 👍

[andreped]

Feel free to tag in @IvanUkhov, if you have any questions regarding the implementation or whatnot.

Might be that we should add support for his implementation to this framework, regardless whether this worked fine in a distributed setting, as it uses the newest Optimizer implementation. Likely we should also move to the new one. It seems to require less code than the old legacy Optimizer.

[tno123]

@IvanUkhov
I am getting a few errors (might be due to old test code):

            opt = CumulativeAdam()

            # compile model
            model.compile(
                optimizer=opt,
                loss=tf.keras.losses.CategoricalCrossentropy(),
                metrics=["acc"],
                run_eagerly=False,
            )

        # train model
        model.fit(
            ds_train,
            epochs=10,
            validation_data=ds_test,
            verbose=1
        )

errors:

  File "/opt/conda/lib/python3.10/site-packages/keras/src/utils/traceback_utils.py", line 65, in error_handler

  File "/opt/conda/lib/python3.10/site-packages/keras/src/engine/training.py", line 1807, in fit

3 root error(s) found.
  (0) INVALID_ARGUMENT:  You must feed a value for placeholder tensor 'cond/else/_1/cond/Placeholder_1' with dtype int32
	 [[{{node cond/Placeholder_1}}]]
	 [[cond/else/_1/cond/StatefulPartitionedCall/cond/output/_307/_218]]
	 [[cond/then/_0/cond/StatefulPartitionedCall/div_no_nan_1/ReadVariableOp_3/_272]]
  (1) INVALID_ARGUMENT:  You must feed a value for placeholder tensor 'cond/else/_1/cond/Placeholder_1' with dtype int32
	 [[{{node cond/Placeholder_1}}]]
	 [[cond/else/_1/cond/StatefulPartitionedCall/cond/output/_307/_218]]
  (2) INVALID_ARGUMENT:  You must feed a value for placeholder tensor 'cond/else/_1/cond/Placeholder_1' with dtype int32
	 [[{{node cond/Placeholder_1}}]]
0 successful operations.
0 derived errors ignored. [Op:__inference_fn_with_cond_42426]

Any ideas what might be causing them? Alternatively, do you have some complete training code that you can share?

Thanks in advance!

[IvanUkhov]

Sure! My use case is distributed training, and I have been following closely this tutorial under Training loop:

https://www.tensorflow.org/tutorials/distribute/custom_training#training_loop

I have not tried with Model.fit, but I will take a look at what happens there.

[andreped]

@tno123 Could you share the notebook/gist/script that reproduces the issue? That way, @IvanUkhov can use it for debugging as well :]

[IvanUkhov]

Just quickly opened this tutorial in Colab and replaced the optimizer:

https://www.tensorflow.org/tutorials/quickstart/beginner

There is no error. Not sure what the difference is with respect to that test setup you mentioned.

[andreped]

Just quickly opened this tutorial in Colab

@IvanUkhov I believe @tno123 tested this in Kaggle, where he has access to two GPUs. Hence, he can test whether using this implementation with multiple GPUs would work out-of-the-box.

I made a very simple Jupyter Notebook, which reproduces the issue. I have made it publicly available at:
https://www.kaggle.com/code/andreped/cumulativeadam-distributed-test

The notebook is setup to use two GPUs already, so you should be able to reproduce the issue if you were to run it yourself.

Feel free to use the notebook as a base to do your own experiments, and feel free to share an updated notebook, when you have gotten something working, or found some interesting observations :] We can use this thread to share experiences.

If you only select one GPU, this script runs fine. The problem is when you want to use multiple GPUs, as highlighted in my notebook. Hence, in CoLab this script should run fine, as you only have access to a single GPU.

[IvanUkhov]

Thank you for the pointers! I have managed to reproduce the problem with the following tutorial on my machine with TensorFlow 2.15 and two GPUs:

https://www.tensorflow.org/guide/distributed_training

Please give a try to the updated version:

https://blog.ivanukhov.com/2024/01/31/gradient-accumulation.html

It works now. Here is the main change:

I missed this section:

https://github.com/keras-team/keras/blob/v2.15.0/keras/optimizers/optimizer.py#L635-L638

[andreped]

It works now. Here is the main change:

@IvanUkhov Sounds great! Could you run some tests using it, @tno123? :]

[tno123]

Nice! I'll test it now 👍

Edit: It seems to work! It is training, but I'll do some more tests just to make sure everything is good.