Summary
Tensor materialization trades the memory with recomputation. Prior tensor materialization methods do not consider the memory fragmentation problem of the memory system used in deep learning frameworks, which makes them evict unnecesary tensors. The authors of this paper proposed a memory-system-aware rematerialization method called Coop to reduce the memory fragmentation. Experiments show that Coop can achieve up to 2x memory saving comparied with prior works.
Weaknesses
1. The requirement of underlying memory allocator might limit the applicability of the proposed method
This paper assumed the memory allocator used by the deep learning systems (discussed in section 2.1). The underlying memory allocator must be able to **merge** the freed chuncks if they are contiguous. There are other kinds of memory allocators (e.g., record the mapping from chunck size to a list of free chuncks with the specific chunk size) that do not have this feature, thus the proposed method can not be (directly) used for deep learning systems with this kind of memory allocators.
2. More discussion on the effectiveness of the the page-table-based memory system is needed
Similar to CPU memory system, the GPU memory system also employed the page table to manage its memory. Thus, we can free the discontiguous memory chunks and allocate a new one with the sum of the sizes of the freed chunks. From the virtual memory's view, the allocated memory is contiguous. Thus, there is no fragmentation problem discussed in this paper, and the prior works can be directly used. The good news is that, since CUDA 11.2, we can directly use the memory pool [1] implemented in cuda runtime/driver to enjoy this feature. Thus, I am interested in whether the prior works have the fragmentation problem if they use the memory pool implemented in cuda?
We can use the following program to validate that the page-table based GPU memory system. In the program, we allocate 3 chuncks of 7 GB called p1, p2, and p3. Then, we free p1 and p3 (they are not contiguous). At last, we allocate a chunk of memory with 14 GB called p4 successfully.
```python
#include <stdio.h>
#include <cuda.h>
#include <unistd.h>
#define CHECK(e) {auto s = e; if (s != cudaSuccess) printf("CUDA error: %s", cudaGetErrorString(s));}
#define GB(x) ((x) * 1024ull * 1024 * 1024)
int main() {
void *p1, *p2, *p3, *p4;
CHECK(cudaMalloc(&p1, GB(7)));
CHECK(cudaMalloc(&p2, GB(7)));
CHECK(cudaMalloc(&p3, GB(7)));
printf("first allocation done\n");
printf("p1=%p\n", p1);
printf("p2=%p\n", p2);
printf("p3=%p\n", p3);
// sleep 5 seconds, we can check the memory usage in `nvidia-smi` during this time
sleep(5);
printf("free p1 and p3, allocate p4\n");
CHECK(cudaFree(p1));
CHECK(cudaFree(p3));
CHECK(cudaMalloc(&p4, GB(14)));
printf("p4=%p\n", p4);
printf("successfully allocated p4\n");
CHECK(cudaFree(p2));
CHECK(cudaFree(p4));
}
```
I can run above program on a RTX 3090 (24GB memory).
```
first allocation done
p1=0x7f4a68000000
p2=0x7f48a8000000
p3=0x7f46e8000000
free p1 and p3, allocate p4
p4=0x7f4528000000
successfully allocated p4
```
[1] https://docs.nvidia.com/cuda/cuda-runtime-api/group__CUDART__MEMORY__POOLS.html
Rating
7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.