Add linear layer and ffn config to enable TransformerEngine layers (with FP8)

[vchiley]

This PR adds a config for Linear Layers and FFN modules which allows the use TransformerEngine's te.Linear and te.LayerNormMLP modules (which have fp8 with amp.fp8).
+ I did a little cleanup

This PR is built on top of #271

In the future, this'll also allow us to add and prototype other linear layers and ffn blocks. Furthermore it enables us to configure TP/SP for the MLP block, in the build_ffn util fn.

AMP FP8 training gets results which are nearly identical to AMP BF16:

but has faster runtime.

Furthermore ffn_config_defaults: ffn_type: te_ln_mlp allows us to use transformer engine's LayerNormMLP layer which has SP and TP support if configured correctly.

[vancoykendall]

I'm excited that this PR is being worked on, mainly because I've been extending the MPTMLP myself in order to try using gated linear units (GEGLU, SwiGLU, etc) and also decided to add a ffn_config in my implementation. I was wondering though, why expansion_ratio is not a part of the ffn_config? From my understanding, expansion_ratio is only used when building the ffn block, so it makes more sense to me for it to be defined in the ffn_config. Also, what does fc stand for in fc_type? It isn't clear to me.

[vchiley]

@vancoykendall

fc stands for fully connected layer

While expansion_ratio should really be part of the ffn_config, moving it into the ffn_config breaks backwards compatibility with models already uploaded to HuggingFace hub and all of our internal checkpoints so we are leaving that config variable alone.

[sashaDoubov]

LGTM