When using AMP, there is a large GPU memory increase on GPU 0. Below is the pseudo code to show my training pipeline. Before running self.scaler.update(), 19G memory is used on each of the 8 GPUs. After the line is done, GPU 0 memory usage will be 29G. This happens in the first training iteration. Is it normal? Or is there any method to distribute the extra memory usage evenly to all the GPUs?
More info:

I am using DistributedDataParallel.
Pretrained checkpoint is loaded to different GPUs.

class Trainer:
    def __init__():
        self.optimizer = ...
        self.scaler = GradScaler()

    def train(input):
        output = model(input)
        loss = get_loss(output)
        self.scaler.scale(loss).backward()
        self.scaler.step(self.optimizer)
        self.scaler.update()
        self.optimizer.zero_grad()