In my case the “mean_all” and “invstd_all” should be weighted average according to different “counts” in GPUs, is it the actual situation?

I think you’re right.
torch.batch_norm_gather_stats_with_counts leads to aten\src\ATen\native\cuda\Normalization.cuh.
The function you’re finding is batch_norm_reduce_statistics_kernel.
In the loop starts at L405, you can see that all statistics are calculated w.r.t. their own count.
 for (int j = 0; j < world_size; j++) {
  scalar_t count = counts[j];
  accscalar_t m = vec_mean[j][i];
  accscalar_t v = accscalar_t(1.0) / (vec_invstd[j][i]);
  v = (v * v - epsilon) * count;
  accscalar_t factor = 1.0 / (n + count);
  var_n += v + (avg - m) * (avg - m) * n * count * factor;
  avg = n * factor * avg + count * factor * m;
  n += count;
}