bug fixed when stacking pre_scales

Azure · Dec 11, 2023 · df68631 · df68631
1 parent 6e58e1e
commit df68631
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/msamp/megatron/optimizer/distrib_optimizer.py b/msamp/megatron/optimizer/distrib_optimizer.py
@@ -566,7 +566,7 @@ def reduce_model_grads(self, args, timers):    # noqa: C901
                     # pre_scales in the partition `data_parallel_rank`
                     pre_scales = [g.meta.pre_scale for g in fp8_grads[data_parallel_rank]]
                     max_elems_per_rank = max(model._grad_buffer_num_params)
-                    pre_scales = torch.cat(pre_scales)
+                    pre_scales = torch.stack(pre_scales)
                     # padding to max_elems_per_rank
                     pad = max_elems_per_rank - pre_scales.numel()
                     pre_scales = F.pad(pre_scales, (0, pad))