Flush GPU memory and run garbage collection. If the flush_compile flag is set, we also clear the everything related to compile cache.
(flush_compile: bool = True)
| 80 | |
| 81 | |
| 82 | def flush_memory(flush_compile: bool = True) -> None: |
| 83 | """Flush GPU memory and run garbage collection. If the flush_compile flag is set, we also clear the everything |
| 84 | related to compile cache.""" |
| 85 | gc.collect() |
| 86 | # If needed, flush everything related to torch.compile |
| 87 | if flush_compile: |
| 88 | # Dynamo resets |
| 89 | torch._dynamo.reset() |
| 90 | torch._dynamo.reset_code_caches() |
| 91 | if hasattr(torch._inductor, "codecache"): |
| 92 | # Clear FX graph cache |
| 93 | if hasattr(torch._inductor.codecache, "FxGraphCache"): |
| 94 | torch._inductor.codecache.FxGraphCache.clear() |
| 95 | # Clear PyCodeCache |
| 96 | if hasattr(torch._inductor.codecache, "PyCodeCache"): |
| 97 | torch._inductor.codecache.PyCodeCache.cache_clear() |
| 98 | # Clear TritonFuture cache (for async compilation) |
| 99 | if hasattr(torch._inductor.codecache, "TritonFuture"): |
| 100 | if hasattr(torch._inductor.codecache.TritonFuture, "_compile_cache"): |
| 101 | torch._inductor.codecache.TritonFuture._compile_cache.clear() |
| 102 | # Clear device cache |
| 103 | if torch.cuda.is_available(): |
| 104 | torch.cuda.empty_cache() |
| 105 | torch.cuda.synchronize() |
| 106 | elif is_torch_xpu_available(): |
| 107 | torch.xpu.empty_cache() |
| 108 | torch.xpu.synchronize() |
| 109 | gc.collect() |
| 110 | |
| 111 | |
| 112 | class BenchmarkStreamer(BaseStreamer): |
no test coverage detected