[quantization] Fix tracing for dynamic quantized LSTM #29331
Closed
Conversation
This file contains hidden or 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
[ghstack-poisoned]
|
@jamesr66a merged this pull request in f17e02f. |
|
Reverted. |
Closes #27954 This fixes the hard-coding of packed parameter values for the dynamic quantized LSTM by orchestrating the following dance: 1) Each variadic parameter on the module has its own Module. That Module defines the `__getstate__` and __setstate__` method s.t. packed weights are properly re-done on model load. 2) Each of these modules is wrapped into a `torch.nn.ModuleList`, s.t. the parameters appear as attributes in the hierarchy. Then, `gatherParametersAndBuffers` (https://github.com/pytorch/pytorch/blob/9c43b16df9dad3dfb4da1efab68d8c88e6437e8f/torch/csrc/jit/tracer.cpp#L285) can see these parameters and create a `Value*` for them in the traced graph. 3) In forward, we need to convert from ModuleList -> Module -> Parameter to a simple TensorList of the parameters. We just use a loop here. In tracing, we simply record a `ListConstruct` with each of the proper parameter values. In scripting, the `ModuleList` is const, so it can be unrolled into the graph and a subsequent `ListConstruct` does its business. The `forward` of the traced LSTM before and after this change are as follows: Before ``` def forward(self, input: Tensor, argument_2: Tuple[Tensor, Tensor]) -> Tuple[Tensor, Tuple[Tensor, Tensor]]: hx, hx0, = argument_2 _0, _1, _2 = torch.quantized_lstm(input, [hx, hx0], [CONSTANTS.c0, CONSTANTS.c1], True, 1, 0., True, False, False, dtype=12, use_dynamic=True) return (_0, (_1, _2)) ``` After ``` def forward(self, input: Tensor, argument_2: Tuple[Tensor, Tensor]) -> Tuple[Tensor, Tuple[Tensor, Tensor]]: _0 = self.cell._all_weight_values _1 = getattr(_0, "0").param _2 = getattr(_0, "1").param hx, hx0, = argument_2 _3, _4, _5 = torch.quantized_lstm(input, [hx, hx0], [_1, _2], True, 1, 0., True, False, False, dtype=12, use_dynamic=True) return (_3, (_4, _5)) ``` Differential Revision: [D18359880](https://our.internmc.facebook.com/intern/diff/D18359880) [ghstack-poisoned]
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.
Stack from ghstack:
Closes #27954
This fixes the hard-coding of packed parameter values for the dynamic quantized LSTM by orchestrating the following dance:
__getstate__and setstate` method s.t. packed weights are properly re-done on model load.torch.nn.ModuleList, s.t. the parameters appear as attributes in the hierarchy. Then,gatherParametersAndBuffers(pytorch/torch/csrc/jit/tracer.cpp
Line 285 in 9c43b16
Value*for them in the traced graph.ListConstructwith each of the proper parameter values. In scripting, theModuleListis const, so it can be unrolled into the graph and a subsequentListConstructdoes its business.The
forwardof the traced LSTM before and after this change are as follows:Before
After
Differential Revision: D18374904