Skip to content

Init distributed torch only if needed #108

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
jeffra merged 7 commits into from
Feb 26, 2020
Merged

Init distributed torch only if needed #108

Show file tree
Hide file tree
Changes from 2 commits
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
7c3653c
add auto-detect to torch dist init
jeffra Feb 24, 2020
7423439
update tests to infer distributed init status
jeffra Feb 26, 2020
3183a1c
prevent crash if dist_init_required is True but already initiliazed
jeffra Feb 26, 2020
51640fc
Merge branch 'master' into jeffra/issue85
jeffra Feb 26, 2020
ca42b63
only init if safe to do so (forgot to add this file in prev commit)
jeffra Feb 26, 2020
5249990
Merge branch 'jeffra/issue85' of github.com:microsoft/DeepSpeed into …
jeffra Feb 26, 2020
7a852fe
slight refactor
jeffra Feb 26, 2020
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 3 additions & 2 deletions deepspeed/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -32,7 +32,7 @@ def initialize(args,
training_data=None,
lr_scheduler=None,
mpu=None,
dist_init_required=True,
dist_init_required=None,
collate_fn=None):
r"""Initialize the DeepSpeed Engine.

Expand All @@ -56,7 +56,8 @@ def initialize(args,
mpu: Optional: A model parallelism unit object that implements
get_{model,data}_parallel_{rank,group,world_size}()

dist_init_required: Optional: Initializes torch.distributed
dist_init_required: Optional: None will auto-initialize torch.distributed if needed,
Comment thread
jeffra marked this conversation as resolved.
otherwise the user can force it to be initialized or not via boolean.

collate_fn: Optional: Merges a list of samples to form a
mini-batch of Tensor(s). Used when using batched loading from a
Expand Down
10 changes: 8 additions & 2 deletions deepspeed/pt/deepspeed_light.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -97,7 +97,7 @@ def __init__(self,
training_data=None,
lr_scheduler=None,
mpu=None,
dist_init_required=True,
dist_init_required=None,
collate_fn=None):
super(DeepSpeedLight, self).__init__()

Expand All @@ -119,8 +119,14 @@ def __init__(self,
self.gradient_average = True
self.warn_unscaled_loss = True

if dist_init_required is None:
dist_init_required = not dist.is_initialized()

self.dist_backend = "nccl"
if dist_init_required:
dist.init_process_group(backend="nccl")
logging.info("Initializing torch distributed with backend: {}".format(
Comment thread
jeffra marked this conversation as resolved.
Outdated
self.dist_backend))
dist.init_process_group(backend=self.dist_backend)

self._do_args_sanity_check(args)
self._configure_with_arguments(args, mpu)
Expand Down
3 changes: 1 addition & 2 deletions tests/unit/test_config.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -142,8 +142,7 @@ def test_deprecated_deepscale_config(tmpdir):
def _test_deprecated_deepscale_config(args, model, hidden_dim):
model, _, _,_ = deepspeed.initialize(args=args,
model=model,
model_parameters=model.parameters(),
dist_init_required=False)
model_parameters=model.parameters())
data_loader = random_dataloader(model=model,
total_samples=5,
hidden_dim=hidden_dim,
Expand Down
12 changes: 4 additions & 8 deletions tests/unit/test_fp16.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -32,8 +32,7 @@ def test_lamb_fp16_basic(tmpdir):
def _test_lamb_fp16_basic(args, model, hidden_dim):
model, _, _,_ = deepspeed.initialize(args=args,
model=model,
model_parameters=model.parameters(),
dist_init_required=False)
model_parameters=model.parameters())
data_loader = random_dataloader(model=model,
total_samples=50,
hidden_dim=hidden_dim,
Expand Down Expand Up @@ -70,8 +69,7 @@ def test_lamb_fp16_empty_grad(tmpdir):
def _test_lamb_fp16_empty_grad(args, model, hidden_dim):
model, _, _,_ = deepspeed.initialize(args=args,
model=model,
model_parameters=model.parameters(),
dist_init_required=False)
model_parameters=model.parameters())
data_loader = random_dataloader(model=model,
total_samples=50,
hidden_dim=hidden_dim,
Expand Down Expand Up @@ -102,8 +100,7 @@ def _test_adamw_fp16_basic(args, model, hidden_dim):
optimizer = torch.optim.AdamW(params=model.parameters())
model, _, _,_ = deepspeed.initialize(args=args,
model=model,
optimizer=optimizer,
dist_init_required=False)
optimizer=optimizer)
data_loader = random_dataloader(model=model,
total_samples=50,
hidden_dim=hidden_dim,
Expand Down Expand Up @@ -134,8 +131,7 @@ def _test_adamw_fp16_empty_grad(args, model, hidden_dim):
optimizer = torch.optim.AdamW(params=model.parameters())
model, _, _,_ = deepspeed.initialize(args=args,
model=model,
optimizer=optimizer,
dist_init_required=False)
optimizer=optimizer)
data_loader = random_dataloader(model=model,
total_samples=50,
hidden_dim=hidden_dim,
Expand Down