Inconsistent values for the same extent `Val*`

[zasdfgbnm]

As demonstrated in #5708, The same extent Val* could have different multiplier. This is a blocker for my PR #5082, because this could lead to PrecomputedValues binding the same Val* with different value for the same instance of FusionExecutorCache.

Any suggestion? cc: @wujingyue

[wujingyue]

I've yet to understand why

Fuser/csrc/multidevice/execution_utils.cpp

Line 73 in 878c3e2

// TODO(#5525): hack for MultiDeviceExecutor. MultiDeviceExecutor looks

can create problems.

I failed to reproduce this reliably. I suspect it has something to do with the use of static: https://github.com/NVIDIA/Fuser/pull/5708/files#diff-140da8af7df6a9ae6ddfbc972a5db1bfdd0efba4b6412902f285b32f3d131347R72. Some Val*s from a previous test might have stayed in the hashtable even though they've been deleted.

[wujingyue]

Repro

$ git fetch origin eval-segment
$ git checkout eval-segment
$ mpirun -np 1 pytest tests/python/multidevice/test_overlap.py -k test_row_parallel_linear_forward --only-mpi -vs
...
E       RuntimeError:  INTERNAL ASSERT FAILED at /opt/pytorch/nvfuser/csrc/evaluator_common.cpp:371, please report a bug with repro script to NVFuser at https://github.com/NVIDIA/Fuser/issues. 
E       Expected isSame(values_[it.first], it.second) . Precomputed values failed to validate.
E       Something unexpected changed between the compilation and execution.
E       6 != 18
E       Exception raised from validate at /opt/pytorch/nvfuser/csrc/evaluator_common.cpp:371 (most recent call first):
...

Root cause

LinearOp::evaluate doesn't respect the stream parallel type in allocation. The LinearOp's output allocation ought to be stream-parallelized but LinearOp::evaluate still produces a tensor of shape [6, 2, 1]. Below is the preseg fusion IR. The problematic LinearOp is the local LinearOp before allreduce.

This problem doesn't show up today because:

1. Without Use ATen ops on meta tensor to compute output shapes and strides for ExprEval segments #5082, we never use ExpressionEvaluator to evaluate a stream-parallel LinearOp.
2. In HostIrEvaluator, stream-parallel linears take the result of ShardByStreams (

   Fuser/tests/python/multidevice/test_overlap.py

   Line 75 in 75731ee

   # T4_l___bfloat[istreamIdx31{3}, ideviceIdx.x33{2}, iS32{( ceilDiv(i0, 3) )}, iS34{4}] (DeviceMesh{0 1}) = ShardByStream(T0_g___bfloat[istreamIdx7{3}, ideviceIdx.x9{2}, iS8{( ceilDiv(i0, 3) )}, iS10{4}] (DeviceMesh{0 1}), stream_index = i535)

   ) so the m dimension is 2 not 6. ShardByStreams are inserted by host_ir/lowering.cpp.

Fusion IR after pre-segmenter optimization passes:
Inputs:
  T0_g___bfloat[istreamIdx7{3}, ideviceIdx.x9{1}, iS8{( ceilDiv(i0, 3) )}, iS10{8}] (DeviceMesh{0})
  T1_g___bfloat[ideviceIdx.x11{1}, iS2{2}, iS12{8}] (DeviceMesh{0})
Outputs:
  T2_g___bfloat[istreamIdx27{3}, rdeviceIdx.x26{1}, iS28{( ceilDiv(i0, 3) )}, iS25{2}] (DeviceMesh{0})

%kernel {
T3_g___bfloat[istreamIdx20{3}, ideviceIdx.x22{1}rf, iS21{( ceilDiv(i0, 3) )}, iS18{2}, rS23{8}rf] (DeviceMesh{0})
   = linear(T0_g___bfloat[istreamIdx7{3}, ideviceIdx.x9{1}, iS8{( ceilDiv(i0, 3) )}, iS10{8}] (DeviceMesh{0}),
            T1_g___bfloat[ideviceIdx.x11{1}, iS2{2}, iS12{8}] (DeviceMesh{0})  )
T2_g___bfloat[istreamIdx27{3}, rdeviceIdx.x26{1}, iS28{( ceilDiv(i0, 3) )}, iS25{2}] (DeviceMesh{0})
   = reduction( T3_g___bfloat[istreamIdx20{3}, ideviceIdx.x22{1}rf, iS21{( ceilDiv(i0, 3) )}, iS18{2}, rS23{8}rf] (DeviceMesh{0}), op = add, initial value = __bfloat(0), allreduce = false )

TransformPrinter : 
T0_g___bfloat[istreamIdx7{3}, ideviceIdx.x9{1}, iS8{( ceilDiv(i0, 3) )}, iS10{8}] (DeviceMesh{0})
 logical domain : (iS0{i0}, iS1{8})
 allocation domain : (iS0{i0}, ideviceIdx.x9{1}, iS10{8})
 contiguity: t t t
  Outer split: iS0{i0} by factor 3 -> istreamIdx7{3}, iS8{( ceilDiv(i0, 3) )}
  Outer split: iS1{8} by factor 1 -> ideviceIdx.x9{1}, iS10{8}
 loop domain : (istreamIdx7{3}, ideviceIdx.x9{1}, iS8{( ceilDiv(i0, 3) )}, iS10{8})
T1_g___bfloat[ideviceIdx.x11{1}, iS2{2}, iS12{8}] (DeviceMesh{0})
 logical domain : (iS2{2}, iS3{8})
 allocation domain : (iS2{2}, ideviceIdx.x11{1}, iS12{8})
 contiguity: t t t
  Outer split: iS3{8} by factor 1 -> ideviceIdx.x11{1}, iS12{8}
 loop domain : (ideviceIdx.x11{1}, iS2{2}, iS12{8})
T3_g___bfloat[istreamIdx20{3}, ideviceIdx.x22{1}rf, iS21{( ceilDiv(i0, 3) )}, iS18{2}, rS23{8}rf] (DeviceMesh{0})
 root domain : (iS17{i0}, iS18{2}, rS19{8}rf)
  Outer split: rS19{8}rf by factor 1 -> ideviceIdx.x22{1}rf, rS23{8}rf
 logical domain : (iS17{i0}, iS18{2}, ideviceIdx.x22{1}rf, rS23{8}rf)
 allocation domain : (istreamIdx20{3}, iS21{( ceilDiv(i0, 3) )}, iS18{2}, ideviceIdx.x22{1}rf, rS23{8}rf)
 contiguity: t t t t n
  Outer split: iS17{i0} by factor 3 -> istreamIdx20{3}, iS21{( ceilDiv(i0, 3) )}
  Outer split: rS19{8}rf by factor 1 -> ideviceIdx.x22{1}rf, rS23{8}rf
 loop domain : (istreamIdx20{3}, ideviceIdx.x22{1}rf, iS21{( ceilDiv(i0, 3) )}, iS18{2}, rS23{8}rf)
T2_g___bfloat[istreamIdx27{3}, rdeviceIdx.x26{1}, iS28{( ceilDiv(i0, 3) )}, iS25{2}] (DeviceMesh{0})
 logical domain : (iS24{i0}, iS25{2}, rdeviceIdx.x26{1})
 allocation domain : (iS24{i0}, iS25{2}, rdeviceIdx.x26{1})
 contiguity: t t n
  Outer split: iS24{i0} by factor 3 -> istreamIdx27{3}, iS28{( ceilDiv(i0, 3) )}
 loop domain : (istreamIdx27{3}, rdeviceIdx.x26{1}, iS28{( ceilDiv(i0, 3) )}, iS25{2})
} // %kernel

[wujingyue]

I'm not sure the best way to fix this.

Option 1: make LinearOp::evaluate support stream-parallel allocation domains. Other ops likely have this problem too, so we'll have to fix many *Op::evaluate. Side question: do meta ops have a similar problem? E.g.

Fuser/csrc/ir/internal_nodes.cpp

Line 2425 in a53ae37

std::vector<PolymorphicValue> ReshapeOp::evaluate(

doesn't check the output allocation domain so shape inference as implemented in #5082 could produce a meta tensor that has the right shape but wrong strides.

Option 2: Let preseg insert LoadStoreOp wherever a switch between stream-parallel allocation and non-stream-parallel allocation happens. For example, in this particular test, we'll expect a LoadStoreOp between the LinearOp and the input activation. This way, the input meta tensor of the LinearOp will be of shape [t/s=2, h*4/d] instead of [t=6, h*4/d]. The benefit is to not have to change every *Op::evaluate.

Option 3: Let shape inference work on global tensors. For example, the output meta tensor of the LinearOp will be of shape [t, h, d]. Then, we have to remember to bind the shape directly to the logical domain without unshardedSizes multiplying t by s, which would lead to t=6 * s=3 = 18. That'll probably lead to a bifurcation in the "binding" logic -- don't call unshardedSizes if during shape inference and/or the tensor is on meta.

[zasdfgbnm]

Thanks @wujingyue for looking into it and for the detailed explanation!

To me, option 1 sounds like the cleanest. Although many *Op::evaluate needs to be changed, the added complexity sounds most manageable. I believe I should proceed with landing #5082 first by making it an enable option. Then I can incrementally enable all tests, in Python and C++, in separate PRs. What do you think? @wujingyue

[wujingyue]

I believe I should proceed with landing #5082 first by making it an enable option.

SGTM.

For the initial PR, I recommend we go simple and broken. I.e., let's not target passing all unit tests with that new code path, but rather do the right thing without hacks like https://github.com/NVIDIA/Fuser/pull/5082/files#r2381242779. What do you think?