[BYO-FT] Improve check for available memory

[Lunderberg]

Previously, any parameters whose allocations was performed by the relax VM would be double-counted.

[masahi]

sorry there is an incredibly embarrassing hack I did to TVM's PooledAllocator based on an assumption that get_used_memory() is only called once:
https://github.com/octoml/tvm/blob/for-mlc-serve-jan12/src/runtime/memory/pooled_allocator.h#L114-L116

This was to change the behavior of PooledAllocator before / after memory profiling, to workaround a limitation of the allocator. I have a branch that integrates a third-party CUDA allocator https://github.com/rapidsai/rmm to replace PooledAllocator, so that we don't have to worry about this issue.

[masahi]

That said, we can make get_used_memory() pure by introducing a global function to toggle the allocation mode. I should have done that but I did the shortcut of abusing get_used_memory() to signal the end of memory profiling.

Such global function is also useful even after we land the allocator based on rmm, since tracking statistics for memory allocation is costly and should only be done during memory profiling.

[Lunderberg]

Ah, got it. So, by calling it twice, I'm invalidating the safety net anyways.

I think in that case, I'd lean toward disabling this check altogether, and replacing it with a compile-time tracking of memory usage. That way, instead of profiling, we could know the maximum size of the live set.

[masahi]

I'm curious if such compile-time estimate of memory usage agrees with the "ground truth" that we collect at runtime, given that the latter is affected by the behavior of the underlying allocator which the memory planner is not aware of. And my impression of Relax's memory planning has been that it allocates much more than necessary, compared to the allocator in torch.

[Lunderberg]

Yeah. I think that increased memory would show up in the compile-time estimation, because a large part of the footprint is due to operations that aren't done in place. @slyubomirsky has been doing a lot of good work for in-place operations, which should give us a path forward on footprint.

[masahi]

any parameters whose allocations was performed by the relax VM

When do we encounter such params? I thought all params are allocated by NDArray::Empty().

[Lunderberg]

When do we encounter such params? I thought all params are allocated by NDArray::Empty().

In the current flow with pre-processed weights, that is correct. All parameters are allocated by NDArray::Empty(), as part of tvm.nd.array.

For the pre-process at initialization flow, which will be used to support bring-your-own-fine-tune, the transform_params is called at runtime, and outputs the parameters to be used. If there is any processing required, transform_params will allocate space for the pre-processed parameters. That allocation is performed by the relax VM, but is later used as a model parameter.

[sunggg]

Hi, @Lunderberg! Thanks for the contribution!
One quick question:

If there is any processing required, transform_params will allocate space for the pre-processed parameters.

Does it mean there will be a dynamic allocation? To guarantee that we wouldn't get OOM during the runtime, we've been very careful about the memory allocation assuming the worst worst scenario. I think we are reaching the point where we would need to revisit this assumption, so your input will be very appreciated :)

[Lunderberg]

Currently, all the allocations are static, but that's entirely due to the weights having static shape. While we could in theory write a transform_params function that is dynamic over the model parameters, but I don't see a need for that at the moment. What would be useful would be to have the transform_params function be dynamic over world_size, as then we could have a single .so file that can be used with different multi-GPU strategies.

I think it would be good to have a PrimFunc that takes any dynamic parameters, and returns the memory footprint for that set of dynamic parameters.