Accessing input/output of unnamed functional layers via hooks

Patricsen: <p>Hey everyone,<br>
I am currently trying to use the Intel FP8 Toolkit and trying to expand it with additional formats. From what I gathered, it builds on the eager PyTorch Quantization method, which uses hooks to modify intermediate layer inputs and outputs.I am interested in static PTQ.<br>
I’ll give a little bit of background, hoping that this will help describing my question precisely:</p>
<p>I really like this approach, as it is easy to comprehend. However, when applying it to transformer models, such as Albert I am facing some issues.</p>
<p>Albert is composed of Attention heads, which can be (simplified) described as:<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/6/1/611ff9805356222f3e9392a2e25d40fbc6022cf3.png" data-download-href="https://discuss.pytorch.org/uploads/default/611ff9805356222f3e9392a2e25d40fbc6022cf3" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/6/1/611ff9805356222f3e9392a2e25d40fbc6022cf3_2_517x299.png" alt="image" data-base62-sha1="dRcOPO05i8IjzwgC1Muiw1hcLDR" width="517" height="299" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/6/1/611ff9805356222f3e9392a2e25d40fbc6022cf3_2_517x299.png, https://discuss.pytorch.org/uploads/default/original/3X/6/1/611ff9805356222f3e9392a2e25d40fbc6022cf3.png 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/6/1/611ff9805356222f3e9392a2e25d40fbc6022cf3.png 2x" data-dominant-color="EDEDEE"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">image</span><span class="informations">691×400 39.3 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>When using the <code>model.named_modules()</code> function on this model, like suggested <a href="https://discuss.pytorch.org/t/how-to-register-forward-hooks-for-each-module/43347/2">here</a> this returns (just an excerpt):</p>
<details>
<summary>
List of module names</summary>
<p>albert.encoder.albert_layer_groups.0.albert_layers.0.attention<br>
albert.encoder.albert_layer_groups.0.albert_layers.0.attention.query<br>
albert.encoder.albert_layer_groups.0.albert_layers.0.attention.key<br>
albert.encoder.albert_layer_groups.0.albert_layers.0.attention.value<br>
albert.encoder.albert_layer_groups.0.albert_layers.0.attention.attention_dropout<br>
albert.encoder.albert_layer_groups.0.albert_layers.0.attention.output_dropout<br>
albert.encoder.albert_layer_groups.0.albert_layers.0.attention.dense</p>
</details>
<details>
<summary>
Code to reprodude</summary>
<pre data-code-wrap="python"><code class="lang-python">from transformers import AlbertForQuestionAnswering

# get model
model = AlbertForQuestionAnswering.from_pretrained("twmkn9/albert-base-v2-squad2")

# list layer names and layer types
for name, module in model.named_modules():
    print(f"{name} | {type(module)}")
</code></pre>
</details>
<p>Now one could use the names of the linear modules and easily start modifiying the inputs and outputs of for example Q,K and V. But what if I wanted to modify intermediate layers, such as the output of the softmax, to quantize the input of the MatMul?</p>
<p>In the file defining Albert <a href="https://github.com/huggingface/transformers/blob/main/src/transformers/models/albert/modeling_albert.py#L342" rel="noopener nofollow ugc"><em>modeling_albert.py</em></a> we can see the softmax layer, defined as a <code>nn.functional.softmax</code>, which results in it not being created as named_module(). At least that is my understanding.</p>
<p><strong>So, in essence, my question is:</strong><br>
Is there a way to access the input/output of intermediate (unnamed) layers like these? If yes, how should this be done?</p>
<p>Thanks a lot for looking into this!</p>

HDCharles: <p>You either have to edit the model to modularize those operations, or you can’t use eager mode for that.</p>

Patricsen: <p>Thanks! Would there be a smart way to it? I would for example change <code>nn.functional.softmax</code> to <code>nn.softmax</code> and then take it from there, I guess that’s the only way?</p>
<p>I guess I am stuck with eager mode, since I want to try dtypes that are not supported by PyTorch yet, so both prototype export and graph methods are out.</p>

Patricsen: <p>Update for other readers that might run into this problem:<br>
I’ve updated the structure of for example the <a href="https://github.com/huggingface/transformers/blob/main/src/transformers/models/albert/modeling_albert.py#L249" rel="noopener nofollow ugc">Albert Attention Class</a> to include the softmax as a module. I added a softmax layer to the <code>__init__()</code>:</p>
<pre data-code-wrap="python"><code class="lang-python">self.softmax1 = nn.Softmax(dim=-1)
</code></pre>
<p>Then I replaced the actual execution in the <code>forward()</code> function with:</p>
<pre data-code-wrap="python"><code class="lang-python"># Original implementation (line 342)
attention_probs = nn.functional.softmax(attention_scores, dim=-1)

# Replaced by:
attention_probs = self.softmax1(attention_scores)
</code></pre>
<p>Although a little finicky, it works like a charm and lists the softmax gets listed with <code>model.named_modules()</code>.</p>