Extendability refactors

mosaicml / llm-foundry

LLM training code for Databricks foundation models

Apache License 2.0

3.83k stars 502 forks source link

Closed dakinggg closed 1 week ago

dakinggg commented 2 weeks ago

This PR includes a few changes for increased extendability of the code:

Fixes some typing
Adds slice_attention_mask to MPTBlock
Removes the need for extra imports in configuration_mpt.py just for HF checkpointing
Factors out the attention config validation
Makes the block class overridable on MPTModel
Adds a registry for transforms that can be applied to the TrainConfig

Loss before and after:

dakinggg commented 1 week ago

@milocress the GPU test is unrelated. It will be fixed by the next composer release (which is why that test isn't marked as required yet)