Repository: microsoft/DeepSpeed
File: deepspeed/runtime/pipe/engine.py (Python)

Code before:
            # XXX: torch 1.6.0 DataLoader will auto convert tuple to list
            elif isinstance(batch[1], (tuple, list)):
                loaded = []
                for x in batch[1]:
                    assert torch.is_tensor(x)
                    x = x.to(self.device).detach()
                    loaded.append(x)
                loaded = tuple(loaded)

            self.pipe_buffers['labels'][buffer_id] = loaded

        if self.wall_clock_breakdown():
            self.timers(BATCH_INPUT_TIMER).stop()

    def _send_tensor_meta(self, buffer, recv_stage):
        """ Communicate metadata about upcoming p2p transfers.

        Metadata is communicated in this order:
            * type (0: tensor, 1: list)
            * num_tensors if type=list
            foreach tensor in buffer:
                * ndims
                * shape
        """
        meta_buffer = torch.empty(TENSOR_META_SIZE, dtype=torch.int32, device=self.device)
        if isinstance(buffer, torch.Tensor):
            meta_buf_list = [
                0,  # type of data (0: tensor, 1: list, 2: tuple)
                self.DTYPE_TO_ID[buffer.dtype],  # dtype
                len(buffer.size())  # ndims
            ]
            meta_buf_list.extend(buffer.size())
            meta_buffer[:len(meta_buf_list)].copy_(torch.tensor(meta_buf_list, dtype=torch.int32))
            p2p.send(meta_buffer, recv_stage)

        elif isinstance(buffer, list):
            assert (False)
            type_tensor = torch.LongTensor(data=[1]).to(self.device)
            p2p.send(type_tensor, recv_stage)
            count_tensor = torch.LongTensor(data=[len(buffer)]).to(self.device)
            p2p.send(count_tensor, recv_stage)
            for tensor in buffer:
                assert isinstance(tensor, torch.Tensor)
                send_shape = torch.LongTensor(data=tensor.size()).to(self.device)
                send_ndims = torch.LongTensor(data=[len(tensor.size())]).to(self.device)
                p2p.send(send_ndims, recv_stage)
                p2p.send(send_shape, recv_stage)
        elif isinstance(buffer, tuple):
            meta_buf_list = [
                2,  # type of data (0: tensor, 1: list, 2: tuple)
                len(buffer)  # num_tensors

Review comment (bug):
How can we guard against overflow if entries into `meta_buffer` is more than `TENSOR_META_SIZE`

Code after:
                loaded = batch[1].to(self.device)
            # XXX: torch 1.6.0 DataLoader will auto convert tuple to list
            elif isinstance(batch[1], (tuple, list)):
                loaded = []
                for x in batch[1]:
                    assert torch.is_tensor(x)
                    x = x.to(self.device).detach()
                    loaded.append(x)
                loaded = tuple(loaded)

            self.pipe_buffers['labels'][buffer_id] = loaded

        if self.wall_clock_breakdown():
            self.timers(BATCH_INPUT_TIMER).stop()

    def _send_tensor_meta(self, buffer, recv_stage):
        """ Communicate metadata about upcoming p2p transfers.

        Metadata is communicated in this order:
            * type (0: tensor, 1: list)
            * num_tensors if type=list
            foreach tensor in buffer:
                * ndims
                * shape
        """
        meta_buffer = torch.empty(TENSOR_META_SIZE, dtype=torch.int32, device=self.device)
        if isinstance(buffer, torch.Tensor):
            meta_buf_list = [
                0,  # type of data (0: tensor, 1: list (unused), 2: tuple)
                self.DTYPE_TO_ID[buffer.dtype],  # dtype
                len(buffer.size())  # ndims
            ]
            meta_buf_list.extend(buffer.size())
            assert len(
                meta_buf_list
            ) <= TENSOR_META_SIZE, f"Buffer for metadata is too small. Current buffer size: {TENSOR_META_SIZE} but required {len(meta_buf_list)}"
            meta_buffer[:len(meta_buf_list)].copy_(torch.tensor(meta_buf_list, dtype=torch.int32))
            p2p.send(meta_buffer, recv_stage)

        elif isinstance(buffer, tuple):
            meta_buf_list = [
                2,  # type of data (0: tensor, 1: list (unused), 2: tuple)
                len(buffer)  # num_tensors
            ]

            for tensor in buffer:
                assert isinstance(tensor, torch.Tensor)
                meta_buf_list.append(self.DTYPE_TO_ID[tensor.dtype])
                meta_buf_list.append(len(tensor.size()))
                meta_buf_list.extend(tensor.size())