Improve testing coverage across all device hw types

[rpsilva-aws]

🐛 Bug

We're currently maintaining test/run_tests.sh (CPU), test/tpu/run_tests.sh (TPU) and test/neuron/run_tests.sh (Trainium). Subsequently, anyone adding a test is required to add the test on all 3 locations, in addition to being given the option to skip (intentionally or not) a specific hardware device type. Therefore, it is evident that not all tests are present in all three files, and there is no tracking visibility if these were deliberately not added, in addition to its reasoning.

Expected behavior

We should move all of the tests under a single location, and use fixtures to mark tests to skip for a specific device / hardware. In addition we can specify and split/group tests to run under a specific runner, when distributing the tests across the CI/CD. The expected behavior is that skipped tests have a reason/motivation that is concisely described or tracked under a bug/feature item.

[ysiraichi]

Thank you for your suggestion.

I would argue that we actually benefit from having different test sets for different XLA backends. That's for a couple of reasons:

• There are some tests that make sense only for specific backends (e.g. due to specific lowerings). For those, IMHO it's better to write the test on the backend-specific file, rather than mixing all of them with @onlyXXXXX decorators
• PyTorch/XLA is something that sits in the middle (e.g. not a front-end, but also not a backend). As far as I understand it, there's just a little backend-specific code. I believe only a handful of tests would be benefitial for all backends

Could you give examples on when you have to add the same tests for these 3 backends? Maybe we could have a file that implements something like PyTorch's device parameterized tests

cc @miladm @tengyifei

[rpsilva-aws]

Thanks for the feedback! Based on our recent Neuron test migration effort, we found that most of the tests from CPU (at least) were compatible with our backend. Several of the tests were commented out as needing RCA to mitigate/resolve any backend limitations, but this gives us room to cover those cases. I believe that most PyTorch/XLA functionality is intended to be consistent and agnostic to XLA backends, and if tests (which granted, can be written to be backend-specific) are disabled or not captured as part of a backend-specific test, we want to make sure that we are aware of these limitations, which are not tracked and lack sufficient visibility - which is the main concern.

A unified test structure with skip decorators (as an example, not necessarily having to use decorators, e.g., @skipCPU, @skipTPU) requiring reason parameters offers several advantages. It provides clear visibility into backend-specific limitations and allows us to track technical debt through known references. When tests are skipped, the reasoning is documented directly in the code, making it easier for teams to understand intentional behavioral differences. If we have backend-specific tests, we can still use @onlyCPU/@onlyTPU decorators. If there are more general topics to capture, we can group these in dedicated testing classes that only test a specific backend, but that again, is a better means for us to document our coverage (we have these for Neuron as well). In my opinion, having the decorators with concise reasoning behind the skip is a delta improvement from having to maintain distinct test files. Note that the standard at the moment is for the developer to copy the test file across all test files. PyTorch's device parameterized tests is an interesting suggestion. I am not particularly in favor of using skip decorators, but mostly in how we can avoid having to maintain difference test files.

From a maintenance perspective, consolidating tests to run with any XLA backend unless granted an exemption, mitigates the risk of tests being accidentally omitted from specific backends. In addition to the better visibility into cross-backend compatibility, it will help us identify areas where we can improve consistency across implementations. IMHO, there's added benefit in developing XLA backend agnostic solutions when possible as an additional criteria. Adding that level of visibility that X/Y/Z doesn't work for a specific backend allows those teams/contributors to add/investigate that gap and benefit from it.

I think we could first survey which tests are XLA backend specific, and see what coverage we might be missing at the moment. I'll see if I can provide the examples, but started with the motivation first.

[rpsilva-aws]

I did a short skim, and I believe we have ~26% overlap of TPU & CPU, and 86% overlap of Neuron & CPU, and around ~26% between Neuron & TPU.

An example I see here is that there's a large number of CPU tests that are not present in TPU (see below) or vice-versa, but it's unclear if they all should not be.

  - scan_examples.py
  - test_async_closures.py
  - test_bridge.py
  - test_callback.py
  - test_cli.py
  - test_collective_ops_tpu.py
  - test_composite.py
  - test_convolution.py
  - test_core_aten_ops.py
  - test_ddp.py
  - test_deprecation.py
  - test_devices.py
  - test_dropout.py
  - test_dtensor_integration.py
  - test_dtensor_integration2.py
  - test_dynamic_plugin_tpu.py
  - test_dynamo_aliasing.py
  - test_dynamo_config.py
  - test_dynamo_graph_dump.py
  - test_dynamo_integrations_util.py
  - test_dynamo_spmd.py
  - test_eager.py
  - test_eager_all_reduce_in_place.py
  - test_eager_with_torch_compile.py
  - test_eager_with_xla_compile.py
  - test_embedding.py
  - test_env_var_mapper.py
  - test_export_fx_passes.py
  - test_exports.py
  - test_fsdp_auto_wrap.py
  - test_fsdp_v2.py
  - test_gmm.py
  - test_gpu_device_detection.py
  - test_grad_checkpoint.py
  - test_graph_input_matcher.py
  - test_hlo_metadata.py
  - test_implicit_broadcasting.py
  - test_indexing.py
  - test_internal_tpu.py
  - test_mesh_service.py
  - test_metrics.py
  - test_mp_all_gather.py
  - test_mp_all_to_all.py
  - test_mp_collective_permute.py
  - test_mp_distributed_mm.py
  - test_mp_early_exit.py
  - test_mp_mesh_reduce.py
  - test_mp_pt_xla_debug.py
  - test_mp_reduce_scatter.py
  - test_mp_replication.py
  - test_mp_save.py
  - test_mp_sync_batch_norm.py
  - test_multi_queries_paged_attention_kernel.py
  - test_multihead_attention.py
  - test_nn.py
  - test_num_output.py
  - test_operations_hlo.py
  - test_ops.py
  - test_pallas_spmd.py
  - test_persistent_cache.py
  - test_pooling.py
  - test_profiler.py
  - test_pt2e_qdq.py
  - test_pt_xla_debug.py
  - test_python_ops.py
  - test_quantized_matmul.py
  - test_ragged_paged_attention_kernel.py
  - test_runtime.py
  - test_runtime_multi_cpu.py
  - test_runtime_multi_gpu.py
  - test_runtime_single_proc_gpu.py
  - test_runtime_tpu.py
  - test_spmd_graph_dump.py
  - test_stablehlo_compile.py
  - test_stablehlo_custom_call.py
  - test_stablehlo_inference.py
  - test_stablehlo_save_load.py
  - test_torch.py
  - test_torch_distributed_bucketed_all_reduce_xla_backend.py
  - test_torch_distributed_fsdp_meta.py
  - test_torch_distributed_xla_backend.py
  - test_torchrun.py
  - test_type_promotion.py
  - test_unbounded_dynamism.py
  - test_view_ops.py
  - test_xla_sharding_hlo.py
  - test_zero1.py
  - train_decoder_only_base.py
  - train_decoder_only_eager.py
  - train_decoder_only_eager_multi_process.py
  - train_decoder_only_eager_spmd_data_parallel.py
  - train_decoder_only_eager_with_compile.py
  - train_decoder_only_fsdp_v2.py
  - train_resnet_amp.py
  - train_resnet_fsdp_auto_wrap.py
  - train_resnet_spmd_data_parallel.py
  - train_resnet_xla_ddp.py

[tengyifei]

Seems that the data has revealed that we're missing a lot of test coverage on TPU and it's worth addressing that, irregardless of if we use decorators or separate files to isolate backend specific tests.

[rpsilva-aws]

I agree. This matches my expectation since I think this is usually unavoidable with the current test structure, but +1 to prioritizing closing the coverage gap. I made it a bit clearer above, but the test files above are not overlapping - meaning either missing on CPU or TPU, mostly the latter.

[ysiraichi]

An example I see here is that there's a large number of CPU tests that are not present in TPU (see below) or vice-versa, but it's unclear if they all should not be.

I think this highlights exactly what I am thinking.

I know that there are some backend-specific code inside PyTorch/XLA, which definitely should be tested. Since they are backend-specific, IMHO it makes more sense to me that they should live inside their own directories in test/.

In the end, I think it mostly depends on what kind of tests we wish to have in this repository:

• If a given program works for, say, XLA:CPU but doesn't work for XLA:TPU (given that their lowering is the same), is this an issue we want to catch in this repository? Or is this an issue for OpenXLA?
• Do we wish to track the feature parity of CPU, GPU, TPU, and Neuron backends inside PyTorch/XLA (even though we don't actually implement any of them)?

Also, if tests are the same for all backends, we could just run them while setting PJRT_DEVICE. Or is there something I'm missing where it would not work?

---

That said, I agree that we shouldn't be copy-and-pasting tests for testing on different backends. Also, if the whole files are meant to be the same, we could just run them with different PJRT_DEVICE set.