Improve torch formatting performance
@vblagoje I tested my proposed approach before posting it here and it worked for me. 

Is it not working in your case because of the SSH protocol? In that case you could try the same approach but using HTTPS:
```
"datasets @ git+https://github.com/huggingface/datasets.git@refs/pull/2505/head",
``` 
**Is your feature request related to a problem? Please describe.**
It would be great, if possible, to further improve read performance of raw encoded datasets and their subsequent conversion to torch tensors. 

A bit more background.  I am working on LM pre-training using HF ecosystem. We use encoded HF Wikipedia and BookCorpus datasets. The training machines are similar to DGX-1 workstations. We use HF trainer torch.distributed training approach on a single machine with 8 GPUs.

The current performance is about 30% slower than NVidia optimized BERT [examples](https://github.com/NVIDIA/DeepLearningExamples/tree/master/PyTorch/LanguageModeling) baseline. Quite a bit of customized code and training loop tricks were used to achieve the baseline performance. It would be great to achieve the same performance while using nothing more than off the shelf HF ecosystem. Perhaps, in the future, with @stas00 work on deepspeed integration, it could even be exceeded. 

**Describe the solution you'd like**
Using profiling tools we've observed that appx. 25% of cumulative run time is spent on data loader next call.
![dataloader_next](https://user-images.githubusercontent.com/458335/121895543-59742a00-ccee-11eb-85fb-f07715e3f1f6.png)

As you can observe most of the data loader next call is spent in HF datasets torch_formatter.py format_batch call. 

Digging a bit deeper into format_batch we can see the following profiler data:
![torch_formatter](https://user-images.githubusercontent.com/458335/121895944-c7b8ec80-ccee-11eb-95d5-5875c5716c30.png)

Once again, a lot of time is spent in pyarrow table conversion to pandas which seems like an intermediary step. Offline @lhoestq told me that this approach was, for some unknown reason, faster than direct to numpy conversion. 

**Describe alternatives you've considered**
I am not familiar with pyarrow and have not yet considered the alternatives to the current approach. 

Most of the online advice around data loader performance improvements revolve around increasing number of workers, using pin memory for copying tensors from host device to gpus but we've already tried these avenues without much performance improvement.  Weights & Biases dashboard for the pre-training task reports CPU utilization of ~ 10%, GPUs are completely saturated (GPU utilization is above 95% on all GPUs), while disk utilization is above 90%. 


Improve torch formatting performance 
 **Is your feature request related to a problem? Please describe.**
It would be great, if possible, to further improve read performance of raw encoded datasets and their subsequent conversion to torch tensors. 

A bit more background.  I am working on LM pre-training using HF ecosystem. We use encoded HF Wikipedia and BookCorpus datasets. The training machines are similar to DGX-1 workstations. We use HF trainer torch.distributed training approach on a single machine with 8 GPUs.

The current performance is about 30% slower than NVidia optimized BERT [examples](https://github.com/NVIDIA/DeepLearningExamples/tree/master/PyTorch/LanguageModeling) baseline. Quite a bit of customized code and training loop tricks were used to achieve the baseline performance. It would be great to achieve the same performance while using nothing more than off the shelf HF ecosystem. Perhaps, in the future, with @stas00 work on deepspeed integration, it could even be exceeded. 

**Describe the solution you'd like**
Using profiling tools we've observed that appx. 25% of cumulative run time is spent on data loader next call.
![dataloader_next](https://user-images.githubusercontent.com/458335/121895543-59742a00-ccee-11eb-85fb-f07715e3f1f6.png)

As you can observe most of the data loader next call is spent in HF datasets torch_formatter.py format_batch call. 

Digging a bit deeper into format_batch we can see the following profiler data:
![torch_formatter](https://user-images.githubusercontent.com/458335/121895944-c7b8ec80-ccee-11eb-95d5-5875c5716c30.png)

Once again, a lot of time is spent in pyarrow table conversion to pandas which seems like an intermediary step. Offline @lhoestq told me that this approach was, for some unknown reason, faster than direct to numpy conversion. 

**Describe alternatives you've considered**
I am not familiar with pyarrow and have not yet considered the alternatives to the current approach. 

Most of the online advice around data loader performance improvements revolve around increasing number of workers, using pin memory for copying tensors from host device to gpus but we've already tried these avenues without much performance improvement.  Weights & Biases dashboard for the pre-training task reports CPU utilization of ~ 10%, GPUs are completely saturated (GPU utilization is above 95% on all GPUs), while disk utilization is above 90%. 

 
 @vblagoje I tested my proposed approach before posting it here and it worked for me. 

Is it not working in your case because of the SSH protocol? In that case you could try the same approach but using HTTPS:
```
"datasets @ git+https://github.com/huggingface/datasets.git@refs/pull/2505/head",
``` 