| 813 | |
| 814 | |
| 815 | def see_memory_usage(message, force=False): |
| 816 | if not force: |
| 817 | return |
| 818 | if dist.is_initialized() and not dist.get_rank() == 0: |
| 819 | return |
| 820 | |
| 821 | # python doesn't do real-time garbage collection so do it explicitly to get the correct RAM reports |
| 822 | gc.collect() |
| 823 | |
| 824 | # Print message except when distributed but not rank 0 |
| 825 | print(message) |
| 826 | print(f"MA {round(get_accelerator().memory_allocated() / (1024 * 1024 * 1024),2 )} GB \ |
| 827 | Max_MA {round(get_accelerator().max_memory_allocated() / (1024 * 1024 * 1024),2)} GB \ |
| 828 | CA {round(torch_memory_reserved() / (1024 * 1024 * 1024),2)} GB \ |
| 829 | Max_CA {round(torch_max_memory_reserved() / (1024 * 1024 * 1024))} GB ") |
| 830 | |
| 831 | vm_stats = psutil.virtual_memory() |
| 832 | used_GB = round(((vm_stats.total - vm_stats.available) / (1024**3)), 2) |
| 833 | print(f'CPU Virtual Memory: used = {used_GB} GB, percent = {vm_stats.percent}%') |
| 834 | |
| 835 | # get the peak memory to report correct data, so reset the counter for the next call |
| 836 | get_accelerator().reset_peak_memory_stats() |
| 837 | |
| 838 | |
| 839 | def call_to_str(base, *args, **kwargs): |