[Enhancements] add a vcude device to help mitigate compile time GPU memory usage

[xinli-git]

The problem

During compilation, some passes needs to replace certain operators with one or more new operators. For example, resolve_variants, subgraph_rewrite. In doing so, hidet relies on imperative_run to generate the new operators. This creates a lot of intermediate tensors that can exceed the actual run-time memory consumption, which makes larger models (such as Llama-7B) unable to compile even for a GPU with 24GB using fp16.

Fix

We introduce a vcuda device that allows GPU tensors to be stored on CPU and only transferred to GPU on demand. We call this vcuda. With this change, any additional GPU memory usage during compilation is off-loaded to CPU.

This might incur a bit compilation overhead when enabled, but given the time-consuming nature of such large models, this compile time increase is negligible.

Now, on RTX3090, this is the GPU memory consumption for running llama test

Status of cuda:0 memory pool
   Allocated: 14081 MiB
        Peak: 14081 MiB
    Reserved: 1196 MiB
      Active: 12884 MiB
Status of cpu memory pool
   Allocated: 4 KiB
        Peak: 26022 MiB
    Reserved: 4 KiB
      Active: 0 Bytes
Status of vcuda:0 memory pool
   Allocated: 0 Bytes
        Peak: 25486 MiB
    Reserved: 0 Bytes
      Active: 0 Bytes

[yaoyaoding]

Hi @xinli-git , good work!

I left some alternative implementation suggestions and let's discuss.

[yaoyaoding]

There are two ways to implement "vcuda" device:

1. [The current version]: we copy the tensor to cuda before we run the operator and copy back to cpu when we finish it.
2. We can also reply on the unified addressing of cuda. (see here and here). Essentially, the unified addressing allows the cuda kernel access the host memory (page-locked) directly. Thus, we do not need to copy to and from the cuda device memory.

I prefer the second way. With it, we do not need to update our run_async function.

[xinli-git]

this is a great suggestion and it works! thank you

[yaoyaoding]

One minor issue is that the cpu tensor allocated by other library (e.g., pytorch and numpy) will not use cudaMallocHost to allocate, and it is not page-locked. Thus, when we convert the tensors from them, we should allocate the tensor by hidet instead of sharing the weights (on the platform that cuda is avaliable).

[xinli-git]

I think this case is fine because if users give a cpu tensor from other library and expect it to work on CUDA device that is already wrong

[xinli-git]

runtime check should catch this

[yaoyaoding]

Make sense.

[yaoyaoding]

Thanks @xinli-git !