[QUESTION] One possible typo in docs/source/distrib_optimizer.md

[wplf]

Your question
Ask a clear and concise question about Megatron-LM.

Is backward below supposed to be forward ?

[Advait251206]

Yes, I believe that's a typo.

The sequence described is:

1. Backward pass computes gradients.
2. Gradients are reduced via reduce-scatter.
3. The optimizer updates the FP32 master parameters (Optimizer.step()).
4. The updated FP16 model parameters are reconstructed via all-gather.
5. The model parameters are copied back into the model tensors.
6. The next iteration begins with the forward pass, not another backward pass.

So the sentence:

"(At this point, grad buffer is ready to be zero'd for the next iteration)"

is describing the state after parameter synchronization has completed. The next computation performed by the training loop is the forward pass, since the gradients have already been consumed by the optimizer step.

Therefore, if the documentation refers to the backward pass immediately following this sequence, it should likely say forward instead.

This also matches the standard training order:

Forward
    ↓
Backward
    ↓
Reduce-scatter gradients
    ↓
Optimizer.step()
    ↓
All-gather updated parameters
    ↓
Forward (next iteration)

So yes, I believe "backward" should be "forward" in that context.

---

If this answer helped or pointed you in the right direction, I'd appreciate it if you could mark it as the accepted answer so it's easier for others with the same issue to find.

Also, if you found my contribution useful, I'd appreciate it if you could check out my GitHub profile, follow me, and star any repositories you find interesting.

GitHub: https://github.com/Advait251206