Fix decomposeLinearWithBias to shard all created tensorviews

csrc/preseg_passes/decompose_reshardings.cpp

@@ -313,7 +313,7 @@ void decomposeRowParallelLinearWithBias(Fusion* fusion) {
     }
 
     auto* without_bias = linear(linear_op->inA(), linear_op->inB());
-    TransformReplay::selfReplay(out->domain(), without_bias->domain());
+    auto* upcast_without_bias = maybeCastOp(DataType::Float, without_bias);
 
     TensorView* broadcasted_bias = [&]() {
       const int64_t rank_after_broadcast = std::ssize(
@@ -327,10 +327,15 @@ void decomposeRowParallelLinearWithBias(Fusion* fusion) {
       return broadcast(linear_op->bias(), is_broadcast_dim);
     }();
 
-    TensorView* new_out =
-        maybeCastOp(out->dtype(), add(without_bias, broadcasted_bias));
-    TransformReplay::selfReplay(out->domain(), new_out->domain());
+    TensorView* with_bias = add(upcast_without_bias, broadcasted_bias);
+    TensorView* new_out = maybeCastOp(out->dtype(), with_bias);
+
     ir_utils::replaceValInAllExprInputsAndFusionOutputs(out, new_out);
+
+    for (TensorView* tv :
+         {without_bias, upcast_without_bias, with_bias, new_out}) {
+      TransformReplay::selfReplay(out->domain(), tv->domain());
+    }
   }
 }
 

csrc/runtime/communication_executor.cpp

@@ -84,7 +84,7 @@ KernelArgumentHolder CommunicationExecutor::run(
         group_id_);
     SegmentProfiler& sprof = FusionProfiler::segment(group_id_);
     sprof.inputBytesAccessed(computeBytes(args));
-    sprof.scheduler(toString(SchedulerType::ExprEval));
+    sprof.scheduler(toString(SchedulerType::Communication));
     sprof.startKernel();
   }
   NVF_ERROR(host_ir_container_, "Need to compile before you can run.");

tests/python/multidevice/test_matmul.py

@@ -6,7 +6,7 @@
 import torch
 
 import nvfuser_direct as nvfuser
-from nvfuser_direct import DataType, FusionDefinition
+from nvfuser_direct import DataType, FusionDefinition, PythonProfiler
 
 
 # Avoid doing this when possible. This test started to exist before nvFuser
@@ -200,47 +200,57 @@ def test_linear_reduce_scatter(multidevice_direct_test):
     e = 768
 
     def _definition(fd: FusionDefinition):
-        inp = fd.define_tensor([-1, -1, d * e])
-        weight = fd.define_tensor([e, d * e])
-        out = fd.ops.linear(inp, weight, None)
+        inp = fd.define_tensor([-1, -1, d * e], dtype=DataType.BFloat16)
+        weight = fd.define_tensor([-1, d * e], dtype=DataType.BFloat16)
+        bias = fd.define_tensor([e], dtype=DataType.BFloat16)
+        out = fd.ops.linear(inp, weight, bias)
         fd.add_output(out)
 
     def _multidevice_schedule(fd: FusionDefinition):
-        inp, weight = fd.fusion.inputs()
+        inp, weight, bias = fd.fusion.inputs()
         (out,) = fd.fusion.outputs()
-        for t in [inp, weight, out]:
-            t.set_device_mesh(mesh)
-            t.split(-1, d, inner_split=False)
-            t.axis(-2).parallelize(nvfuser.ParallelType.mesh_x)
+        bias.set_device_mesh(mesh)
+        for tv in [inp, weight, out]:
+            tv.set_device_mesh(mesh)
+            tv.split(-1, d, inner_split=False)
+            tv.axis(-2).parallelize(nvfuser.ParallelType.mesh_x)
 
         # Scatter
         out.split(1, d, inner_split=False)
         out.axis(1).parallelize(nvfuser.ParallelType.mesh_x)
 
     torch.cuda.set_device(multidevice_direct_test.local_rank)
 
-    # set b=1 as a temporary fix for the test to pass.
-    # TODO: set b>1 once reduce scatter is fixed.
-    b, s = 2, 1024
-    unsharded_inp = torch.randn(b, s, d * e)
-    unsharded_weight = torch.randn(e, d * e)
-
+    b, s = 2, 8
+    unsharded_inp = torch.randint(-2, 3, (b, s, d * e)).to(torch.bfloat16)
+    unsharded_weight = torch.randint(-2, 3, (e, d * e)).to(torch.bfloat16)
+    bias = torch.randint(-2, 3, (e,)).to(torch.bfloat16)
     inp = multidevice_direct_test.shard_tensor(unsharded_inp, -1, mesh)
     weight = multidevice_direct_test.shard_tensor(unsharded_weight, -1, mesh)
 
     with FusionDefinition() as fd:
         _definition(fd)
         _multidevice_schedule(fd)
 
-    (out,) = fd.execute([inp, weight])
+    with PythonProfiler() as prof:
+        (out,) = fd.execute([inp, weight, bias.cuda()])
 
-    unsharded_out = torch.nn.functional.linear(unsharded_inp, unsharded_weight, None)
-    # rtol is the same as the default for fp32. atol is slightly increased.
+    # Only one reduce scatter kernel should be scheduled.
+    assert (
+        len(
+            [
+                kp
+                for kp in prof.profile.kernel_profiles
+                if kp.scheduler == "communication"
+            ]
+        )
+        == 1
+    )
+
+    unsharded_out = torch.nn.functional.linear(unsharded_inp, unsharded_weight, bias)
     torch.testing.assert_close(
         out,
         multidevice_direct_test.shard_tensor(unsharded_out, 1, mesh),
-        rtol=1.3e-6,
-        atol=1e-3,
     )