使用distributed optimzer时grad_norm计算准确度的疑问

[chivychao]

Megatron-LLaMA/megatron/optimizer/distrib_optimizer.py

Lines 926 to 939 in 25306de

	# Scale grad buffers by '1 / data_parallel_world_size'.
	for model in self.models:
	for dtype, gbuf in model._grad_buffers.items():
	gbuf.data /= data_parallel_world_size
	# Reduce-scatter all grads.
	gbuf_view_items = self.get_model_grad_buffer_dp_views()
	for index, (model_index, dtype, gbuf, gbuf_views) \
	in enumerate(gbuf_view_items):
	torch.distributed._reduce_scatter_base(
	gbuf_views[data_parallel_rank],
	gbuf,
	group=data_parallel_group,
	)

这里执行的应该是使得dp组内每个成员只获得自己维护的那一部分参数梯度的求和吧？

但这样做的话，在后面optimizer.step()中计算的grad_norm是不是就不是很准确了？

因为我看grad_norm计算的时候是dp组内每个成员把自己那部分模型的所有param的所有梯度都平方加和了，但是每个成员的grad只有一部分完成了dp组内求和，这样求出来的grad_norm感觉是错的。

请问是否确实存在这样的问题呢？

[li-yi-dong]

Megatron-LLaMA/megatron/optimizer/clip_grads.py

Line 92 in 25306de

torch.distributed.all_reduce(total_norm,

可以看看这段代码