[docs][pytorch] Add examples for compiling with external weights. #18658
Conversation
ScottTodd
added
documentation
| @@ -322,6 +322,132 @@ their values independently at runtime. | |||
| self.value = new_value | |||
| ``` | |||
|
|
|||
| #### :octicons-code-16: Separating parameters during compilation | |||
|
|
|||
| In practical scenarios, it is usually unfeasible to embed model parameters | |||
There was a problem hiding this comment.
nit: infeasible is more common.
nit: instead of "In practical scenarios", maybe put "For large models", since this is the practical scenario where this applies.
| Safetensors can be created directly from a parameters dictionary created from | ||
| randomly generated torch tensors. | ||
|
|
||
| ```python | ||
| # Create a params dictionary. Note that the keys here match | ||
| # LinearModule's attributes. We will use the saved safetensor | ||
| # file for use from the command line. | ||
| from safetensors.torch import save_file | ||
|
|
||
| wt = torch.randn(4,3) | ||
| bias = torch.randn(3) | ||
| params = { "weight": wt, "bias": bias } | ||
| save_file(params, "params.safetensors") | ||
| ``` |
There was a problem hiding this comment.
I'm finding this section misleading from the example above. The linear_module already has specific weight and bias parameters. Why are we storing new random ones?
There was a problem hiding this comment.
Just as an example! We might not always have the class that actually represents the model, but just some representation of the graph. This is a demonstration of how, in such a case, one can load the weights separately.
There was a problem hiding this comment.
@zjgarvey I changed the examples to use the modules parameters.
| # tensors from the IR module, allowing them to be loaded | ||
| # at runtime. Symbolic references to these parameters are | ||
| # still retained in the IR. | ||
| aot.externalize_module_parameters(linear_module) |
There was a problem hiding this comment.
Where does this externalize the weights to? Does the file have to be called params.safetensors in the PWD or something?
There was a problem hiding this comment.
Reading further on, it looks like you specify the params file when invoking iree-run-module. I suppose the symbolic reference is to the name of the parameter in the safetensors dictionary you saved earlier?
There was a problem hiding this comment.
No! It can be any path you specify. Here it is this file in PWD.
There was a problem hiding this comment.
Reading further on, it looks like you specify the params file when invoking iree-run-module. I suppose the symbolic reference is to the name of the parameter in the safetensors dictionary you saved earlier?
Yes, and under the scope named "model".
| Note here that the `--parameters` flag has `model=` following it immediately. | ||
| This simply specifies the scope of the parameters, and is reflected in the | ||
| compiled IR. |
There was a problem hiding this comment.
I'm not sure I understand this. In what situations are parameters not model wide?
There was a problem hiding this comment.
model is a scope here. It could actually be anything. If is even possible that the biases could be under a scope called "b", and weights under a scope called "a". In that case, they would be in different indices, and model= would be then b= and a= respectively.
|
Thanks for doing this. Once landed, can someone make sure to add a test to iree-turbine which does this exact thing? |
|
@stellaraccident Sure! I'd be happy to do that. |
vinayakdsci
force-pushed
the
docs-st-examples
branch
from
d3cc3a7
to
0d91d52
Compare
Progress on #18564.
Adds examples to the PyTorch guide, showing how to externalize module parameters, and load them at runtime, both through command line (
iree-run-module) and through the iree-runtime Python API (usingParameterIndex).