About offload stage3 source code learning problems

[lzy-edu]

Hello everyone, I want to learn the stage3 part of zerooffload when learning the source code of deepspeed, but I can't find the scheduling process code of the gradient between cpu and gpu, please help me if you know

[jomayeri]

The gradients are not offloaded in ZeRO. The only parts that can be offloaded are the optimizer states and the parameters.

[tjruwase]

@lzy-edu, see

DeepSpeed/deepspeed/runtime/zero/stage3.py

Line 1463 in fc4e733

def partition_grads(self, params_to_release: List[Parameter], grad_partitions: List[Tensor]) -> None:

[lzy-edu]

@lzy-edu, see

DeepSpeed/deepspeed/runtime/zero/stage3.py

Line 1463 in fc4e733

def partition_grads(self, params_to_release: List[Parameter], grad_partitions: List[Tensor]) -> None:

Thank you for your answer. I would also like to ask you a question about the initial parameter partitioning. When I enabled zerooffload3 during initialization, will all parameters be unloaded to the cpu first? Where is this part