Refactor quantization.py to use TSP #2522
Conversation
tejapulagam
requested review from
mruberry,
lantiga and
t-vi
as code owners
for more information, see https://pre-commit.ci
There was a problem hiding this comment.
Looks great, thank you @tejapulagam .
I wonder if we can process the additional args a tad more elegantly, other than that it seems great.
| if psym.shape != csym.shape or psym.dtype != csym.dtype | ||
| } | ||
|
|
||
| new_computation_trace = trace_with_replaced_proxy_metadata(computation_trace, computation_proxy_map) |
There was a problem hiding this comment.
Do we still need this btw or could we absorb it into the QuantizationProcessor?
There was a problem hiding this comment.
I think it might be better to keep them separate because the TraceSubstitutionProcessor is only concerned with symbol-by-symbol replacement, i tried to experiment with it and ran into more errors.
thunder/transforms/quantization.py
Outdated
| ) | ||
|
|
||
| # Add new compute inputs to the trace args before processing | ||
| new_computation_trace.args = (*new_computation_trace.args, *new_compute_inputs) |
There was a problem hiding this comment.
I wonder if we can do this before the processing and also add the new bindings.
thunder/transforms/quantization.py
Outdated
| else: | ||
| new_computation_trace.bound_symbols.append(bsym.from_bsym()) | ||
| # Insert the new bindings at the beginning of the trace | ||
| new_computation_trace.bound_symbols = new_bindings + new_computation_trace.bound_symbols |
There was a problem hiding this comment.
I would prefer to keep the order of the bindings aligned with the arg order.
There was a problem hiding this comment.
not sure if this is what you had in mind
for more information, see https://pre-commit.ci
…m/tejapulagam/lightning-thunder into refactor/traceSubstitionProcessor
for more information, see https://pre-commit.ci
What does this PR do?
Refactors thunder/transforms/quantization.py to build the computation trace using TraceSubstitutionProcessor instead of manual bound-symbol rewriting.
Test test_networks.py::test_quantization is still passing.