""" | |
Nanotron training script. | |
Usage: | |
``` | |
export CUDA_DEVICE_MAX_CONNECTIONS=1 # important for some distributed operations | |
torchrun --nproc_per_node=8 run_train.py --config-file config_tiny_mistral.yaml | |
``` | |
""" | |
import argparse | |
from config_tiny_mistral import MistralConfig | |
from dataloader import get_dataloader | |
from modeling_mistral import MistralForTraining | |
from nanotron.trainer import DistributedTrainer | |
def get_args(): | |
parser = argparse.ArgumentParser() | |
parser.add_argument("--config-file", type=str, required=True, help="Path to the YAML or python config file") | |
return parser.parse_args() | |
if __name__ == "__main__": | |
args = get_args() | |
config_file = args.config_file | |
# Load trainer and data | |
trainer = DistributedTrainer(config_file, model_class=MistralForTraining, model_config_class=MistralConfig) | |
dataloader = get_dataloader(trainer) | |
# Train | |
trainer.train(dataloader) | |