File size: 3,833 Bytes
d9115cb | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 | 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:
# Perform special initialization
# We scale the loss by the world size for proper initialization
opt = initalization_optimizer(model.parameters(), lr=initalization_optimizer_lr)
loss = loss * world_size
else:
# Use the default training loop
opt = model.optimizers()
opt.zero_grad()
model.manual_backward(loss)
opt.step()
|