2021 the huggingface team licensed under the apache license version 2 0 the license you may not use this file except in compliance with the license you may obtain a copy of the license at http www apache orglicenseslicense2 0 unless required by applicable law or agreed to in writing software distributed under the license is distributed on an as is basis without warranties or conditions of any kind either express or implied see the license for the specific language governing permissions and limitations under the license todo should be moved to utils after refactoring of sagemakertrainer get the sagemaker specific mp parameters from smpoptions variable parse it and check the field partitions is included it is required for model parallel get the sagemaker specific framework parameters from mpioptions variable parse it and check the field sagemakerdistributeddataparallelenabled lastly check if the smdistributed module is present if ngpu is 1 we ll use nn dataparallel if you only want to use a specific subset of gpus use cudavisibledevices0 explicitly set cuda to the first index 0 cuda device otherwise setdevice will trigger an error that a device index is missing index 0 takes into account the gpus available in the environment so cudavisibledevices1 2 with cuda 0 will use the first gpu in that env i e gpu1 sometimes the line in the postinit has not been run before we end up here so just checking we re not at the default value here we ll use torch distributed initializes the distributed backend which will take care of synchronizing nodesgpus 2021 the huggingface team licensed under the apache license version 2 0 the license you may not use this file except in compliance with the license you may obtain a copy of the license at http www apache org licenses license 2 0 unless required by applicable law or agreed to in writing software distributed under the license is distributed on an as is basis without warranties or conditions of any kind either express or implied see the license for the specific language governing permissions and limitations under the license todo should be moved to utils after refactoring of sagemakertrainer get the sagemaker specific mp parameters from smp_options variable parse it and check the field partitions is included it is required for model parallel get the sagemaker specific framework parameters from mpi_options variable parse it and check the field sagemaker_distributed_dataparallel_enabled lastly check if the smdistributed module is present noqa f401 if n_gpu is 1 we ll use nn dataparallel if you only want to use a specific subset of gpus use cuda_visible_devices 0 explicitly set cuda to the first index 0 cuda device otherwise set_device will trigger an error that a device index is missing index 0 takes into account the gpus available in the environment so cuda_visible_devices 1 2 with cuda 0 will use the first gpu in that env i e gpu 1 sometimes the line in the postinit has not been run before we end up here so just checking we re not at the default value here we ll use torch distributed initializes the distributed backend which will take care of synchronizing nodes gpus
import importlib.util
import json
import os
import warnings
from dataclasses import dataclass, field

import torch

from ..training_args import TrainingArguments
from ..utils import cached_property, is_sagemaker_dp_enabled, logging


logger = logging.get_logger(__name__)




def is_sagemaker_model_parallel_available():
    
    smp_options = os.getenv("SM_HP_MP_PARAMETERS", "{}")
    try:
        
        smp_options = json.loads(smp_options)
        if "partitions" not in smp_options:
            return False
    except json.JSONDecodeError:
        return False

    
    mpi_options = os.getenv("SM_FRAMEWORK_PARAMS", "{}")
    try:
        
        mpi_options = json.loads(mpi_options)
        if not mpi_options.get("sagemaker_mpi_enabled", False):
            return False
    except json.JSONDecodeError:
        return False
    
    return importlib.util.find_spec("smdistributed") is not None


if is_sagemaker_model_parallel_available():
    import smdistributed.modelparallel.torch as smp

    smp.init()


@dataclass
class SageMakerTrainingArguments(TrainingArguments):
    mp_parameters: str = field(
        default="",
        metadata={"help": "Used by the SageMaker launcher to send mp-specific args. Ignored in SageMakerTrainer"},
    )

    def __post_init__(self):
        super().__post_init__()
        warnings.warn(
            "`SageMakerTrainingArguments` is deprecated and will be removed in v5 of Transformers. You can use "
            "`TrainingArguments` instead.",
            FutureWarning,
        )

    @cached_property
    def _setup_devices(self) -> "torch.device":
        logger.info("PyTorch: setting up devices")
        if torch.distributed.is_available() and torch.distributed.is_initialized() and self.local_rank == -1:
            logger.warning(
                "torch.distributed process group is initialized, but local_rank == -1. "
                "In order to use Torch DDP, launch your script with `python -m torch.distributed.launch"
            )
        if self.no_cuda:
            device = torch.device("cpu")
            self._n_gpu = 0
        elif is_sagemaker_model_parallel_available():
            local_rank = smp.local_rank()
            device = torch.device("cuda", local_rank)
            self._n_gpu = 1
        elif is_sagemaker_dp_enabled():
            import smdistributed.dataparallel.torch.torch_smddp  

            torch.distributed.init_process_group(backend="smddp", timeout=self.ddp_timeout_delta)
            self.local_rank = int(os.getenv("SMDATAPARALLEL_LOCAL_RANK"))
            device = torch.device("cuda", self.local_rank)
            self._n_gpu = 1
        elif self.local_rank == -1:
            
            
            
            
            
            
            device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
            
            
            self._n_gpu = torch.cuda.device_count()
        else:
            
            
            if not torch.distributed.is_initialized():
                torch.distributed.init_process_group(backend="nccl", timeout=self.ddp_timeout_delta)
            device = torch.device("cuda", self.local_rank)
            self._n_gpu = 1

        if device.type == "cuda":
            torch.cuda.set_device(device)

        return device

    @property
    def world_size(self):
        if is_sagemaker_model_parallel_available():
            return smp.dp_size()

        return super().world_size

    @property
    def place_model_on_device(self):
        return not is_sagemaker_model_parallel_available()

    @property
    def _no_sync_in_gradient_accumulation(self):
        return False