rapidsai / crossfit

Metric calculation library
Apache License 2.0
2 stars 5 forks source link

Ensure Memory Estimation works with RMM backed memory manager #74

Open VibhuJawa opened 1 month ago

VibhuJawa commented 1 month ago

Currently our memory estimation works only with torch's default memory allocator, we should make sure it works with RMM allocator plugged in too.

Like ensure below works for RMM backed Pytorch allocator too. https://github.com/rapidsai/crossfit/blob/6c213c96b3366298c5681c8d6f5bce72c2dda34f/crossfit/backend/torch/hf/memory_curve_utils.py#L27-L89

VibhuJawa commented 1 month ago

Raised Pytorch issue here: https://github.com/pytorch/pytorch/issues/133280

VibhuJawa commented 1 month ago

Another related issue is: https://github.com/pytorch/pytorch/issues/133281