Skip to content
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

[fix] [FSDP] Make _get_default_cuda_device more robust to modules without params #606

Merged
merged 2 commits into from
Apr 14, 2021
Merged

[fix] [FSDP] Make _get_default_cuda_device more robust to modules without params #606

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
dfe2f67
Make _get_default_cuda_device more robust to modules without params
Apr 14, 2021
75cd3d3
Fix auto_wrap docstring
Apr 14, 2021
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
13 changes: 8 additions & 5 deletions fairscale/nn/data_parallel/fully_sharded_data_parallel.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1540,11 +1540,14 @@ def _print_r0(self, msg: str, restart: bool = False) -> None:

def _get_default_cuda_device(module: nn.Module) -> torch.device:
"""Try to infer CUDA device from module parameters."""
compute_device = next(module.parameters()).device
if compute_device.type != "cuda":
# Fall back to current CUDA device.
compute_device = torch.device("cuda")
return compute_device
try:
compute_device = next(module.parameters()).device
if compute_device.type == "cuda":
return compute_device
except StopIteration:
pass
# Fall back to current CUDA device
return torch.device("cuda")
min-xu-ai marked this conversation as resolved.
Show resolved Hide resolved


@torch.no_grad()
Expand Down
8 changes: 5 additions & 3 deletions fairscale/nn/wrap/auto_wrap.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -88,9 +88,11 @@ def enable_wrap(auto_wrap_policy: Optional[Callable] = None, **wrapper_kwargs: A
with enable_wrap(**params):
# Wraps layer in FSDP by default if within context
self.l1 = wrap(torch.nn.Linear(5, 5))
# Wraps children modules based on a different min_num_params
my_auto_wrap_policy = functools.partial(auto_wrap_policy, min_num_params=1e7)
self.l2 = auto_wrap(TransformerBlock(), shuold_wrap=my_auto_wrap_policy)
self.l2 = auto_wrap(
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks!

TransformerBlock(),
# Wraps children modules based on a different min_num_params
auto_wrap_policy=functools.partial(default_auto_wrap_policy, min_num_params=1e7)
)

Args:
auto_wrap_policy (Callable, Optional):
Expand Down