Add a parallel_mode property to TrainingArguments (#8877)
* Add a `distributed_env` property to TrainingArguments * Change name * Address comment
This commit is contained in:
@@ -465,6 +465,27 @@ class TrainingArguments:
|
||||
"""
|
||||
return self._setup_devices[1]
|
||||
|
||||
@property
|
||||
@torch_required
|
||||
def parallel_mode(self):
|
||||
"""
|
||||
The current mode used for parallelism if multiple GPUs/TPU cores are available. One of:
|
||||
|
||||
- :obj:`ParallelMode.NOT_PARALLEL`: no parallelism (CPU or one GPU).
|
||||
- :obj:`ParallelMode.NOT_DISTRIBUTED`: several GPUs in one single process (uses :obj:`torch.nn.DataParallel`).
|
||||
- :obj:`ParallelMode.DISTRIBUTED`: several GPUs, each ahving its own process (uses
|
||||
:obj:`torch.nn.DistributedDataParallel`).
|
||||
- :obj:`ParallelMode.TPU`: several TPU cores.
|
||||
"""
|
||||
if is_torch_tpu_available():
|
||||
return ParallelMode.TPU
|
||||
elif self.local_rank != -1:
|
||||
return ParallelMode.DISTRIBUTED
|
||||
elif self.n_gpu > 1:
|
||||
return ParallelMode.NOT_DISTRIBUTED
|
||||
else:
|
||||
return ParallelMode.NOT_PARALLEL
|
||||
|
||||
def to_dict(self):
|
||||
"""
|
||||
Serializes this instance while replace `Enum` by their values (for JSON serialization support).
|
||||
@@ -493,3 +514,10 @@ class TrainingArguments:
|
||||
valid_types.append(torch.Tensor)
|
||||
|
||||
return {k: v if type(v) in valid_types else str(v) for k, v in d.items()}
|
||||
|
||||
|
||||
class ParallelMode(Enum):
|
||||
NOT_PARALLEL = "not_parallel"
|
||||
NOT_DISTRIBUTED = "not_distributed"
|
||||
DISTRIBUTED = "distributed"
|
||||
TPU = "tpu"
|
||||
|
||||
Reference in New Issue
Block a user