Add AWS Neuron torchrun support (#20806)
* Add XLA torchrun support * Clarify that currently DDP doesn't work with torch.distributed XLA backend yet * Enable DDP with torchrun and XLA (now available in PT-XLA 1.13) * Add check for AWS Neuron availability and AWS Neuron specific compiler flag * Change the new test's name to TestTrainerDistributedNeuronCore * Remove "assert" and replace raised exception * Remove compiler flag as it is optional. If needed, will be another PR. * Use TORCHELASTIC_RUN_ID to determine whether torchrun is used
This commit is contained in:
@@ -83,6 +83,7 @@ from .utils import (
|
||||
is_torch_available,
|
||||
is_torch_bf16_cpu_available,
|
||||
is_torch_bf16_gpu_available,
|
||||
is_torch_neuroncore_available,
|
||||
is_torch_tensorrt_fx_available,
|
||||
is_torch_tf32_available,
|
||||
is_torch_tpu_available,
|
||||
@@ -500,6 +501,15 @@ def require_torch_tpu(test_case):
|
||||
return unittest.skipUnless(is_torch_tpu_available(check_device=False), "test requires PyTorch TPU")(test_case)
|
||||
|
||||
|
||||
def require_torch_neuroncore(test_case):
|
||||
"""
|
||||
Decorator marking a test that requires NeuronCore (in PyTorch).
|
||||
"""
|
||||
return unittest.skipUnless(is_torch_neuroncore_available(check_device=False), "test requires PyTorch NeuronCore")(
|
||||
test_case
|
||||
)
|
||||
|
||||
|
||||
if is_torch_available():
|
||||
# Set env var CUDA_VISIBLE_DEVICES="" to force cpu-mode
|
||||
import torch
|
||||
|
||||
Reference in New Issue
Block a user