facebookresearch / facebookresearch/fairscale

[FSDP] Can we access parameter views when using flatten_parameters=True?

Open
#430 9 comments 0 reactions 0 assignees View on GitHub
FSDP
Dominant language
Python
Stars
3.4k
Forks
293
PR merge metrics
No merged PRs in 30d

Description

## ❓ Questions and Help

This should explain the case:

```python
import torch
from fairscale.nn.data_parallel import FullyShardedDataParallel
import os

os.environ['MASTER_ADDR'] = 'localhost'
os.environ['MASTER_PORT'] = '1337'
torch.distributed.init_process_group("gloo", rank=0, world_size=1)

class Model(torch.nn.Module):
def __init__(self):
super().__init__()
self.layer = torch.nn.Linear(5, 5)

model = FullyShardedDataParallel(Model(), flatten_parameters=False)

# prints parameter list of module
print([p for p in model.module.layer.parameters()])

model = FullyShardedDataParallel(Model(), flatten_parameters=True)

# prints []
print([p for p in model.module.layer.parameters()])

# Throws an error: ValueError: optimizer got an empty parameter list
optimizer = torch.optim.SGD(model.layer.parameters(), lr=1e-3)

```

When `flatten_parameters=True` we remove the parameter as we have migrated it to a contiguous buffer, but this means when we call `.parameters()` on specific modules (in the case we only want to wrap certain parts of the model with optimizers) this can not be done.

Any remedy to this problem? We were experimenting with the possibility of using views to replace this functionality however this doesn't return a parameter I think. Alternatively, we could tell the users if they run into issues like above, to turn off `flatten_parameters`.

Contributor guide

Open the contributing guide

Research direction

Start with the supplied FSDP reproducer and inspect the behavior of FullyShardedDataParallel when flatten_parameters is true, especially module.parameters() and optimizer construction. Determine whether parameter views can support access to specific wrapped modules, or whether the documented resolution should be flatten_parameters=False; done requires an agreed behavior and a regression test.

Written by the indexing model from the issue text.

Assessment

Tech stack
python, pytorch
Domain
distributed-systems, machine-learning
Issue type
Feature
Difficulty
5/5
Estimated time
Over a week
Activity status
Stale
Clarity
Needs clarification
Newbie friendliness
25/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.