checkpointing is not still covering all cases 
## Environment info
<!-- You can run the command `transformers-cli env` and copy-and-paste its output below.
     Don't forget to fill out the missing fields in that output! -->

- `transformers` version: 4.5.0
- Platform: linux
- Python version: 3.8 
- PyTorch version (GPU?): 1.8 
- Tensorflow version (GPU?): 1.8 
- Using GPU in script?: - 
- Using distributed or parallel set-up in script?: - 

### Who can help
<!-- Your issue will be replied to more quickly if you can figure out the right person to tag with @
 If you know how to use git blame, that is the easiest way, otherwise, here is a rough guide of **who to tag**.
 Please tag fewer than 3 people.

Models:

- albert, bert, xlm: @LysandreJik
- blenderbot, bart, marian, pegasus, encoderdecoder,  t5: @patrickvonplaten, @patil-suraj
- longformer, reformer, transfoxl, xlnet: @patrickvonplaten
- fsmt: @stas00
- funnel: @sgugger
- gpt2: @patrickvonplaten, @LysandreJik
- rag: @patrickvonplaten, @lhoestq
- tensorflow: @Rocketknight1

Library:

- benchmarks: @patrickvonplaten
- deepspeed: @stas00
- ray/raytune: @richardliaw, @amogkam
- text generation: @patrickvonplaten
- tokenizers: @LysandreJik
- trainer: @sgugger
- pipelines: @LysandreJik

Documentation: @sgugger

Model hub:

- for issues with a model report at https://discuss.huggingface.co/ and tag the model's creator.

HF projects:

- datasets: [different repo](https://github.com/huggingface/datasets)
- rust tokenizers: [different repo](https://github.com/huggingface/tokenizers)

Examples:

- maintained examples (not research project or legacy): @sgugger, @patil-suraj
- research_projects/bert-loses-patience: @JetRunner
- research_projects/distillation: @VictorSanh

 -->
@sgugger 

## Information
Hi
sometime ago, you improved the checkopinting in huggingface repo, but this is not covering all cases, here is one example, let assume a user is wrapping a model into a class as below:
```
model = .... // one of huggingface model here 
# lets wrap it
model = intrinsic_dimension_said(model, intrinsic_dim, training_args.output_dir, set())
```

I put the wrap class for completeness, but feel free to ignore it, this can be any class:

```
class IntrinsicDimensionLight:
    def __init__(self, module: nn.Module, intrinsic_dimension: int,  output_dir, str_filter: Set[str] = set(), said=False, random_seed=1997):
        torch.manual_seed(random_seed)
        np.random.seed(random_seed)

        self.initial_value_path = os.path.join(output_dir, "initial_value")
        self.fastfood_params_path = os.path.join(output_dir, "fastfood_params")


        self.name_base_localname = []
        self.initial_value = dict()
        self.fastfood_params = {}
        self.said = said
        self.said_size = len(list(module.named_parameters()))
        if self.said:
            assert intrinsic_dimension > self.said_size
            intrinsic_dimension -= self.said_size

        self.intrinsic_parameter = nn.Parameter(
            torch.zeros((intrinsic_dimension)).cpu())
        module.register_parameter(
            "intrinsic_parameter", self.intrinsic_parameter)
        setattr(module, "intrinsic_parameter", self.intrinsic_parameter)

        length = 0
        for name, param in module.named_parameters():
            if param.requires_grad and all([x not in name for x in str_filter]):
                length += 1
                self.initial_value[name] = v0 = (
                    param.clone().detach().requires_grad_(False).to(self.intrinsic_parameter.device)
                )
                DD = np.prod(v0.size())
                self.fastfood_params[name] = fastfood_vars(
                   DD, self.intrinsic_parameter.device)
                base, localname = module, name
                while "." in localname:
                    prefix, localname = localname.split(".", 1)
                    base = base.__getattr__(prefix)
                self.name_base_localname.append((name, base, localname))
                if "intrinsic_parameter" not in name:
                    param.requires_grad_(False)
        if said:
            self.intrinsic_parameter_said = nn.Parameter(
                torch.ones((length)).cpu())
            module.register_parameter(
                "intrinsic_parameter_said", self.intrinsic_parameter_said)
            setattr(module, "intrinsic_parameter_said",
                    self.intrinsic_parameter_said)

        # If this is created before, here we save it and here it loads it.
        if not self.is_projection_params_saved():
            self.save_required_params()
        self.load_required_params()

    def is_projection_params_saved(self):
        return  os.path.isfile(self.fastfood_params_path) and\
                os.path.isfile(self.initial_value_path)

    def load_required_params(self):
        # check and if intrinsic porjection mats exists load them.
        if self.is_projection_params_saved():
            self.fastfood_params = torch.load(self.fastfood_params_path)
            self.initial_value = torch.load(self.initial_value_path)
        
    def save_required_params(self):
        # Saves the generates projection params.
        torch.save(self.initial_value, self.initial_value_path)
        torch.save(self.fastfood_params, self.fastfood_params_path)

    def move_to(self, x_tuple, target):
        if isinstance(x_tuple, torch.Tensor):
            return x_tuple.to(target)
        a = []
        for x in x_tuple:
            if isinstance(x, torch.Tensor):
                a.append(x.to(target))
            else:
                a.append(x)
        return tuple(a)

    def requires_to(self, x_tuple, target):
        if isinstance(x_tuple, torch.Tensor):
            x_tuple.requires_grad_(target)
        for x in x_tuple:
            if isinstance(x, torch.Tensor):
                x.requires_grad_(target)

    def fastfood_vars_requires_grad_(self, requires_grad):
        for item in self.fastfood_params.items():
            self.requires_to(item, requires_grad)

    def __call__(self, module, inputs):
        index = 0
        with torch.enable_grad():
            for name, base, localname in self.name_base_localname:
                if localname == "intrinsic_parameter":
                    continue
                self.initial_value[name] = self.initial_value[name].to(
                    getattr(base, localname))
                device_dtype = getattr(base, localname).dtype
                init_shape = self.initial_value[name].size()

                DD = np.prod(init_shape)
                self.fastfood_params[name] = self.move_to(
                    self.fastfood_params[name], module.intrinsic_parameter.device)

                # Fastfood transform te replace dence P
                ray = fastfood_torched(module.intrinsic_parameter, DD, self.fastfood_params[name]).view(
                    init_shape
                )
                if self.said:
                    ray = ray * self.intrinsic_parameter_said[index]
                param = (self.initial_value[name] + ray).to(device_dtype)
                delattr(base, localname)
                setattr(base, localname, param)
                index += 1

    @staticmethod
    def apply(module, intrinsic_dimension, output_dir, str_filter=set(), said=False):
        for k, hook in module._forward_pre_hooks.items():
            if isinstance(hook, IntrinsicDimensionLight) and hook.name == name:
                raise RuntimeError("Cannot register two intrinsic dimension hooks on "
                                   "the same parameter {}".format(name))

        fn = IntrinsicDimensionLight(
            module, intrinsic_dimension, output_dir, str_filter, said)
        module.register_forward_pre_hook(fn)
        return fn

    @staticmethod
    def apply_with_tensor(module, intrinsic_vector, str_filter=set()):
        assert isinstance(intrinsic_vector,
                          torch.Tensor) and intrinsic_vector.ndim == 1

        for k, hook in module._forward_pre_hooks.items():
            if isinstance(hook, IntrinsicDimensionLight) and hook.name == name:
                raise RuntimeError("Cannot register two intrinsic dimension hooks on "
                                   "the same parameter {}".format(name))

        fn = IntrinsicDimensionLight(
            module, intrinsic_vector.size(0), str_filter, False)
        fn.intrinsic_parameter = intrinsic_vector
        module.register_forward_pre_hook(fn)
        return fn


def intrinsic_dimension(module, intrinsic_dimension,  output_dir, str_filter):
    IntrinsicDimensionLight.apply(
        module, intrinsic_dimension,  output_dir, str_filter, False)
    return module


def intrinsic_dimension_said(module, intrinsic_dimension,  output_dir, str_filter):
    IntrinsicDimensionLight.apply(
        module, intrinsic_dimension,  output_dir, str_filter, True)
    return module
```

Now if you see the model after reloading in trainer, this is without this wrapper class, could you also cover this case that user might wrap a model into a class?

thanks 

## Expected behavior

wrapped model needs to be reloaded fine in checkpoint reloading