Memory grows due to keeping losses on device

[carmocca]

If logging is disabled (or very infrequent), the memory usage slowly grows because the max and average loss is kept in a list on-device: https://github.com/pytorch/torchtitan/blob/main/train.py#L353-L354

The training loop should offload these tensors to the CPU right after their aggregation is finished. Especially because the logging prints will do that anyways under the hood

[tianyu-l]

Thanks for raising this issue, @carmocca !

It looks to me that what's kept on device is https://github.com/pytorch/torchtitan/blob/main/train.py#L330 (the max and average you mentioned are on CPU?).

I think the reason we keep it is because we don't want call .item() (which incurs synchronization between CPU and GPU) unless hitting a log step. I do agree that if logging is disabled / infrequent, this overhead is unnecessary. Although, may I ask what's the use case where you'd log too infrequently so that this overhead becomes unacceptable?

[carmocca]

One case would be debugging or testing: you don't want to dump logs to TensorBoard but then your memory will grow until OOM. The code currently expects that you eventually log

Another is small and fast models where the step time is small and it's excessive to log that often. There you might notice the ragged look in the memory metrics.

Also, I don't personally like the default behaviour of smoothing the loss over the last log_freq steps. I find it misleading when compared to log_freq=1. My expectation was that log_freq > 1 would simply log the loss at that modulo step. But that might just be me.

[tianyu-l]

My expectation was that log_freq > 1 would simply log the loss at that modulo step.

Hmm makes sense to me. Please feel free to change the behavior with a PR. Looks it'll address both issues.