| 408 | |
| 409 | @torch.no_grad() |
| 410 | def step(self, parameters: Iterable[torch.nn.Parameter]): |
| 411 | if isinstance(parameters, torch.nn.Module): |
| 412 | deprecation_message = ( |
| 413 | "Passing a `torch.nn.Module` to `ExponentialMovingAverage.step` is deprecated. " |
| 414 | "Please pass the parameters of the module instead." |
| 415 | ) |
| 416 | deprecate( |
| 417 | "passing a `torch.nn.Module` to `ExponentialMovingAverage.step`", |
| 418 | "1.0.0", |
| 419 | deprecation_message, |
| 420 | standard_warn=False, |
| 421 | ) |
| 422 | parameters = parameters.parameters() |
| 423 | |
| 424 | parameters = list(parameters) |
| 425 | |
| 426 | self.optimization_step += 1 |
| 427 | |
| 428 | # Compute the decay factor for the exponential moving average. |
| 429 | decay = self.get_decay(self.optimization_step) |
| 430 | self.cur_decay_value = decay |
| 431 | one_minus_decay = 1 - decay |
| 432 | |
| 433 | context_manager = contextlib.nullcontext |
| 434 | if is_transformers_available() and transformers.integrations.deepspeed.is_deepspeed_zero3_enabled(): |
| 435 | import deepspeed |
| 436 | |
| 437 | if self.foreach: |
| 438 | if is_transformers_available() and transformers.integrations.deepspeed.is_deepspeed_zero3_enabled(): |
| 439 | context_manager = deepspeed.zero.GatheredParameters(parameters, modifier_rank=None) |
| 440 | |
| 441 | with context_manager(): |
| 442 | params_grad = [param for param in parameters if param.requires_grad] |
| 443 | s_params_grad = [ |
| 444 | s_param for s_param, param in zip(self.shadow_params, parameters) if param.requires_grad |
| 445 | ] |
| 446 | |
| 447 | if len(params_grad) < len(parameters): |
| 448 | torch._foreach_copy_( |
| 449 | [s_param for s_param, param in zip(self.shadow_params, parameters) if not param.requires_grad], |
| 450 | [param for param in parameters if not param.requires_grad], |
| 451 | non_blocking=True, |
| 452 | ) |
| 453 | |
| 454 | torch._foreach_sub_( |
| 455 | s_params_grad, torch._foreach_sub(s_params_grad, params_grad), alpha=one_minus_decay |
| 456 | ) |
| 457 | |
| 458 | else: |
| 459 | for s_param, param in zip(self.shadow_params, parameters): |
| 460 | if is_transformers_available() and transformers.integrations.deepspeed.is_deepspeed_zero3_enabled(): |
| 461 | context_manager = deepspeed.zero.GatheredParameters(param, modifier_rank=None) |
| 462 | |
| 463 | with context_manager(): |
| 464 | if param.requires_grad: |
| 465 | s_param.sub_(one_minus_decay * (s_param - param)) |
| 466 | else: |
| 467 | s_param.copy_(param) |