| 146 | |
| 147 | |
| 148 | class SyncMultiGPUTrainerReplicated(SingleCostTrainer): |
| 149 | |
| 150 | __doc__ = SyncMultiGPUReplicatedBuilder.__doc__ + """ |
| 151 | |
| 152 | Attributes: |
| 153 | devices (list[int]): List of GPU ids. |
| 154 | |
| 155 | BROADCAST_EVERY_EPOCH (bool): |
| 156 | Whether to broadcast the variables every epoch. |
| 157 | Theoretically this is a no-op (because the variables |
| 158 | are supposed to be in-sync). |
| 159 | But this cheap operation may help prevent |
| 160 | certain numerical issues in practice. |
| 161 | |
| 162 | Note that in cases such as BatchNorm, the variables may not be in sync: |
| 163 | e.g., non-master worker may not maintain EMAs. |
| 164 | |
| 165 | For benchmark, disable this option. |
| 166 | """ |
| 167 | |
| 168 | @map_arg(gpus=_int_to_range) |
| 169 | def __init__(self, gpus, average=True, mode=None): |
| 170 | """ |
| 171 | Args: |
| 172 | gpus (int or [int]): list of GPU ids. |
| 173 | average (bool): whether to average or sum gradients. |
| 174 | mode (str or None): Gradient aggregation mode. |
| 175 | Supported values: ['nccl', 'hierarchical', 'cpu', 'gpu']. |
| 176 | These modes may differ in speed. |
| 177 | Default to pick automatically by heuristics. |
| 178 | "hierarchical" mode was designed for DGX-like 8-GPU machines. |
| 179 | """ |
| 180 | self.devices = gpus |
| 181 | if mode is not None: |
| 182 | mode = mode.lower() |
| 183 | |
| 184 | # Heuristics about mode selection: |
| 185 | if mode == 'hierarchical' and len(gpus) != 8: |
| 186 | logger.warn("mode='hierarchical' requires 8 GPUs. Will fallback to default mode.") |
| 187 | mode = None |
| 188 | if mode is None: |
| 189 | if len(gpus) == 8: |
| 190 | mode = 'hierarchical' |
| 191 | else: |
| 192 | # https://github.com/tensorflow/tensorflow/issues/41539 |
| 193 | mode = 'nccl' if get_tf_version_tuple() < (1, 15) else 'gpu' |
| 194 | if mode == 'cpu' and get_tf_version_tuple() >= (2, 0): |
| 195 | # cpu mode causes the entire model to get located on cpu |
| 196 | mode = 'gpu' |
| 197 | if mode == 'nccl' and get_tf_version_tuple() >= (1, 15): |
| 198 | logger.warning( |
| 199 | "NCCL in TensorFlow has a serious bug that is likely to trigger in TF>=1.15. " |
| 200 | "Try 'mode=None' to use a better default mode.") |
| 201 | |
| 202 | self._builder = SyncMultiGPUReplicatedBuilder(gpus, average, mode) |
| 203 | self.BROADCAST_EVERY_EPOCH = True |
| 204 | |
| 205 | super(SyncMultiGPUTrainerReplicated, self).__init__() |
no outgoing calls
no test coverage detected