Discrepancy between FusedMLP and Mlp

[netw0rkf10w]

Hello,

It seems to me that FusedMLP uses the tanh approximation of GELU by default (and the exact version is not available), whereas Mlp uses the exact variant. This will cause numerical mismatches of the two variants when using them with default arguments.

I would suggest to set the same approximate variant for Mlp by default:

class Mlp(nn.Module):
    def __init__(self, in_features, hidden_features=None, out_features=None, activation=nn.GELU(approximate='tanh'),
                 return_residual=False, device=None, dtype=None):

The implementation of some models would need minor adjustments as well, for example for ViT, act_layer = act_layer or nn.GELU should be replaced by

act_layer = act_layer or nn.GELU(approximate='tanh')

With these changes, setting fused_mlp=True or fused_mlp=False would change only the running time and not numerical results. I propose this because it has the same spirit as Flash Attention: use_flash_attention=True or use_flash_attention=False should give the same results.

Please let me know what you think.