There are some small buffers in my model, but they are just some constant matrix, the synchronization might take some time.
Another interesting speedup also confuses me a lot, in the Megatron-LM (GitHub - NVIDIA/Megatron-LM: Ongoing research training transformer language models at scale, including: BERT & GPT-2) package, the model will be explicitly converted into the Float16, by calling something like model.half() before training the model with fp16 enabled, and by doing so the backward time will be largely reduced (e.g. the time takes to finish loss.backward()). I think the backward time should be agnostic to whether the model is float16 or not since the gradients are fp16 when the fp16 is enabled. Do you have  any hint about why this may happen?