[Code scan] Aggregate pretraining DDP validation metrics across ranks

[njzjz]

This issue is a result of a Codex global repository scan.

Summary

In pretraining DDP evaluation, each rank evaluates its validation shard, but rank 0 chooses best-checkpoint and early-stopping decisions from its local metrics. The metric reducer multiplies bsz by world size, but loss sums and sample counts are not all-reduced or gathered before checkpoint decisions.

Code references

unimol_tools/unimol_tools/pretrain/trainer.py

Lines 296 to 321 in 4596596

	if save_every > 0 and self.global_step % save_every == 0:
	val_metrics = None
	if self.valid_dataloader is not None:
	val_metrics = self.evaluate(self.global_step, log=self.rank == 0)
	if self.rank == 0:
	self._save_checkpoint(epoch, self.global_step, 'checkpoint_last.ckpt')
	if self.valid_dataloader is not None and val_metrics is not None:
	curr_loss = val_metrics.get('loss', float('inf'))
	if curr_loss < self.best_loss:
	self.best_loss = curr_loss
	self.no_improve_steps = 0
	logger.info(
	f"New best model at step {self.global_step} with loss {curr_loss:.4f}"
	)
	self._save_checkpoint(epoch, self.global_step, 'checkpoint_best.ckpt')
	else:
	self.no_improve_steps += 1
	if (
	self.patience >= 0
	and self.no_improve_steps > self.patience
	):
	logger.info(
	f"Early stopping triggered at step {self.global_step}"
	)
	stop_training = True
	self._save_checkpoint(epoch, self.global_step)

unimol_tools/unimol_tools/pretrain/trainer.py

Lines 400 to 417 in 4596596

	def evaluate(self, step, log: bool = False):
	self.model.eval()
	logging_infos = []
	with torch.no_grad():
	for batch in self.valid_dataloader:
	net_input, net_target = self.decorate_batch(batch)
	loss, sample_size, logging_info = self.loss_fn(
	self.model, net_input, net_target
	)
	logging_infos.append(logging_info)
	metrics = self.reduce_metrics(
	logging_infos,
	writer=self.writer,
	logger=logger if log else None,
	step=step,
	split="valid",
	unit="Step",
	)

unimol_tools/unimol_tools/pretrain/trainer.py

Lines 442 to 458 in 4596596

	def reduce_metrics(
	self,
	logging_outputs,
	writer=None,
	logger=None,
	step=None,
	split="train",
	unit="Epoch",
	epoch=None,
	inner_step=None,
	inner_total=None,
	):
	metrics_mean = self.loss_fn.reduce_metrics(logging_outputs, split=split)
	if split.startswith("train") and self.fp16 and self.scaler is not None:
	metrics_mean["loss_scale"] = self.scaler.loss_scale
	if "bsz" in metrics_mean:
	metrics_mean["bsz"] *= self.world_size

Impact

Best-checkpoint selection and early stopping can depend on only rank 0 validation samples, not the full validation set. This can select the wrong checkpoint or stop too early/late in distributed pretraining.

Suggested fix

All-reduce metric numerators and counts, or gather logging outputs from all ranks, before reduce_metrics and before checkpoint or early-stop decisions.