| 65 | |
| 66 | |
| 67 | class CommsLogger: |
| 68 | |
| 69 | def __init__(self): |
| 70 | from deepspeed.comm.constants import COMMS_LOGGER_VERBOSE_DEFAULT, COMMS_LOGGER_DEBUG_DEFAULT, COMMS_LOGGER_PROF_OPS_DEFAULT, COMMS_LOGGER_PROF_ALL_DEFAULT, COMMS_LOGGER_ENABLED_DEFAULT |
| 71 | self.comms_dict = {} |
| 72 | self.verbose = COMMS_LOGGER_VERBOSE_DEFAULT |
| 73 | self.debug = COMMS_LOGGER_DEBUG_DEFAULT |
| 74 | self.prof_ops = COMMS_LOGGER_PROF_OPS_DEFAULT |
| 75 | self.prof_all = COMMS_LOGGER_PROF_ALL_DEFAULT |
| 76 | self.enabled = COMMS_LOGGER_ENABLED_DEFAULT |
| 77 | |
| 78 | def configure(self, comms_config): |
| 79 | self.enabled = comms_config.comms_logger_enabled |
| 80 | if self.enabled: |
| 81 | self.verbose = comms_config.comms_logger.verbose |
| 82 | self.debug = comms_config.comms_logger.debug |
| 83 | self.prof_ops = comms_config.comms_logger.prof_ops |
| 84 | self.prof_all = comms_config.comms_logger.prof_all |
| 85 | |
| 86 | # There are three settings for the op profiler: |
| 87 | # - Global profiling (profile all comms) |
| 88 | # - Op-type profiling (e.g. profile all all_reduce comms) |
| 89 | # - Op profiling (e.g. profile a specific all_reduce op) |
| 90 | def start_profiling_comms(self): |
| 91 | self.prof_all = True |
| 92 | |
| 93 | def stop_profiling_comms(self): |
| 94 | self.prof_all = False |
| 95 | |
| 96 | # E.g. start_profiling_op('all_reduce') |
| 97 | def start_profiling_op(self, op_name_list): |
| 98 | self.prof_ops = list(set(self.prof_ops) | set(op_name_list)) |
| 99 | |
| 100 | def stop_profiling_op(self, op_name_list): |
| 101 | self.prof_ops = [op for op in self.prof_ops if op not in op_name_list] |
| 102 | |
| 103 | # Add log entry |
| 104 | def append(self, raw_name, record_name, latency, msg_size): |
| 105 | algbw, busbw = calc_bw_log(raw_name, msg_size, latency) |
| 106 | if record_name in self.comms_dict.keys(): |
| 107 | # If this comm_op has already been logged with this message size, just add to existing record |
| 108 | if msg_size in self.comms_dict[record_name].keys(): |
| 109 | self.comms_dict[record_name][msg_size][0] += 1 |
| 110 | self.comms_dict[record_name][msg_size][1].append(latency) |
| 111 | self.comms_dict[record_name][msg_size][2].append(algbw) |
| 112 | self.comms_dict[record_name][msg_size][3].append(busbw) |
| 113 | # If this is a new message size for this comm_op, add new record under existing comm_op |
| 114 | else: |
| 115 | self.comms_dict[record_name][msg_size] = [1, [latency], [algbw], [busbw]] |
| 116 | else: |
| 117 | # Create entirely new record |
| 118 | self.comms_dict[record_name] = {msg_size: [1, [latency], [algbw], [busbw]]} |
| 119 | # If verbose, print every comm op |
| 120 | # TODO: Add to tensorboard |
| 121 | if self.verbose: |
| 122 | log_str = f"comm op: {record_name} | time (ms): {latency:.2f} | msg size: {convert_size(msg_size)} | algbw (Gbps): {algbw:.2f} | busbw (Gbps): {busbw:.2f}" |
| 123 | log_dist(log_str, [0]) |
| 124 |
no outgoing calls