[ROCm] bfloat16 enablement

[liligwu]

Enable bfloat16 on ROCm

[netlify]

✅ Deploy Preview for pytorch-fbgemm-docs canceled.

Name	Link
🔨 Latest commit	2cde382
🔍 Latest deploy log	https://app.netlify.com/sites/pytorch-fbgemm-docs/deploys/63e1837a882c1a0008455cf0

[amathews-amd]

cc : @jianyuh @shintaro-iwasaki

[q10]

LGTM; importing to Phabricator

[facebook-github-bot]

@q10 has imported this pull request. If you are a Meta employee, you can view this diff on Phabricator.

[facebook-github-bot]

@q10 has imported this pull request. If you are a Meta employee, you can view this diff on Phabricator.

[q10]

@liligwu Sorry to bug you one more time. Could I trouble you to do 2 things to keep our build system happe:

1. Rebase this PR branch to top of latest main
2. Apply approximately the following diff to keep the internal linter happy:

@@ -1398,12 +1398,13 @@
 };

 #ifdef __HIP_PLATFORM_HCC__
-using __nv_bfloat16=hip_bfloat16;
+using __nv_bfloat16 = hip_bfloat16;

 typedef struct __align__(4) {
   uint16_t x;
   uint16_t y;
-} __nv_bfloat162_raw;
+}
+__nv_bfloat162_raw;

 struct __align__(4) __nv_bfloat162 {
   __nv_bfloat16 x;
@@ -1479,12 +1480,12 @@
 #ifdef __HIP_PLATFORM_HCC__
 // the descriptions of __float2bfloat16 and __float2bfloat16_rn are identical
 // https://docs.nvidia.com/cuda/cuda-math-api/group__CUDA__MATH____BFLOAT16__MISC.html#group__CUDA__MATH____BFLOAT16__MISC
-static __host__ __device__ __nv_bfloat16 __float2bfloat16(float f){
+static __host__ __device__ __nv_bfloat16 __float2bfloat16(float f) {
   __nv_bfloat16 output;
   return output.round_to_bfloat16(f);
 }

-static __host__ __device__ __nv_bfloat16 __float2bfloat16_rn(float f){
+static __host__ __device__ __nv_bfloat16 __float2bfloat16_rn(float f) {
   __nv_bfloat16 output;
   return output.round_to_bfloat16(f);
 }

After this, we should be good to merge. Thanks for your patience!

[liligwu]

@q10 , it's done.

[liligwu]

@q10 , @shintaro-iwasaki . Could you please close this PR if everything looks good? Thank you.

[facebook-github-bot]

@q10 has imported this pull request. If you are a Meta employee, you can view this diff on Phabricator.

[q10]

@q10 , @shintaro-iwasaki . Could you please close this PR if everything looks good? Thank you.

Yep, was making sure the GitHub builds were passing

[liligwu]

@q10 , @shintaro-iwasaki . Could you please close this PR if everything looks good? Thank you.

Yep, was making sure the GitHub builds were passing

Regarding the ROCm CI failure

2023-02-07T19:53:19.5087315Z test_gpu (main.TableBatchedEmbeddingsTest) ... Memory access fault by GPU node-2 (Agent handle: 0x55e30ef37ad0) on address 0x7f38a5a82000. Reason: Page not present or supervisor privilege.
2023-02-07T19:53:19.7516263Z /workspace/fbgemm-private-jenkins/.jenkins/rocm/build_and_test.sh: line 49: 19939 Aborted (core dumped) python batched_unary_embeddings_test.py --verbose
2023-02-07T19:53:20.3188419Z ##[error]Process completed with exit code 134.

I opened an issue #1559

I don't know if it was observed in CI, but it failed in a CUDA container I tested.

[q10]

@q10 , @shintaro-iwasaki . Could you please close this PR if everything looks good? Thank you.

Yep, was making sure the GitHub builds were passing

Regarding the ROCm CI failure

2023-02-07T19:53:19.5087315Z test_gpu (main.TableBatchedEmbeddingsTest) ... Memory access fault by GPU node-2 (Agent handle: 0x55e30ef37ad0) on address 0x7f38a5a82000. Reason: Page not present or supervisor privilege.
2023-02-07T19:53:19.7516263Z /workspace/fbgemm-private-jenkins/.jenkins/rocm/build_and_test.sh: line 49: 19939 Aborted (core dumped) python batched_unary_embeddings_test.py --verbose
2023-02-07T19:53:20.3188419Z ##[error]Process completed with exit code 134.

I opened an issue #1559

I don't know if it was observed in CI, but it failed in a CUDA container I tested.

Thanks for filing the issue. The issue is observed in CI, but it is very unreliable, as one of the earlier builds in this PR had previously passed prior to your last lint commit. I suspect it is an issue with the instance that the job was running on. We will investigate this further after this PR is merged.

[facebook-github-bot]

@q10 merged this pull request in abefe30.