Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[context] support lazy init of module #1088

Merged
merged 2 commits into from
Jun 10, 2022

Conversation

FrankLeeeee
Copy link
Contributor

The current ColoTensor does has not supported lazy initialization yet. This PR proposes a method for lazy initialization and it's currently experimental and subject to future changes.

In this PR, a context is implemented for lazy initialization, this context will do the following things.

  1. hijack the module initialization so that all modules are created as meta tensors. This can have two benefits 1) super fast module initialization 2) make the module traceable for future integration with torch.fx
  2. hijack the weight initialization calls and put them into cache
  3. materialize the meta tensors to real-value tensors and apply weight initialization when reset_parameters is called.

@FrankLeeeee FrankLeeeee merged commit bad5d4c into hpcaitech:main Jun 10, 2022
@FrankLeeeee FrankLeeeee deleted the feature/lazy-init-ctx branch June 13, 2022 02:07
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

2 participants