DID loop split on reshaped IDs

[Priya2698]

This PR updates propagateReshapeTransform to support DID loop split.

When the loop split is on the iterdomains being reshaped, the logical reshaped iterdomain is no longer present in the loop domain since it is split. In this case, we check if there is a sharded loop ID and compare the logical reshape iterdomain to the producer of this DID split.

[Priya2698]

!test

[github-actions]

Review updated until commit 73d237a

Description

• Added support for DID loop split in propagateReshapeTransform

• Enhanced propagateReshapeTransform to handle reshaped iterdomains

• Added a new test case for transform propagation with reshape

---

Changes walkthrough 📝

	Relevant files
Enhancement	utils.cpp Handle DID loop split in propagateReshapeTransform              csrc/scheduler/utils.cpp Added logic to check if logical ID is directly in the loop domain Added handling for sharded loop ID and DID split producer comparison Reordered reshape dimensions to the front of the domain +21/-1
Tests	test_multidevice_sharding.cpp Add test for transform propagation with reshape                    tests/cpp/test_multidevice_sharding.cpp Added a new test case TransformPropagatorWithReshape Demonstrates transform propagation with reshape and DID loop split +60/-0

---

PR Reviewer Guide 🔍

Here are some key observations to aid the review process:

🧪 PR contains tests

⚡ Recommended focus areas for review Error Handling The error message in NVF_ERROR could be more descriptive, especially regarding the context of the DID loop split and reshaped IDs. "Require ", logical_id, " is in the active domain of ", tv->toString(), " for view propagation."); Code Clarity The logic for handling the DID loop split and reshaped IDs could be refactored for better readability and maintainability. // Check if logical ID is directly in the loop domain auto find_it = std::find( tv->getLoopDomain().begin(), tv->getLoopDomain().end(), logical_id); // If not found directly and there is a sharded loop ID, // check if the logical ID is the same as the producer of the DID split. if (find_it == tv->getLoopDomain().end()) { int64_t sharded_axis = getShardedLoopAxis(tv, ParallelType::DIDx); if (sharded_axis != -1) { // Get the split operation that created the DIDx dimension auto split = dynamic_cast<Split*>( tv->getLoopDomain().at(sharded_axis)->definition()); if (split != nullptr && split->in() == logical_id) { find_it = std::find( tv->getLoopDomain().begin(), tv->getLoopDomain().end(), split->inner()); } } } NVF_ERROR( Test Coverage Ensure that the new test case covers all edge cases and potential failure scenarios related to DID loop split and reshaped IDs. TEST_F(MultiDeviceTest, TransformPropagatorWithReshape) { auto fusion = std::make_unique<Fusion>(); FusionGuard fg(fusion.get()); const int d = communicator_->size(); const int64_t b = 2, s = 2, h = 4, e = 3; TensorView* in = makeContigConcreteTensor( {b, s, d * h * e}); // in: loop domain: {b, s, d*h*e} TensorView* out = reshape( in, {b, s, d * h * e}, {b, s, d * h, e}); // out: loop domain: {b, s, d*h, e} fusion->addInput(in); fusion->addOutput(out); auto mesh = DeviceMesh::createForNumDevices(d); // Propagate transform from reshaped output to input. // Without this propagation, the two DID axes on `in` and `out` will not be // mapped in together in ID model. This causes scheduling to fail due to // resharding. TransformPropagator propagator_c2p(out); MaxLogicalDomainInfoSpanningTree(out).traverse(&propagator_c2p); // in: loop domain: {b, s, d*h, e} after transform propagation // Loop split and parallelize input in->setDeviceMesh(mesh); in->split(-2, d, /*inner_split=*/false); in->axis(-3)->parallelize(ParallelType::DIDx); // in: loop domain: {b, s, DIDx{d}, h, e} // Propagate DID loop split to output TransformPropagator propagator_p2c(in); MaxLogicalDomainInfoSpanningTree(in).traverse(&propagator_p2c); // out: loop domain: {b, s, d, h, e} after transform propagation // Parallelize output scheduler_utils::parallelizeAllLike( in, /*pos=*/-1, /*selected_tv=*/{out}); // out: loop domain: {b, s, DIDx{d}, h, e} after parallelization in->setAllocationDomain(in->getLoopDomain(), true); out->setAllocationDomain(out->getLoopDomain(), true); FusionExecutorCache executor_cache(std::move(fusion)); at::Tensor in_tensor = at::randn({b, s, h * e}, tensor_options); at::Tensor out_tensor = executor_cache.runFusionWithInputs({in_tensor})[0]; testValidate( executor_cache.fusion(), {out_tensor}, {in_tensor}, {in_tensor.view({b, s, h, e})}, __LINE__, __FILE__); }

[Priya2698]

!test

[wujingyue]

@naoyam I'm quite confused by this pre-existing logic before I can understand this PR. Why is it necessary to move reshape dimensions to the front of the loop domain? It can cause conflict with the pre-existing assumption that DIDs have to be the front as well.

[naoyam]

I think this is related to the propagation done at line 2279. IIRC, it propagates the outermost N dimensions, where N is old2new.size() in this case. Since here we just want to propagate the transformations related to the rfactor, this is how we limit the propagation.

We can probably just reorder tv back after line 2279.

[naoyam]

I think I see what the below part is trying to do and why, which seems to make sense, but can you expand the comment and elaborate a little more?

[wujingyue]

This bears several assumptions that will break in a foreseeable future.

With context parallelism, the sequence dimension s will be split into [tp, iDIDy{cp}, s/tp/cp], so the code below won't be able to find tp and s/tp/cp. Similarly, with overlapping, the sequence dimension s will be split into [sp, iDID{tp}, s/sp/tp] where the sp is the stream parallelization factor. See this test for the idea.

I understand this change does fix some narrow cases that we care about at this very moment, but I'll have to think more about how to fix the broader issue...

[wujingyue]

(I still haven't given up on improving ID model.)

If we have to do graph traversal like this, we may want to do it in a place where the logic can be generalized and reused (and therefore ID model). At this moment, there are two use cases:

1. splitting reshape: [h]=>[d, h/d] and [h]=>[d,a/d,h/a]
2. merging reshape: [a,h/a]=>[d,a/d,h/a] and [a,h/a]=>[h]=>[d,h/d]
   We want ID model to map the ds in both cases so these reshapes won't be considered resharding.

How much harder is it to make ID model support these cases than working around using reshape transformation? I suspect the latter has a bigger blast radius because the former is local to ID model and the latter changes TensorViews.

[Priya2698]

I realized the same limitation for DID loop split on slice:

  auto fusion = std::make_unique<Fusion>();
  FusionGuard fg(fusion.get());

  const int d = communicator_->size();
  const int64_t b = 2, s = 2, h = 4;

  TensorView* in = makeContigConcreteTensor(
      {b, s, 3 * d * h});
  TensorView* out = slice(
      in,
      {0, 0, 0},
      {b, s, d * h});

  fusion->addInput(in);
  fusion->addOutput(out);

  auto mesh = DeviceMesh::createForNumDevices(d);
  for (auto tv: {in, out}) {
    tv->setDeviceMesh(mesh);
    tv->split(-1, d, /*inner_split=*/false);
    tv->axis(-2)->parallelize(ParallelType::DIDx);
    tv->setAllocationDomain(tv->getLoopDomain(), true);
  }

I was trying to manually handle the case of SliceOp in hasDifferentShardings but it would make certain assumptions about the parallelization patterns and can easily break.

[Priya2698]

I understand this change does fix some narrow cases that we care about at this very moment, but I'll have to think more about how to fix the broader issue...

Yes, I agree. I wanted to add an example to demonstrate how reshapes can be loop split but it certainly does not cover all the cases.

[wujingyue]

I still haven't given up on improving ID model.

(Sorry -- I wish I knew more about IdModel to be more constructive.)

Another use case to consider is manual sharding -- the user wants to manually shard a subset of TVs to improve perf when our sharding propagation is suboptimal.

They may well annotate [b,s,h]=>(reshape)=>[b,s,a,h/a] as follows

in: [b,s,h] => [b,s,d,h/d]
out: [b,s,a,h/a] => [b,s,d,a/d,h/a]

and expect nvFuser to recognize this reshape is local. In this case, it's hard to replay the reshape on the input because h there is already split by d.

[wujingyue]

While I understand your intention, I don't know the implications of doing this for TransformPropagator. E.g.

root=[b, s, h], logical=[b, s, a, h/a], loop=[b, s, d, a/d, h/a]

This will move a/d and h/a to the front so the new loop domain becomes [a/d, h/a, b, s, d] and later ask TransformPropagator to replay at replayed_pos_ 2. What is TransformPropagator supposed to do with that? The first two loop IDs (a/d and h/a) don't even form a split in this TV.

cc @naoyam

[Priya2698]

I see your point. In the case of reshape with DID loop split, we have already propagated the reshape upwards, so the TransformPropagator only reorders the axis when called later. In the absence of the earlier reshape propagation before the loop split, the behavior could be erroneous since they don't form a split.

Although, since the reshape has already been propagated, and, as @naoyam mentioned above, the tv is reordered back, maybe this propagation can be skipped altogether.
Let me think about it more and see what the schedulers expect from this propagateReshapeTransform.

However, this may not work for the manual sharding case you mentioned in the above comment.