Stream weights to the GPU when loading a model #7994
Merged
Conversation
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
mattsoulanille
force-pushed
the
weights_stream
branch
from
bf9c6fa
to
af89cf0
Compare
mattsoulanille
force-pushed
the
weights_stream
branch
from
5c93720
to
ed74c0e
Compare
pyu10055
approved these changes
Nov 16, 2023
tfjs-core/src/io/http.ts
Outdated
| return fetchURLs; | ||
| } | ||
|
|
||
| // private get loadOptions(): LoadOptions { |
| * Whether to stream the model directly to the backend or cache all its | ||
| * weights on CPU. Useful for large models. | ||
| */ | ||
| streamWeights?: boolean; |
There was a problem hiding this comment.
this flag is the same name as the function streamWeights?: () => ReadableStream<ArrayBuffer>?
There was a problem hiding this comment.
Yes, but they're part of different interfaces. One is on the ModelArtifacts and is the actual function that, when called, will start streaming the weights. It's internal, and users shouldn't really need to access it.
The other is on the LoadOptions interface, and it's used to configure whether the model loader should stream the weights or load them normally. It's exposed to the user.
I can rename the one on ModelArtifacts to something like startWeightsStream or streamWeightsData.
This was referenced May 31, 2024
Closed
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
When downloading model weight data, slice it into weight tensors and push them to the GPU eagerly. This avoids storing an extra copy of the weights on CPU, allowing for larger models (1.3B to possibly ~6.7B or larger) to be loaded without causing a V8 OOM crash.
When streaming the weights, check
CPU_HANDOFF_SIZE_THRESHOLDorWEBGPU_CPU_HANDOFF_SIZE_THRESHOLDto determine whether the weight should be sent to GPU or remain on CPU.This feature is guarded by the
streamWeightsoption inLoadOptions. Since most of TFJS's graph model saving relies on the CPU copy of the model, model saving is disabled when the model was streamed (i.e. it will throw an error since the weights ArrayBuffer is missing).To see the logs from the Cloud Build CI, please join either our discussion or announcement mailing list.