[AMD][ROCm] Improve support of AMD

[k-artem]

The patch delivers several fixes for building issues for CUDA part of DeepSpeed library.
Percentage of passed unit tests improved(tested on RDNA hardware, gfx110x and gfx12x) Before:
collected 5298 items / 15 skipped
2773 failed, 862 passed, 1665 skipped, 13 errors
After:
collected 5851 items / 11 skipped
4187 failed, 1373 passed, 292 skipped, 10 errors

Regarding testing of fp_quantizer(DS_BUILD_FP_QUANTIZER) via tests/unit/ops/fp_quantizer/test_fp_quant.py, this test depends on QPyTorch which should be patched before run on AMD, please apply Tiiiger/QPyTorch#71

[k-artem]

@hwchen2017 kindly ask for review after fixed your comments.

[loadams]

@k-artem - is this ready for final review? @hwchen2017 - any remaining review requests?

[hwchen2017]

Can you share the error message you get on AMD GPU and explain why these changes can fix issues? It can help us better understand this PR. Thanks!

[hwchen2017]

The change here seems unnecessary because this ops only runs on cuda.

[k-artem]

Yep, currently no support, can be removed or we can keep for proper cast, because both cuda and hip declares 1st arg as const void *. please share your opinion.

[hwchen2017]

Let's remove it for now.

[k-artem]

ok

[hwchen2017]

what's the purpose of changing the file extension here? Usually .cpp extension is used for the api definition.

[k-artem]

if bf16 support enabled it should be build using nvcc/hipcc not regular compiler which by default uses for files with cpp extensions in cpp_exetension.py. And actually, as I see it's cuda essentially.

[hwchen2017]

The binding code may or may not contain cuda code. That's why people use .cpp by default. Let's make it consistent in the repo.

[k-artem]

Yep, let me add more details why I did rename, in case of HIP and BF16_AVAILABLE these files include cuda staff that cannot be compiled by regular compiler which selected in torch/utils/cpp_extension.py (_is_cuda_file) based on file's extension, example:

     In file included from /opt/rocm/include/hip/amd_detail/amd_warp_functions.h:27,
                       from /opt/rocm/include/hip/amd_detail/amd_hip_bf16.h:113,
                       from /opt/rocm/include/hip/hip_bf16.h:29,
                       from /myworkspace/DeepSpeed/csrc/includes/ds_kernel_utils_hip.h:18,
                       from /myworkspace/DeepSpeed/csrc/includes/quantization_hip.h:10,
                       from csrc/quantization/pt_binding_hip.cpp:11

and as result __builtin_amdgcn_* cannot be found.

[hwchen2017]

Can you move the change to https://github.com/deepspeedai/DeepSpeed/blob/master/op_builder/builder.py#L613?

[k-artem]

ok

[k-artem]

hi @hwchen2017 @loadams Apologies for the delay in this PR. I've updated the code according to the last set of comments. Please review the changes. I've enabled bf16 library-wide; however, I've disabled it for transformer_inference. This is because we need to explicitly declare the bf16 type inside pt_binding.cpp, which is currently not possible due to a limitation I previously mentioned. For some extensions, such as fp_quantizer, I was able to resolve the issue using a forward declaration. Thank you in advance for your review. Also I'd appreciate you share any acceptable ideas about ways to enable bf16 for transformer_inference on the library side.