Add gradient variance logging as a callback

[mariannearr]

• Gradient variance computed as per Block Diffusion eq. 10, where the variance is taken over gradient vectors from accumulation_steps consecutive gradient updates

• Gradients are collected after gradient accumulation and before the optimizer step at after_train_batch