Suppress reset_parameters of torch.nn.Linear,Conv2d... inside no_init_weights

[YouJiacheng]

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

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

[github-actions]

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 are likely to be ignored.

[jph00]

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

[LysandreJik]

May be interesting to you @ArthurZucker

[ArthurZucker]

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

[Chillee]

Why not just do

with torch.device('meta'):
     model_init()

?

[YouJiacheng]

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.

pytorch/pytorch#97951 (comment)

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

[ArthurZucker]

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

[jph00]

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.

[YouJiacheng]

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.

[ArthurZucker]

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.

[pacman100]

Hello, please also see this comment #26258 (comment)

[ArthurZucker]

Fixed by #27709 🤗