init

transformers/utils/set_cuda_devices_for_ci.py

@@ -0,0 +1,26 @@
+"""A simple script to set flexibly CUDA_VISIBLE_DEVICES in GitHub Actions CI workflow files."""
+
+import argparse
+import os
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--test_folder",
+        type=str,
+        default=None,
+        help="The test folder name of the model being tested. For example, `models/cohere`.",
+    )
+    args = parser.parse_args()
+
+    # `test_eager_matches_sdpa_generate` for `cohere` needs a lot of GPU memory!
+    # This depends on the runners. At this moment we are targeting our AWS CI runners.
+    if args.test_folder == "models/cohere":
+        cuda_visible_devices = "0,1,2,3"
+    elif "CUDA_VISIBLE_DEVICES" in os.environ:
+        cuda_visible_devices = os.environ.get("CUDA_VISIBLE_DEVICES")
+    else:
+        cuda_visible_devices = "0"
+
+    print(cuda_visible_devices)