Log global batch size in tokens

[tomiock]

we need this metric to normalize the end to end step time and make accurate reports for scaling

[tomiock]

torchtitan/torchtitan/train.py

Line 391 in 7d17a92

ntokens_batch = labels.numel()

here the batch number of tokens are calculated.

maybe an estimation can be made by max_seq_len and batch_size

[tomiock]

there are some ways to make an average of one metrics across the data parallel dims

torchtitan/torchtitan/train.py

Lines 508 to 521 in 7d17a92

	if parallel_dims.dp_cp_enabled:
	loss = loss.detach()
	ft_pg = self.ft_manager.loss_sync_pg
	global_avg_loss, global_max_loss, global_ntokens_seen = (
	dist_utils.dist_mean(loss, parallel_dims.world_mesh["dp_cp"], ft_pg),
	dist_utils.dist_max(loss, parallel_dims.world_mesh["dp_cp"], ft_pg),
	dist_utils.dist_sum(
	torch.tensor(
	self.ntokens_seen, dtype=torch.int64, device=self.device
	),
	parallel_dims.world_mesh["dp_cp"],
	ft_pg,
	),
	)