Merge branch 'main' into add-s3-uploader

Author: eliebak

src/nanotron/models/llama.py

@@ -165,7 +165,7 @@ def __init__(
             bias=False,
             async_communication=tp_linear_async_communication and tp_mode is TensorParallelLinearMode.REDUCE_SCATTER,
         )
-        self.split_silu_mul = torch.compile(GLUActivation(config.hidden_act))
+        self.split_silu_mul = GLUActivation(config.hidden_act)
 
     def forward(self, hidden_states):  # [seq_length, batch_size, hidden_dim]
         merged_states = self.gate_up_proj(hidden_states)

src/nanotron/trainer.py

@@ -723,8 +723,7 @@ def _load_model_checkpoint(self, model: NanotronModel) -> NanotronModel:
                 )
             reloaded_from_checkpoint=True
         if not reloaded_from_checkpoint:
-            # TODO @eliebak add s3 support also here
-            log_rank("No checkpoint path provided.", logger=logger, level=logging.INFO)
+            log_rank("No checkpoint path provided.", logger=logger, level=logging.INFO, rank=0)
             if isinstance(self.config.model.init_method, ExistingCheckpointInit):
                 # Initialize model from an pretrained model checkpoint (without optimizer, lr_scheduler...)
                 self.param_shard_metadata = load_weights(