Fix, pipeline model with moe cause error when send grad

[wukong1992]

Fix, pipeline model with moe, send and recv grad use same buffer will cause error when moe loss not in position 1. Solution, activation and grad use diff buffer. BUG DETAIL: in forward， the grad_buf will allocate hidden_states，some
tensors(required_grad is false), moe_loss；in backward，only send act grad（floating point），so only hidden_states and moe loss's grads will be send, and forward and backward use same grad_buf. This will cause send and recv grad use diff dtype, and cause numerical error。

[wukong1992]

@tohtana Could you review the merge requests?

[loadams]

Hi @wukong1992 - looks like your changes are failing the formatting check - could you either run the pre-commit formatter locally or take a look at the errors? It looks like your changes mean we can remove an import and flake8 caught that.

[tohtana]

Hi @wukong1992, Thank you for this great fix, and sorry for my late response.

I wonder if some behaviors might change with this PR other than solving the issue. _allocate_or_extend_buffers reuses the buffer when possible. Your change seems to allocate memory every time. Honestly, I'm not sure how much it affects the performance, but afraid that it can increase memory fragmentation. Can you share any thought on this?