`RuntimeError: tensors must be contiguous` when predicting GPTJForClassification trainer
Author: Dahoas
Pull request: no

### System Info

- `transformers` version: 4.21.2
- Platform: Linux-5.15.0-1023-aws-x86_64-with-glibc2.29
- Python version: 3.8.10
- Huggingface_hub version: 0.10.0
- PyTorch version (GPU?): 1.12.1+cu113 (True)
- Tensorflow version (GPU?): not installed (NA)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed
- Using GPU in script?: 1
- Using distributed or parallel set-up in script?: huggingface transformers deepspeed

### Who can help?

@sgugger @stas00 

### Information

- [ ] The official example scripts
- [X] My own modified scripts

### Tasks

- [ ] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [X] My own task or dataset (give details below)

### Reproduction

```python
import os
import torch
from torch.utils.data import Dataset, random_split
from transformers import AutoTokenizer, TrainingArguments, Trainer, AutoModelForCausalLM, IntervalStrategy, AutoModel, AutoConfig, PreTrainedModel, AutoModelForSequenceClassification
import json
import deepspeed
import argparse
from datasets import load_dataset
import wandb
from tqdm import tqdm


class PairwiseEvalDataset(Dataset):
    def __init__(self, pairs, tokenizer, max_length):
        self.input_ids = []
        self.attn_masks = []

        for pair in tqdm(pairs):
            prompt = pair["prompt"]
            chosen, rejected = pair["chosen"], pair["rejected"]
            tok_chosen = tokenizer(prompt + chosen + "
", return_tensors="pt")["input_ids"]
            tok_rejected = tokenizer(prompt + rejected + "
", return_tensors="pt")["input_ids"]
            # Reject data with num tokens > max_length
            if tok_chosen.shape[-1] <= max_length and tok_rejected.shape[-1] <= max_length:
                chosen_encodings_dict = tokenizer(prompt + chosen + '
', truncation=True,
                                        max_length=max_length, padding="max_length", return_tensors="pt")
                rejected_encodings_dict = tokenizer(prompt + rejected + '
', truncation=True,
                                        max_length=max_length, padding="max_length", return_tensors="pt")
                # First append chosen then rejected
                self.input_ids.append(chosen_encodings_dict['input_ids'])
                self.attn_masks.append(chosen_encodings_dict['attention_mask'])
                self.input_ids.append(rejected_encodings_dict['input_ids'])
                self.attn_masks.append(rejected_encodings_dict['attention_mask'])

    def __len__(self):
        return len(self.input_ids)

    def __getitem__(self, idx):
        return self.input_ids[idx], self.attn_masks[idx]

def pairwise_data_collator(data):
    if len(data[0]) == 4:
        return {'input_ids': torch.cat([f[0] for f in data] + [f[2] for f in data]),
                'attention_mask': torch.cat([f[1] for f in data] + [f[3] for f in data])}
    elif len(data[0]) == 2:
        return {'input_ids': torch.cat([f[0] for f in data]),
                'attention_mask': torch.cat([f[1] for f in data])}
    else:
        raise ValueError("Invalid data format")

class PairwiseTrainer(Trainer):
    def compute_loss(self, model, inputs, return_outputs=False):
        # forward pass
        PAD_ID = model.PAD_ID
        assert len(inputs["input_ids"].shape) == 2
        bs = inputs["input_ids"].shape[0] // 2
        chosen = inputs["input_ids"][:bs]
        rejected = inputs["input_ids"][bs:]
        rewards = model(**inputs).logits
        chosen_rewards = rewards[:bs]
        rejected_rewards = rewards[bs:]
        loss = -torch.log(torch.sigmoid(chosen_rewards - rejected_rewards)).mean()
        return (loss, outputs) if return_outputs else loss

def make_rm(model_name):
    config = AutoConfig.from_pretrained(model_name)
    config.num_labels = 1
    reward_model = AutoModelForSequenceClassification.from_config(config)
    return reward_model

tokenizer = AutoTokenizer.from_pretrained("EleutherAI/gpt-j-6B")
tokenizer.pad_token = tokenizer.eos_token
model = make_rm("Dahoas/gptj-sft-static")

data = load_dataset("Dahoas/rm-static")
max_length = 1024
eval_dataset = PairwiseEvalDataset(data["test"], tokenizer, max_length=max_length)

train_args = TrainingArguments(output_dir=".", per_device_eval_batch_size=1)
trainer = PairwiseTrainer(model=model, args=train_args, train_dataset=eval_dataset, data_collator=pairwise_data_collator)

# TODO(dahoas): Unsure how to compute metrics in trainer for non-classification task
preds = torch.tensor(trainer.predict(eval_dataset)[0])
```

with ds_config

```yaml
{
	"train_batch_size": "auto",
	"fp16": {
	  "enabled": "auto",
	  "min_loss_scale": 1,
	  "loss_scale_window": 1000,
	  "hysteresis": 2,
	  "initial_scale_power": 32
	},
	"bf16": {
		"enabled": "auto"
	},
	"zero_optimization": {
	  "stage": 3,
	  "offload_param": {
		"device": "none"
	  },
	  "offload_optimizer": {
		"device": "none"
	  },
	  "allgather_partitions": true,
	  "allgather_bucket_size": 5e8,
	  "contiguous_gradients": true
	},
	"optimizer": {
	  "type": "AdamW",
	  "params": {
		"lr": "auto",
		"betas": [
		  0.9,
		  0.999
		],
		"eps": 1e-08
	  }
	},
	"scheduler": {
	  "type": "WarmupLR",
	  "params": {
		"warmup_min_lr": 0,
		"warmup_max_lr": "auto",
		"warmup_num_steps": 100
	  }
	}
  }
```

Launch with `deepspeed --num_gpus 1 test.py --deepspeed ../configs/ds_configs/ds_config_gpt_j_z3.json`

I get the error `RuntimeError: Tensors must be contiguous`. The script runs as expected when replacing `gptj` with `gpt2`. I am using 1 A100 40gb gpu. Thank you for any insight.

### Expected behavior

trainer.predict should infer without error

Comments:
Trying to run your example, @Dahoas, I get:
```
Traceback (most recent call last):
  File "contiguous.py", line 8, in <module>
    from rm_datasets import PairwiseDataset, PairwiseEvalDataset, pairwise_data_collator
ModuleNotFoundError: No module named 'rm_datasets'
```

fixed that by removing:

```
from rm_datasets import PairwiseDataset, PairwiseEvalDataset, pairwise_data_collator
```
and now it fails on:

```
Traceback (most recent call last):
  File "contiguous.py", line 10, in <module>
    from utils import freeze_bottom_causal_layers, load_yaml, make_rm
ModuleNotFoundError: No module named 'utils'
```

Sorry about that, I forgot to remove some extraneous imports. I've done so now. 

Thank you for the corrections, @Dahoas  -  I'm able to reproduce the issue. Thank you for that.

Will follow up again once I get a chance to investigate the issue.

The missing from report full traceback is:

```
***** Running Prediction *****
  Num examples = 10206
  Batch size = 1
Traceback (most recent call last):
  File "contiguous.py", line 83, in <module>
    preds = torch.tensor(trainer.predict(eval_dataset)[0])
  File "/mnt/nvme0/code/huggingface/transformers-master-2/src/transformers/trainer.py", line 2894, in predict
    output = eval_loop(
  File "/mnt/nvme0/code/huggingface/transformers-master-2/src/transformers/trainer.py", line 3024, in evaluation_loop
    logits = self._nested_gather(logits)
  File "/mnt/nvme0/code/huggingface/transformers-master-2/src/transformers/trainer.py", line 3140, in _nested_gather
    tensors = distributed_concat(tensors)
  File "/mnt/nvme0/code/huggingface/transformers-master-2/src/transformers/trainer_pt_utils.py", line 191, in distributed_concat
    return type(tensor)(distributed_concat(t, num_total_examples) for t in tensor)
  File "/mnt/nvme0/code/huggingface/transformers-master-2/src/transformers/trainer_pt_utils.py", line 191, in <genexpr>
    return type(tensor)(distributed_concat(t, num_total_examples) for t in tensor)
  File "/mnt/nvme0/code/huggingface/transformers-master-2/src/transformers/trainer_pt_utils.py", line 191, in distributed_concat
    return type(tensor)(distributed_concat(t, num_total_examples) for t in tensor)
  File "/mnt/nvme0/code/huggingface/transformers-master-2/src/transformers/trainer_pt_utils.py", line 191, in <genexpr>
    return type(tensor)(distributed_concat(t, num_total_examples) for t in tensor)
  File "/mnt/nvme0/code/huggingface/transformers-master-2/src/transformers/trainer_pt_utils.py", line 191, in distributed_concat
    return type(tensor)(distributed_concat(t, num_total_examples) for t in tensor)
  File "/mnt/nvme0/code/huggingface/transformers-master-2/src/transformers/trainer_pt_utils.py", line 191, in <genexpr>
    return type(tensor)(distributed_concat(t, num_total_examples) for t in tensor)
  File "/mnt/nvme0/code/huggingface/transformers-master-2/src/transformers/trainer_pt_utils.py", line 194, in distributed_concat
    dist.all_gather(output_tensors, tensor)
  File "/home/stas/anaconda3/envs/py38-pt113/lib/python3.8/site-packages/torch/distributed/distributed_c10d.py", line 2275, in all_gather
    work = default_pg.allgather([tensor_list], [tensor])
RuntimeError: Tensors must be contiguous
```

@Dahoas, this should fix the problem: https://github.com/huggingface/transformers/pull/20951

Thank you for making it super easy for us to identify the problem!

Excellent thank you very much!