Misc CI fixes (#4813)

Summary:
X-link: facebookresearch/FBGEMM#1842

- Increased timeout for Nova genai aarch64 jobs

- Disable `test_backward_sgd_writeback` on GitHub + ROCm for now bc it is flaky

Pull Request resolved: #4813

Reviewed By: cthi

Differential Revision: D81604499

Pulled By: q10

fbshipit-source-id: 5f67627945b4235e3e750c6838f41305fe39d038

Author: q10

.github/workflows/build_wheels_genai_linux_aarch64.yml

@@ -88,4 +88,4 @@ jobs:
       trigger-event: ${{ github.event_name }}
       architecture: aarch64
       setup-miniconda: false
-      timeout: 150
+      timeout: 210

fbgemm_gpu/test/tbe/training/backward_none_test.py

@@ -48,13 +48,15 @@
         additional_decorators,
         gpu_unavailable,
         optests,
+        running_on_github,
         TEST_WITH_ROCM,
     )
 else:
     from fbgemm_gpu.test.test_utils import (
         additional_decorators,
         gpu_unavailable,
         optests,
+        running_on_github,
         TEST_WITH_ROCM,
     )
 VERBOSITY: Verbosity = Verbosity.verbose
@@ -84,7 +86,10 @@
 @optests.generate_opcheck_tests(fast=True, additional_decorators=additional_decorators)
 class BackwardNoneTest(unittest.TestCase):
     @unittest.skipIf(*gpu_unavailable)
-    @unittest.skipIf(torch.version.hip is not None, "Test is flaky on ROCm")
+    @unittest.skipIf(
+        running_on_github and torch.version.hip is not None,
+        "Test is flaky on GitHub + ROCm",
+    )
     @given(
         T=st.integers(min_value=1, max_value=5),
         D=st.integers(min_value=2, max_value=256),

fbgemm_gpu/test/tbe/training/backward_sgd_test.py

@@ -48,6 +48,7 @@
         additional_decorators,
         gpu_unavailable,
         optests,
+        running_on_github,
         TEST_WITH_ROCM,
         use_cpu_strategy,
     )
@@ -56,6 +57,7 @@
         additional_decorators,
         gpu_unavailable,
         optests,
+        running_on_github,
         TEST_WITH_ROCM,
         use_cpu_strategy,
     )
@@ -540,6 +542,10 @@ def test_backward_sgd_really_long_segments(  # noqa C901
             SparseType.FP32,  # output_dtype
         )
 
+    @unittest.skipIf(
+        running_on_github and torch.version.hip is not None,
+        "Test is flaky on GitHub + ROCm",
+    )
     @given(
         T=st.integers(min_value=1, max_value=3),
         D=st.integers(min_value=2, max_value=256),