Get gradients of parameters

[thomas0809]

I am using ZeRO optimization (stage 2) to train a model. I am trying to analyze the gradients, but all the parameters in model_engine.module.parameters() have no .grad. Where are the gradients located?

[tjruwase]

@thomas0809, can you give more details for this issue. For example, at what point in the training process are you examining the parameters for gradients?

[thomas0809]

I am doing it right after model_engine.backward(loss).
I found .grad exists for parameters in optimizer.optimizer.param_groups. However, I haven't figured out how to map the gradients to the model. I am looking for something like

for n, p in model.named_parameters():
    p.grad

Not sure if it is possible.

[tjruwase]

Got it. Your observation is due to the gradient memory optimization of ZeRO stage 2. This optimization frees up fp16 gradient memory after writing the gradients out of the fp32 gradient copies in the optimizer. So optimizer.optimizer.param_groups is the reliable way to examine gradient values.

[Moldoteck]

Got it. Your observation is due to the gradient memory optimization of ZeRO stage 2. This optimization frees up fp16 gradient memory after writing the gradients out of the fp32 gradient copies in the optimizer. So optimizer.optimizer.param_groups is the reliable way to examine gradient values.
@tjruwase I have similar problem for my use-case
Is there a reason why printing: print(model.optimizer.optimizer.param_groups) shows this?
[{'params': [], 'lr': 5.490483162518434e-05, 'bias_correction': True, 'betas': (0.9, 0.999), 'eps': 1e-08, 'weight_decay': 0.0001, 'initial_lr': 0, 'step': 15}]
In other words the params field is empty

[tjruwase]

@Moldoteck, yes for zero-3 we need to add support for extracting gradients #1098.

[Moldoteck]

@Moldoteck, yes for zero-3 we need to add support for extracting gradients #1098.

Understand. Thank you. So if i will not use zero_optimization at all and only use activation_checkpointing, it should work?

[tjruwase]

Sorry for the inconvenience. You should be able to use zero stage 1 and activation_checkpointing.

[Moldoteck]

Thank you. I found out the problem for me. I wanted to extract value.grad from params field, but it was emptied when calling model.step(). After moving model.step() after my calculations, everything worked as expected. Thank you!