How to compute per sample gradient in deepspeed

[vermouthdky]

Hi there, I was wondering how can I do per sample gradient computing in deepspeed for multi-gpu training?

It seems like model_engine.backward() automatically do all_reduce, and use deepspeed.utils.safe_get_full_grad() will get the grads after gradient synchronisation. But how can I get the local grads per sample after backward.

Thanks in advance for the kind help!