Upload folder using huggingface_hub

README.md

@@ -16,4 +16,4 @@ python config_tiny_mistral.py
 # Run training
 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
-```
+```

config_tiny_mistral.py

@@ -6,6 +6,8 @@ python config_tiny_mistral.py
 ```
 """
 import os
+from dataclasses import dataclass
+from typing import Optional
 
 from nanotron.config import (
     CheckpointsArgs,
@@ -23,8 +25,6 @@ from nanotron.config import (
     TokensArgs,
 )
 from nanotron.logging import human_format
-from dataclasses import dataclass
-from typing import Optional
 
 
 @dataclass
@@ -58,6 +58,7 @@ class MistralConfig:
         if self.num_key_value_heads is None:
             self.num_key_value_heads = self.num_attention_heads
 
+
 model_config = MistralConfig(
     # Config for a tiny model model with 1.62M parameters
     bos_token_id=1,

dataloader.py

@@ -1,3 +1,4 @@
+from nanotron import logging
 from nanotron.config import (
     PretrainDatasetsArgs,
 )
@@ -13,7 +14,6 @@ from nanotron.trainer import DistributedTrainer
 from nanotron.utils import (
     main_rank_first,
 )
-from nanotron import logging
 
 try:
     from huggingface_hub import __version__ as hf_hub_version

modeling_mistral.py

@@ -23,16 +23,13 @@ from flash_attn.flash_attn_interface import (
     flash_attn_with_kvcache,
 )
 from flash_attn.layers.rotary import RotaryEmbedding as FlashRotaryEmbedding
-from torch import nn
-from transformers import MistralConfig
-from transformers.activations import ACT2FN
-
 from nanotron import distributed as dist
 from nanotron import logging
 from nanotron.config import ParallelismArgs, RecomputeGranularity
-from nanotron.nn.layer_norm import TritonRMSNorm
+from nanotron.generation.generate_store import AttachableStore
 from nanotron.logging import log_rank
 from nanotron.models import NanotronModel
+from nanotron.nn.layer_norm import TritonRMSNorm
 from nanotron.parallel import ParallelContext
 from nanotron.parallel.parameters import NanotronParameter
 from nanotron.parallel.pipeline_parallel.block import (
@@ -49,7 +46,9 @@ from nanotron.parallel.tensor_parallel.nn import (
 )
 from nanotron.random import RandomStates
 from nanotron.utils import checkpoint_method
-from nanotron.generation.generate_store import AttachableStore
+from torch import nn
+from transformers import MistralConfig
+from transformers.activations import ACT2FN
 
 logger = logging.get_logger(__name__)
 
@@ -852,6 +851,7 @@ class MistralForTraining(NanotronModel):
     ):
         super().__init__()
         import warnings
+
         warnings.warn("This is just a Llama Model, not a Mistral one for demo purpose. Please fix implementation")
         self.model = MistralModel(config=config, parallel_context=parallel_context, parallel_config=parallel_config)
         self.loss = PipelineBlock(
@@ -1120,4 +1120,4 @@ def get_flops(
     else:
         raise ValueError("recompute_granularity must be one of 'full' or 'selective'")
 
-    return model_flops, hardware_flops
+    return model_flops, hardware_flops

run_train.py

@@ -9,11 +9,10 @@ torchrun --nproc_per_node=8 run_train.py --config-file config_tiny_mistral.yaml
 """
 import argparse
 
-from modeling_mistral import MistralForTraining
+from config_tiny_mistral import MistralConfig
 from dataloader import get_dataloader
+from modeling_mistral import MistralForTraining
 from nanotron.trainer import DistributedTrainer
-from config_tiny_mistral import MistralConfig
-
 
 
 def get_args():