I am calling backward on computed reward which is calculated in the following fashion:
For each training sample in the batch, I will have to first decode n complete sequences (n = beam_size), evaluate them based on a metric to calculate reward(or loss) for back-propagation.

For example, for each sample in the batch, decode n sequences so that output of decoder is of size: (batch_size, n, max_seq_len). The metric can evaluate sequences only after they are complete (i.e. only after end_token is sampled). Reward is computed for a total of batch_size*n sequences.


If using batch_size = 8, and beam size = 32. I am sampling 8*32 sequences from the decoder, calculating their (1) log probability and (2) reward, using Reinforce Trick to calculate loss which is reward*log_probability before calling backward. Here is the loss -


The loss is the expectation of reward for decoded sequences, r(.) is the reward for complete sequences, y_t are hidden states sample from the decoder and pi is decoder model.
In this training regime, I have no GT to match predictions with at each time-step. I can only back-propagate after I have sampled complete sequence. Thus saving predictions for all sequences for all time-steps is saving their computation graphs which I think is incrementally increasing GPU memory consumption. Right?
Solutions tried: I have tried gradient accumulation where I am splitting batch_size into sub-batches, calling loss.backward for each of them before calling optimiser.step().
Question: Is there anything that can be done to minimise such huge memory consumption. Let me know if you need specific code excerpt to debug.
Paper for reference https://openreview.net/pdf?id=H1Xw62kRZ