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

2023-07-26T17:44:00

众所周知,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

当前页面是本站的「Baidu MIP」版。发表评论请点击:完整版 »