def _test_dtypes(cfg: Dict,
autocast,
in_dtype,
p_dtype,
loss_dtype,
reduce_dtype,
rank,
group,
expected_buffer_type=None):
# Patch torch.distributed.reduce_scatter to check the dtype of the reduction
orig_reduce_scatter = torch.distributed.reduce_scatter
model: nn.Module = DeviceAndTypeCheckModule(
expected_input_dtype=in_dtype,
expected_param_dtype=p_dtype,
expected_loss_dtype=loss_dtype,
expected_buffer_dtype=expected_buffer_type,
)
def _reduce_scatter(output, input_list, **kwargs):
for tensor in input_list:
model._check("reduce_scatter.dtype",
tensor.dtype,
expected=reduce_dtype)
return orig_reduce_scatter(output, input_list, **kwargs)
with mock.patch("torch.distributed.reduce_scatter",
new=_reduce_scatter):
model = FullyShardedDataParallel(model, group, **cfg).cuda()
device = next(model.parameters()).device
x = torch.rand(2, 5).to(device)
with torch.cuda.amp.autocast(enabled=autocast):
loss = model(x)
loss.backward()
def _test_func(rank, world_size, tempfile_name, unused):
result = dist_init(rank, world_size, tempfile_name, unused)
assert result, "Dist init failed"
# Keep initialization deterministic.
torch.manual_seed(0)
model = FullyShardedDataParallel(SimpleModuleWithCheckpointing().cuda())
optim = torch.optim.SGD(model.parameters(), lr=0.001, momentum=0.9)
# Collect parameter sizes to ensure these stay consistent through the steps below.
expected_param_shapes = {
name: tuple(param.shape)
for name, param in model.named_parameters()
}
# For clarity, this is what `expected_param_shapes` should look like depending on world size:
assert expected_param_shapes == {
"_fsdp_wrapped_module.flat_param_0": (12, ),
"_fsdp_wrapped_module._fpw_module.ffn.1._fsdp_wrapped_module.flat_param_0":
(6, ),
}, expected_param_shapes
torch.manual_seed(1 + rank)
# Train for a step.
_train_step(model, optim, expected_param_shapes)
# Now do an eval step.
_eval_step(model, optim, expected_param_shapes)
# And finally do another train step.
_train_step(model, optim, expected_param_shapes)
teardown()
def _test_func(rank, world_size, fsdp_config, tempfile_name, unused):
result = dist_init(rank, world_size, tempfile_name, unused)
assert result, "Dist init failed"
assert isinstance(fsdp_config, dict), str(fsdp_config)
class Model(Module):
def __init__(self):
super().__init__()
self.inner = FSDP(Linear(4, 4), **fsdp_config)
self.outer = Linear(4, 5)
def forward(self, x):
# Forward twice.
i = self.inner(x)
j = self.inner(x)
return self.outer(i + j)
model = FSDP(Model(), **fsdp_config).cuda()
optim = SGD(model.parameters(), lr=0.1)
for _ in range(3):
in_data = torch.rand(64, 4).cuda()
in_data.requires_grad = True
out = model(in_data)
out.sum().backward()
optim.step()
optim.zero_grad()
model.assert_state(TrainingState.IDLE)
teardown()
def _test_identical_outputs_eval(
cls,
model_init_fn,
config,
rank,
group,
num_steps=2,
use_cuda=True,
lr=0.01,
ref_ddp_fn=None,
):
if config.get("mixed_precision", False):
autocast = True
# Force the compute dtype to be torch.float32 so that we get
# identical results as PyTorch DDP when using autocast. Note that
# this will cause the all-gather to happen in FP32, which is slower
# than necessary in most cases.
config["compute_dtype"] = torch.float32
else:
autocast = False
# Establish reference behavior with PyTorch DDP (+ optionally autocast).
model = model_init_fn(group=group, wrapper_config=None).cuda()
if ref_ddp_fn is None:
model = nn.parallel.DistributedDataParallel(model,
device_ids=[rank],
output_device=rank,
process_group=group)
else:
model = ref_ddp_fn(model, group)
ref_loss = cls._eval_with_config(model, autocast)
ref_state_dict = model.module.state_dict()
if config.get("cpu_offload", False):
for k in ref_state_dict.keys():
ref_state_dict[k] = ref_state_dict[k].cpu()
# Confirm we get the same behavior using FullyShardedDataParallel.
if config.get("ssd_offload", False):
config["offload_config"] = OffloadConfig(
offload_type="ssd_offload")
del config["ssd_offload"]
model = FullyShardedDataParallel(
model_init_fn(group=group, wrapper_config=config), group, **config)
if not model.ssd_offload and not model.move_params_to_cpu:
if use_cuda:
model = model.cuda()
else:
assert next(model.parameters()).device == torch.device("cpu")
shard_loss = cls._eval_with_config(model, autocast)
try:
torch.testing.assert_allclose(ref_loss, shard_loss)
except (AssertionError, RuntimeError) as e:
raise Exception(
f"FullyShardedDataParallel didn't match PyTorch DDP using config: {config}\n\n {e}"
)
if config.get("flatten_parameters", True):
metadata = model.local_metadata_dict()
assert isinstance(metadata, dict)
def _test_func(rank, world_size, model, fsdp_config, tempfile_name, unused,
test_case):
result = dist_init(rank, world_size, tempfile_name, unused)
assert result, "Dist init failed"
if test_case["assert_ref_out"]:
with torch.no_grad():
weight = model.weight.T.clone().cuda()
v = torch.Tensor(test_case["inputs"][0][rank]).cuda()
ref_out = torch.matmul(v, weight)
model.to("cuda")
assert isinstance(fsdp_config, dict), str(fsdp_config)
model = FSDP(model, **fsdp_config)
optim = SGD(model.parameters(), lr=0.1)
inputs = test_case["inputs"]
assert len(inputs) == 1 or not test_case["assert_ref_out"]
assert len(inputs[0]) >= world_size
for in_data in inputs:
in_data = Tensor(in_data[rank]).cuda()
out = model(in_data)
out.sum().backward()
optim.step()
optim.zero_grad()
if test_case["assert_ref_out"]:
torch.testing.assert_allclose(ref_out, out)
model.assert_state(TrainingState.IDLE)
teardown()
def _distributed_worker(gpu_id, world_size, with_fsdp, freezing_method,
tempfile_name, unused, rank_0_output, expected_state):
torch.cuda.set_device(gpu_id)
rank = gpu_id
result = dist_init(rank, world_size, tempfile_name, unused)
assert result, "Dist init failed"
torch.manual_seed(0)
torch.backends.cudnn.deterministic = True
batch = torch.randn(size=(2, 3, 224, 224)).cuda()
model = _create_model(with_fsdp)
model = model.cuda()
# freezing the trunk using requires_grad.
assert freezing_method in ["requires_grad", "grad_to_none"]
if freezing_method == "requires_grad":
for param in model.trunk.parameters():
param.requires_grad = False
if with_fsdp:
model = FSDP(model)
else:
model = DistributedDataParallel(model, device_ids=[gpu_id])
if gpu_id == 0:
print(model)
target = torch.LongTensor([0, 1]).cuda()
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=0.1, momentum=0.9)
for iteration in range(3):
out = model(batch)
fake_loss = criterion(out, target)
print("Loss", iteration, ":", fake_loss.item())
optimizer.zero_grad()
fake_loss.backward()
if freezing_method == "grad_to_none":
for param in model.trunk.parameters():
param.grad = None
optimizer.step()
if with_fsdp:
fsdp_state = model.state_dict()
# Move tensors to CPU to compare numerics.
for k, v in fsdp_state.items():
fsdp_state[k] = v.cpu()
assert objects_are_equal(expected_state,
fsdp_state,
raise_exception=True)
elif rank == 0:
state_after = model.module.cpu().state_dict()
torch.save(state_after, rank_0_output)
teardown()
def _test_func(rank, world_size, model, fsdp_config, tempfile_name, unused,
test_case):
result = dist_init(rank, world_size, tempfile_name, unused)
assert result, "Dist init failed"
my_lr = 0.1
device = torch.device("cuda")
if fsdp_config.get("mixed_precision", False):
dtype = torch.float16
fsdp_config["fp32_reduce_scatter"] = True
else:
dtype = torch.float32
if test_case["assert_ref_out"]:
with torch.no_grad():
# Compute one iteration local output.
fp32_weight = model.weight.T.clone().to(device)
weight = fp32_weight.to(dtype)
v = torch.Tensor(test_case["inputs"][0][rank]).to(device, dtype)
ref_forward_output_my_rank = torch.matmul(v, weight)
# Compute one iteration global weight update.
v = torch.Tensor(test_case["inputs"][0][:world_size]).to(
device, dtype)
grad = v.float().sum(0).repeat(weight.shape[0], 1).div(world_size)
ref_weight_out = fp32_weight - grad.T * my_lr
assert ref_weight_out.dtype == torch.float32
model.to(
device) # not dtype, since FSDP will manage mixed precision internally
assert isinstance(fsdp_config, dict), str(fsdp_config)
model = FSDP(model, **fsdp_config)
optim = SGD(model.parameters(), lr=my_lr)
inputs = test_case["inputs"]
assert len(inputs) == 1 or not test_case["assert_ref_out"]
assert len(inputs[0]) >= world_size
for in_data in inputs:
in_data = Tensor(in_data[rank]).to(device, dtype)
out = model(in_data)
out.float().sum().backward()
optim.step()
optim.zero_grad()
if test_case["assert_ref_out"]:
with model.summon_full_params():
weight_out = model.module.weight.data.T.clone()
# make sure we can do more fwd/bwd
loss = model(in_data)
loss.sum().backward()
if test_case["assert_ref_out"]:
torch.testing.assert_allclose(ref_forward_output_my_rank, out)
torch.testing.assert_allclose(ref_weight_out, weight_out)
model.assert_state(TrainingState.IDLE)
teardown()
def _test_func(rank, world_size, tempfile_name, unused, flatten,
mixed_precision, amp_context, half_input, fsdp_wrap_ckpt):
result = dist_init(rank, world_size, tempfile_name, unused)
assert result, "Dist init failed"
# Keep initialization deterministic.
torch.manual_seed(0)
model = FSDP(
SimpleModuleWithCheckpointing(flatten, mixed_precision,
fsdp_wrap_ckpt).cuda(),
flatten_parameters=flatten,
mixed_precision=mixed_precision,
)
optim = torch.optim.SGD(model.parameters(), lr=0.001, momentum=0.9)
# Collect parameter sizes to ensure these stay consistent through the steps below.
expected_param_shapes = {
name: tuple(param.shape)
for name, param in model.named_parameters()
}
# For clarity, this is what `expected_param_shapes` should look like depending on world size:
if not flatten:
assert expected_param_shapes == {
"ffn.0.weight": (5, ),
"ffn.0.bias": (2, ),
"ffn.1.weight": (5, ),
"ffn.1.bias": (2, ),
"ffn.2.weight": (5, ),
"ffn.2.bias": (2, ),
}
else:
assert expected_param_shapes == {
"_fsdp_wrapped_module.flat_param_0": (12, ),
"_fsdp_wrapped_module._fpw_module.ffn.1._fsdp_wrapped_module.flat_param_0":
(6, ),
}, expected_param_shapes
torch.manual_seed(1 + rank)
# Train for a step.
_train_step(model, optim, expected_param_shapes, amp_context,
mixed_precision, half_input)
# Now do an eval step.
_eval_step(model, optim, expected_param_shapes, amp_context,
mixed_precision, half_input)
# And finally do another train step.
_train_step(model, optim, expected_param_shapes, amp_context,
mixed_precision, half_input)
teardown()
def test_input_type(temp_files, fsdp_config, input_cls):
"""Test FSDP with input being a list or a dict, only single GPU."""
if torch_version() < (1, 7, 0):
# This test runs multiple test cases in a single process. On 1.6.0 it
# throw an error like this:
# RuntimeError: Container is already initialized! Cannot initialize it twice!
pytest.skip(
"older pytorch doesn't work well with single process dist_init multiple times"
)
result = dist_init(rank=0,
world_size=1,
filename=temp_files[0],
filename_rpc=temp_files[1])
assert result, "Dist init failed"
assert isinstance(fsdp_config, dict), str(fsdp_config)
class Model(Module):
def __init__(self):
super().__init__()
self.layer = Linear(4, 4)
def forward(self, input):
if isinstance(input, list):
input = input[0]
else:
assert isinstance(input, dict), input
input = input["in"]
return self.layer(input)
model = FSDP(Model(), **fsdp_config).cuda()
optim = SGD(model.parameters(), lr=0.1)
for _ in range(5):
in_data = torch.rand(64, 4).cuda()
in_data.requires_grad = True
if input_cls is list:
in_data = [in_data]
else:
assert input_cls is dict
in_data = {"in": in_data}
out = model(in_data)
out.sum().backward()
optim.step()
optim.zero_grad()
model.assert_state(TrainingState.IDLE)
teardown()
def _dist_worker(rank, world_size, files, wrap_middle, test_fn):
# Get data from files.
file1, file2, sd_before, sd_after, in_data = files
sd_before = torch.load(
sd_before, map_location=lambda storage, loc: storage.cuda(rank))
if test_fn == "train":
sd_after = torch.load(
sd_after, map_location=lambda storage, loc: storage.cuda(rank))
in_data = torch.load(in_data,
map_location=lambda storage, loc: storage.cuda(rank))
result = dist_init(rank=rank,
world_size=world_size,
filename=file1,
filename_rpc=file2)
assert result, "Dist init failed"
fsdp_model = FSDP(
# To debug: first make with_fsdp=False (no inner wrapping) work, then enable inner wrapping
# and make that work.
Model(with_fsdp=True, wrap_middle=wrap_middle),
flatten_parameters=test_fn == "optim_state",
mixed_precision=False,
compute_dtype=torch.float16,
)
fsdp_model.load_state_dict(sd_before)
if test_fn == "train":
_train(fsdp_model, in_data)
objects_are_equal(sd_after,
fsdp_model.state_dict(),
raise_exception=True)
elif test_fn == "eval":
_eval(fsdp_model, in_data)
elif test_fn == "optim_state":
optim = SGD(fsdp_model.parameters(), lr=0.1)
for _ in range(3):
out = fsdp_model(in_data)
out.backward()
optim.step()
sd = fsdp_model.gather_full_optim_state_dict(optim)
if rank == 0:
# There should 8 momentum buffers in the state.
assert len(sd["state"].keys()) == 8
else:
assert sd is None, "only rank 0 should have the optim state"
else:
assert 0, f"invalid test_fn {test_fn}"
teardown()
def test_it(fsdp_config, input_cls):
"""Test FSDP with input being a list or a dict, only single GPU."""
if torch_version() < (1, 6, 0):
pytest.skip("older pytorch doesn't support reduce_scatter")
# Random port in case the next test run quickly, same port would cause conflict.
os.environ["MASTER_ADDR"] = "localhost"
os.environ["MASTER_PORT"] = str(random.randint(2000, 3000))
torch.distributed.init_process_group(backend="nccl", rank=0, world_size=1)
try:
assert isinstance(fsdp_config, dict), str(fsdp_config)
class Model(Module):
def __init__(self):
super().__init__()
self.layer = Linear(4, 4)
def forward(self, input):
if isinstance(input, list):
input = input[0]
else:
assert isinstance(input, dict), input
input = input["in"]
return self.layer(input)
model = FSDP(Model(), **fsdp_config).cuda()
optim = SGD(model.parameters(), lr=0.1)
for _ in range(5):
in_data = torch.rand(64, 4).cuda()
in_data.requires_grad = True
if input_cls is list:
in_data = [in_data]
else:
assert input_cls is dict
in_data = {"in": in_data}
out = model(in_data)
out.sum().backward()
optim.step()
optim.zero_grad()
model.assert_state(TrainingState.IDLE)
finally:
# Clean-up is important or the next test in this file may fail to init the PG.
torch.distributed.destroy_process_group()
del os.environ["MASTER_ADDR"]
del os.environ["MASTER_PORT"]
def _distributed_worker(
gpu_id: int, with_fsdp: bool, sync_file: str, result_file: str
):
torch.cuda.set_device(gpu_id)
dist.init_process_group(
backend="nccl", init_method="file://" + sync_file, world_size=2, rank=gpu_id
)
# Create the inputs
torch.manual_seed(0)
torch.backends.cudnn.deterministic = True
batch = torch.randn(size=(8, 3, 224, 224)).cuda()
# Create a fake model based on SWAV blocks
config = TestRegnetFSDP._create_config(with_fsdp)
model = build_model(config["MODEL"], config["OPTIMIZER"])
model = model.cuda()
if with_fsdp:
model = FSDP(model)
else:
model = DistributedDataParallel(model, device_ids=[gpu_id])
criterion = SwAVLoss(loss_config=config["LOSS"]["swav_loss"])
optimizer = optim.SGD(model.parameters(), lr=1e-2)
# Run a few iterations and collect the losses
losses = []
for iteration in range(5):
out = model(batch)
loss = criterion(out[0], torch.tensor(0.0).cuda())
if gpu_id == 0:
losses.append(loss.item())
optimizer.zero_grad()
loss.backward()
if iteration <= 2:
for name, param in model.named_parameters():
if "prototypes" in name:
param.grad = None
optimizer.step()
# Store the losses in a file to compare several methods
if gpu_id == 0:
with open(result_file, "wb") as f:
pickle.dump(losses, f)
def _test_func(rank, world_size, model, fsdp_config, tempfile_name, unused,
test_case):
result = dist_init(rank, world_size, tempfile_name, unused)
assert result, "Dist init failed"
my_lr = 0.1
if test_case["assert_ref_out"]:
with torch.no_grad():
# Compute one iteration local output.
weight = model.weight.T.clone().cuda()
v = torch.Tensor(test_case["inputs"][0][rank]).cuda()
ref_forward_output_my_rank = torch.matmul(v, weight)
# Compute one iteration global weight update.
v = torch.Tensor(test_case["inputs"][0][:world_size]).cuda()
grad = v.sum(0).repeat(weight.shape[0], 1).div(world_size)
ref_weight_out = weight - grad.T * my_lr
model.to("cuda")
assert isinstance(fsdp_config, dict), str(fsdp_config)
model = FSDP(model, **fsdp_config)
optim = SGD(model.parameters(), lr=my_lr)
inputs = test_case["inputs"]
assert len(inputs) == 1 or not test_case["assert_ref_out"]
assert len(inputs[0]) >= world_size
for in_data in inputs:
in_data = Tensor(in_data[rank]).cuda()
out = model(in_data)
out.sum().backward()
optim.step()
optim.zero_grad()
if test_case["assert_ref_out"]:
with model.summon_full_params():
weight_out = model.module.weight.data.T.clone()
# make sure we can do more fwd/bwd
loss = model(in_data)
loss.sum().backward()
if test_case["assert_ref_out"]:
torch.testing.assert_allclose(ref_forward_output_my_rank, out)
torch.testing.assert_allclose(ref_weight_out, weight_out)
model.assert_state(TrainingState.IDLE)
teardown()
def _test_identical_outputs(
cls, model_init_fn, config, rank, group, num_steps=2, use_cuda=True, lr=0.01, ref_ddp_fn=None, norm_type=2,
):
if config["mixed_precision"]:
autocast = True
# Force the compute dtype to be torch.float32 so that we get
# identical results as PyTorch DDP when using autocast. Note that
# this will cause the all-gather to happen in FP32, which is slower
# than necessary in most cases.
config["compute_dtype"] = torch.float32
else:
autocast = False
# Establish reference behavior with PyTorch DDP (+ optionally autocast).
model = model_init_fn(group=group, wrapper_config=None).cuda()
if ref_ddp_fn is None:
model = nn.parallel.DistributedDataParallel(
model, device_ids=[rank], output_device=rank, process_group=group
)
else:
model = ref_ddp_fn(model, group)
ref_loss = cls._train_for_several_steps(model, num_steps, autocast, lr=lr, norm_type=norm_type)
ref_state_dict = model.module.state_dict()
# Confirm we get the same behavior using FullyShardedDataParallel.
model = FullyShardedDataParallel(model_init_fn(group=group, wrapper_config=config), group, **config)
if use_cuda:
model = model.cuda()
else:
assert next(model.parameters()).device == torch.device("cpu")
shard_loss = cls._train_for_several_steps(model, num_steps, autocast, lr=lr, norm_type=norm_type)
shard_state_dict = model.state_dict()
try:
torch.testing.assert_allclose(ref_loss, shard_loss)
assert objects_are_equal(ref_state_dict, shard_state_dict, raise_exception=True)
except (AssertionError, RuntimeError) as e:
raise Exception(f"FullyShardedDataParallel didn't match PyTorch DDP using config: {config}\n\n {e}")
def _test_func(rank, world_size, fsdp_config, tempfile_name, unused):
result = dist_init(rank, world_size, tempfile_name, unused)
assert result, "Dist init failed"
assert isinstance(fsdp_config, dict), str(fsdp_config)
class InnerModel(Module):
def __init__(self):
super().__init__()
self.layers = Sequential(FSDP(Linear(5, 5), **fsdp_config), )
def forward(self, x):
return self.layers(x)
inner_model = InnerModel()
model = FSDP(inner_model, **fsdp_config).cuda()
optim = SGD(model.parameters(), lr=0.1)
for i in range(3):
input = torch.rand((1, 5), dtype=torch.float).cuda()
input.requires_grad = True
output = model(input)
output.sum().backward()
optim.step()
optim.zero_grad()
input = torch.rand((1, 5), dtype=torch.float).cuda()
output = model(input)
model.assert_state(TrainingState.IDLE)
# second time to rewrap the inner model
rewrapped_model = FSDP(inner_model, **fsdp_config).cuda()
rewrapped_output = rewrapped_model(input)
assert torch.allclose(output, rewrapped_output)
teardown()
def main(local_rank, *args):
torch.backends.cudnn.benchmark = True
init_method = "tcp://%s:%s" % ("0.0.0.0", "9999")
torch.distributed.init_process_group(backend="nccl",
rank=local_rank,
world_size=8,
init_method=init_method)
print("[Train]: Time = %s, Initialized Dist Process for Rank = %s" %
(get_time_string(), local_rank))
device = torch.device(
f'cuda:{local_rank}') # Unique only on individual node.
torch.cuda.set_device(device)
torch.cuda.set_device(device)
fsdp_params = dict(mixed_precision=True,
flatten_parameters=True,
bucket_cap_mb=25,
reshard_after_forward=False,
fp32_reduce_scatter=False,
cpu_offload=False,
move_grads_to_cpu=False,
process_group=torch.distributed.group.WORLD)
with enable_wrap(wrapper_cls=FullyShardedDDP, **fsdp_params):
nn_model = nn.Sequential(
nn.Linear(200, 200),
wrap(
checkpoint_wrapper(nn.Sequential(
nn.Linear(200, 200), nn.Linear(200, 200),
wrap(
checkpoint_wrapper(nn.Linear(200, 200),
offload_to_cpu=True)),
checkpoint_wrapper(nn.GELU(), offload_to_cpu=True),
nn.Linear(200, 200)),
offload_to_cpu=True)),
checkpoint_wrapper(nn.GELU(), offload_to_cpu=True),
nn.LayerNorm(200, eps=1e-7), nn.Linear(200, 64)).cuda()
model = FullyShardedDDP(nn_model, **fsdp_params)
optimizer = torch.optim.AdamW(model.parameters(),
lr=1e-4,
eps=1e-7,
weight_decay=1e-2,
betas=(0.9, 0.99))
optimizer.zero_grad(set_to_none=True)
for i in range(1000):
optimizer.zero_grad(set_to_none=True)
fake_inputs = torch.randn(32, 200, device=device)
fake_labels = torch.randn(32, 64, device=device)
outputs = model(fake_inputs)
loss = ((outputs - fake_labels)**2).mean()
loss.backward()
model.clip_grad_norm_(1.0)
optimizer.step()
if i % 100 == 0:
print("Loss = %s, rank = %s" % (loss.item(), local_rank))
state_dict = model.state_dict()
nn_model = nn.Sequential(
nn.Linear(200, 200),
nn.Sequential(nn.Linear(200, 200), nn.Linear(200, 200),
nn.Linear(200, 200),
checkpoint_wrapper(nn.GELU(), offload_to_cpu=True),
nn.Linear(200, 200)),
checkpoint_wrapper(nn.GELU(), offload_to_cpu=True),
nn.LayerNorm(200, eps=1e-7), nn.Linear(200, 64)).cuda()
nn_model.load_state_dict(state_dict)
print("[Train]: Time = %s, Trainable Params = %s" %
(get_time_string(), numel(nn_model) / 1_000_000))
def _distributed_worker(
gpu_id,
world_size,
with_model2,
with_sync_bn,
with_fsdp,
with_checkpoint,
files,
mixed_precision,
flatten,
wrap_bn,
fp32_reduce_scatter,
bucket_cap_mb,
):
filename, filename_rpc = files[:2]
filename_loss = files[2:]
torch.cuda.set_device(gpu_id)
rank = gpu_id
result = dist_init(rank, world_size, filename, filename_rpc)
assert result, "Dist init failed"
# use False below to debug since error msg is not as good with cudnn.
torch.backends.cudnn.enabled = True
# these make things deterministic.
torch.manual_seed(0)
torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = False
# Ensure we have multiple forward passes.
batch = [
torch.randn(size=(2, 3, 16, 16)).cuda(),
torch.randn(size=(2, 3, 9, 9)).cuda(),
torch.randn(size=(2, 3, 9, 9)).cuda(),
]
if mixed_precision and not with_fsdp:
batch = [x.half() for x in batch]
model = _create_model(
with_model2,
with_sync_bn,
with_fsdp,
with_checkpoint,
mixed_precision,
flatten,
wrap_bn,
fp32_reduce_scatter,
bucket_cap_mb,
)
model = model.cuda()
if with_fsdp:
model = FSDP(
model,
flatten_parameters=flatten,
mixed_precision=mixed_precision,
compute_dtype=torch.float32,
fp32_reduce_scatter=fp32_reduce_scatter,
bucket_cap_mb=bucket_cap_mb,
)
model.set_gradient_divide_factors(1.0, 2.0, True)
no_sync_context = contextlib.suppress()
else:
# With DDP, we need no_sync and manual gradient reduction below because
# it can't handle multiple forward pass + checkpointing otherwise.
model = DistributedDataParallel(model, device_ids=[gpu_id])
no_sync_context = model.no_sync()
mp_context = contextlib.suppress()
if mixed_precision:
mp_context = torch.cuda.amp.autocast(enabled=True)
if gpu_id == 0:
print(model)
target = torch.tensor([0, 1, 2, 3, 4, 5], dtype=torch.long).cuda()
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=0.1, momentum=0.9)
losses = {}
i = 0
with no_sync_context:
for iteration in range(3):
with mp_context:
out = model(batch)
loss = criterion(out, target)
print("Loss", iteration, ":", loss.item())
losses[f"iter_{i}"] = loss
i += 1
optimizer.zero_grad()
loss.backward()
# Manual grad reduction, no autocast.
if not with_fsdp:
for p in model.parameters():
dist.all_reduce(p.grad.data)
p.grad.data.div_(2.0)
# Stepping, no autocast
optimizer.step()
# Due to dist.all_reduce code block above with ddp.no_sync, we seem to hit a bug
# in DDP where tensor.cpu() and torch.save() calls both hang. FSDP is not affected.
# Therefore, we have to compare losses here instead of states.
with open(filename_loss[rank], "wb") as f:
pickle.dump(losses, f)
teardown()