Introduce explicit parameter mapping

[jonatanklosko]

Closes #71.

Advantages:

• This removes a lot of duplication between models. For example, now all of the models use the same attention implementation, and even the whole transformer block.
• Certain models looked very differently, even though they were for the most part the same, that's because their reference implementations applied the same concepts, but implemented them on their own (naming, structuring and nesting the layers). Now the models are much more consistent, both in terms of layer nesting and naming, which makes it much easier to see the similarities and differences.
• We no longer need to follow an artificial layers nesting, that follows from the class composition in the Python implementations. For example, when building a model with head, we don't need to prefix the base model with something like bert., we just add the head layers.
• We have full control over the parameters, which makes us much less coupled to the original implementation details. For example the original GPT2 implementation uses conv1d layers (in terms of implementation those are essentially transposed dense layers) for attention. The parameter mapping has enough flexibility that we actually slice the conv1d kernels into dense layer kernels and use the same attention implementation as all other models.

Disadvantages:

• A single file doesn't contain all of the model. This already wasn't the case and I don't consider this a major issue. Also, we could still be doing that and use consistent naming/implementation, but again I don't think we should.
• Each model needs explicit mapping to hf/transformers parameters. This also implies we use different parameter names than hf/transformers.

[seanmor5]

@jonatanklosko I've just skimmed a bit, but my initial thoughts are that this is fantastic! It should also make adding new models significantly easier, and in general is much more manageable for us to maintain.

I will give a more in-depth review later

[seanmor5]

Also, on requiring mappings, I don't think that's a big deal, HF actually requires mappings for non-PT models anyway, and we mostly end up hiding that fact from the user anyway. It's a consequence of depending on an external community for models, but I don't think it's a bad one

[josevalim]

Fantastic work!

[seanmor5]

I wonder if Axon should just make the switch to use weight everywhere like PyTorch. It would make the conversions a bit easier, though you'd still have to pay attention to doing the transpose in some cases

[jonatanklosko]

I don't think we should be changing just because of PyTorch in this case, flax does use kernel. Parameter naming in particular is not much hassle, especially that in case we don't need to transpose, we just always look for weight instead of kernel by default.

[seanmor5]

When is offset not s64?

[jonatanklosko]

Offset comes from the cache, which we pass as model input and get back as model output, so Axon would automatically convert it to float either way.

[seanmor5]

Ahh makes sense, I have opened this issue in Axon to fix this: elixir-nx/axon#464