Flushes the memory of the current device and, if the flush_compile flag is True, all data related to torch.compile.
(flush_compile: bool = True)
| 76 | |
| 77 | # Helper functions |
| 78 | def flush_memory(flush_compile: bool = True) -> None: |
| 79 | """Flushes the memory of the current device and, if the flush_compile flag is True, all data related to |
| 80 | torch.compile.""" |
| 81 | gc.collect() |
| 82 | # If needed, flush everything related to torch.compile |
| 83 | if flush_compile: |
| 84 | # Dynamo resets |
| 85 | torch._dynamo.reset() |
| 86 | torch._dynamo.reset_code_caches() |
| 87 | if hasattr(torch._inductor, "codecache"): |
| 88 | # Clear FX graph cache |
| 89 | if hasattr(torch._inductor.codecache, "FxGraphCache"): |
| 90 | torch._inductor.codecache.FxGraphCache.clear() |
| 91 | # Clear PyCodeCache |
| 92 | if hasattr(torch._inductor.codecache, "PyCodeCache"): |
| 93 | torch._inductor.codecache.PyCodeCache.cache_clear() |
| 94 | # Clear TritonFuture cache (for async compilation) |
| 95 | if hasattr(torch._inductor.codecache, "TritonFuture"): |
| 96 | if hasattr(torch._inductor.codecache.TritonFuture, "_compile_cache"): |
| 97 | torch._inductor.codecache.TritonFuture._compile_cache.clear() |
| 98 | # Clear CUDA cache |
| 99 | if torch.cuda.is_available(): |
| 100 | torch.cuda.empty_cache() |
| 101 | torch.cuda.synchronize() |
| 102 | elif torch.xpu.is_available(): |
| 103 | torch.xpu.empty_cache() |
| 104 | torch.xpu.synchronize() |
| 105 | gc.collect() |
| 106 | |
| 107 | |
| 108 | def get_tokenizer_and_model( |
no test coverage detected