python3Packages.torchao, python3Packages.torchtune: disable tests that fail on darwin, adopt (#471546)

pkgs/development/python-modules/executorch/default.nix

@@ -200,6 +200,8 @@ buildPythonPackage rec {
     badPlatforms = [
       # Many tests segfault. Supporting this platform will need additional work
       "aarch64-linux"
+      # Compilation errors on Darwin
+      lib.systems.inspect.patterns.isDarwin
     ];
   };
 }

pkgs/development/python-modules/torchao/default.nix

@@ -168,7 +168,7 @@ buildPythonPackage rec {
     "test_save_load_int8woqtensors_0_cpu"
     "test_save_load_int8woqtensors_1_cpu"
   ]
-  ++ lib.optionals (stdenv.hostPlatform.isDarwin) [
+  ++ lib.optionals stdenv.hostPlatform.isDarwin [
     # AssertionError: Scalars are not equal!
     "test_comm"
     "test_fsdp2"
@@ -176,8 +176,7 @@ buildPythonPackage rec {
     "test_precompute_bitnet_scale"
     "test_qlora_fsdp2"
     "test_uneven_shard"
-  ]
-  ++ lib.optionals (stdenv.hostPlatform.isDarwin && stdenv.hostPlatform.isAarch64) [
+
     # RuntimeError: No packed_weights_format was selected
     "TestIntxOpaqueTensor"
     "test_accuracy_kleidiai"
@@ -185,15 +184,31 @@ buildPythonPackage rec {
     # RuntimeError: quantized engine NoQEngine is not supported
     "test_smooth_linear_cpu"
     "test_smooth_linear_edge_cases"
+
+    # TypeError: Trying to convert Float8_e4m3fn to the MPS backend but it does not have support for that dtype.
+    "test_subclass_slice_subclass2_shape0_device_mps"
+    "test_subclass_slice_subclass2_shape1_device_mps"
+    # torch._inductor.exc.InductorError: KeyError: torch.float8_e4m3fn
+    "test_optim_default_dtype_bf16_optim_name_AdamFp8_device_mps"
+    "test_optim_smoke_optim_name_AdamWFp8_bfloat16_device_mps"
+    "test_optim_smoke_optim_name_AdamWFp8_float32_device_mps"
+    "test_param_groups_optim_name_AdamFp8_device_mps"
+    "test_subclass_slice_subclass0_shape0_device_mps"
+    "test_optim_smoke_optim_name_AdamFp8_bfloat16_device_mps"
+    "test_optim_smoke_optim_name_AdamFp8_float32_device_mps"
+    "test_subclass_slice_subclass0_shape1_device_mps"
   ]
   ++ lib.optionals (stdenv.hostPlatform.isDarwin && stdenv.hostPlatform.isx86_64) [
     # Flaky: [gw0] node down: keyboard-interrupt
     "test_int8_weight_only_quant_with_freeze_0_cpu"
     "test_int8_weight_only_quant_with_freeze_1_cpu"
     "test_int8_weight_only_quant_with_freeze_2_cpu"
+
+    # Illegal instruction in subclass_4bit.py::dequantize
+    "test_subclass_slice"
   ];
 
-  disabledTestPaths = lib.optionals (stdenv.hostPlatform.isDarwin && stdenv.hostPlatform.isAarch64) [
+  disabledTestPaths = lib.optionals stdenv.hostPlatform.isDarwin [
     # Require unpackaged 'coremltools'
     "test/prototype/test_groupwise_lowbit_weight_lut_quantizer.py"
 
@@ -208,6 +223,9 @@ buildPythonPackage rec {
     homepage = "https://github.com/pytorch/ao";
     changelog = "https://github.com/pytorch/ao/releases/tag/v${version}";
     license = lib.licenses.bsd3;
-    maintainers = with lib.maintainers; [ GaetanLepage ];
+    maintainers = with lib.maintainers; [
+      GaetanLepage
+      sarahec
+    ];
   };
 }

pkgs/development/python-modules/torchtune/default.nix

@@ -85,6 +85,9 @@ buildPythonPackage rec {
     writableTmpDirAsHomeHook
   ];
 
+  # Exclude `regression` which depends on a specific llama model and `recipies` which are sample code
+  enabledTestPaths = [ "tests/torchtune" ];
+
   disabledTests = [
     # AssertionError (tensors are not equal)
     "test_stop_tokens"
@@ -111,6 +114,21 @@ buildPythonPackage rec {
   ++ lib.optionals (stdenv.hostPlatform.isLinux && stdenv.hostPlatform.isAarch64) [
     # Fatal Python error: Segmentation fault
     "test_forward_gqa"
+  ]
+  ++ lib.optionals stdenv.hostPlatform.isDarwin [
+    # tests/torchtune/training/test_distributed.py
+    "test_init_from_env_no_dup"
+    "test_init_from_env_dup"
+  ];
+
+  disabledTestPaths = lib.optionals stdenv.hostPlatform.isDarwin [
+    # fail due to floating-point precision differences
+    "tests/torchtune/models/flux/test_flux_autoencoder.py::TestFluxAutoencoder::test_encode"
+    "tests/torchtune/modules/peft/test_dora.py::TestDoRALinear::test_qdora_parity[True-dtype1]"
+    "tests/torchtune/modules/peft/test_lora.py::TestLoRALinear::test_qlora_parity[True-dtype1]"
+
+    # hangs
+    "tests/torchtune/utils"
   ];
 
   meta = {
@@ -118,6 +136,9 @@ buildPythonPackage rec {
     homepage = "https://github.com/meta-pytorch/torchtune";
     changelog = "https://github.com/meta-pytorch/torchtune/releases/tag/${src.tag}";
     license = lib.licenses.bsd3;
-    maintainers = with lib.maintainers; [ GaetanLepage ];
+    maintainers = with lib.maintainers; [
+      GaetanLepage
+      sarahec
+    ];
   };
 }