Fix bug in scatter

[samnordmann]

Fixes a subtle bug, exposed by #2168

[samnordmann]

I find this bug really counter-intuitive!

I thought it was already checked here that the buffers were contiguous...

[samnordmann]

!build --dist

[cowanmeg]

I find this bug really counter-intuitive!

I thought it was already checked here that the buffers were contiguous...

Huh...I don't know how this didn't come up earlier...

[cowanmeg]

We have assumed that the input tensor is contiguous when lowering comms (probably should have added this contiguous() call before), but our tests have all have contiguous aten inputs, so I'm not certain where the non-contiguity got introduced...

[wujingyue]

Code before #2172 doesn't .contiguous() either, so it has likely existed for a while: https://github.com/NVIDIA/Fuser/pull/2172/files/0cf8e27ac816d7d0751df4e64ec1cb7f442ff442#diff-7312205b6b6416d47bc9ad0e0bc815836e7353d4cbbad00dd88196c2e17baff1L122-L123

[wujingyue]

We have assumed that the input tensor is contiguous when lowering comms (probably should have added this contiguous() call before), but our tests have all have contiguous aten inputs, so I'm not certain where the non-contiguity got introduced...

This is exposed by #2168. But the root cause I believe is that insertReshardings doesn't set allocation domain properly. IIRC, you and @jjsjann123 noticed this potential problem in another PR, but we never get a chance to fix this properly.

Below are fusion IR for the two stages. Although Set.Permute is inserted, because the output of the first stage doesn't have an allocation domain, Set.Permute is allowed to produce a non-contiguous tensor (in this case, strides=[3, 12, 1]) for speed.

Inputs:
  T0_g[ iS0{3}, iS1{4}, iS2{3}, iS3{5} ] (DeviceMesh{1, }), float
Outputs:
  T4_g[ iS15{4}, iS14{3}, iS16{3} ] (DeviceMesh{1, }), float

%kernel_math {
T1_l[ iS4{3}, iS5{4}, iS6{3}, rS7{5} ] (DeviceMesh{1, })
   = reduction( T0_g[ iS0{3}, iS1{4}, iS2{3}, iS3{5} ] (DeviceMesh{1, }), op = add, initial value = float(0), allreduce = false )
T4_g[ iS15{4}, iS14{3}, iS16{3} ] (DeviceMesh{1, })
   = Set.Permute( T1_l[ iS4{3}, iS5{4}, iS6{3}, rS7{5} ] (DeviceMesh{1, }), cache_op=Streaming )
}

[4, 3, 3]
[3, 12, 1]
Inputs:
  T5_g[ ideviceIdx.x17{4}, iS18{3}, iS19{3} ] (DeviceMesh{0, 1, 2, 3, }), float
Outputs:
  T3_g[ iS11{3}, ideviceIdx.x12{4}, iS13{3} ] (DeviceMesh{0, 1, 2, 3, }), float

%kernel_math {
T6_l[ iS21{3}, ideviceIdx.x20{4}, iS22{3} ] (DeviceMesh{0, 1, 2, 3, })
   = Set.Permute( T5_g[ ideviceIdx.x17{4}, iS18{3}, iS19{3} ] (DeviceMesh{0, 1, 2, 3, }), cache_op=Streaming )
T3_g[ iS11{3}, ideviceIdx.x12{4}, iS13{3} ] (DeviceMesh{0, 1, 2, 3, })
   = T6_l[ iS21{3}, ideviceIdx.x20{4}, iS22{3} ] (DeviceMesh{0, 1, 2, 3, })
   + T6_l[ iS21{3}, ideviceIdx.x20{4}, iS22{3} ] (DeviceMesh{0, 1, 2, 3, });
}

I think our options are:

1. Merge this PR, which now sounds like a bandaid because it leaves performance on the table. In the above example, doing reduction+permute in one kernel is likely fast than in two kernels.
2. Fix insertResharding to set allocation domains properly. This sounds like the right fix.
3. Revert Allocation order refactor #2168, which isn't the root cause but is unfortunately the trigger.

I suspect option 2 will take a while and it's a bad idea to leave CI broken, so I'll cross that out. Option 3 is safest because #2168 could triggered other failure cases that we are not aware just yet. Option 1 is suboptimal, but as long as we (👀 @cowanmeg) fix the root cause soon, we should be fine. Wdyt?

[jjsjann123]

Oops. sorry about that. 😛

@samnordmann ignore the thunder tests. there's something with transformer_engine.

[cowanmeg]

Yes, I agree option 2 is the correct fix. Plus we need to set allocation domain for DID parallelism on the leaf domain to work. Hopefully, I can find some time soon to work on this!