riaz.somc/transformers
1
0
Fork 0
mirror of https://github.com/huggingface/transformers.git synced 2025-08-01 18:51:14 +06:00
Code Issues Actions 3 Packages Projects Releases Wiki Activity

Disable DDP for neuron (#21953)

Browse Source 
Disable DDp for neuron

Co-authored-by: EC2 Default User <ec2-user@ip-172-31-42-72.us-west-2.compute.internal>
This commit is contained in:
aws-sangeetha 2023-03-06 06:33:44 -08:00 committed by GitHub
parent bc33fbf956
commit 0bb17295f0
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 3 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
src/transformers/trainer.py
View File

@ -148,6 +148,7 @@ from .utils import (
is_sagemaker_dp_enabled, is_sagemaker_dp_enabled,
is_sagemaker_mp_enabled, is_sagemaker_mp_enabled,
is_torch_compile_available, is_torch_compile_available,
is_torch_neuroncore_available,
is_torch_tpu_available, is_torch_tpu_available,
logging, logging,
) )
@ -1537,6 +1538,8 @@ class Trainer:
if self.args.ddp_bucket_cap_mb is not None: if self.args.ddp_bucket_cap_mb is not None:
kwargs["bucket_cap_mb"] = self.args.ddp_bucket_cap_mb kwargs["bucket_cap_mb"] = self.args.ddp_bucket_cap_mb
if is_torch_neuroncore_available:
return model
model = nn.parallel.DistributedDataParallel( model = nn.parallel.DistributedDataParallel(
model, model,
device_ids=[self.args.local_rank] if self.args._n_gpu != 0 else None, device_ids=[self.args.local_rank] if self.args._n_gpu != 0 else None,