AttributeError: 'GPTJModel' object has no attribute 'first_device'
### System Info

```
transformers.__version__ # 4.28.0.dev0
torch.__version__ # 2.0.0+cu117
python # Python 3.7.12
```

### Who can help?

@sgugger @younesbelkada

### Information

- [ ] The official example scripts
- [X] My own modified scripts

### Tasks

- [X] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [ ] My own task or dataset (give details below)

### Reproduction

The followig gist uses `model_name = 'facebook/opt-2.7b'`, which seems working as expected. 

[Gist-model-parallel](https://gist.github.com/innat/e6c4826382641f640cc91def95026ad3)

But for model like `'EleutherAI/gpt-j-6b'` or `gpt2`, it gives error. 

```
AttributeError: 'GPTJModel' object has no attribute 'first_device'
```

Full logs

```
---------------------------------------------------------------------------
AttributeError                            Traceback (most recent call last)
Cell In[14], line 20
     12 trainer = Trainer(
     13     model=model,
     14     args=training_args,
     15     data_collator=data_collator,
     16     train_dataset=train_dataset,
     17 )
     19 model.config.use_cache = False 
---> 20 trainer.train()

File /opt/conda/envs/gpt_neox/lib/python3.9/site-packages/transformers/trainer.py:1639, in Trainer.train(self, resume_from_checkpoint, trial, ignore_keys_for_eval, **kwargs)
   1634     self.model_wrapped = self.model
   1636 inner_training_loop = find_executable_batch_size(
   1637     self._inner_training_loop, self._train_batch_size, args.auto_find_batch_size
   1638 )
-> 1639 return inner_training_loop(
   1640     args=args,
   1641     resume_from_checkpoint=resume_from_checkpoint,
   1642     trial=trial,
   1643     ignore_keys_for_eval=ignore_keys_for_eval,
   1644 )

File /opt/conda/envs/gpt_neox/lib/python3.9/site-packages/transformers/trainer.py:1906, in Trainer._inner_training_loop(self, batch_size, args, resume_from_checkpoint, trial, ignore_keys_for_eval)
   1904         tr_loss_step = self.training_step(model, inputs)
   1905 else:
-> 1906     tr_loss_step = self.training_step(model, inputs)
   1908 if (
   1909     args.logging_nan_inf_filter
   1910     and not is_torch_tpu_available()
   1911     and (torch.isnan(tr_loss_step) or torch.isinf(tr_loss_step))
   1912 ):
   1913     # if loss is nan or inf simply add the average of previous logged losses
   1914     tr_loss += tr_loss / (1 + self.state.global_step - self._globalstep_last_logged)

File /opt/conda/envs/gpt_neox/lib/python3.9/site-packages/transformers/trainer.py:2652, in Trainer.training_step(self, model, inputs)
   2649     return loss_mb.reduce_mean().detach().to(self.args.device)
   2651 with self.compute_loss_context_manager():
-> 2652     loss = self.compute_loss(model, inputs)
   2654 if self.args.n_gpu > 1:
   2655     loss = loss.mean()  # mean() to average on multi-gpu parallel training

File /opt/conda/envs/gpt_neox/lib/python3.9/site-packages/transformers/trainer.py:2684, in Trainer.compute_loss(self, model, inputs, return_outputs)
   2682 else:
   2683     labels = None
-> 2684 outputs = model(**inputs)
   2685 # Save past state if it exists
   2686 # TODO: this needs to be fixed and made cleaner later.
   2687 if self.args.past_index >= 0:

File /opt/conda/envs/gpt_neox/lib/python3.9/site-packages/torch/nn/modules/module.py:1501, in Module._call_impl(self, *args, **kwargs)
   1496 # If we don't have any hooks, we want to skip the rest of the logic in
   1497 # this function, and just call forward.
   1498 if not (self._backward_hooks or self._backward_pre_hooks or self._forward_hooks or self._forward_pre_hooks
   1499         or _global_backward_pre_hooks or _global_backward_hooks
   1500         or _global_forward_hooks or _global_forward_pre_hooks):
-> 1501     return forward_call(*args, **kwargs)
   1502 # Do not call functions when jit is used
   1503 full_backward_hooks, non_full_backward_hooks = [], []

File /opt/conda/envs/gpt_neox/lib/python3.9/site-packages/accelerate/hooks.py:165, in add_hook_to_module.<locals>.new_forward(*args, **kwargs)
    163         output = old_forward(*args, **kwargs)
    164 else:
--> 165     output = old_forward(*args, **kwargs)
    166 return module._hf_hook.post_forward(module, output)

File /opt/conda/envs/gpt_neox/lib/python3.9/site-packages/transformers/models/gptj/modeling_gptj.py:869, in GPTJForCausalLM.forward(self, input_ids, past_key_values, attention_mask, token_type_ids, position_ids, head_mask, inputs_embeds, labels, use_cache, output_attentions, output_hidden_states, return_dict)
    867 # Set device for model parallelism
    868 if self.model_parallel:
--> 869     torch.cuda.set_device(self.transformer.first_device)
    870     hidden_states = hidden_states.to(self.lm_head.weight.device)
    872 # make sure sampling in fp16 works correctly and
    873 # compute loss in fp32 to match with mesh-tf version
    874 # https://github.com/EleutherAI/gpt-neo/blob/89ce74164da2fb16179106f54e2269b5da8db333/models/gpt2/gpt2.py#L179

File /opt/conda/envs/gpt_neox/lib/python3.9/site-packages/torch/nn/modules/module.py:1614, in Module.__getattr__(self, name)
   1612     if name in modules:
   1613         return modules[name]
-> 1614 raise AttributeError("'{}' object has no attribute '{}'".format(
   1615     type(self).__name__, name))

AttributeError: 'GPTJModel' object has no attribute 'first_device'
```

### Expected behavior

Couldn't interpret the problem (`no attribute 'first_device'`), otherwise, it's expected to work same as other model.
This comes from your hack of setting the attributes of the model in cell 6. This then makes the model want to try to use the old model parallel API which crashes since you didn't really use it ;-)<|||||>Ah, I see. Sorry, it was bit confusing. As mentioned, model `fb/opt` worked. Also `abeja/gpt-neox-japanese-2.7b` worked either. Is there any easy fix for the newer API?<|||||>It will work with any model that does not implement the `parallelize` API. As for fixes, the issue you originally psoted on will fix the models with head if needed, and the Trainer has been fixed as @younesbelkada mentioned, so you shouldn't need this hack anymore.<|||||>This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md) are likely to be ignored.