Question on fine-grained synchronization in MiCS

[learning-chip]

In the MiCS paper, Section 4, "Fine-grained synchronization", it is stated that

MiCS follows the good practice in existing works, e.g., BytePS [26], that leverages much finer-grained wait_event, wait_stream and record_stream operations for synchronization, which allow us to maintain the relative order of computation and communication operations in different streams.

However, in the implementation #2964, I cannot find invocation of wait_event, wait_stream or record_stream. Is it not implemented yet? How large is the performance impact? Thanks!

[tjruwase]

@zarzen, FYI

[zarzen]

Thanks @tjruwase for mentioning.

@learning-chip Thanks for being interested in our work, the latest DeepSpeed implementation for parameter gathering has already included these fine-grained event control. You can view it as implemented already.

[learning-chip]

Thanks! Just to confirm: do you mean this reduce_and_partition_stream.wait_stream inside DeepSpeedZeroOptimizer_Stage3:

DeepSpeed/deepspeed/runtime/zero/stage3.py

Line 1085 in 7e8bcc0

self.reduce_and_partition_stream.wait_stream(get_accelerator().default_stream())

which is then inherited by class MiCS_Optimizer?

[zarzen]

Yes it is