解决APU平台上PyTorch的显存不足问题
众所周知,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)