empty Cache after logps_per_token

[shirinyamani]

Shouldn't we torch.cuda.empty_cache() after computing the logprobs in [here] (

trl/trl/trainer/grpo_trainer.py

Line 430 in 801582e

def get_per_token_logps(model, input_ids, num_logits_to_keep):

) also after loss computation?

[qgallouedec]

I usually refer to https://discuss.pytorch.org/t/about-torch-cuda-empty-cache/34232/2 when I consider emptying the cache. In my understanding, when you leave the function, the internal variables no longer exist so the tensor stop being reference and the underlying memory is freed.

But I might be wrong.

[shirinyamani]

right, but we wipe the cache in same scenario in ppo ,

I am not sure either!

[qgallouedec]

It's probably worth doing some tests/profiling with a toy example to see the diff