[Code scan] Guarantee masked atoms or skip empty pretraining masked losses

[njzjz]

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

Summary

Pretraining computes the number of masked atoms as int(sz * mask_prob + random). For small molecules this can be zero. The loss then computes mean token, coordinate, and distance losses over empty tensors, which can produce NaN and poison training.

Code references

unimol_tools/unimol_tools/pretrain/dataset.py

Lines 267 to 274 in 4596596

	num_mask = int(sz * mask_prob + np.random.rand())
	mask_idc = np.random.choice(sz, num_mask, replace=False)
	mask = np.full(sz, fill_value=False)
	mask[mask_idc] = True
	tgt_tokens = np.full(sz, dictionary.pad())
	tgt_tokens[mask] = src_tokens[mask]
	tgt_tokens = torch.from_numpy(tgt_tokens).long()

unimol_tools/unimol_tools/pretrain/loss.py

Lines 44 to 68 in 4596596

	masked_tokens = tgt_tokens.ne(self.padding_idx)
	(
	logits_encoder,
	encoder_distance,
	encoder_coord,
	x_norm,
	delta_encoder_pair_rep_norm,
	) = model(**net_input, encoder_masked_tokens=masked_tokens)
	target = tgt_tokens
	if masked_tokens is not None:
	target = target[masked_tokens]
	masked_token_loss = F.nll_loss(
	F.log_softmax(logits_encoder, dim=-1, dtype=torch.float32),
	target,
	ignore_index=self.padding_idx,
	reduction="mean",
	)
	masked_pred = logits_encoder.argmax(dim=-1)
	masked_hit = (masked_pred == target).long().sum()
	masked_cnt = masked_tokens.long().sum()
	loss = masked_token_loss * self.masked_token_loss

unimol_tools/unimol_tools/pretrain/loss.py

Lines 79 to 99 in 4596596

	if encoder_coord is not None:
	coord_target = tgt_coordinates
	masked_coord_loss = F.smooth_l1_loss(
	encoder_coord[masked_tokens].view(-1, 3).float(),
	coord_target[masked_tokens].view(-1, 3),
	reduction="mean",
	beta=1.0,
	)
	loss = loss + masked_coord_loss * self.masked_coord_loss
	logging_output["masked_coord_loss"] = masked_coord_loss.data
	if encoder_distance is not None:
	masked_dist_loss = self.cal_dist_loss(
	encoder_distance,
	tgt_distance,
	net_input["src_tokens"],
	masked_tokens,
	normalize=True,
	)
	loss = loss + masked_dist_loss * self.masked_dist_loss
	logging_output["masked_dist_loss"] = masked_dist_loss.data

Impact

A batch containing molecules with zero selected masked atoms can generate NaN loss values. Once NaNs enter gradients or checkpoints, the pretraining run may become unrecoverable.

Suggested fix

When mask_prob > 0, clamp num_mask to at least one and at most the atom count. Alternatively, explicitly skip masked objectives for samples or batches with zero masked atoms and avoid mean reductions over empty tensors.