Suppress `reset_parameters` of `torch.nn.Linear,Conv2d...` inside `no_init_weights`
Author: YouJiacheng
Pull request: no

### Feature request

`torch.nn.Linear,Conv2d...` will call `self.reset_parameters()` inside their `__init__`.
I'd like to make `reset_parameters` be a no-op inside `no_init_weights` context manager.

### Motivation

`no_init_weights` is used in `from_pretrained` to speed up loading large models.
However, torch-built-in modules like `torch.nn.Linear` are heavily used in models of `transformers`, while its weights initialization cannot be disabled by `no_init_weights`.
And in the doc string of `no_init_weights`, it should "globally disable weight initialization".

### Your contribution

possible implementation
```python
class SupportsResetParameters(Protocol):
    def reset_parameters(self): ...

@contextmanager
def no_init(module_classes: Iterable[Type[SupportsResetParameters]]):
    saved = {m: vars(m).get('reset_parameters') for m in module_classes}
    def no_op(_): pass
    for m in saved: m.reset_parameters = no_op # Iterable can only be safely iterated through once
    try:
        yield
    finally:
        for m, init in saved.items():
            del m.reset_parameters
            if init is not None:
                m.reset_parameters = init

TORCH_BUILT_IN_MODULES = [nn.Linear, nn.Conv2d, ...]

@contextmanager
def no_init_weights():
    """
    Context manager to globally disable weight initialization to speed up loading large models.
    """
    global _init_weights
    saved = _init_weights
    _init_weights = False
    try:
        with no_init(TORCH_BUILT_IN_MODULES):
            yield
    finally:
        _init_weights = saved
```

Comments:
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.

@younesbelkada @pacman100 Can I suggest reopening this issue - it's pretty important IMO and it's hitting lots of people.

May be interesting to you @ArthurZucker 

Yes! I'll take this one, makes sense 😉 

Why not just do
```
with torch.device('meta'):
     model_init()
```
?

Yes, meta device or faketensor is the correct and recommended choice for deferring/skipping initialization.

https://pytorch.org/torchdistx/latest/fake_tensor_and_deferred_init.html

https://huggingface.co/blog/accelerate-large-models

But `with torch.device('meta')` requires pytorch 2.0 IIUC.

https://github.com/pytorch/pytorch/issues/97951#issuecomment-1490417946

TBH, I didn't know meta device when I posted this issue (2022). I knew meta device this year.

Oups sorry for the delay! Have not forgotten about this 😉 

> Yes, meta device or faketensor is the correct and recommended choice for deferring/skipping initialization.

meta device doesn't solve the problem, because buffers (e.g sin/cos in llama2) don't get initialized in that case.

> > Yes, meta device or faketensor is the correct and recommended choice for deferring/skipping initialization.
> 
> meta device doesn't solve the problem, because buffers (e.g sin/cos in llama2) don't get initialized in that case.

You are right, meta device would skip all initializations, builtin or user-defined, parameters or buffers.
But as long as those buffers are not marked as persistent=False, it can be loaded from the checkpoint.


We have a lot of buffers with persistant = False. I answered on the other issue but I’ll most probably go about this with skipping initialization for all layers but the ones that are missing in the checkpoints.


Hello, please also see this comment https://github.com/huggingface/transformers/issues/26258#issuecomment-1829210754

Fixed by #27709 🤗