目前在大模型模型训练及accelerate训练框架下时经常会使用到的cuda函数:
import gcdef reset_memroy(device):gc.collect()torch.cuda.empty_cache()torch.cuda.reset_peak_memory_stats(device)torch.cuda.reset_accumulated_memory_stats(device)
- torch.cuda.empty_cache()
torch.cuda.empty_cache() 是 PyTorch 中的一个函数,用于释放未被使用的 GPU 内存缓存。具体来说,它会清空 PyTorch 内部缓存中已经分配但目前未被使用的内存块。这个函数并不会释放正在使用中的内存或减少当前 PyTorch 程序中的内存占用,而是清理那些 PyTorch 内部缓存中暂时未使用的内存。
在训练深度学习模型时,PyTorch 会自动管理内存以避免频繁的内存分配和释放,从而提高性能。调用 torch.cuda.empty_cache() 可以帮助释放一些内存,尤其是在内存紧张的情况下。但它不会影响到 GPU 上当前活跃的内存分配,也不会对模型训练的结果产生影响。
- torch.cuda.reset_peak_memory_stats()
torch.cuda.reset_peak_memory_stats() 用于重置 GPU 内存的峰值使用统计信息。具体来说,它会清除从上次调用此函数以来的最大 GPU 内存使用量的记录。在 PyTorch 中,可以通过 torch.cuda.max_memory_allocated() 和 torch.cuda.max_memory_cached() 等函数获取 GPU 内存的最大使用情况。调用 reset_peak_memory_stats() 后,这些统计信息将被重置为当前的内存使用状态,因此可以重新开始监测内存峰值的变化。
- torch.cuda.reset_accumulated_memory_stats()
torch.cuda.reset_accumulated_memory_stats() 是 PyTorch 中用于重置 GPU 内存统计信息的函数。具体来说,它会清除从上次调用此函数以来的所有 CUDA 内存统计信息,包括 GPU 上分配的内存、已释放的内存、内存碎片等。这对于调试和性能分析非常有用,尤其是在需要重新开始内存统计的情况下。此函数并不会释放任何已分配的内存,而是仅仅清除与内存使用情况相关的累积统计数据。