Add Multi-head Latent Attention (DeepSeekv2)

[simoneangarano]

This pull request contains the modifications made to the codebase in the litgpt repository to add support for Multi-Head Latent Attention (MLA) block from DeepSeekV2.

Changes

• Configuration: Added latent_attention: Optional[bool] = False parameter to the configuration file to enable the MLA block.
• MLA module: Implemented the MLA module as a separate component in the litgpt codebase.
• KVCacheCompressed: Added support for the KVCacheCompressed class to store the key-value pairs for the MLA block.
• Model: Modified the GPT class to include the MLA block as an alternative component based on the configuration parameter latent_attention.
• Training: Updated the training script to support the MLA block and added support for training with the new configuration file config_hub/pretrain/cfg.yaml.

Usage

• Configuration: Set the latent_attention parameter to True in the configuration file to enable the MLA block.
• Training: Run the training script with the updated configuration file.
  litgpt pretrain --config config_hub/pretrain/cfg.yaml
• Inference: Use the trained model for inference as follows:
  litgpt generate out/pretrain/mla/final/

[Borda]

@simoneangarano, mind checking the failing tests? :)

[ysjprojects]

@Borda @simoneangarano I can look into the failing tests if that's okay with everyone.

[Borda]

@simoneangarano I can look into the failing tests if that's okay with everyone.

@ysjprojects that would be great!