Feat: Various Configurable Initializations

[flxst]

What does this PR do?

This PR implements the following weight initializations (see https://arxiv.org/abs/2312.16903):

• plain (= all weights normally distributed)
• scaled (= same as plain, but narrower distribution for projection weights W0 & W2)
• scaled embed (= same as scaled, but wider distribution for embedding)

A weight initialisation component is introduced which modifies the model weights in place (see #168 for more details)

General Changes

• Components and factories for plain, scaled and scaled_embed initialisation.
• in GPT2 model training configs, the standard deviation std can now be set to the string auto (in which case it will equal sqrt(2/(5*hidden_dim)), see e.g. https://arxiv.org/abs/2312.16903)
• The CoCa model, which previously used a hardcoded, (probably not entirely correct) scaled initialization (see Limited and potentially incorrect weight initialization for CoCa model #165), can now only use plain initialization

Breaking Changes

• All training configs require an additional component for initialization of the raw model (i.e. the model with random weights), as shown here.

Checklist before submitting final PR

• My PR is minimal and addresses one issue / enhancement in isolation
• I have merged the latest version of the target branch into this feature branch
• I have reviewed my own code w.r.t. correct implementation, missing type hints, proper documentation, etc.
• I have run a sample config for model training
• I have fixed all failing tests (python tests/tests.py)

[le1nux]

Nice work! I think the implementation is correct (checked with the paper). I only had 2-3 questions where things were not clear to me. Regarding the integration into modalities, I think we need more iterations fixing some issues about the dependencies and coupling.

I'm not the biggest fan of having the initialisation in the parent class of the models (i.e., NNModel).
The reason is that the weight initialisation is pretty much dependent on the concrete model implementation. For instance, the gpt2 model must have the c_proj parameters initialised in certain ways, whereas CoCa for instance has different named parameters. Currently, we do string matching in the parent class which introduces low level, model-specific dependencies in the parent class. Similarly, this can be the case for modules, e.g., a custom linear layer that we introduce in a particular model. Also in this case, we would need custom initialisation in the parent class. For each of these special cases we would have to modify the parent class, strongly coupling dependencies.

We could resolve this inverse dependency by introducing an generic WeightInitializer class, that initializes the weights of a model in a generic, configurable way. The WeightInitializer class would be passed to the constructor of the concrete model. The model would then call something like weight_initalizer.init_weights(self, weight_init_config). Basically, the strategy pattern that modifies the calling object (here, the model) in place.

We should have a generic WeightInitializer covering the general cases. For specific models, we can also introduce new WeightInitializers that are specific to a certain model. These WeightInitializers should be instantiable as part of the hierarchical instantiation, as we do for the other components.

The config for the model and WeightInitializer would look like this:

weight_initializer:
   <weight init config...>

model:
  component_key: model
  variant_key: gpt2
  config:
    n_layer: 2
    n_head_q: 8
    n_head_kv: 4
    ffn_hidden: 128
    weight_initializer:
              instance_key: weight_initializer
              pass_type: BY_REFERENCE

Additionally, I left a bunch of mostly minor comments regarding some questions and ideas.

[le1nux]

LGTM :)