# Remove causal mask from Llama decoder

**URL:** <https://discuss.huggingface.co/t/remove-causal-mask-from-llama-decoder/111506>\
**Category:** Intermediate\
**Created:** [October 12, 2024, 5:11am UTC](https://discuss.huggingface.co/t/remove-causal-mask-from-llama-decoder/111506 "2024-10-12T05:11:50Z")\
**Posts on this page:** 6\
**Page:** 1

<div class="post-metadata">

**Author:** ![abhinavkulkarni](https://sea2.discourse-cdn.com/hellohellohello/user_avatar/discuss.huggingface.co/abhinavkulkarni/32/32750_2.png) [@abhinavkulkarni](https://discuss.huggingface.co/u/abhinavkulkarni)\
**Post date:** [October 12, 2024, 5:11am UTC](https://discuss.huggingface.co/t/remove-causal-mask-from-llama-decoder/111506/1 "2024-10-12T05:11:50Z")

</div>

Hi,

I want to train the recently released smaller Llama 3.2 models ([link](https://huggingface.co/meta-llama/Llama-3.2-1B)) for an NER task.

However, I would like to remove the causal LM triangular mask during training and inference. How do I go about this?

First of all, I thought the mask was automatically generated based on `model.config.is_decoder` in `get_extended_attention_mask` ([link](https://github.com/huggingface/transformers/blob/617b21273a349bd3a94e2b3bfb83f8089f45749b/src/transformers/modeling_utils.py#L1138)). If so, simply setting this to `False` should enable bidrectional attention.

However, `model.config.is_decoder` is `False` for `AutoModelForCausalLM` (as well as for `AutoModelForTokenClassification`).

```python
from transformers import AutoTokenizer, AutoModelForCausalLM

model_checkpoint = "meta-llama/Llama-3.2-1B"

model = AutoModelForCausalLM.from_pretrained(model_checkpoint)
assert model.config.is_decoder == False

model = AutoModelForTokenClassification.from_pretrained(model_checkpoint)
assert model.config.is_decoder == False

```

How can this be the case? The flag should be `True` for `AutoModelForCausalLM`.

---

<div class="post-metadata">

**Author:** ![not-lain](https://sea2.discourse-cdn.com/hellohellohello/user_avatar/discuss.huggingface.co/not-lain/32/59092_2.png) [@not-lain](https://discuss.huggingface.co/u/not-lain)\
**Post date:** [October 12, 2024, 4:21pm UTC](https://discuss.huggingface.co/t/remove-causal-mask-from-llama-decoder/111506/2 "2024-10-12T16:21:20Z")

</div>

Hi @abhinavkulkarni  
I might propose the following patch, let me know if this works for you.

```python
class MyModel(AutoModelForCausalLM) : 
   def get_extended_attention_mask(self, **kwargs,*args): #let's override the method here
       pass

```

---

<div class="post-metadata">

**Author:** ![SuperSonnix71](https://sea2.discourse-cdn.com/hellohellohello/user_avatar/discuss.huggingface.co/supersonnix71/32/50662_2.png) [@SuperSonnix71](https://discuss.huggingface.co/u/SuperSonnix71)\
**Post date:** [October 21, 2024, 9:54pm UTC](https://discuss.huggingface.co/t/remove-causal-mask-from-llama-decoder/111506/3 "2024-10-21T21:54:26Z")

</div>

setting `model.config.is_decoder` to `False` is not going to work since the flag is already `False` for both `AutoModelForCausalLM` and `AutoModelForTokenClassification`. Exactly as the previous answer you can override the ‘ get\_extended\_attention\_mask’.

But may I ask why you want to change the casual to bidirectional? NER task? Then why not using BERT?  
Also, you could try using LLM2Vec.

---

<div class="post-metadata">

**Author:** ![alex-hh](https://avatars.discourse-cdn.com/v4/letter/a/c6cbf5/32.png) [@alex-hh](https://discuss.huggingface.co/u/alex-hh)\
**Post date:** [October 21, 2024, 11:43pm UTC](https://discuss.huggingface.co/t/remove-causal-mask-from-llama-decoder/111506/4 "2024-10-21T23:43:08Z")

</div>

You need to override LlamaModel.\_update\_causal\_mask, or pass your own 4d attention bias as the attention mask argument, with shape 1, 1, L, L, where mask[0,0,i,j] = 0 unless i or j is padding in which case it equals minus some large number

---

<div class="post-metadata">

**Author:** ![abhinavkulkarni](https://sea2.discourse-cdn.com/hellohellohello/user_avatar/discuss.huggingface.co/abhinavkulkarni/32/32750_2.png) [@abhinavkulkarni](https://discuss.huggingface.co/u/abhinavkulkarni)\
**Post date:** [October 22, 2024, 3:59am UTC](https://discuss.huggingface.co/t/remove-causal-mask-from-llama-decoder/111506/5 "2024-10-22T03:59:48Z")

</div>

I was able to create a full attention mask as follows:

```python
class LlamaFullAttentionModel(LlamaModel):
    def _update_causal_mask(
        self, *args, **kwargs
    ):
        causal_mask = self._update_causal_mask(*args, **kwargs)
        full_attention_mask = torch.tril(causal_mask)
        return full_attention_mask

```

---

<div class="post-metadata">

**Author:** ![system](https://us1.discourse-cdn.com/hellohellohello/original/2X/d/de4155eb4aa4108ecb32a1389d7cc37ae69f88b7.png) [@system](https://discuss.huggingface.co/u/system)\
**Post date:** [October 22, 2024, 3:59pm UTC](https://discuss.huggingface.co/t/remove-causal-mask-from-llama-decoder/111506/6 "2024-10-22T15:59:48Z")

</div>

This topic was automatically closed 12 hours after the last reply. New replies are no longer allowed.