Clarification for gradient memory in mixed precision training about ml-engineering

[stas00]

There are 2 ways that I know of to measure peak GPU memory:

• using pynvml to take frequent snapshots (imprecise, but good enough) - that's what ipyexperiments does
• https://pytorch.org/docs/stable/generated/torch.cuda.max_memory_allocated.html - (but there is only one counter, so if more than application uses it and resets the counter - the results are unpredictable.)

[stas00]

There are a few places where the Deepspeed implementation of ZeRO upcasts grads to fp32:

1. it does it before stepping as it tries not to meddle with the optim API, so any optimizer can be used. It'd have to meddle with the optim API to perform the upcasting on the fly.
2. If you use communication_data_type=fp32, to avoid the lossy reduction of grads, it'll also upcast to fp32, though only doing it to the amount of params fitting into the current reduction bucket (so typically far less than all params).

So there are at least those 2 cases where 4 bytes are used. This might change in the future.

Now those upcasts can really be considered temp memory usage as they use peak memory, but practically this is 4 bytes per param.

I don't know how Apex does it.

How do you measure memory usage, are you taking into the account the peak memory? Which would have led to OOM if there weren't enough spare.

You can try this package in a jupyter notebook https://github.com/stas00/ipyexperiments if you want automated reports of peak and normal memory usages. The package is very trivial to setup and you will just need to load your code into a jupyter notebook, which sometimes can be more difficult to setup. But you could use even something as simple as https://github.com/Syllo/nvtop to get the peak memory usage.

[SumanthRH]

Got it! I understand now. Part of the issue was that I did not understand torch.autocast properly. After inspecting training step with mixed precision, it's clear that the full FP32 gradients are stored in HBM and this matches the memory consumption seen as well.

If you use communication_data_type=fp32, to avoid the lossy reduction of grads, it'll also upcast to fp32, though only doing it to the amount of params fitting into the current reduction bucket (so typically far less than all params).

Ah interesting. I believe I'll probably need to dig into the code to actually understand this fully :)

How do you measure memory usage, are you taking into the account the peak memory? Which would have led to OOM if there weren't enough spare.

Hmm I was using pynvml to measure current memory usage. You're right that peak memory is not accounted for here. ipyexperiments looks great! I'll check it out!

Thank you so much for the detailed reply!