Simple GPU memory report.
(name)
| 90 | |
| 91 | |
| 92 | def report_memory(name): |
| 93 | """Simple GPU memory report.""" |
| 94 | mega_bytes = 1024.0 * 1024.0 |
| 95 | string = name + " memory (MB)" |
| 96 | string += " | allocated: {}".format(torch.cuda.memory_allocated() / mega_bytes) |
| 97 | string += " | max allocated: {}".format( |
| 98 | torch.cuda.max_memory_allocated() / mega_bytes |
| 99 | ) |
| 100 | string += " | reserved: {}".format(torch.cuda.memory_reserved() / mega_bytes) |
| 101 | string += " | max reserved: {}".format( |
| 102 | torch.cuda.max_memory_reserved() / mega_bytes |
| 103 | ) |
| 104 | if mpu.get_data_parallel_rank() == 0: |
| 105 | print("[Rank {}] {}".format(torch.distributed.get_rank(), string), flush=True) |
| 106 | |
| 107 | |
| 108 | def print_params_min_max_norm(optimizer, iteration): |