[ROCm] support torch._C._set_sm_carveout_experimental - Parity with Nvidia

[functionstackx]

🐛 Describe the bug

Hi @hliuca

On Nvidia, they support torch._C._set_sm_carveout_experimental for better compute-comms overlapping. this is useful during bwd pass of DDP and fwd/bwd pass of FSDP to ensure there is enough available SM/CUs for the rccl comms kernels to not be blocked by compute kernels that use up all the SM/CUs

Furthermore, it is useful towards benchmarking real world GEMMs that occurs in the backwards pass when the GEMM is unavailable to take up all the available SM/CUs due to rccl comms kernels occupying some of the SM/CUs

related to #147966

I was looking into implementing this myself but it seems like it isn't as simple as calling hipblasLtMatmulDescSetAttribute as it requires changes to hipblaslt itself since unlike cublasLtMatmulDescSetAttribute, HIPBLASLT_MATMUL_DESC_CU_COUNT_TARGET is not an option for hipblasLtMatmulDescSetAttribute function which takes in enum of hipblasLtMatmulDescAttributes_t at least according to the AMD docs

https://rocm.docs.amd.com/projects/hipBLASLt/en/latest/datatypes.html#_CPPv431hipblasLtMatmulDescAttributes_t

computeDesc.setAttribute<int32_t>(
        CUBLASLT_MATMUL_DESC_SM_COUNT_TARGET,
        at::cuda::getCurrentDeviceProperties()->multiProcessorCount -
            at::globalContext()._SMCarveout_EXPERIMENTAL().value());
  }

Versions

any rocm torch version

cc @jeffdaily @sunway513 @jithunnair-amd @pruthvistony @ROCmSupport @dllehr-amd @jataylo @hongxiayang @naromero77amd

[hliuca]

@OrenLeung we have created internal jira to solve this. Thank you.

For reference: SWDEV-522108

cc @jeffdaily @pruthvistony

[jeffdaily]

This is still very experimental but attempted a solution here #149466.

[functionstackx]

interesting solution creating a bitmask to mask out CUs with hipExtStreamCreateWithCUMask. seems like it would the kernel from CU index [0, value) even if [0, value) already has occupied with a kernel but [value, max_CU) is idle.

i wonder what would the steps be to implement this in hipblaslt directly be. not too familar with hipblaslt myself and unfortunately cublaslt is closed source so unable to look at an example to pattern match from

[hliuca]

In hipBLASLt, I think the Gemm performance is tuned to use all CUs. We may modify some code when calling hipBLASLt and add some CUMask to force hipBLASLt to use some of the CUs, but we may see performance impact.

[functionstackx]

@hliuca in real world inference & training, even if CU carveout isn't explicitly set, the GEMM would never be using up all the CUs anyways due to compute-comms overlapping.

[hliuca]

Thank you @OrenLeung nice explanation and image. I was thinking if we needed to consider this in Gemm performance optimization.

[functionstackx]

I think so. iirc the way that cublas GEMMs r tuned internally is that they tune it to the common amounts of "available SMs" . they look at how many SMs something like a mlperf nccl takes and and tune according such that the out of box real world GEMM performance is properly tuned to an realistic SM count

https://github.com/mlcommons/training_results_v4.0/blob/main/NVIDIA/benchmarks/gpt3/implementations/nemo/conf/nccl/custom_communicator_cta.yaml

[hliuca]

I will pass this to BLAS team. Thank you @OrenLeung

[functionstackx]

thanks @hliuca

[functionstackx]

u can see this in other gemm libraries too such as deepgemm where they have deep_gemm.set_num_sms to set the number of SMs for the compute to be used during development to be the same amount as during prod

https://github.com/deepseek-ai/DeepGEMM?tab=readme-ov-file#utilities

https://github.com/deepseek-ai/DeepGEMM/blob/3b3783d06cd4d06ac4ba048633e604151d1ee535/deep_gemm/jit_kernels/gemm.py#L160