     } else {
       if (local_shard) {
         if (use_hierarchical_allreduce) {
-          NCCLAllReduceWithScale(fp32_grad,
-                                 fp32_sum_grad,
-                                 fp32_numel,
-                                 nranks / num_devices,
-                                 external_comm,
-                                 stream,
-                                 dev_ctx);
           NCCLReduceScatterWithScale(
-              fp32_sum_grad,
+              fp32_grad,
               fp32_sum_grad + local_rank * fp32_numel_each_device,
               fp32_numel_each_device,
               num_devices,
               local_comm,
               stream,
               dev_ctx);
-
-          NCCLAllReduceWithScale(fp16_grad,
-                                 fp16_sum_grad,
-                                 fp16_numel,
-                                 nranks / num_devices,
-                                 external_comm,
-                                 stream,
-                                 dev_ctx);
+          NCCLAllReduceWithScale(
+              fp32_sum_grad + local_rank * fp32_numel_each_device,
+              fp32_sum_grad + local_rank * fp32_numel_each_device,
+              fp32_numel_each_device,
+              nranks / num_devices,
+              external_comm,
+              stream,
+              dev_ctx);
+
           NCCLReduceScatterWithScale(
-              fp16_sum_grad,
+              fp16_grad,
               fp16_sum_grad + local_rank * fp16_numel_each_device,
               fp16_numel_each_device,
               num_devices,
               local_comm,
+              stream,
+              dev_ctx);
+          NCCLAllReduceWithScale(
+              fp16_sum_grad + local_rank * fp16_numel_each_device,
+              fp16_sum_grad + local_rank * fp16_numel_each_device,
+              fp16_numel_each_device,
+              nranks / num_devices,
+              external_comm,
               stream,
               dev_ctx);
         } else {

        // (3) Do ReduceScatter with scale
        if (local_shard) {
          if (use_hierarchical_allreduce) {
            NCCLAllReduceWithScale(fp32_grad,
                                   fp32_sum_grad,
                                   fp32_numel,
                                   nranks / num_devices,
                                   external_comm,
                                   stream,
                                   dev_ctx,
                                   fp32_scale);
            NCCLReduceScatterWithScale(
                fp32_sum_grad,
                fp32_sum_grad + local_rank * fp32_numel_each_device,
                fp32_numel_each_device,
                num_devices,
                local_comm,
                stream,
                dev_ctx);

            NCCLAllReduceWithScale(fp16_grad,
                                   fp16_sum_grad,
                                   fp16_numel,
                                   nranks / num_devices,
                                   external_comm,
                                   stream,
                                   dev_ctx,
                                   fp16_scale);
            NCCLReduceScatterWithScale(
                fp16_sum_grad,
                fp16_sum_grad + local_rank * fp16_numel_each_device,
                fp16_numel_each_device,
                num_devices,
                local_comm,
                stream,
                dev_ctx);
          } else {

        // (1) ReduceScater first
        if (local_shard) {
          if (use_hierarchical_allreduce) {
            NCCLReduceScatterWithScale(
                fp32_grad,
                fp32_sum_grad + local_rank * fp32_numel_each_device,
                fp32_numel_each_device,
                num_devices,
                local_comm,
                stream,
                dev_ctx);
            NCCLAllReduceWithScale(
                fp32_sum_grad + local_rank * fp32_numel_each_device,
                fp32_sum_grad + local_rank * fp32_numel_each_device,
                fp32_numel_each_device,
                nranks / num_devices,
                external_comm,
                stream,
                dev_ctx);

            NCCLReduceScatterWithScale(
                fp16_grad,
                fp16_sum_grad + local_rank * fp16_numel_each_device,
                fp16_numel_each_device,
                num_devices,
                local_comm,
                stream,
                dev_ctx);
            NCCLAllReduceWithScale(
                fp16_sum_grad + local_rank * fp16_numel_each_device,
                fp16_sum_grad + local_rank * fp16_numel_each_device,
                fp16_numel_each_device,
                nranks / num_devices,
                external_comm,
                stream,
                dev_ctx);
          } else {
