<h1 dir="auto"><g-emoji class="g-emoji" alias="rocket" fallback-src="https://github.githubassets.com/images/icons/emoji/unicode/1f680.png">🚀</g-emoji> Feature request</h1>
<p dir="auto">Pytorch supports mimicking quantization errors while training the models.<br>
Here is the <a href="https://pytorch.org/tutorials/recipes/quantization.html#quantization-aware-training" rel="nofollow">tutorial</a> on this. For our NLP transformers, it requires a "fake quantization" operation to be done on the embeddings. I found this <a href="https://github.com/IntelLabs/nlp-architect/blob/0f00215dcaf81f8a9b296035834310d77015f085/nlp_architect/models/transformers/quantized_bert.py">repository</a> converting BERT to support this.</p>
<h2 dir="auto">Motivation</h2>
<p dir="auto">I think quantization aware fine-tuning (if it works) will help a lot of use-cases where dynamic quantization alone doesn't suffice in maintaining the performance of the quantized model. Supporting it out of the box will remove the duplication of model code in end use cases.</p>
<h2 dir="auto">Your contribution</h2>
<p dir="auto">I can work on this ASAP. Would appreciate initial thoughts on what a the MVP for it would be, any thoughts on the API (should we take in a "qat" boolean in config?), any pitfalls that I should be aware of, etc.</p>
<h2 dir="auto">Environment info</h2>
<ul dir="auto">
<li><code class="notranslate">transformers</code> version: 4.17.0.dev0</li>
<li>Platform: Linux-5.13.0-27-generic-x86_64-with-glibc2.10</li>
<li>Python version: 3.8.0</li>
<li>PyTorch version (GPU?): 1.10.1 (True)</li>
<li>Tensorflow version (GPU?): not installed (NA)</li>
<li>Flax version (CPU?/GPU?/TPU?): not installed (NA)</li>
<li>Jax version: not installed</li>
<li>JaxLib version: not installed</li>
<li>Using GPU in script?: yes</li>
<li>Using distributed or parallel set-up in script?: yes (deepspeed)</li>
<li>Note: I installed DeepSpeed from source</li>
</ul>
<h3 dir="auto">Who can help</h3>
<p dir="auto">Models:<br>
(I'm actually trying to use T0pp but T5 is close enough)</p>
<ul dir="auto">
<li>T5, BART, Marian, Pegasus, EncoderDecoder: <a class="user-mention notranslate" data-hovercard-type="user" data-hovercard-url="/users/patrickvonplaten/hovercard" data-octo-click="hovercard-link-click" data-octo-dimensions="link_type:self" href="https://github.com/patrickvonplaten">@patrickvonplaten</a></li>
</ul>
<p dir="auto">Library:</p>
<ul dir="auto">
<li>Deepspeed: <a class="user-mention notranslate" data-hovercard-type="user" data-hovercard-url="/users/stas00/hovercard" data-octo-click="hovercard-link-click" data-octo-dimensions="link_type:self" href="https://github.com/stas00">@stas00</a></li>
<li>Text generation: <a class="user-mention notranslate" data-hovercard-type="user" data-hovercard-url="/users/patrickvonplaten/hovercard" data-octo-click="hovercard-link-click" data-octo-dimensions="link_type:self" href="https://github.com/patrickvonplaten">@patrickvonplaten</a> <a class="user-mention notranslate" data-hovercard-type="user" data-hovercard-url="/users/Narsil/hovercard" data-octo-click="hovercard-link-click" data-octo-dimensions="link_type:self" href="https://github.com/Narsil">@Narsil</a></li>
</ul>
<h2 dir="auto">Information</h2>
<p dir="auto">Model I am using (Bert, XLNet ...): T0pp / T0_3B</p>
<p dir="auto">The problem arises when using:</p>
<ul class="contains-task-list">
<li class="task-list-item"><input type="checkbox" id="" disabled="" class="task-list-item-checkbox"> the official example scripts: (give details below)</li>
<li class="task-list-item"><input type="checkbox" id="" disabled="" class="task-list-item-checkbox" checked=""> my own modified scripts: (give details below)</li>
</ul>
<p dir="auto">The tasks I am working on is:</p>
<ul class="contains-task-list">
<li class="task-list-item"><input type="checkbox" id="" disabled="" class="task-list-item-checkbox"> an official GLUE/SQUaD task: (give the name)</li>
<li class="task-list-item"><input type="checkbox" id="" disabled="" class="task-list-item-checkbox" checked=""> my own task or dataset: (give details below)</li>
</ul>
<h2 dir="auto">To reproduce</h2>
<p dir="auto">I want to load T0pp across 2 24GB GPUs and only run inference. I know Deepspeed wit zeRO stage 3 is the way to go for this from reading documentation. I am following the HuggingFace example <a href="https://huggingface.co/docs/transformers/v4.16.1/en/main_classes/deepspeed#deepspeed-non-trainer-integration" rel="nofollow">here</a> to use Deepspeed without a <code class="notranslate">Trainer</code> object.</p>
<p dir="auto">The error I get is</p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="[2022-01-28 18:36:41,193] [INFO] [partition_parameters.py:456:__exit__] finished initializing model with 2.85B parameters
Traceback (most recent call last):
  File &quot;multi_gpu_T0pp.py&quot;, line 26, in &lt;module&gt;
    engine = deepspeed.initialize(model=model, config_params=ds_config)
AttributeError: module 'transformers.deepspeed' has no attribute 'initialize'"><pre class="notranslate"><code class="notranslate">[2022-01-28 18:36:41,193] [INFO] [partition_parameters.py:456:__exit__] finished initializing model with 2.85B parameters
Traceback (most recent call last):
  File "multi_gpu_T0pp.py", line 26, in &lt;module&gt;
    engine = deepspeed.initialize(model=model, config_params=ds_config)
AttributeError: module 'transformers.deepspeed' has no attribute 'initialize'
</code></pre></div>
<p dir="auto">My code:</p>
<p dir="auto">Run with <code class="notranslate">CUDA_VISIBLE_DEVICES="0,1" deepspeed &lt;script.py&gt;</code></p>
<div class="highlight highlight-source-python notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="&quot;&quot;&quot;
Example code to load a PyTorch model across GPUs
&quot;&quot;&quot;
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
from transformers.deepspeed import HfDeepSpeedConfig
from transformers import deepspeed
import pandas as pd
import torch
import pdb
import os

seed = 42
torch.manual_seed(seed)

ds_config = {
    &quot;fp16&quot;: {
        &quot;enabled&quot;: &quot;auto&quot;,
        &quot;loss_scale&quot;: 0,
        &quot;loss_scale_window&quot;: 1000,
        &quot;initial_scale_power&quot;: 16,
        &quot;hysteresis&quot;: 2,
        &quot;min_loss_scale&quot;: 1
    },
    &quot;zero_optimization&quot;: {
        &quot;stage&quot;: 3,
        &quot;overlap_comm&quot;: true,
        &quot;contiguous_gradients&quot;: true,
        &quot;sub_group_size&quot;: 1e9,
        &quot;reduce_bucket_size&quot;: &quot;auto&quot;,
        &quot;stage3_prefetch_bucket_size&quot;: &quot;auto&quot;,
        &quot;stage3_param_persistence_threshold&quot;: &quot;auto&quot;,
        &quot;stage3_max_live_parameters&quot;: 1e9,
        &quot;stage3_max_reuse_distance&quot;: 1e9,
        &quot;stage3_gather_fp16_weights_on_model_save&quot;: true
    },
    &quot;gradient_accumulation_steps&quot;: 1,
    &quot;gradient_clipping&quot;: 0,
    &quot;steps_per_print&quot;: 2000,
    &quot;train_batch_size&quot;: 2,
    &quot;train_micro_batch_size_per_gpu&quot;: 1,
    &quot;wall_clock_breakdown&quot;: false
}

if __name__ == &quot;__main__&quot;:
    # must run before instantiating the model
    # ds_config is deepspeed config object or path to the file
    dschf = HfDeepSpeedConfig(ds_config)  # keep this object alive

    model_name = &quot;bigscience/T0_3B&quot;
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    model = AutoModelForSeq2SeqLM.from_pretrained(model_name)

    engine = deepspeed.initialize(model=model, config_params=ds_config)

    inputs = tokenizer.encode(
        &quot;Is this review positive or negative? Review: this is the best cast iron skillet you will ever buy&quot;,
        return_tensors=&quot;pt&quot;)
    outputs = model.generate(inputs)
    print(tokenizer.decode(outputs[0]))"><pre class="notranslate"><span class="pl-s">"""</span>
<span class="pl-s">Example code to load a PyTorch model across GPUs</span>
<span class="pl-s">"""</span>
<span class="pl-k">from</span> <span class="pl-s1">transformers</span> <span class="pl-k">import</span> <span class="pl-v">AutoTokenizer</span>, <span class="pl-v">AutoModelForSeq2SeqLM</span>
<span class="pl-k">from</span> <span class="pl-s1">transformers</span>.<span class="pl-s1">deepspeed</span> <span class="pl-k">import</span> <span class="pl-v">HfDeepSpeedConfig</span>
<span class="pl-k">from</span> <span class="pl-s1">transformers</span> <span class="pl-k">import</span> <span class="pl-s1">deepspeed</span>
<span class="pl-k">import</span> <span class="pl-s1">pandas</span> <span class="pl-k">as</span> <span class="pl-s1">pd</span>
<span class="pl-k">import</span> <span class="pl-s1">torch</span>
<span class="pl-k">import</span> <span class="pl-s1">pdb</span>
<span class="pl-k">import</span> <span class="pl-s1">os</span>

<span class="pl-s1">seed</span> <span class="pl-c1">=</span> <span class="pl-c1">42</span>
<span class="pl-s1">torch</span>.<span class="pl-en">manual_seed</span>(<span class="pl-s1">seed</span>)

<span class="pl-s1">ds_config</span> <span class="pl-c1">=</span> {
    <span class="pl-s">"fp16"</span>: {
        <span class="pl-s">"enabled"</span>: <span class="pl-s">"auto"</span>,
        <span class="pl-s">"loss_scale"</span>: <span class="pl-c1">0</span>,
        <span class="pl-s">"loss_scale_window"</span>: <span class="pl-c1">1000</span>,
        <span class="pl-s">"initial_scale_power"</span>: <span class="pl-c1">16</span>,
        <span class="pl-s">"hysteresis"</span>: <span class="pl-c1">2</span>,
        <span class="pl-s">"min_loss_scale"</span>: <span class="pl-c1">1</span>
    },
    <span class="pl-s">"zero_optimization"</span>: {
        <span class="pl-s">"stage"</span>: <span class="pl-c1">3</span>,
        <span class="pl-s">"overlap_comm"</span>: <span class="pl-s1">true</span>,
        <span class="pl-s">"contiguous_gradients"</span>: <span class="pl-s1">true</span>,
        <span class="pl-s">"sub_group_size"</span>: <span class="pl-c1">1e9</span>,
        <span class="pl-s">"reduce_bucket_size"</span>: <span class="pl-s">"auto"</span>,
        <span class="pl-s">"stage3_prefetch_bucket_size"</span>: <span class="pl-s">"auto"</span>,
        <span class="pl-s">"stage3_param_persistence_threshold"</span>: <span class="pl-s">"auto"</span>,
        <span class="pl-s">"stage3_max_live_parameters"</span>: <span class="pl-c1">1e9</span>,
        <span class="pl-s">"stage3_max_reuse_distance"</span>: <span class="pl-c1">1e9</span>,
        <span class="pl-s">"stage3_gather_fp16_weights_on_model_save"</span>: <span class="pl-s1">true</span>
    },
    <span class="pl-s">"gradient_accumulation_steps"</span>: <span class="pl-c1">1</span>,
    <span class="pl-s">"gradient_clipping"</span>: <span class="pl-c1">0</span>,
    <span class="pl-s">"steps_per_print"</span>: <span class="pl-c1">2000</span>,
    <span class="pl-s">"train_batch_size"</span>: <span class="pl-c1">2</span>,
    <span class="pl-s">"train_micro_batch_size_per_gpu"</span>: <span class="pl-c1">1</span>,
    <span class="pl-s">"wall_clock_breakdown"</span>: <span class="pl-s1">false</span>
}

<span class="pl-k">if</span> <span class="pl-s1">__name__</span> <span class="pl-c1">==</span> <span class="pl-s">"__main__"</span>:
    <span class="pl-c"># must run before instantiating the model</span>
    <span class="pl-c"># ds_config is deepspeed config object or path to the file</span>
    <span class="pl-s1">dschf</span> <span class="pl-c1">=</span> <span class="pl-v">HfDeepSpeedConfig</span>(<span class="pl-s1">ds_config</span>)  <span class="pl-c"># keep this object alive</span>

    <span class="pl-s1">model_name</span> <span class="pl-c1">=</span> <span class="pl-s">"bigscience/T0_3B"</span>
    <span class="pl-s1">tokenizer</span> <span class="pl-c1">=</span> <span class="pl-v">AutoTokenizer</span>.<span class="pl-en">from_pretrained</span>(<span class="pl-s1">model_name</span>)
    <span class="pl-s1">model</span> <span class="pl-c1">=</span> <span class="pl-v">AutoModelForSeq2SeqLM</span>.<span class="pl-en">from_pretrained</span>(<span class="pl-s1">model_name</span>)

    <span class="pl-s1">engine</span> <span class="pl-c1">=</span> <span class="pl-s1">deepspeed</span>.<span class="pl-en">initialize</span>(<span class="pl-s1">model</span><span class="pl-c1">=</span><span class="pl-s1">model</span>, <span class="pl-s1">config_params</span><span class="pl-c1">=</span><span class="pl-s1">ds_config</span>)

    <span class="pl-s1">inputs</span> <span class="pl-c1">=</span> <span class="pl-s1">tokenizer</span>.<span class="pl-en">encode</span>(
        <span class="pl-s">"Is this review positive or negative? Review: this is the best cast iron skillet you will ever buy"</span>,
        <span class="pl-s1">return_tensors</span><span class="pl-c1">=</span><span class="pl-s">"pt"</span>)
    <span class="pl-s1">outputs</span> <span class="pl-c1">=</span> <span class="pl-s1">model</span>.<span class="pl-en">generate</span>(<span class="pl-s1">inputs</span>)
    <span class="pl-en">print</span>(<span class="pl-s1">tokenizer</span>.<span class="pl-en">decode</span>(<span class="pl-s1">outputs</span>[<span class="pl-c1">0</span>]))</pre></div>
<h2 dir="auto">Expected behavior</h2>
<p dir="auto">T0pp (or T0_3B) to load across 2 GPUs, generate an answer, and then quit.</p>