def __init__(self, group, wrapper_config, checkpoint_act=False):
super().__init__(group, wrapper_config)
self.group = group
# "expert" params are different on each rank
torch.manual_seed(42 + group.rank())
expert = nn.Linear(16, 4)
for p in expert.parameters():
p.expert = True
# everything else is shared
torch.manual_seed(0)
shared = nn.Linear(4, 16)
if checkpoint_act:
expert = checkpoint_wrapper(expert)
shared = checkpoint_wrapper(shared)
if wrapper_config is not None:
# we create a process group of size 1 for the expert params
expert_group = torch.distributed.new_group([group.rank()])
expert = FullyShardedDataParallel(expert, expert_group,
**wrapper_config)
shared = FullyShardedDataParallel(shared, group, **wrapper_config)
self.module = nn.Sequential(nn.Linear(8, 4), shared, expert,
nn.Linear(4, 8))
def __init__(self, multiout=False, checkpoint_config=0):
super().__init__()
torch.manual_seed(0) # make sure weights are deterministic.
self.multiout = multiout
self.conv1 = nn.Sequential(nn.Conv2d(1, 5, 3), nn.ReLU(), nn.Conv2d(5, 5, 3))
self.conv2 = nn.Sequential(nn.Conv2d(3, 5, 3), nn.ReLU(), nn.Conv2d(5, 5, 3))
assert 0 <= checkpoint_config <= 3
if checkpoint_config & 1:
self.conv1 = checkpoint_wrapper(self.conv1)
if checkpoint_config & (1 << 1):
self.conv2 = checkpoint_wrapper(self.conv2)
def __init__(self,
group,
wrapper_config,
checkpoint_act=False,
delay_before_free_ms=0):
super().__init__(group, wrapper_config)
self.group = group
self.delay_before_free_ms = delay_before_free_ms
# "expert" params are different on each rank
torch.manual_seed(42 + group.rank())
d_expert = 23
d_shared = 12
d_input = 8
expert = nn.Linear(d_expert, d_shared)
self.num_expert_params = sum([p.numel() for p in expert.parameters()])
for p in expert.parameters():
p.expert = True
# everything else is shared
torch.manual_seed(0)
shared = nn.Linear(d_shared, d_expert)
if checkpoint_act:
expert = checkpoint_wrapper(expert)
shared = checkpoint_wrapper(shared)
if wrapper_config is not None:
# we create a process group of size 1 for the expert params
expert_group = torch.distributed.new_group(
[group.rank()]) # world size 1 means no shard
expert = FullyShardedDataParallel(expert, expert_group,
**wrapper_config)
shared = FullyShardedDataParallel(shared, group, **wrapper_config)
self.module = nn.Sequential(nn.Linear(d_input, d_shared), shared,
expert, nn.Linear(d_shared, d_input))
def __init__(self, use_pytorch_checkpoint=False, use_fairscale_checkpoint=False, **kwargs):
super().__init__()
torch.manual_seed(0) # make sure weights are deterministic.
assert not (
use_pytorch_checkpoint and use_fairscale_checkpoint
), "Cannot use both pytorch and fairscale checkpointing mechanisms."
self.use_pytorch_checkpoint = use_pytorch_checkpoint
self.ffn = nn.Sequential(
nn.Linear(32, 128),
# add a Dropout layer to test RNG save/restore
nn.Dropout(p=0.5),
nn.Linear(128, 32),
)
if use_fairscale_checkpoint:
self.ffn = checkpoint_wrapper(self.ffn, **kwargs)
self.out = nn.Linear(32, 1)
def __init__(self,
use_pytorch_checkpoint=False,
use_fairseq_checkpoint=False,
**kwargs):
super().__init__()
torch.manual_seed(0)
self.use_pytorch_checkpoint = use_pytorch_checkpoint
self.ffn = nn.Sequential(
nn.Linear(32, 128),
# add a Dropout layer to test RNG save/restore
nn.Dropout(p=0.5),
nn.Linear(128, 32),
)
if use_fairseq_checkpoint:
self.ffn = checkpoint_wrapper(self.ffn, **kwargs)
self.out = nn.Linear(32, 1)
def __init__(self, enable_checkpoint=False, cpu_offload=False):
super().__init__()
torch.manual_seed(0) # make sure weights are deterministic.
# These numbers are picked to show cpu_offload memory saving.
# Inner (recomputed) activation sizes need to be just right
# to show the benefit.
self.layers = nn.Sequential(
nn.Sequential(nn.Linear(4, 4), nn.Linear(4, 4), nn.Linear(4, 8)),
nn.Sequential(nn.Linear(8, 4), nn.Linear(4, 4), nn.Linear(4, 4)),
nn.Sequential(nn.Linear(4, 6), nn.Linear(6, 8), nn.Linear(8, 2)),
)
if enable_checkpoint:
for i, layer in enumerate(self.layers):
# Only middle layer needs to have offloading
self.layers[i] = checkpoint_wrapper(layer, cpu_offload if i == 1 else False)
def get_model(norm_type, checkpointed, mixed_precision):
assert norm_type in NORM_TYPES, norm_type
assert checkpointed in [True, False], checkpointed
assert mixed_precision in MP_TYPES
model = Sequential(Linear(3, 2), norm_type(2))
if mixed_precision == "fp16":
# Set param.data and buffers as fp16
for p in model.parameters():
p.data = p.data.half()
for m in model:
for n, b in m.named_buffers():
setattr(m, n, b.half())
elif mixed_precision == "call_half":
model.half()
if checkpointed:
model = checkpoint_wrapper(model)
return model