| import torch
|
| from pytorch_lightning import LightningModule
|
|
|
|
|
| def scale_loss_by_world_size_for_initialization_training_loop(
|
| model: LightningModule,
|
| loss: torch.Tensor,
|
| world_size: int,
|
| is_initialized: bool = True,
|
| initalization_optimizer_lr: float = 0.5,
|
| initalization_optimizer: torch.optim.Optimizer = torch.optim.SGD,
|
| ):
|
| """
|
| Training loop that scales the loss by the number of GPUs used for training during
|
| model initialization.
|
|
|
| This is used for training models with DDP that require special initialization
|
| computed based on the data, such as KMeans. To use this function to initialize a
|
| model:
|
|
|
| 1. Ensure that the model parameters are set to zero at the very start of training.
|
|
|
| 2. On one device, compute the desired initial parameters based on the data, and
|
| compute the loss as the squared distance between the model's parameters and the
|
| desired initial parameters. We use only one device to compute the initialization
|
| because different devices may have different data, and the average initialization
|
| across devices may not be a reasonable initialization.
|
|
|
| 3. Set the loss on other devices to zero.
|
|
|
| 4. Feed the loss to this training loop with the `is_initialized=False`. This will
|
| tell this function to scale the loss by the world size, which is necessary for
|
| the following reason:
|
|
|
| Without loss rescaling, the losses are:
|
| - Device 1: || model.parameters - desired_initial_parameters||^2
|
| - Device 2-world_size: 0
|
| The gradients are (recalling that model.parameters are all zero):
|
| - Device 1: -2 * desired_initial_parameters
|
| - Device 2-world_size: 0
|
| Recall that DDP will average the gradients across all devices. At this point, the
|
| average gradient is:
|
| -2 * desired_initial_parameters / world_size
|
| which means that with SGD optimizer with learning rate 0.5, the model's parameters
|
| will be `desired_initial_parameters / world_size` after the update. To avoid this,
|
| we scale the loss by the world size so that the average gradient across all devices
|
| is -2 * desired_initial_parameters, . This means that the model's parameters will be
|
| updated to the desired initial parameters in one step on all devices via DDP.
|
| Note that this is only used for the initialization step. After the model is
|
| initialized, we can use the default training loop without loss rescaling because the
|
| gradients on all devices are non-trivial.
|
|
|
| 5. Ensure that the optimizer is SGD with learning rate 0.5 so that the model's
|
| parameters are updated to the desired initial parameters in one step on all devices
|
| via DDP.
|
|
|
| 6. After the model is initialized in this step, set the `is_initialized` flag to
|
| `True` to use the default training loop, as now we can compute losses on all devices
|
| and no longer need to scale the loss by the world size.
|
|
|
| Args:
|
| trainer: The trainer object.
|
| loss: The loss value. If the model is not yet initialized, this loss
|
| should be the squared distance between the model's parameters, which should
|
| have value zero at this point, and the desired initial parameters.
|
| world_size: The number of GPUs used for training.
|
| is_initialized: A boolean indicating whether the model is already initialized.
|
| """
|
| if not is_initialized:
|
|
|
|
|
| opt = initalization_optimizer(model.parameters(), lr=initalization_optimizer_lr)
|
| loss = loss * world_size
|
| else:
|
|
|
| opt = model.optimizers()
|
| opt.zero_grad()
|
| model.manual_backward(loss)
|
| opt.step()
|
|
|
