[CUTLASS] Add parallel split-k support to wgrad

[masahi]

Building on #10177, this adds parallel split-k support to wgrad.

@comaniac @Laurawly @junrushao1994 @vinx13 @YuchenJin @hwu36 @manishucsd

Split-k is described in https://github.com/NVIDIA/cutlass/blob/master/media/docs/efficient_gemm.md#parallelized-reductions.
This is my first experience using split-k in cutlass or any other API. Wgrad is particularly interesting for split-k since the implicit gemm K dimension is really large in wgrad (N * P * Q where P and Q are the output H and W). Without split-k, wgrad on large spatial inputs is extremely slow.

For now, I'm not trying anything smart to pick the split-k parameter, instead we ask users to provide possible candidates. I tuned over [1, 4, 8, 16, 32, 64] below and that already showed excellent performance. The benchmark code is here.

Benchmark result against cuDNN. Note that currently there are non-trivial difference in cuDNN and TVM + cutlass outputs, especially for the larger batch size. I didn't find anything obviously wrong in the generated code and I gave up fixing accuracy difference at some point. Also note that difference is not due to parallel-split-k, even in a normal case the results were different (and actually improved after split-k lol).

• Batch size 8
• Batch size 256

The result showed that cutlass winning across the board (Profiler time vs cuDNN columns, but again, the results do not match exactly). However, there is a serious problem when cutlass wgrad + split-k kernels are called from TVM (TVM + CUTLASS column): Split-k requires large workspace, and the space requirement grows linearly with split-k-slices parameter. Right now we naively allocate the workspace on every cutlass kernel call on each run, while for cuDNN we have a simple workspace memory reuse mechanism implemented in (together with a thread local storage)

tvm/src/runtime/contrib/cudnn/cudnn_utils.cc

Lines 153 to 161 in 211291f

	void ConvEntry::UpdateWorkspace(const size_t wsize) {
	if (workspace_size < wsize) {
	if (workspace != nullptr) {
	CleanWorkspace();
	}
	workspace_size = wsize;
	workspace = cuda_api->AllocWorkspace(device, workspace_size);
	}
	}

. So during benchmarking, we run compiled TVM + CUTLASS 100 times, and each time we are allocating the same workspace 🤦‍♂️ while cuDNN allocates only once during algo selection outside of the profiling loop. This is causing huge perf penalty in TVM + cutlass results above.

I attempted adding a simple workspace memory management in https://github.com/masahi/tvm/compare/cutlass-split-k...masahi:cutlass-workspace?expand=1, it kind of works in terms of the expected perf improvement. However, I'm getting segfault and other strange issues. I'm a bit confused as to what the right behavior should be for a thread local memory manager in the context of JIT- generated and compiled multiple translation units. Let me know if you have any thoughts on this issue.

Known issues and TODO

• Accuracy alignment with cuDNN
• Figure out workspace memory reuse
• Split-k parameter selection strategy
• Support split-k in GEMM and other conv2d kind

[hwu36]

If you want to investigate accuracy issue, i suggest you compare both cutlass and cudnn with a naive fp64 or fp32 version.

[jroesch]

cc @mbs-octoml interesting example for perf work

[manishucsd]

Hi Masa, This is amazing progress. Some questions on the known issues:

• Accuracy alignment with cuDNN
  Can you share the size that has accuracy issues. Can you repro the accuracy issue in profiler?
• Figure out workspace memory reuse
  Both cuDNN and CUTLASS offers similar get_workspace_size(...) API. Thus, I believe this part should be similar.
• Split-k parameter selection strategy
  As we discussed in an another thread, we run sweeps to find the best split. You can cut down the sweep on k by using a simple analytic model.

[masahi]

Hi Manish,

Can you share the size that has accuracy issues. Can you repro the accuracy issue in profiler?

The benchmark result I linked above show accuracy difference in the last two columns. Most workload have some differences, except for some deeper layers in batch = 8 which showed exact match. It seems deeper layers, those having small spatial size and large channels, have generally less accuracy problems. The differences become much bigger for batch = 256. So it kind of works but not quite, it is very hard to debug. The profiler in cutlass doesn't report any accuracy problem, which is another mystery. It could be TVM's use of cuDNN wgrad having some issues.

Both cuDNN and CUTLASS offers similar get_workspace_size(...) API. Thus, I believe this part should be similar.

The issue is memory reuse across multiple calls. The way we integrate cuDNN and cutlass are significantly different. I tried to apply a similar memory management strategy we use for cuDNN to the JIT-generated cutlass, but as I said above I'm having strange issues.

As we discussed in an another thread, we run sweeps to find the best split. You can cut down the sweep on k by using a simple analytic model.

Yes, I haven't grokked your note in that thread. I just tried a dumb strategy in my benchmark and it already shows good performance. I didn't pursue perf improvement further, since the accuracy problem was more concerning.

[manishucsd]

On accuracy, floating point additions are not associative. The change the order can change the result. Parallel reduction does change the order of accumulation over GEMM-K (NPQ). Thus, some change between runs is expected. I don't have a guidance on what threshold to set in checking relative error.

I would take Haicheng's suggestions here and follow:

If you want to investigate accuracy issue, i suggest you compare both cutlass and cudnn with a naive fp64 or fp32 version.
Run FP32 wgrad with no split-k and compare both cutlass and cudnn against this golden reference.

CUTLASS profiler uses integer input to initialize tensors and matrices. This is to make the error checking easier. You can also use the CUTLASS profiler approach to make sure there are no functional error, i.e., try the operation on integer input.

[masahi]

Actually, accuracy difference was there even before I added parallel split-k to wgrad. And that the result got closer to cuDNN after adding split-k. So I believe the issue is not in parallel reduction, there is something off elsewhere. I have seen some workload where cuDNN uses cutlass's wgrad and reduction kernel, even in that case there was difference. Probably I should look at how TVM is using cuDNN wgrad first.

I haven't applied fp32 wgrad on large inputs, for small ones we used in the unit test, the result looked good. We also have an option of comparing against TVM native results, which I only looked briefly.

CUTLASS profiler uses integer input to initialize tensors and matrices. This is to make the error checking easier. You can also use the CUTLASS profiler approach to make sure there are no functional error, i.e., try the operation on integer input.

That's very interesting... I didn't know that. I can definitely try, thanks.