(self, buffer_id)
| 720 | return batch |
| 721 | |
| 722 | def _exec_forward_pass(self, buffer_id): |
| 723 | self.tput_timer.start() |
| 724 | |
| 725 | if isinstance(self.pipe_buffers['inputs'][buffer_id], tuple): |
| 726 | inputs = tuple(t.clone() for t in self.pipe_buffers['inputs'][buffer_id]) |
| 727 | else: |
| 728 | inputs = self.pipe_buffers['inputs'][buffer_id].clone() |
| 729 | |
| 730 | # collect the partitioned input from the previous stage |
| 731 | if self.is_pipe_partitioned and not self.is_first_stage(): |
| 732 | if self.pipe_partition_input_meta_cache is None: |
| 733 | self.pipe_partition_input_meta_cache = inputs[0].to('cpu') |
| 734 | part_input = PartitionedTensor.from_meta(meta=self.pipe_partition_input_meta_cache, |
| 735 | local_part=inputs[1], |
| 736 | group=self.grid.get_slice_parallel_group()) |
| 737 | |
| 738 | inputs = (part_input.full(), *inputs[2:]) |
| 739 | inputs[0].requires_grad = True |
| 740 | # skip mask |
| 741 | #inputs[1].requires_grad = True |
| 742 | part_input = None |
| 743 | inputs = inputs[0] if len(inputs) == 1 else inputs |
| 744 | self.pipe_buffers['inputs'][buffer_id] = inputs |
| 745 | |
| 746 | # inputs has no gradient because it is from a cloned tensor |
| 747 | outputs = super().forward(inputs) |
| 748 | |
| 749 | # Reset activation checkpointing buffers. |
| 750 | # Need to call this between evaluation iterations |
| 751 | if not self.module.training: |
| 752 | ds_checkpointing.reset() |
| 753 | |
| 754 | # Partition the outputs if we are not the last stage |
| 755 | if self.is_pipe_partitioned and not self.is_last_stage(): |
| 756 | if isinstance(outputs, tuple): |
| 757 | first_output = outputs[0] |
| 758 | # TODO: Improve pipe partitioning to pass multiple tensors that require grads |
| 759 | assert all([torch.is_tensor(elt) and elt.requires_grad is False for elt in outputs[1:]]) |
| 760 | outputs_tail = outputs[1:] |
| 761 | elif torch.is_tensor(outputs): |
| 762 | first_output = outputs |
| 763 | outputs_tail = [] |
| 764 | else: |
| 765 | raise ValueError("expecting a tensor or a tuple of tensors") |
| 766 | part = PartitionedTensor(tensor=first_output, group=self.grid.get_slice_parallel_group()) |
| 767 | # Clear the large output data, but save the computation graph |
| 768 | first_output.data = torch.zeros(1, device=first_output.data.device) |
| 769 | self.pipe_buffers['output_tensors'][buffer_id] = first_output |
| 770 | # Inject the partitioned tensor into the output before sending |
| 771 | outputs = (part.to_meta(), part.data(), *outputs_tail) |
| 772 | part = None |
| 773 | |
| 774 | self.pipe_buffers['outputs'][buffer_id] = outputs |
| 775 | |
| 776 | # Optionally compute loss on the last device |
| 777 | if self.is_last_stage(): |
| 778 | if self._compute_loss and self.module.loss_fn is not None: |
| 779 | labels = self.pipe_buffers['labels'][buffer_id] |
nothing calls this directly
no test coverage detected