Memory debug for MPS & report method

CAIIVS / chuchichaestli

Where you find all the state-of-the-art cooking utensils (salt, pepper, gradient descent... the usual).

GNU General Public License v3.0

4 stars 0 forks source link

Memory debug for MPS & report method #22

Closed phdenzel closed 5 months ago

phdenzel commented 5 months ago

Added a few more lines in the debug module:

visitor.Visitor.report: instances can now use .report() to retrieve a memory report printout of visited modules
moved a few blocks to the base Visitor class to avoid code duplication
changes should work for the MPSMemoryAllocationVisitor class
changes untested on the CudaMemoryStatsVisitor class
unit tests now use CUDA/MPS only if available

bil-y commented 5 months ago

Thanks for the additions. The code looks good to me, will test it tomorrow.

phdenzel commented 5 months ago

Note: I had problems running multiple test functions with MPS... I had to cache the MPSMemoryAllocationVisitor instance and access it via pytest.fixture, otherwise the 2nd instance wouldn't fetch any stats. If you have an idea why this happens, let me know...

bil-y commented 5 months ago

Note: I had problems running multiple test functions with MPS... I had to cache the MPSMemoryAllocationVisitor instance and access it via pytest.fixture, otherwise the 2nd instance wouldn't fetch any stats. If you have an idea why this happens, let me know...

I think that's related to how PyTorch hooks work. I have also noticed similar issues in other projects.