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, 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 init_distributed_data_parallel_model(self):
"""
Initialize FSDP if needed.
This method overloads the ClassificationTask class's method from ClassyVision.
"""
if not is_distributed_training_run():
return
# Make sure default cuda device is set. TODO (Min): we should ensure FSDP can
# be enabled for 1-GPU as well, but the use case there is likely different.
# I.e. perhaps we use it for cpu_offloading.
assert get_cuda_device_index(
) > -1, "Distributed training not setup correctly"
# The model might be already wrapped by FSDP internally. Check regnet_fsdp.py.
# Here, we wrap it at the outer most level.
fsdp_config = self.config["MODEL"]["FSDP_CONFIG"]
if is_primary():
logging.info(f"Using FSDP, config: {fsdp_config}")
# First, wrap the head's prototype_i layers if it is SWAV.
# TODO (Min): make this more general for different models, which may have multiple
# heads.
if len(self.base_model.heads) != 1:
raise ValueError(
f"FSDP only support 1 head, not {len(self.base_model.heads)} heads"
)
head0 = self.base_model.heads[0]
if isinstance(head0, SwAVPrototypesHead):
# This is important for convergence!
#
# Since we "normalize" this layer in the update hook, we need to keep its
# weights in full precision. It is output is going into the loss and used
# for clustering, so we need to have that in full precision as well.
fp_fsdp_config = fsdp_config.copy()
fp_fsdp_config["flatten_parameters"] = False
fp_fsdp_config["mixed_precision"] = False
fp_fsdp_config["fp32_reduce_scatter"] = False
for j in range(head0.nmb_heads):
module = getattr(head0, "prototypes" + str(j))
module = FSDP(module=module, **fp_fsdp_config)
setattr(head0, "prototypes" + str(j), module)
head0 = FSDP(module=head0, **fsdp_config)
self.base_model.heads[0] = head0
# Init the head properly since the weights are potentially initialized on different
# ranks with different seeds. We first summon the full params from all workers.
# Then, within that context, we set a fixed random seed so that all workers init the
# weights the same way. Finally, we reset the layer's weights using reset_parameters().
#
# TODO (Min): This will go away once we have a way to sync from rank 0.
with head0.summon_full_params():
with set_torch_seed(self.config["SEED_VALUE"]):
for m in head0.modules():
if isinstance(m, Linear):
m.reset_parameters()
head0._reset_lazy_init()
head0.prototypes0._reset_lazy_init()
# TODO (Min): We can load checkpoint, but it ends up setting the trunk's _is_root
# flag to true. We need to set it back to None here.
# Also, right now, the head's weight is only partially loaded from the checkpoint
# because we dump the checkpoint after the head if wrapped, but loading it before
# it is wrapped.
# For very big models, we need re-work the checkpoint logic because we don't have
# enough memory to load the entire model on one node. We need to use local_state_dict()
# API to load checkpoint shards.
for module in self.base_model.trunk.modules():
if isinstance(module, FSDP):
module._is_root = None
# Then, wrap the whole model. We replace the base_model since it is used
# when checkpoint is taken.
self.base_model = FSDP(module=self.base_model, **fsdp_config)
self.distributed_model = self.base_model