How to use make_private() with a model that registers a buffer?

username: fsommers
cooked: <p>Specifically, I’m trying to use Opacus with the the transformers implementation of the LayoutLMv3 model. When run a forward pass, I receive the following error message:</p>
<p><code>NotImplementedError: [NotImplementedError("Model contains a trainable layer that Opacus doesn't currently support(model:LayoutLMv3Model....</code></p>
<p>Upon looking at the source of this error, it comes from this check:</p>
<aside class="onebox githubblob" data-onebox-src="https://github.com/pytorch/opacus/blob/32a465bfa980777e60a763616c83429785aa2ac6/opacus/grad_sample/grad_sample_module.py#L467-L478">
  <header class="source">

      <a href="https://github.com/pytorch/opacus/blob/32a465bfa980777e60a763616c83429785aa2ac6/opacus/grad_sample/grad_sample_module.py#L467-L478" target="_blank" rel="noopener nofollow ugc">github.com</a>
  </header>

  <article class="onebox-body">
    <h4><a href="https://github.com/pytorch/opacus/blob/32a465bfa980777e60a763616c83429785aa2ac6/opacus/grad_sample/grad_sample_module.py#L467-L478" target="_blank" rel="noopener nofollow ugc">pytorch/opacus/blob/32a465bfa980777e60a763616c83429785aa2ac6/opacus/grad_sample/grad_sample_module.py#L467-L478</a></h4>



    <pre class="onebox"><code class="lang-py">
      <ol class="start lines" start="467" style="counter-reset: li-counter 466 ;">
          <li>    NotImplementedError(</li>
          <li>        f"Model contains a trainable layer "</li>
          <li>        f"that Opacus doesn't currently support({m_name}:{m}). "</li>
          <li>        f"Please implement and register grad sampler for this layer. "</li>
          <li>        f"(See opacus.grad_sample.utils.register_grad_sampler)"</li>
          <li>    )</li>
          <li>    for m_name, m in trainable_modules(module)</li>
          <li>    # With functorch, all modules are trainable</li>
          <li>    # We still want to avoid module that have buffers (e.g. BatchNorm)</li>
          <li>    # as the buffers are not private</li>
          <li>    if len(list(m.buffers())) &gt; 0</li>
          <li>]</li>
      </ol>
    </code></pre>



  </article>

  <div class="onebox-metadata">
    
    
  </div>

  <div style="clear: both"></div>
</aside>

<p>Indeed, the LayoutLMv3 model does register a buffer:</p>
<aside class="onebox githubblob" data-onebox-src="https://github.com/huggingface/transformers/blob/4fdf58afb72b0754da30037fc800b6044e7d9c99/src/transformers/models/layoutlmv3/modeling_layoutlmv3.py#L246">
  <header class="source">

      <a href="https://github.com/huggingface/transformers/blob/4fdf58afb72b0754da30037fc800b6044e7d9c99/src/transformers/models/layoutlmv3/modeling_layoutlmv3.py#L246" target="_blank" rel="noopener nofollow ugc">github.com</a>
  </header>

  <article class="onebox-body">
    <h4><a href="https://github.com/huggingface/transformers/blob/4fdf58afb72b0754da30037fc800b6044e7d9c99/src/transformers/models/layoutlmv3/modeling_layoutlmv3.py#L246" target="_blank" rel="noopener nofollow ugc">huggingface/transformers/blob/4fdf58afb72b0754da30037fc800b6044e7d9c99/src/transformers/models/layoutlmv3/modeling_layoutlmv3.py#L246</a></h4>



    <pre class="onebox"><code class="lang-py">
      <ol class="start lines" start="236" style="counter-reset: li-counter 235 ;">
          <li></li>
          <li>def __init__(self, config):</li>
          <li>    super().__init__()</li>
          <li>    self.word_embeddings = nn.Embedding(config.vocab_size, config.hidden_size, padding_idx=config.pad_token_id)</li>
          <li>    self.token_type_embeddings = nn.Embedding(config.type_vocab_size, config.hidden_size)</li>
          <li></li>
          <li>    self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)</li>
          <li>    self.dropout = nn.Dropout(config.hidden_dropout_prob)</li>
          <li></li>
          <li>    # position_ids (1, len position emb) is contiguous in memory and exported when serialized</li>
          <li class="selected">    self.register_buffer(</li>
          <li>        "position_ids", torch.arange(config.max_position_embeddings).expand((1, -1)), persistent=False</li>
          <li>    )</li>
          <li></li>
          <li>    self.padding_idx = config.pad_token_id</li>
          <li>    self.position_embeddings = nn.Embedding(</li>
          <li>        config.max_position_embeddings, config.hidden_size, padding_idx=self.padding_idx</li>
          <li>    )</li>
          <li></li>
          <li>    self.x_position_embeddings = nn.Embedding(config.max_2d_position_embeddings, config.coordinate_size)</li>
          <li>    self.y_position_embeddings = nn.Embedding(config.max_2d_position_embeddings, config.coordinate_size)</li>
      </ol>
    </code></pre>



  </article>

  <div class="onebox-metadata">
    
    
  </div>

  <div style="clear: both"></div>
</aside>

<p>What is the recommended way to use Opacus in this situation?</p>

username: HuanyuZhang
cooked: <p>I have not looked at your model clearly. But if you are sure that the buffer you have will not lead to a privacy leakage (unlike batch normalization), feel free to just comment it out in the code.</p>