RNG output depends on segment execution order

[jacobhinkle]

I found this when debugging #4973. That PR reorders segment execution order sometimes. When the segments include RNG methods (rand, randn, etc.) the values we compute depend on the segment execution order. This might not be a problem in practice but we do seem to be testing that this is the expected behavior: bin/test_rng --gtest_filter=RNGTest.Normal is an example that runs 4 separate rng-based kernels and reordering the segments causes the test to fail.

Fuser/tests/cpp/test_rng.cpp

Lines 357 to 367 in c490444

	fusion->addInput(size_val);
	fusion->addInput(mean);
	fusion->addInput(std);
	TensorView* tv0 = normal({size_val}, mean, std, DataType::Float);
	TensorView* tv1 = normal({size_val}, mean, std, DataType::Double);
	TensorView* tv2 = randn({size_val}, DataType::Double);
	TensorView* tv3 = randn_like(tv2);
	fusion->addOutput(tv0);
	fusion->addOutput(tv1);
	fusion->addOutput(tv2);
	fusion->addOutput(tv3);

To repro, try making this change then running bin/test_rng --gtest_filter=RNGTest.Normal:

diff --git a/csrc/fusion_segmenter.cpp b/csrc/fusion_segmenter.cpp
index 22bf34d30..f08ff9395 100644
--- a/csrc/fusion_segmenter.cpp
+++ b/csrc/fusion_segmenter.cpp
@@ -5369,6 +5369,12 @@ RuntimeWorkSpace prepareRuntimeOrder(const SegmentedFusion& segmented_fusion) {
   }
 
   runtime_workspace.group_run_order = toposort(segmented_fusion.groups());
+  runtime_workspace.group_run_order = {
+    runtime_workspace.group_run_order[3],
+    runtime_workspace.group_run_order[2],
+    runtime_workspace.group_run_order[1],
+    runtime_workspace.group_run_order[0],
+  };
 
   return runtime_workspace;
 }

Possible fixes:

• We could try and enforce execution order of rng factory functions when sorting segments. It might still be possible to reorder RNG methods within a segment.
• We could introduce a preseg pass where we set the seed and offset of each RNG method based on the definition order, using expr->name() to determine definition order. This is not ideal because it means we would be baking in the random seed.

In any case, I think we should clearly document the circumstances in which we expect pytorch and nvFuser output to match.

[naoyam]

How cumbersome to do the first fix?

[jacobhinkle]

How cumbersome to do the first fix?

I think it's medium cumbersome. We'd want to go through each segment in prepareRuntimeOrder and determine whether it has any methods that need a random seed. For those that do, we'd need to define dependencies in a chain. We could do this by just choosing one of the earlier segment's outputs to act like an input of the later segment.

I think this would be sufficient to enforce the ordering in cases where it's possible. But there are cases where it is not possible. Say we had the following:

addInput(t0)
t1 = randn(sh)
t2 = randn(sh)
t3 = add(t0, t2)
t4 = segment_set(t3)
t5 = mul(t1, t4)
addOutput(t5)

This is getting segmented into two segments, but the segment where we compute t1 is dependent on the segment where we compute t2 so there's no way for us to compute those segments in the "proper" order.

[naoyam]

t1 doesn't seem to be used. Is it intentional?

[jacobhinkle]

t1 doesn't seem to be used. Is it intentional?

Sorry that was a typo. I wanted to just show that t1 is earlier in definition order but later in topological order.

[jacobhinkle]

We have had a similar discussion previously: #557. In that case the issue was that our random values do not match PyTorch's exactly because the torch implementation depends on the parallelization of the kernel whereas nvFuser's rng depend only on the linear logical index of the element being generated so it is independent of parallelization. This issue is slightly different because it concerns the order of selection of seeds and offsets to provide to each kernel. In cases where these seeds and offsets are important, for instance when we need to save dropout seeds and offsets for recomputation in the backward pass, I suspect that we will receive them and handle them explicitly so this issue will be irrelevant.

Fuser/csrc/rng.cpp

Lines 21 to 33 in b67ea0c

	// Note [CUDA graph capture and RNG seed and offset]
	// This is how we handle RNG seeds and offsets in PyTorch. In PyTorch,
	// depending on whether we are running under CUDA graph capture or not, the
	// seed and offset are stored differently. When not under CUDA graph capture,
	// the seed and offset are stored on the host. When under CUDA graph capture,
	// the seed and offset are stored on the device, and there are device pointers
	// to these values. The pointers themselves are stored on the host. For seed,
	// because all random ops within the same CUDA graph always use the same seed,
	// so just dereferencing the pointer is enough. For offset, because each
	// random op has its own offset, so we need to add the offset value on the
	// host to the dereferenced offset pointer. In this ways, the offset value on
	// the host is the intra-cuda-graph offset, and the offset value on the device
	// is the base offset for this entire CUDA graph.

Fuser/csrc/runtime/executor.cpp

Lines 956 to 973 in b67ea0c

	KernelArgumentHolder resolveRNGSeed(
	const kir::Kernel* kernel,
	KernelArgumentHolder& args) {
	ExpressionEvaluator expr_eval;
	KernelArgumentHolder resolved_args;
	resolved_args.reserve(args.size());
	int64_t arg_idx = 0;
	for (auto param : kernel->parameters()) {
	if (param->definition() &&
	param->definition()->isA<kir::GetRNGSeedAndOffsetFromHost>()) {
	resolved_args.push(expr_eval.evaluate(param));
	} else {
	resolved_args.push(args[arg_idx++]);
	}
	}
	return resolved_args;
	}
	} // namespace

If we want to address this, I would imagine a solution would involve evaluating some of these parameters in a prespecified order reflecting the definition order of the RNG ops in the original fusion definition.

[naoyam]

Yeah, it seems we would need to change the fusion segmenter such that the order is preserved. For example, in the above case, if it got segmented into:

Segment 1:
t1 = randn(sh)

Segment 2:
t2 = randn(sh)
t3 = add(t0, t2)
t4 = segment_set(t3)

Segment 3:
t5 = mul(t1, t4)

Then, there's a valid ordering.

At this moment, I think documenting the behavior would be enough.