# Skip optimizer update when gradient norm is large with Accelerate gradient accumulation

**URL:** <https://discuss.huggingface.co/t/skip-optimizer-update-when-gradient-norm-is-large-with-accelerate-gradient-accumulation/61869>\
**Category:** 🤗Accelerate\
**Created:** [November 10, 2023, 11:48pm UTC](https://discuss.huggingface.co/t/skip-optimizer-update-when-gradient-norm-is-large-with-accelerate-gradient-accumulation/61869 "2023-11-10T23:48:50Z")\
**Posts on this page:** 1\
**Page:** 1

<div class="post-metadata">

**Author:** ![wren93](https://avatars.discourse-cdn.com/v4/letter/w/f1d935/32.png) [@wren93](https://discuss.huggingface.co/u/wren93)\
**Post date:** [November 10, 2023, 11:48pm UTC](https://discuss.huggingface.co/t/skip-optimizer-update-when-gradient-norm-is-large-with-accelerate-gradient-accumulation/61869/1 "2023-11-10T23:48:50Z")

</div>

I’m trying to finetune a stable diffusion model using the example finetuning code in huggingface diffusers ([https://github.com/huggingface/diffusers/blob/main/examples/text\_to\_image/train\_text\_to\_image.py](https://github.com/huggingface/diffusers/blob/main/examples/text_to_image/train_text_to_image.py)) The main training loop uses accelerate gradient accumulation and is similar to this pseudo code example:

```auto
global_step = 0.0
train_loss = 0.0
for step, batch in enumerate(train_dataloader):
    with accelerator.accumulate(unet):
        # steps for generating model inputs (noisy latents, timesteps, etc.)
        ...
        # model forward
        model_pred = unet(noisy_latents, timesteps, encoder_hidden_states).sample
        loss = F.mse_loss(model_pred.float(), target.float(), reduction="mean")
    
        # Gather the losses across all processes for logging (if we use distributed training).
        avg_loss = accelerator.gather(loss.repeat(args.train_batch_size)).mean()
        train_loss += avg_loss.item() / args.gradient_accumulation_steps
    
        # Backpropagate
        accelerator.backward(loss)
        if accelerator.sync_gradients:
            grad_norm = accelerator.clip_grad_norm_(unet.parameters(), args.max_grad_norm)
        optimizer.step()
        lr_scheduler.step()
        optimizer.zero_grad()

    if accelerator.sync_gradients:
        global_step += 1
        train_loss = 0.0

```

I noticed that the gradient norm (stored in the grad\_norm variable) sometimes explodes and becomes a very large value for certain samples in my dataset. For some reason this will cause the loss to diverge even after applying gradient clipping. I would thus like to define a threshold and if the grad norm is above the threshold, I would like to discard this entire batch (that has a batch size equal to the total batch size after applying gradient accumulation) and don’t do any gradient update. I wonder if I can do something like this:

```auto
if accelerator.sync_gradients:
    grad_norm = accelerator.clip_grad_norm_(unet.parameters(), args.max_grad_norm)
    if grad_norm.item() > grad_norm_threshold:
        lr_scheduler.step()
        optimizer.zero_grad()
        continue
optimizer.step()
lr_scheduler.step()
optimizer.zero_grad()

```

I’m not sure if this can correctly skip the update of the current batch. If not, what should I do here? Thanks in advance.