def _test_state_dict_device(self, config, rank, group, pure_fp16=False, **model_kwargs):
model = TransformerWithSharedParams(group, **model_kwargs)
if pure_fp16:
assert not config["mixed_precision"]
model = model.half()
fsdp_model = FSDP(model, group, **config)
if not config["cpu_offload"]:
fsdp_model = fsdp_model.cuda()
autocast = fsdp_model.mixed_precision or pure_fp16
self._train_for_several_steps(fsdp_model, 1, autocast)
sd = fsdp_model.state_dict()
sd_device = config.get("state_dict_device")
for k, v in sd.items():
if config["cpu_offload"] or (sd_device is not None and sd_device.type == "cpu"):
assert v.device.type == "cpu", v.device.type
else:
assert v.device.type == "cuda", v.device.type
expected_dtype = torch.float16 if pure_fp16 else torch.float32
for k, v in sd.items():
if not torch.is_floating_point(v):
continue
assert v.dtype == expected_dtype, f"{v.dtype} != {expected_dtype}"
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_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_state_dict_device(self,
config,
rank,
group,
pure_fp16=False,
**model_kwargs):
model = TransformerWithSharedParams(group, **model_kwargs)
if pure_fp16:
assert not config["mixed_precision"]
model = model.half()
fsdp_model = FullyShardedDataParallel(model, group, **config)
if not config["cpu_offload"]:
fsdp_model = fsdp_model.cuda()
autocast = fsdp_model.mixed_precision or pure_fp16
self._train_for_several_steps(fsdp_model, 1, autocast)
sd = fsdp_model.state_dict()
sd_device = config.get("state_dict_device")
for k, v in sd.items():
if config["cpu_offload"] or (sd_device is not None
and sd_device.type == "cpu"):
assert v.device.type == "cpu", v.device.type
else:
assert v.device.type == "cuda", v.device.type
expected_dtype = torch.float16 if pure_fp16 else torch.float32
buffers = {
k.replace("_fsdp_wrapped_module.", "").replace("_fpw_module.", "")
for k, _ in fsdp_model.named_buffers()
}
for k, v in sd.items():
if not torch.is_floating_point(v):
continue
if k in buffers:
assert v.dtype == fsdp_model.buffer_dtype, f"{v.dtype} != {fsdp_model.buffer_dtype}"
else:
assert v.dtype == expected_dtype, f"{v.dtype} != {expected_dtype}"
def _test_named_params(self, rank, group, config):
# Get the named parameters before wrapping.
before_wrap_model = TransformerWithSharedParams(group)
before_wrap_params = before_wrap_model.named_parameters()
with tempfile.TemporaryDirectory() as current_tempdir:
if config["ssd_offload"]:
config["offload_config"] = OffloadConfig(
offload_type="ssd_offload",
ssd_filepath_dir=current_tempdir)
del config["ssd_offload"]
model = FullyShardedDataParallel(before_wrap_model, **config)
print(f"model.ssd_offload {model.ssd_offload}")
if not model.ssd_offload and not model.move_params_to_cpu:
model = model.cuda()
self._eval_with_config(model, autocast=config["mixed_precision"])
# Get the named parameters after wrapping to compare.
after_wrap_params = model.named_parameters()
if not config.get("flatten_parameters", False):
for before_nm, after_nm in zip(before_wrap_params,
after_wrap_params):
assert before_nm[0] == after_nm[0]
else:
named_params_flat = [p for p in after_wrap_params][0][0]
assert "flat_param_0" in named_params_flat
after_wrap_params = model.named_parameters()
for before_nm, after_nm_original in zip(before_wrap_params,
after_wrap_params):
assert before_nm[0] == after_nm_original[0]
torch.testing.assert_allclose(before_nm[1].shape,
after_nm_original[1].shape)