Transformer fix

[JayYun98]

Proposed changes

Changes

• Changed order of layernorm and ffn according to tensorflow and pytorch implementations. (norm_first=False is default in pytroch as well as transformer)

• Added Encoder, Decoder with Dropout

Parts want to discuss

• Do you want to use Transformer with Dropout as main class?

• Do you like to add ffn as function or classes?

• Do you like to add norm_first parameter so that you could choose to do layer_norm first or later?

• Tensorflow do not do dropout with ffn while pytorch do. Which one should I refer to?

• I guess you did not add dropout since you would like to benchmark transformer without dropout for all jax, torch, tensorflow, mlx. I also had trouble adding dropout for jax.

Testing

• All Encoder, Decoder (w and w/o dropout) are tested with mlx-exmaples/transformer_lm

Checklist

Put an x in the boxes that apply.

• I have read the CONTRIBUTING document
• I have run pre-commit run --all-files to format my code / installed pre-commit prior to committing changes
• I have added tests that prove my fix is effective or that my feature works
• I have updated the necessary documentation (if needed)

[awni]

Do you want to use Transformer with Dropout as main class?

Let's just make dropout a kwarg parameter to the Transformer* which defaults to 0. Don't need to implement seperate classes.

Do you like to add ffn as function or classes?

Let's keep it the way it is implemented in __call__

Do you like to add norm_first parameter so that you could choose to do layer_norm first or later?

I think we should have this. And maybe default it to True as that seems to be a more common setup?

Tensorflow do not do dropout with ffn while pytorch do. Which one should I refer to?

Seems reasonable to follow PyTorch in that case. Add a comment about it when describing the dropout argument.

[awni]

Thanks for the update! I responded to your questions. Please make the corresponding changes and then we can add this!

[awni]

Also do you mind adding a docstring to the Transformer class?

[awni]

@Jyun1998 are you still planning on updating this?

[JayYun98]

@Jyun1998 are you still planning on updating this?

Yes! I will finish it within this week with your comments :)

[JayYun98]

Hi @awni ! Merry Chirstmas :)

Added docstring, norm_first, and also activation to replace mx.maximum at ffn layer.

There are few points to further discuss.

• I have setted norm_first=False, which is common according to pytorch_docstring

• by adding norm_first, codes got way extensive long. according to pytorch_implementation, they further break down encoder/decoer_layer with private function _sa_block, _ff_block. Do you want to keep it lastest commit or use _sa_block, _ff_block?

• default dropout ratio is 0 where Dropout call is implemented with self._p_1 == 1, so Dropout(0) will do nothing.

[awni]

Looks great, just a couple minor comments!

[awni]

Add defaults for these as well in the docstring.

[awni]

I have setted norm_first=False, which is common according to pytorch_docstring

Ok with me

by adding norm_first, codes got way extensive long.

I left a comment that might work for shortening this. I don't think you need separate modules for these.

default dropout ratio is 0 where Dropout call is implemented with self._p_1 == 1, so Dropout(0) will do nothing.

👍

[awni]

Also could you rebase? There is a merge conflict

[JayYun98]

@awni Thanks for the review!

• I added defaults to doctring.

• for the shortening the code, below will do the same but the code become quite not straight forward, so better left it as it is.

def __call__(self, x, mask):
    y = self.ln1(x) if self.norm_first else x
    y = self.attention(y, y, y, mask)
    y = self.dropout1(y)
    x = x + y  
    x = self.ln1(x) if not self.norm_first else x  

    y = self.ln2(x) if self.norm_first else x
    y = self.linear1(y)  
    y = self.activation(y)
    y = self.dropout2(y)
    y = self.linear2(y)
    x = x + y  
    x = self.ln2(x) if not self.norm_first else x  

    return x

[JayYun98]

Also rebased :)

[awni]

Looks great, thanks for adding!! I pushed a few changes to the doc strings and included stuff in the docs as well!

[JayYun98]

Looks great, thanks for adding!! I pushed a few changes to the doc strings and included stuff in the docs as well!

Thanks for the review :)