params_in_partition in DS stage 3

[jrplatin]

Hi!

I am looking into the deepspeed/runtime/zero/stage3.py code, and I noticed that params_in_partition is initialized but never updated, whereas this same variable is populated in deepspeed/runtime/zero/stage_1_and_2.py: is there a reason for this? How can I get the parameters that will be updated by that specific process directly?

Thanks!