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

Support loading large model weights #7610

Merged
merged 9 commits into from
Apr 24, 2023

Conversation

mattsoulanille
Copy link
Member

This is just for the lm branch. See #7609 for the real PR.

To see the logs from the Cloud Build CI, please join either our discussion or announcement mailing list.

qjia7 and others added 9 commits April 20, 2023 23:31
…#7598)

The loadWeights function loads weights in 4MB chunks and then concatenates them into a single large ArrayBuffer. That ArrayBuffer is used for splitting the weights data back up into tensors. Allocating large ArrayBuffers (3.5GB) can be unstable on Chrome, so this PR avoids this allocation, instead slicing the weights out of the chunks manually.

The implementation wraps the array of weights (stored as ArrayBuffer[]) in a new CompositeArrayBuffer class. This class implements slice by copying the desired range out of the buffer(s) that it overlaps with.
@mattsoulanille mattsoulanille merged commit a6033bb into tensorflow:lm Apr 24, 2023
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.

3 participants