gradient_accumulation_steps in trainer_tf
    part-1:
    self.total_train_batch_size = self.args.train_batch_size * self.args.gradient_accumulation_steps
    ds = (
            self.train_dataset.repeat()
            .shuffle(self.num_train_examples, seed=self.args.seed)
            .batch(self.total_train_batch_size, drop_remainder=self.args.dataloader_drop_last)
            .prefetch(tf.data.experimental.AUTOTUNE)
        )
     part-2:
     for _ in tf.range(self.args.gradient_accumulation_steps):
                reduced_features = {
                    k: ft[: self.args.train_batch_size // self.args.n_replicas] for k, ft in features.items()
                }
                reduced_labels = labels[: self.args.train_batch_size // self.args.n_replicas]

                self.training_step(reduced_features, reduced_labels)

                features = {
                    k: tf.concat(
                        [ft[self.args.train_batch_size // self.args.n_replicas :], reduced_features[k]],
                        axis=0,
                    )
                    for k, ft in features.items()
                }

                labels = tf.concat(
                    [labels[self.args.train_batch_size // self.args.n_replicas :], reduced_labels], axis=0
                )
 
the implementation of gradient_accumulation seems not friendly to users who have small gpu memory.

What costs the most is the gradient computation, storing few predictions is ok general. I can run a sequence classification training with a batch of 32 of 128 sequence length and an acummulation of 3 with a 8GB GPU.

Did you encounter during your experiments a memory issue? If yes, let me know and I will look at it.<|||||>This issue has been automatically marked as stale because it has not had recent activity. It will be closed if no further activity occurs. Thank you for your contributions.
