From f2236d1dc1f06abd2b768ed433acf7b2951252ad Mon Sep 17 00:00:00 2001 From: Giuseppe Franco Date: Tue, 21 Jan 2025 12:10:21 +0100 Subject: [PATCH] Fix (brevitas_examples/llm): equalized module + fx compatibility (#1164) --- src/brevitas_examples/llm/main.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/src/brevitas_examples/llm/main.py b/src/brevitas_examples/llm/main.py index 8c4fe1968..7567f1727 100644 --- a/src/brevitas_examples/llm/main.py +++ b/src/brevitas_examples/llm/main.py @@ -397,9 +397,13 @@ def quantize_llm(args): input_quant_format=args.input_quant_format, quantize_embedding=False) if not args.quantize_last_layer: + # Dynamo tracing changes the name of the modules, thus we need this workaround to pick + # up the last module. if require_fx: last_node = [node for node in model.graph.nodes if node.op == 'call_module'][-1] last_module = get_module(model, last_node.target) + # In case we have layerwise rotation/equalization, we need to pick the wrapped module + last_module = last_module.layer if hasattr(last_module, 'layer') else last_module last_layer_kwargs = layer_map[type(last_module)][1] prev_weight_quant = deepcopy(last_layer_kwargs['weight_quant']) prev_input_quant = deepcopy(last_layer_kwargs['input_quant'])