[Python] How can memory fragmentation be further reduced when torch expandable_segments is...

Recently, I've been considering writing a memory manager to reduce GPU memory fragmentation during LLM training. I’ve looked into PyTorch’s CUDACachingAllocator implementation and noticed that it also supports GPU Virtual Memory Management (VMM).

However, I haven’t yet come up with a better way to further improve it or design a specialized memory manager tailored for the LLM pre-training setting based on its current structure.I’d love to hear any suggestions or ideas you might have.

Continue reading...