Trying to implement Transformer-XL using PyTorch Lightning
Author: moksh-enf
Pull request: no

### System Info

```shell
- `transformers` version: 4.22.2
- Platform: Linux-5.15.0-50-generic-x86_64-with-glibc2.29
- Python version: 3.8.10
- Huggingface_hub version: 0.10.0
- PyTorch version (GPU?): 1.11.0+cu102 (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?: Yes
- Using distributed or parallel set-up in script?: Yes
```


### 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

Code: 

```
import itertools
import math
import os
import sys
import time

# third party libraries
import datasets
import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
import pytorch_lightning as pl
from transformers import TransfoXLModel, TransfoXLTokenizer, TransfoXLConfig

class LitMemTransformerLMWT103(pl.LightningModule):
    """Memory Transformer for Attentive Language Models on WikikText 103"""

    def __init__(self, model_name_or_path, **kwargs) -> None:
        """Initialize the model"""
        super().__init__(**kwargs)
        self.configuration = TransfoXLConfig()
        self.tokenizer = TransfoXLTokenizer.from_pretrained("transfo-xl-wt103")
        self.model = TransfoXLModel.from_pretrained(model_name_or_path, config=self.configuration)

    def forward(self, **inputs):
        return self.model(**inputs)

    def training_step(self, batch):
        outputs = self(**batch)
        loss = outputs[0]
        return loss

    def train_dataloader(self) -> DataLoader:
        dm = datasets.load_dataset('wikitext', 'wikitext-103-v1')
        return DataLoader(dm['train'], batch_size=64, num_workers=self.config.cpu_count)

    def configure_optimizers(self) -> torch.optim.Adam:
        """Use the Adam torch optimizer."""
        optimizer = torch.optim.Adam(self.model.parameters(), lr=self.config.learning_rate)
        return optimizer
```

Error:
```
Traceback (most recent call last):
  File "/opt/moksh/venv/lib/python3.8/site-packages/systest/src/lib/agent/mixin/torchtrain.py", line 471, in <module>
    torch_train()  # pylint: disable=no-value-for-parameter
  File "/opt/moksh/venv/lib/python3.8/site-packages/click/core.py", line 1130, in __call__
    return self.main(*args, **kwargs)
  File "/opt/moksh/venv/lib/python3.8/site-packages/click/core.py", line 1055, in main
    rv = self.invoke(ctx)
  File "/opt/moksh/venv/lib/python3.8/site-packages/click/core.py", line 1404, in invoke
    return ctx.invoke(self.callback, **ctx.params)
  File "/opt/moksh/venv/lib/python3.8/site-packages/click/core.py", line 760, in invoke
    return __callback(*args, **kwargs)
  File "/opt/moksh/venv/lib/python3.8/site-packages/systest/src/lib/agent/mixin/torchtrain.py", line 460, in torch_train
    trained_module, profile = train(params)
  File "/opt/moksh/venv/lib/python3.8/site-packages/systest/src/lib/agent/mixin/torchtrain.py", line 346, in train
    trainer.fit(module)
  File "/opt/moksh/venv/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 770, in fit
    self._call_and_handle_interrupt(
  File "/opt/moksh/venv/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 723, in _call_and_handle_interrupt
    return trainer_fn(*args, **kwargs)
  File "/opt/moksh/venv/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 811, in _fit_impl
    results = self._run(model, ckpt_path=self.ckpt_path)
  File "/opt/moksh/venv/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 1236, in _run
    results = self._run_stage()
  File "/opt/moksh/venv/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 1323, in _run_stage
    return self._run_train()
  File "/opt/moksh/venv/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 1353, in _run_train
    self.fit_loop.run()
  File "/opt/moksh/venv/lib/python3.8/site-packages/pytorch_lightning/loops/base.py", line 204, in run
    self.advance(*args, **kwargs)
  File "/opt/moksh/venv/lib/python3.8/site-packages/pytorch_lightning/loops/fit_loop.py", line 266, in advance
    self._outputs = self.epoch_loop.run(self._data_fetcher)
  File "/opt/moksh/venv/lib/python3.8/site-packages/pytorch_lightning/loops/base.py", line 204, in run
    self.advance(*args, **kwargs)
  File "/opt/moksh/venv/lib/python3.8/site-packages/pytorch_lightning/loops/epoch/training_epoch_loop.py", line 208, in advance
    batch_output = self.batch_loop.run(batch, batch_idx)
  File "/opt/moksh/venv/lib/python3.8/site-packages/pytorch_lightning/loops/base.py", line 204, in run
    self.advance(*args, **kwargs)
  File "/opt/moksh/venv/lib/python3.8/site-packages/pytorch_lightning/loops/batch/training_batch_loop.py", line 88, in advance
    outputs = self.optimizer_loop.run(split_batch, optimizers, batch_idx)
  File "/opt/moksh/venv/lib/python3.8/site-packages/pytorch_lightning/loops/base.py", line 204, in run
    self.advance(*args, **kwargs)
  File "/opt/moksh/venv/lib/python3.8/site-packages/pytorch_lightning/loops/optimization/optimizer_loop.py", line 203, in advance
    result = self._run_optimization(
  File "/opt/moksh/venv/lib/python3.8/site-packages/pytorch_lightning/loops/optimization/optimizer_loop.py", line 256, in _run_optimization
    self._optimizer_step(optimizer, opt_idx, batch_idx, closure)
  File "/opt/moksh/venv/lib/python3.8/site-packages/pytorch_lightning/loops/optimization/optimizer_loop.py", line 369, in _optimizer_step
    self.trainer._call_lightning_module_hook(
  File "/opt/moksh/venv/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 1595, in _call_lightning_module_hook
    output = fn(*args, **kwargs)
  File "/opt/moksh/venv/lib/python3.8/site-packages/pytorch_lightning/core/lightning.py", line 1646, in optimizer_step
    optimizer.step(closure=optimizer_closure)
  File "/opt/moksh/venv/lib/python3.8/site-packages/pytorch_lightning/core/optimizer.py", line 168, in step
    step_output = self._strategy.optimizer_step(self._optimizer, self._optimizer_idx, closure, **kwargs)
  File "/opt/moksh/venv/lib/python3.8/site-packages/pytorch_lightning/strategies/ddp.py", line 286, in optimizer_step
    optimizer_output = super().optimizer_step(optimizer, opt_idx, closure, model, **kwargs)
  File "/opt/moksh/venv/lib/python3.8/site-packages/pytorch_lightning/strategies/strategy.py", line 193, in optimizer_step
    return self.precision_plugin.optimizer_step(model, optimizer, opt_idx, closure, **kwargs)
  File "/opt/moksh/venv/lib/python3.8/site-packages/pytorch_lightning/plugins/precision/precision_plugin.py", line 155, in optimizer_step
    return optimizer.step(closure=closure, **kwargs)
  File "/opt/moksh/venv/lib/python3.8/site-packages/torch/optim/optimizer.py", line 88, in wrapper
    return func(*args, **kwargs)
  File "/opt/moksh/venv/lib/python3.8/site-packages/torch/autograd/grad_mode.py", line 27, in decorate_context
    return func(*args, **kwargs)
  File "/opt/moksh/venv/lib/python3.8/site-packages/torch/optim/adam.py", line 100, in step
    loss = closure()
  File "/opt/moksh/venv/lib/python3.8/site-packages/pytorch_lightning/plugins/precision/precision_plugin.py", line 140, in _wrap_closure
    closure_result = closure()
  File "/opt/moksh/venv/lib/python3.8/site-packages/pytorch_lightning/loops/optimization/optimizer_loop.py", line 148, in __call__
    self._result = self.closure(*args, **kwargs)
  File "/opt/moksh/venv/lib/python3.8/site-packages/pytorch_lightning/loops/optimization/optimizer_loop.py", line 134, in closure
    step_output = self._step_fn()
  File "/opt/moksh/venv/lib/python3.8/site-packages/pytorch_lightning/loops/optimization/optimizer_loop.py", line 427, in _training_step
    training_step_output = self.trainer._call_strategy_hook("training_step", *step_kwargs.values())
  File "/opt/moksh/venv/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 1765, in _call_strategy_hook
    output = fn(*args, **kwargs)
  File "/opt/moksh/venv/lib/python3.8/site-packages/pytorch_lightning/strategies/ddp.py", line 349, in training_step
    return self.model(*args, **kwargs)
  File "/opt/moksh/venv/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1110, in _call_impl
    return forward_call(*input, **kwargs)
  File "/opt/moksh/venv/lib/python3.8/site-packages/torch/nn/parallel/distributed.py", line 963, in forward
    output = self.module(*inputs[0], **kwargs[0])
  File "/opt/moksh/venv/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1110, in _call_impl
    return forward_call(*input, **kwargs)
  File "/opt/moksh/venv/lib/python3.8/site-packages/pytorch_lightning/overrides/base.py", line 82, in forward
    output = self.module.training_step(*inputs, **kwargs)
  File "/opt/moksh/venv/lib/python3.8/site-packages/systest/src/lib/agent/mixin/ml/lit_memxfmr_wt103.py", line 35, in training_step
    outputs = self(**batch)
  File "/opt/moksh/venv/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1110, in _call_impl
    return forward_call(*input, **kwargs)
  File "/opt/moksh/venv/lib/python3.8/site-packages/systest/src/lib/agent/mixin/ml/lit_memxfmr_wt103.py", line 32, in forward
    return self.model(**inputs)
  File "/opt/moksh/venv/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1128, in _call_impl
    result = forward_call(*input, **kwargs)
TypeError: forward() got an unexpected keyword argument 'text'
Traceback (most recent call last):
  File "/opt/moksh/venv/lib/python3.8/site-packages/systest/src/lib/agent/mixin/torchtrain.py", line 471, in <module>
    torch_train()  # pylint: disable=no-value-for-parameter
  File "/opt/moksh/venv/lib/python3.8/site-packages/click/core.py", line 1130, in __call__
    return self.main(*args, **kwargs)
  File "/opt/moksh/venv/lib/python3.8/site-packages/click/core.py", line 1055, in main
    rv = self.invoke(ctx)
  File "/opt/moksh/venv/lib/python3.8/site-packages/click/core.py", line 1404, in invoke
    return ctx.invoke(self.callback, **ctx.params)
  File "/opt/moksh/venv/lib/python3.8/site-packages/click/core.py", line 760, in invoke
    return __callback(*args, **kwargs)
  File "/opt/moksh/venv/lib/python3.8/site-packages/systest/src/lib/agent/mixin/torchtrain.py", line 460, in torch_train
    trained_module, profile = train(params)
  File "/opt/moksh/venv/lib/python3.8/site-packages/systest/src/lib/agent/mixin/torchtrain.py", line 346, in train
    trainer.fit(module)
  File "/opt/moksh/venv/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 770, in fit
    self._call_and_handle_interrupt(
  File "/opt/moksh/venv/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 721, in _call_and_handle_interrupt
    return self.strategy.launcher.launch(trainer_fn, *args, trainer=self, **kwargs)
  File "/opt/moksh/venv/lib/python3.8/site-packages/pytorch_lightning/strategies/launchers/subprocess_script.py", line 93, in launch
    return function(*args, **kwargs)
  File "/opt/moksh/venv/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 811, in _fit_impl
    results = self._run(model, ckpt_path=self.ckpt_path)
  File "/opt/moksh/venv/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 1236, in _run
    results = self._run_stage()
  File "/opt/moksh/venv/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 1323, in _run_stage
    return self._run_train()
  File "/opt/moksh/venv/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 1353, in _run_train
    self.fit_loop.run()
  File "/opt/moksh/venv/lib/python3.8/site-packages/pytorch_lightning/loops/base.py", line 204, in run
    self.advance(*args, **kwargs)
  File "/opt/moksh/venv/lib/python3.8/site-packages/pytorch_lightning/loops/fit_loop.py", line 266, in advance
    self._outputs = self.epoch_loop.run(self._data_fetcher)
  File "/opt/moksh/venv/lib/python3.8/site-packages/pytorch_lightning/loops/base.py", line 204, in run
    self.advance(*args, **kwargs)
  File "/opt/moksh/venv/lib/python3.8/site-packages/pytorch_lightning/loops/epoch/training_epoch_loop.py", line 208, in advance
    batch_output = self.batch_loop.run(batch, batch_idx)
  File "/opt/moksh/venv/lib/python3.8/site-packages/pytorch_lightning/loops/base.py", line 204, in run
    self.advance(*args, **kwargs)
  File "/opt/moksh/venv/lib/python3.8/site-packages/pytorch_lightning/loops/batch/training_batch_loop.py", line 88, in advance
    outputs = self.optimizer_loop.run(split_batch, optimizers, batch_idx)
  File "/opt/moksh/venv/lib/python3.8/site-packages/pytorch_lightning/loops/base.py", line 204, in run
    self.advance(*args, **kwargs)
  File "/opt/moksh/venv/lib/python3.8/site-packages/pytorch_lightning/loops/optimization/optimizer_loop.py", line 203, in advance
    result = self._run_optimization(
  File "/opt/moksh/venv/lib/python3.8/site-packages/pytorch_lightning/loops/optimization/optimizer_loop.py", line 256, in _run_optimization
    self._optimizer_step(optimizer, opt_idx, batch_idx, closure)
  File "/opt/moksh/venv/lib/python3.8/site-packages/pytorch_lightning/loops/optimization/optimizer_loop.py", line 369, in _optimizer_step
    self.trainer._call_lightning_module_hook(
  File "/opt/moksh/venv/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 1595, in _call_lightning_module_hook
    output = fn(*args, **kwargs)
  File "/opt/moksh/venv/lib/python3.8/site-packages/pytorch_lightning/core/lightning.py", line 1646, in optimizer_step
    optimizer.step(closure=optimizer_closure)
  File "/opt/moksh/venv/lib/python3.8/site-packages/pytorch_lightning/core/optimizer.py", line 168, in step
    step_output = self._strategy.optimizer_step(self._optimizer, self._optimizer_idx, closure, **kwargs)
  File "/opt/moksh/venv/lib/python3.8/site-packages/pytorch_lightning/strategies/ddp.py", line 286, in optimizer_step
    optimizer_output = super().optimizer_step(optimizer, opt_idx, closure, model, **kwargs)
  File "/opt/moksh/venv/lib/python3.8/site-packages/pytorch_lightning/strategies/strategy.py", line 193, in optimizer_step
    return self.precision_plugin.optimizer_step(model, optimizer, opt_idx, closure, **kwargs)
  File "/opt/moksh/venv/lib/python3.8/site-packages/pytorch_lightning/plugins/precision/precision_plugin.py", line 155, in optimizer_step
    return optimizer.step(closure=closure, **kwargs)
  File "/opt/moksh/venv/lib/python3.8/site-packages/torch/optim/optimizer.py", line 88, in wrapper
    return func(*args, **kwargs)
  File "/opt/moksh/venv/lib/python3.8/site-packages/torch/autograd/grad_mode.py", line 27, in decorate_context
    return func(*args, **kwargs)
  File "/opt/moksh/venv/lib/python3.8/site-packages/torch/optim/adam.py", line 100, in step
    loss = closure()
  File "/opt/moksh/venv/lib/python3.8/site-packages/pytorch_lightning/plugins/precision/precision_plugin.py", line 140, in _wrap_closure
    closure_result = closure()
  File "/opt/moksh/venv/lib/python3.8/site-packages/pytorch_lightning/loops/optimization/optimizer_loop.py", line 148, in __call__
    self._result = self.closure(*args, **kwargs)
  File "/opt/moksh/venv/lib/python3.8/site-packages/pytorch_lightning/loops/optimization/optimizer_loop.py", line 134, in closure
    step_output = self._step_fn()
  File "/opt/moksh/venv/lib/python3.8/site-packages/pytorch_lightning/loops/optimization/optimizer_loop.py", line 427, in _training_step
    training_step_output = self.trainer._call_strategy_hook("training_step", *step_kwargs.values())
  File "/opt/moksh/venv/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 1765, in _call_strategy_hook
    output = fn(*args, **kwargs)
  File "/opt/moksh/venv/lib/python3.8/site-packages/pytorch_lightning/strategies/ddp.py", line 349, in training_step
    return self.model(*args, **kwargs)
  File "/opt/moksh/venv/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1110, in _call_impl
    return forward_call(*input, **kwargs)
  File "/opt/moksh/venv/lib/python3.8/site-packages/torch/nn/parallel/distributed.py", line 963, in forward
    output = self.module(*inputs[0], **kwargs[0])
  File "/opt/moksh/venv/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1110, in _call_impl
    return forward_call(*input, **kwargs)
  File "/opt/moksh/venv/lib/python3.8/site-packages/pytorch_lightning/overrides/base.py", line 82, in forward
    output = self.module.training_step(*inputs, **kwargs)
  File "/opt/moksh/venv/lib/python3.8/site-packages/systest/src/lib/agent/mixin/ml/lit_memxfmr_wt103.py", line 35, in training_step
    outputs = self(**batch)
  File "/opt/moksh/venv/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1110, in _call_impl
    return forward_call(*input, **kwargs)
  File "/opt/moksh/venv/lib/python3.8/site-packages/systest/src/lib/agent/mixin/ml/lit_memxfmr_wt103.py", line 32, in forward
    return self.model(**inputs)
  File "/opt/moksh/venv/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1128, in _call_impl
    result = forward_call(*input, **kwargs)
TypeError: forward() got an unexpected keyword argument 'text'
Exception ignored in: <function Profiler.__del__ at 0x7fa584551670>
Traceback (most recent call last):
  File "/opt/moksh/venv/lib/python3.8/site-packages/pytorch_lightning/profiler/profiler.py", line 170, in __del__
  File "/opt/moksh/venv/lib/python3.8/site-packages/pytorch_lightning/profiler/pytorch.py", line 509, in teardown
  File "/opt/moksh/venv/lib/python3.8/site-packages/pytorch_lightning/profiler/pytorch.py", line 493, in _delete_profilers
  File "/opt/moksh/venv/lib/python3.8/site-packages/torch/profiler/profiler.py", line 456, in __exit__
  File "/opt/moksh/venv/lib/python3.8/site-packages/torch/profiler/profiler.py", line 467, in stop
  File "/opt/moksh/venv/lib/python3.8/site-packages/torch/profiler/profiler.py", line 493, in _transit_action
  File "/opt/moksh/venv/lib/python3.8/site-packages/torch/profiler/profiler.py", line 98, in start_trace
  File "/opt/moksh/venv/lib/python3.8/site-packages/torch/profiler/profiler.py", line 175, in _get_distributed_info
ImportError: sys.meta_path is None, Python is likely shutting down
Exception ignored in: <function Profiler.__del__ at 0x7fa05a25b670>
Traceback (most recent call last):
  File "/opt/moksh/venv/lib/python3.8/site-packages/pytorch_lightning/profiler/profiler.py", line 170, in __del__
  File "/opt/moksh/venv/lib/python3.8/site-packages/pytorch_lightning/profiler/pytorch.py", line 509, in teardown
  File "/opt/moksh/venv/lib/python3.8/site-packages/pytorch_lightning/profiler/pytorch.py", line 493, in _delete_profilers
  File "/opt/moksh/venv/lib/python3.8/site-packages/torch/profiler/profiler.py", line 456, in __exit__
  File "/opt/moksh/venv/lib/python3.8/site-packages/torch/profiler/profiler.py", line 467, in stop
  File "/opt/moksh/venv/lib/python3.8/site-packages/torch/profiler/profiler.py", line 493, in _transit_action
  File "/opt/moksh/venv/lib/python3.8/site-packages/torch/profiler/profiler.py", line 98, in start_trace
  File "/opt/moksh/venv/lib/python3.8/site-packages/torch/profiler/profiler.py", line 175, in _get_distributed_info
ImportError: sys.meta_path is None, Python is likely shutting down
Epoch 0:   0%|          | 0/14074 [00:04<?, ?it/s]
```


### Expected behavior

```shell
I expect the model to train.
```


### Checklist

- [X] I have read the migration guide in the readme. ([pytorch-transformers](https://github.com/huggingface/transformers#migrating-from-pytorch-transformers-to-transformers); [pytorch-pretrained-bert](https://github.com/huggingface/transformers#migrating-from-pytorch-pretrained-bert-to-transformers))
- [X] I checked if a related official extension example runs on my machine.

Comments:
Hi @moksh-enf. Could you please share the code snippet on how you are passing input while training, and also share some data examples (if possible)?

> Hi @moksh-enf. Could you please share the code snippet on how you are passing input while training, and also share some data examples (if possible)?

Hi @uahmad235 , the input is passed from a config file, where as the input data I believe is being downloaded within the code.

> Hi @moksh-enf. Could you please share the code snippet on how you are passing input while training, and also share some data examples (if possible)?

Hey @uahmad235, inputs are being passed via a config file. The data that I am using is from the datasets maintained by hugging face (wikitext-103-v1)

Hi @moksh-enf . For more context, can you post the code where you are using the `LitMemTransformerLMWT103` class ?

> Hi @moksh-enf . For more context, can you post the code where you are using the `LitMemTransformerLMWT103` class ?

Hey @uahmad235, I am passing a params file which checks if the module name and chooses the appropriate model

![image](https://user-images.githubusercontent.com/105234138/199767348-e6942409-e8c8-49d7-9d85-36931a113a86.png)

These are all the arguments that go into the params file

![image](https://user-images.githubusercontent.com/105234138/199768073-b0038980-4856-4a41-bf80-03dfbf254292.png)




Hi @moksh-enf. Could you also attach the config file you are passing? Sorry for the late response.

> Hi @moksh-enf. Could you also attach the config file you are passing? Sorry for the late response.

Hey @uahmad235 do you mean you want to know the data that is being passed to the model?

How does that help with the error?

I believe that the error is in the forward pass of the PyTorch code.

I mean the params file. Based on the information you have provided, I am unable to regenerate this issue.

> I mean the params file. Based on the information you have provided, I am unable to regenerate this issue.

Are you saying you are able to run the model from hugging face?

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.