Send tensor to next rank in pipeline (forward send).
(output_tensor, timers=None)
| 184 | |
| 185 | |
| 186 | def send_forward(output_tensor, timers=None): |
| 187 | """Send tensor to next rank in pipeline (forward send).""" |
| 188 | if not mpu.is_pipeline_last_stage(): |
| 189 | if timers is not None: |
| 190 | timers("forward-send").start() |
| 191 | _communicate( |
| 192 | tensor_send_next=output_tensor, |
| 193 | tensor_send_prev=None, |
| 194 | recv_prev=False, |
| 195 | recv_next=False, |
| 196 | ) |
| 197 | if timers is not None: |
| 198 | timers("forward-send").stop() |
| 199 | |
| 200 | |
| 201 | def send_backward(input_tensor_grad, timers=None): |
nothing calls this directly
no test coverage detected