ammo.torch.quantization

Sbisseb: <p>I want to use the quantization on my custom LLM model, in my my <code>quantization.py</code></p>
<pre><code class="lang-auto">import torch
from torch.utils.data import DataLoader, TensorDataset
import ammo.torch.quantization as atq
import numpy as np

config = atq.FP8_DEFAULT_CFG
model = torch.load("model.pt")
calibration_Data = torch.from_numpy(np.load("data.npy"))
calib_dataloader= DataLoader(calibration_Data, batch_size=10, shuffle=False)
def calibrate_loop():
    for data_c in calib_dataloader:
        model(input_1=data_c) 
with torch.no_grad():
    atq.quantize(model, config, forward_loop=calibrate_loop)
</code></pre>
<p>I got the following output/error, any ideas how to fix that?</p>
<pre><code class="lang-auto">Replaced 54 modules to quantized modules
Traceback (most recent call last):
File "quantization.py", line 22, in
atq.quantize(model, config, forward_loop=calibrate_loop)
File "/lib/python3.10/dist-packages/ammo/torch/quantization/model_quant.py", line 112, in quantize
calibrate(model, config["algorithm"], forward_loop=forward_loop)
File "ammo/torch/quantization/model_calib.py", line 63, in ammo.torch.quantization.model_calib.calibrate
File "/lib/python3.10/dist-packages/torch/utils/_contextlib.py", line 115, in decorate_context
return func(*args, **kwargs)
File "ammo/torch/quantization/model_calib.py", line 73, in ammo.torch.quantization.model_calib.max_calibrate
File "quantization.py", line 19, in calibrate_loop
model(input_1=data_c)
File "/lib/python3.10/dist-packages/torch/fx/graph_module.py", line 736, in call_wrapped
return self._wrapped_call(self, *args, **kwargs)
File "/lib/python3.10/dist-packages/torch/fx/graph_module.py", line 315, in **call**
raise e
File "/lib/python3.10/dist-packages/torch/fx/graph_module.py", line 302, in **call**
return super(self.cls, obj).**call**(*args, **kwargs)
File "/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1510, in _wrapped_call_impl
return self._call_impl(*args, **kwargs)
File "/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1519, in _call_impl
return forward_call(*args, **kwargs)
File "&lt;eval_with_key&gt;.1 from &lt;eval_with_key&gt;.0:10 in forward", line 1625, in forward
File "/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1510, in _wrapped_call_impl
return self._call_impl(*args, **kwargs)
File "/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1519, in _call_impl
return forward_call(*args, **kwargs)
File "/lib/python3.10/dist-packages/onnx2torch/node_converters/reduce.py", line 158, in forward
return torch.sum(input_tensor, dim=self._axes, keepdim=self._keepdims)
TypeError: sum() received an invalid combination of arguments - got (Tensor, keepdim=int, dim=list), but expected one of:

* (Tensor input, *, torch.dtype dtype)
didn't match because some of the keywords were incorrect: keepdim, dim
* (Tensor input, tuple of ints dim, bool keepdim, *, torch.dtype dtype, Tensor out)
* (Tensor input, tuple of names dim, bool keepdim, *, torch.dtype dtype, Tensor out)
</code></pre>
<h1><a name="system-info-1" class="anchor" href="#system-info-1"></a>System Info</h1>
<ul>
<li>Nvidia GPUs H100</li>
<li>python3.10</li>
<li>ammo version 0.7.3</li>
<li>torch 2.2.0a0+81ea7a4</li>
</ul>

ptrblck: <p>Based on the error message it seems the <code>keepdim</code> argument was wrongly passed as an <code>int</code>. If the error comes from the library, could you post an issue in the corresponding GitHub repository?</p>

Sbisseb: <p>where do you mean <code>keepdim</code> argument should be passed here? It can’t be passed in the <code>calibration_loop</code>, if passed then will get:</p>
<pre><code class="lang-auto">/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1519, in _call_impl
    return forward_call(*args, **kwargs)
TypeError: forward() got an unexpected keyword argument 'keepdim'
</code></pre>

ptrblck: <p>From the stacktrace:</p>
<pre data-code-wrap="python"><code class="lang-python">File "/lib/python3.10/dist-packages/onnx2torch/node_converters/reduce.py", line 158, in forward
return torch.sum(input_tensor, dim=self._axes, keepdim=self._keepdims)
TypeError: sum() received an invalid combination of arguments - got (Tensor, keepdim=int, dim=list), but expected one of:
</code></pre>
<p>the <code>torch.sum</code> operation fails as <code>self._keepdims</code> is wrong coming from <code>onnx2torch</code>.</p>

Sbisseb: <p>would you have any idea how to fix that?</p>

ptrblck: <p>It seems this is a known issue and tracked <a href="https://github.com/ENOT-AutoDL/onnx2torch/issues/173">here</a>. However, I don’t see any movement in the issue so you might need to fix it manually by converting the passed <code>int</code> to a <code>bool</code> in <code>/lib/python3.10/dist-packages/onnx2torch/node_converters/reduce.py</code>.</p>

Sbisseb: <p>Many thanks <a class="mention" href="/u/ptrblck">@ptrblck</a> . I am also using <code>return torch.sum(input_tensor, dim=self._axes, keepdim=bool(self._keepdims))</code> but everytime I make a change then a lot of other scripts need to be adapted. For instance I am now getting issues in <code>onnx2torch/node_converters/gather.py</code> in the slicing I am not sure why</p>
<pre><code class="lang-auto">return input_tensor[slice_for_take]
IndexError: index 4 is out of bounds for dimension 0 with size 3
</code></pre>
<p>Would you have any idea how to fix that?</p>