Unverified Commit 0bb17295 authored by aws-sangeetha's avatar aws-sangeetha Committed by GitHub
Browse files

Disable DDP for neuron (#21953)



Disable DDp for neuron
Co-authored-by: default avatarEC2 Default User <ec2-user@ip-172-31-42-72.us-west-2.compute.internal>
parent bc33fbf9
...@@ -148,6 +148,7 @@ from .utils import ( ...@@ -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: ...@@ -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,
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment