[QUESTION] Question about ZeroBubbleVPipeScheduler schedule W

[duanjunwen]

Hello,
I had some questions while reading ZeroBubbleVPipeScheduler, especially schedule W. (As the title)
Can you plz explain exactly how this part is implemented ?

In detailed:

1. Where it strips the weight grad ？
2. Where is the weight grad calculated? (in schedule B, backward_step ?)

[huanggx-sea]

The weight grad computation is stored in WeightGradStore and executed when scheduled.