Update src/nanotron/trainer.py

Co-authored-by: AleHC <[email protected]>
swiss-ai · Nov 22, 2024 · 8659582 · 8659582
1 parent ace4d33
commit 8659582
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/src/nanotron/trainer.py b/src/nanotron/trainer.py
@@ -520,8 +520,10 @@ def train(
                     self.valid_metadata.last_train_step = self.iteration_step
                     self.valid_metadata.data_stages[self.valid_metadata.last_stage_idx].consumed_train_samples += self.global_batch_size
 
-                    log_entries = [LogItem("validation_loss_avg", loss_avg, "human_format")]
-                    self.loggerwriter.add_scalars_from_list(log_entries, self.iteration_step)
+                     if dist.get_rank(self.parallel_context.world_pg) in self.logger_ranks:
+                        log_entries = [LogItem("validation_loss_avg", loss_avg, "human_format")]
+                        self.loggerwriter.add_scalars_from_list(log_entries, self.iteration_step)
+
 
                     # NOTE: only one rank writes to wandb
                     if dist.get_rank(self.parallel_context.world_pg) == self.logger_ranks[0] and wandb is not None: