Fix gradient accumulation related unbound errors for tracking gradnorm

[a-r-r-o-w]

Fixes #41. @Yuancheng-Xu would you mind giving this a review?

[sayakpaul]

Left a small comment.

[sayakpaul]

Oh I thought we always calculate it?

cogvideox-factory/training/cogvideox_image_to_video_lora.py

Line 763 in 6c00cf0

gradient_norm_before_clip = get_gradient_norm(transformer.parameters())

edit: Oh or not when gradient accumulation steps > 1?

[Yuancheng-Xu]

The current fix looks fine to me.

However I have a question: Why not just put all the logging inside if accelerator.sync_gradients:?

What I mean is to change the original code

if accelerator.sync_gradients:
   progress_bar.update(1)
   global_step += 1
   <some code...>
   
last_lr = lr_scheduler.get_last_lr()[0] if lr_scheduler is not None else args.learning_rate
logs = {
                "loss": loss.detach().item(),
                "lr": last_lr,
                "gradient_norm_before_clip": gradient_norm_before_clip,
                "gradient_norm_after_clip": gradient_norm_after_clip,
            }
progress_bar.set_postfix(**logs)
accelerator.log(logs, step=global_step)

to this modified one

if accelerator.sync_gradients:
   progress_bar.update(1)
   global_step += 1
   <some code...>

   last_lr = lr_scheduler.get_last_lr()[0] if lr_scheduler is not None else args.learning_rate
   logs = {
                "loss": loss.detach().item(),
                "lr": last_lr,
                "gradient_norm_before_clip": gradient_norm_before_clip,
                "gradient_norm_after_clip": gradient_norm_after_clip,
            }
    progress_bar.set_postfix(**logs)
    accelerator.log(logs, step=global_step)

It seems that since the logging only happens when accelerator.log(logs, step=global_step), it makes more sense to only log when global_step+=1.

I am not so sure about this, since I also saw the original implementation in other diffuser implementation too, like here

[sayakpaul]

Not a bad idea but I guess it's more a structural preference and logical convenience honestly.