解决APU平台上PyTorch的显存不足问题

July 26, 2023 默认分类

众所周知，ROCm体系问题并不算少，生态也不丰富，更不对APU平台进行官方支持。PyTorch支持了HIP，但是在APU平台上，HIP只认识BIOS中分配的VRAM，而不认识主机可用内存，但对于核显来说，两者没有实际区别，主机可用内存也应该认为是显存。这对于APU平台是极为不合理的。PyTorch提供了显存分配器覆盖机制，除了调试，还能用来修改显存来源。

以下代码将显存分配器的目标替换为主机内存（GTT），这一方式可绕开VRAM限制。

alloc.c：

#define __HIP_PLATFORM_AMD__
#include <sys/types.h>
#include <hip/hip_runtime.h>
// Compile with hipcc alloc.cc -o alloc.so -shared -fPIC
void* gtt_alloc(ssize_t size, int device, hipStream_t stream) {
   void *ptr = NULL;
   hipHostMalloc(&ptr,size,0);
   return ptr;
}

void gtt_free(void* ptr, ssize_t size, int device, hipStream_t stream) {
    hipHostFree(ptr);
}

然后，将以下代码放在使用PyTorch的代码的前面部分（在执行任何实际操作之前）：

new_alloc = torch.cuda.memory.CUDAPluggableAllocator('<alloc.so的路径>','gtt_alloc','gtt_free');
torch.cuda.memory.change_current_allocator(new_alloc)

Github：https://github.com/pomoke/torch-apu-helper