已合并
ShardedGradScale achieves alignment with GPU #36150
我应该是一阵风创建于 5月20日
ShardedGradScale achieves alignment with GPU #36150
已合并
共 1 个文件变更+23-13
| @@ -72,6 +72,8 @@ class _ShardedGradScaler(GradScaler): | |||
| 72 | See :class:`GradScaler` for explanation of scaling/unscaling and more use cases. | 72 | See :class:`GradScaler` for explanation of scaling/unscaling and more use cases. |
| 73 | 73 | ||
| 74 | Args: | 74 | Args: |
| 75 | + device (str, optional, default="npu"): target device used to synchronize CPU offloaded | ||
| 76 | + inf/nan checks across ranks. | ||
| 75 | init_scale (float, optional, default=2.**16): Initial scale factor. | 77 | init_scale (float, optional, default=2.**16): Initial scale factor. |
| 76 | growth_factor (float, optional, default=2.0): Factor by which the scale is multiplied during | 78 | growth_factor (float, optional, default=2.0): Factor by which the scale is multiplied during |
| 77 | :meth:`update` if no inf/NaN gradients occur for ``growth_interval`` consecutive iterations. | 79 | :meth:`update` if no inf/NaN gradients occur for ``growth_interval`` consecutive iterations. |
| @@ -88,6 +90,7 @@ class _ShardedGradScaler(GradScaler): | |||
| 88 | 90 | ||
| 89 | def __init__( | 91 | def __init__( |
| 90 | self, | 92 | self, |
| 93 | + device: str = "npu", | ||
| 91 | init_scale: float = 2.0 ** 16, | 94 | init_scale: float = 2.0 ** 16, |
| 92 | backoff_factor: float = 0.5, | 95 | backoff_factor: float = 0.5, |
| 93 | growth_factor: float = 2.0, | 96 | growth_factor: float = 2.0, |
| @@ -102,6 +105,7 @@ class _ShardedGradScaler(GradScaler): | |||
| 102 | growth_interval=growth_interval, | 105 | growth_interval=growth_interval, |
| 103 | enabled=enabled, | 106 | enabled=enabled, |
| 104 | ) | 107 | ) |
| 108 | + self._device = device | ||
| 105 | if self._enabled: | 109 | if self._enabled: |
| 106 | self.process_group = process_group | 110 | self.process_group = process_group |
| 107 | self._per_optimizer_states = defaultdict(_refresh_per_optimizer_state) | 111 | self._per_optimizer_states = defaultdict(_refresh_per_optimizer_state) |
| @@ -280,22 +284,28 @@ class _ShardedGradScaler(GradScaler): | |||
| 280 | # Synchronize the detected inf across the ranks | 284 | # Synchronize the detected inf across the ranks |
| 281 | optimizer_state = self._per_optimizer_states[id(optimizer)] | 285 | optimizer_state = self._per_optimizer_states[id(optimizer)] |
| 282 | works = [] | 286 | works = [] |
| 287 | + found_inf_on_cpus = [] | ||
| 288 | + found_inf_on_devices = [] | ||
| 283 | 289 | ||
| 284 | for found_inf in optimizer_state["found_inf_per_device"].values(): | 290 | for found_inf in optimizer_state["found_inf_per_device"].values(): |
| 285 | - if found_inf.device.type == "cpu": | 291 | + if self._device != "cpu" and found_inf.device.type == "cpu": |
| 286 | - found_inf_npu = found_inf.to(self._scale.device) | 292 | + found_inf_on_cpus.append(found_inf) |
| 287 | - work = dist.all_reduce(found_inf_npu, async_op=True, group=self.process_group) | 293 | + found_inf_on_device = found_inf.to(self._device) |
| 288 | - works.append((work, found_inf, found_inf_npu)) | 294 | + found_inf_on_devices.append(found_inf_on_device) |
| 295 | + works.append( | ||
| 296 | + dist.all_reduce( | ||
| 297 | + found_inf_on_device, async_op=True, group=self.process_group | ||
| 298 | + ) | ||
| 299 | + ) | ||
| 289 | else: | 300 | else: |
| 290 | - works.append((dist.all_reduce(found_inf, async_op=True, group=self.process_group), None, None)) | 301 | + works.append( |
| 291 | - | 302 | + dist.all_reduce(found_inf, async_op=True, group=self.process_group) |
| 292 | - for item in works: | 303 | + ) |
| 293 | - if item[1] is not None: | 304 | + for work in works: |
| 294 | - work, found_inf_cpu, found_inf_npu = item | 305 | + work.wait() |
| 295 | - work.wait() | 306 | + if found_inf_on_cpus: |
| 296 | - found_inf_cpu.copy_(found_inf_npu) | 307 | + for cpu_t, dev_t in zip(found_inf_on_cpus, found_inf_on_devices): |
| 297 | - else: | 308 | + cpu_t.copy_(dev_t) |
| 298 | - item[0].wait() | ||
| 299 | 309 | ||
| 300 | def step(self, optimizer: SGD, *args, **kwargs) -> Optional[float]: | 310 | def step(self, optimizer: SGD, *args, **kwargs) -> Optional[float]: |
| 301 | return super().step(optimizer, *args, **kwargs) | 311 | return super().step(optimizer, *args, **kwargs) |