(self, buffer_id)
| 1081 | self.timers(PIPE_SEND_OUTPUT_TIMER).stop() |
| 1082 | |
| 1083 | def _exec_send_grads(self, buffer_id): |
| 1084 | if self.wall_clock_breakdown(): |
| 1085 | self.timers(PIPE_SEND_GRAD_TIMER).start() |
| 1086 | |
| 1087 | inputs = self.pipe_buffers['inputs'][buffer_id] |
| 1088 | |
| 1089 | # Partition the gradient |
| 1090 | if self.is_grad_partitioned: |
| 1091 | if isinstance(inputs, tuple): |
| 1092 | first_input = inputs[0] |
| 1093 | assert all([torch.is_tensor(elt) for elt in inputs[1:]]) |
| 1094 | inputs_grad_tail = [elt.grad for elt in inputs[1:]] |
| 1095 | elif torch.is_tensor(inputs): |
| 1096 | first_input = inputs |
| 1097 | inputs_grad_tail = [] |
| 1098 | else: |
| 1099 | raise ValueError("expecting a tensor or a tuple of tensors") |
| 1100 | assert torch.is_tensor(first_input) |
| 1101 | part = PartitionedTensor(tensor=first_input.grad, group=self.grid.get_slice_parallel_group()) |
| 1102 | |
| 1103 | inputs = (part.to_meta(), part.data(), *inputs_grad_tail) |
| 1104 | |
| 1105 | # XXX Terrible hack |
| 1106 | # Drop the attention mask from the input buffer here. It does not have |
| 1107 | # a grad that needs to be communicated. We free the buffer immediately |
| 1108 | # after, so no need to restore it. The receiver also has a hack that skips |
| 1109 | # the recv. This is because NCCL does not let us send torch.BoolTensor :-(. |
| 1110 | if self.has_attention_mask or self.has_bool_tensors: |
| 1111 | inputs = list(inputs) |
| 1112 | inputs.pop() |
| 1113 | inputs = tuple(inputs) |
| 1114 | |
| 1115 | if isinstance(inputs, torch.Tensor): |
| 1116 | assert inputs.grad is not None |
| 1117 | p2p.send(inputs.grad, self.prev_stage) |
| 1118 | else: |
| 1119 | # XXX terrible hacky branch |
| 1120 | if self.is_grad_partitioned: |
| 1121 | # First two sends are partitioned gradient |
| 1122 | p2p.send(inputs[0], self.prev_stage) |
| 1123 | p2p.send(inputs[1], self.prev_stage) |
| 1124 | else: |
| 1125 | for idx, buffer in enumerate(inputs): |
| 1126 | # Skip tensors that will not produce a grad |
| 1127 | if not buffer.is_floating_point(): |
| 1128 | assert buffer.grad is None |
| 1129 | continue |
| 1130 | assert buffer.grad is not None |
| 1131 | p2p.send(buffer.grad, self.prev_stage) |
| 1132 | |
| 1133 | # We can free up the input buffer now |
| 1134 | self.pipe_buffers['inputs'][buffer_id] = None |
| 1135 | |
| 1136 | if self.wall_clock_breakdown(): |
| 1137 | self.timers(PIPE_SEND_GRAD_TIMER).stop() |
| 1138 | |
| 1139 | def _exec_recv_activations(self, buffer_id): |
| 1140 | if self.wall_clock_breakdown(): |
nothing calls this directly
no test coverage detected