Fix support for custom modeling

huggingface / nanotron

Minimalistic large language model 3D-parallelism training

Apache License 2.0

1.14k stars 107 forks source link

Closed NouamaneTazi closed 8 months ago

NouamaneTazi commented 8 months ago

Steps to use your custom model in nanotron

Update the MistralConfig class in config_tiny_mistral.py to match your model's configuration
Update the MistralForTraining class in modeling_mistral.py to match your model's architecture

Pass the previous to the DistributedTrainer class in run_train.py:

trainer = DistributedTrainer(config_file, model_class=MistralForTraining, model_config_class=MistralConfig)

thomwolf commented 8 months ago

Simple and nice! looks good to merge