fairscale_sharded_data_parallel_with_oss_optimizer_quickstart.py

1import torch
2import torch.nn as nn
3from fairscale.nn.data_parallel import ShardedDataParallel as ShardedDDP
4from fairscale.optim.oss import OSS
5from fairscale.nn.wrap import wrap
6
7# 1. Define your model
8model = nn.Sequential(nn.Linear(32, 32), nn.ReLU(), nn.Linear(32, 10))
9
10# 2. Define your optimizer with OSS (Optimizer State Sharding)
11# OSS is a requirement/companion for ShardedDDP to save memory
12base_optimizer = torch.optim.SGD
13optimizer = OSS(params=model.parameters(), optim=base_optimizer, lr=1e-2)
14
15# 3. Wrap the model with ShardedDDP
16# In a real scenario, you would have initialized torch.distributed first
17sharded_model = ShardedDDP(model, optimizer)
18
19# 4. Standard training loop
20input_data = torch.randn(16, 32)
21output = sharded_model(input_data)
22loss = output.sum()
23loss.backward()
24optimizer.step()