Repository: microsoft/DeepSpeed

Issue:
[BUG]Assertion error while using pipeline parallelism
**Describe the bug**
Recently I'm working on pipeline parallelism to train large model, and I found when my dataset returns two tensors as input, an error occurs:
```bash
Traceback (most recent call last):
  File "codegen_pipeline_deepspeed.py", line 108, in <module>
    loss = model_engine.train_batch()
  File "/nvme/xingshuhao.dispatch/anaconda3/envs/deepspeed/lib/python3.8/site-packages/deepspeed/runtime/pipe/engine.py", line 346, in train_batch
    self._exec_schedule(sched)
  File "/nvme/xingshuhao.dispatch/anaconda3/envs/deepspeed/lib/python3.8/site-packages/deepspeed/runtime/pipe/engine.py", line 1376, in _exec_schedule
    self._exec_instr(**cmd.kwargs)
  File "/nvme/xingshuhao.dispatch/anaconda3/envs/deepspeed/lib/python3.8/site-packages/deepspeed/runtime/pipe/engine.py", line 798, in _exec_load_micro_batch
    assert isinstance(batch[0], tuple)
AssertionError
```

This is my demo dataset:
```python
class MyDataset(Dataset):
    def __init__(self, nums, tokenizer):
        super(MyDataset, self).__init__()
        self.nums = nums
        self.raw = "# this function prints hello world"
        self.data = [tokenizer(self.raw)['input_ids'] for _ in range(nums)]
        self.labels = [tokenizer(self.raw)['input_ids'] for _ in range(nums)]

    def __len__(self):
         return self.nums
    
    def __getitem__(self, index):
        # multiple input tensors and label just fot test
        return (self.data[index], self.labels[index]), self.labels[index]
```

I printed the type of `batch[0]` in the traceback, and I found it is of type `list`.

Is it a bug or is there anyway to make it work since `attetion_masks`, `token_type_ids` may be used as input as well? I have already tried `RepeatingLoader`, `torch.utils.data.DataLoader` and initializing deepspeed with this dataset but this error still occurs.

Below is my code:
```python
import time

import torch
from torch import nn
from torch.optim import AdamW
from torch.utils.data import Dataset, DataLoader
from transformers import AutoTokenizer, AutoModelForCausalLM, AutoConfig
import deepspeed
from deepspeed.pipe import PipelineModule

from pipeline_model_deepspeed import CodeGenForCasualMaskPipe

DEEPSPEED_CONFIG = \
{
    'fp16': {
        'enabled': True,
        },
    'optimizer': {
        'type': 'AdamW',
        'params': {
            'lr': 1e-05, 'betas': [0.9, 0.999],
            'eps': 1e-08, 'weight_decay': 0.0
        }
     },
    'scheduler': {'type': 'WarmupLR', 'params': {'warmup_min_lr': 0, 'warmup_max_lr': 1e-05, 'warmup_num_steps': 100}},
    'zero_optimization': {
        'stage': 1,
        'offload_optimizer': {'device': 'cpu', 'pin_memory': False},
        # 'offload_param': {'device': 'cpu', 'pin_memory': False},
    },
    # 'train_batch_size': 12,
    'train_micro_batch_size_per_gpu': 2,
    'gradient_accumulation_steps': 2,
}

class MyDataset(Dataset):
    def __init__(self, nums, tokenizer):
        super(MyDataset, self).__init__()
        self.nums = nums
        self.raw = "# this function prints hello world"
        self.data = [tokenizer(self.raw)['input_ids'] for _ in range(nums)]
        self.labels = [tokenizer(self.raw)['input_ids'] for _ in range(nums)]

    def __len__(self):
         return self.nums
    
    def __getitem__(self, index):
        # multiple input tensors and label just fot test
        return (self.data[index], self.labels[index]), 0

def loss_fn(lm_logits, labels):
    # Shift so that tokens < n predict n
    shift_logits = lm_logits[..., :-1, :].contiguous()
    shift_labels = labels[..., 1:].contiguous()
    # Flatten the tokens
    loss_fct = nn.CrossEntropyLoss()
    loss = loss_fct(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1))

    return loss

MODEL_NAME = 'Salesforce/codegen-350M-mono'
epochs = 20
deepspeed.init_distributed()

print(f'initializing model: {MODEL_NAME}')

start = time.time()
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)

model = PipelineModule(
    CodeGenForCasualMaskPipe.from_pretrained(MODEL_NAME).to_layers(),
    num_stages=3, loss_fn=loss_fn
)

print(f"model loaded: {time.time() - start}s")
print('initializing deepspeed')

dataset = MyDataset(10, tokenizer)

model_parameters = list(filter(lambda p: p.requires_grad, model.parameters()))
optimizer = AdamW(model_parameters, lr=1e-5)
model_engine, optimizer, train_dataloader, _ = deepspeed.initialize(
	config=DEEPSPEED_CONFIG, model=model, model_parameters=model_parameters,
    optimizer=optimizer, training_data=dataset)

torch.cuda.empty_cache()

print('starting training')

for step in range(epochs):
    loss = model_engine.train_batch()

    print(f'{step} {loss:8.3f}')
```

deepspeed version 0.8.1, pytorch 1.13.1


Discussion:
`assert isinstance(mylist, (list,tuple))` should work.
Any reason why its limited to tuple only?

What `pipeline_model_deepspeed` by the way? @x54-729 
> `assert isinstance(mylist, (list,tuple))` should work. Any reason why its limited to tuple only?
> 
> What `pipeline_model_deepspeed` by the way? @x54-729

Yeah if I change my source code to `assert isinstance(mylist, (list, tuple))`, my program works well, so in my opinion it might be a bug.

I define my pipeline model `CodeGenForCasualMaskPipe` in `pipeline_model_deepspeed.py` .  I can post it if needed. @satpalsr 

Resolving patch:
diff --git a/deepspeed/runtime/pipe/engine.py b/deepspeed/runtime/pipe/engine.py
--- a/deepspeed/runtime/pipe/engine.py
+++ b/deepspeed/runtime/pipe/engine.py
@@ -795,7 +795,7 @@ def _exec_load_micro_batch(self, buffer_id):
                 loaded = batch[0].clone().to(self.device).detach()
                 loaded.requires_grad = loaded.is_floating_point()
             else:
-                assert isinstance(batch[0], tuple)
+                assert isinstance(batch[0], (tuple, list))
                 # Assume list or tuple
                 loaded = []
                 for x in batch[0]:


Test patch: