Edit train_dataloader in callbacks
Author: arshadshk
Pull request: no
Labels: name: Feature request

In certain scenarios, it becomes necessary to modify the train_dataloader after each epoch. This can be managed through callbacks since the train dataloader is passed to the callback function. However, it's worth noting that the dataloader for training is not a variable within the base Trainer class. Instead, an [independent variable called train_dataloader](https://github.com/huggingface/transformers/blob/v4.32.1/src/transformers/trainer.py#L1569) is created at the beginning of training. Unfortunately, when we attempt to make changes to this dataloader within callbacks, it doesn't seem to have the desired effect.

Here's a code example illustrating this issue:
```python
class CustomCallback(TrainerCallback):
    def on_epoch_begin(self, args, state, control, model, tokenizer, **kwargs):
        kwargs['train_dataloader'] = new_dataloader
```
The problem appears to be that even though we try to update train_dataloader within the callback, it doesn't affect the dataloader as expected.

Please share a reference to the data loader utilized for training in the callback, rather than a duplicate, to allow for modifications to the data loader within callbacks. Currently, there is no other method available to make changes to the data loader after each epoch or after a certain number of steps.

To make the train_dataloader editable by callbacks, it would be beneficial to associate it directly with the Trainer class, rather than creating it as an independent variable. This would enable callbacks to modify it effectively, as the current implementation creates a copy that is not affected by changes made within callbacks.


### Who can help?

@muellerzr @pacman100 @Narsil 

### 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
class CustomCallback(TrainerCallback):
    def on_epoch_begin(self, args, state, control, model, tokenizer, **kwargs):
        kwargs['train_dataloader'] = new_dataloader
```

### Expected behavior

The dataloader used for training should be replaced by the new_dataloader passed in callbacks

Comments:
What expectation of modifying the dataloader are you expecting? `Trainer` uses it's own `get_train_dataloader`, `get_eval_dataloader`, and `get_test_dataloader` functions and has for a while, so what specific modifications to the dataloader you are looking for would be good to know, else this may require some substantial rewrites of a few areas. 

Most notably: dataloaders are sent through `accelerator.prepare`, so the original dataloader is one modified for distributed systems, and really shouldn't be touched much. 

@muellerzr  for sometimes, we might want to use https://github.com/imoneoi/multipack_sampler/blob/master/multipack_sampler.py#L94 for efficient reasoning 

So we care about a custom sampler to use instead here? 

> So we care about a custom sampler to use instead here?

Yeah. I might need to open a new feature request instead of putting things here

A new feature request would be good @dumpmemory, as that's very different. Honestly I'd open that on the Accelerate side probably since it's what those dataloaders use in DDP. @arshadshk can you provide more details on what you want to modify in the dataloader? We have a lot of preprocessing etc we need to do on the dataloader, which is why we keep it that way in general. More clarity on what you are specifically wanting to do here with it would be great. 

We also wish that we can have this new feature.
We found that we may not be able to modify the `IterableDataset` passed to Trainer then create the train dataloader.
We probably need to update it on every epoch or the train loss would get suddenly dropped and increase every new epoch.
Please let me know if there is a solution for this.