SUMIN/HuggingFace_transformer
1
0
Fork 0
Code Issues Pull Requests Actions 7 Packages Projects Releases Wiki Activity

torch.distributed group initialization for torch_neuron disabled when optimum-neuron is installed (#22728)

Browse Source 
* Make the process group initialization not happen if optimum_neuron is installed

* Add warning

* Remove list and added warning
This commit is contained in:
Michael Benayoun
2023-04-12 17:42:50 +02:00
committed by GitHub
parent 1306b7d3ae
commit 10fab90fe2
2 changed files with 24 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files

29
src/transformers/training_args.py
View File

@@ -54,8 +54,13 @@ from .utils import (
logging,
requires_backends,
)
from .utils.import_utils import is_optimum_neuron_available
logger = logging.get_logger(__name__)
log_levels = logging.get_log_levels_dict().copy()
trainer_log_levels = dict(**log_levels, passive=-1)
if is_torch_available():
import torch
import torch.distributed as dist
@@ -67,12 +72,23 @@ if is_torch_neuroncore_available(check_device=False):
# torchrun support
# https://github.com/pytorch/xla/pull/3609
if os.environ.get("TORCHELASTIC_RUN_ID"):
import torch_xla.distributed.xla_backend as xbn
if is_optimum_neuron_available():
logger.info(
"Make sure that you are performing the training with the TrainiumTrainer from optimum[neuron], this "
"will fail otherwise."
)
else:
logger.warning(
"Please use the TrainiumTrainer from optimum[neuron] instead of the Transformers library to perform "
"training on AWS Trainium instances. More information here: "
"https://github.com/huggingface/optimum-neuron"
)
import torch_xla.distributed.xla_backend as xbn
if not isinstance(torch.distributed.group.WORLD, xbn.ProcessGroupXla):
torch.distributed.init_process_group(backend="xla")
if not isinstance(torch.distributed.group.WORLD, xbn.ProcessGroupXla):
raise AssertionError("Failed to initialize torch.distributed process group using XLA backend.")
torch.distributed.init_process_group(backend="xla")
if not isinstance(torch.distributed.group.WORLD, xbn.ProcessGroupXla):
raise AssertionError("Failed to initialize torch.distributed process group using XLA backend.")
if is_sagemaker_mp_enabled():
@@ -81,11 +97,6 @@ if is_sagemaker_mp_enabled():
smp.init()
logger = logging.get_logger(__name__)
log_levels = logging.get_log_levels_dict().copy()
trainer_log_levels = dict(**log_levels, passive=-1)
def default_logdir() -> str:
"""
Same default as PyTorch