🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/167178

• 📄 Preview Python docs built from this PR
• 📄 Preview C++ docs built from this PR
• ❓ Need help or want to give feedback on the CI? Visit the bot commands wiki

Note: Links to docs will display an error until the docs builds have been completed.

✅ You can merge normally! (1 Unrelated Failure)

As of commit a88ea4f with merge base dc00842 ():

This comment was automatically generated by Dr. CI and updates every 15 minutes.

The committers listed above are authorized under a signed CLA.

• ✅ login: garrett361 / name: Garrett Goon (019e636, 0353eac, 6c015d6, 8862e11, 9810b16, a88ea4f, a90a8d5, b971070, cb6a1d9, e967d2c, ef9dcf2)

Thanks a lot! Sounds right to me. I'll let @weifengpy review / stamp.

Yes, a test would be appreciated.

good catch!

triggering CI - waiting for signals

Thanks @weifengpy ! Figuring out CLA stuff with work, then will add a test.

@weifengpy still waiting on CLA stuff, but started looking at where I'd add a test. I'm not seeing gradient_divide_factor tested anywhere in the code base. Am I missing a test somewhere? Closest I can find is in TestFullyShardCollectiveOps where it's mentioned but left as None.

Should I add a test here? Seems like I'd have to add a decent bit of infra code to get a proper test set up.

Should I add a test here? Seems like I'd have to add a decent bit of infra code to get a proper test set up.

adding a unit test would be great! I was mentioning CI to make sure it does not break shard world size 2+, but having unit test on 1 is better

cc @anshul-si for a bug fix in world size 1

I found some other issues in the code and tests related to this topic:

1. There is a current edge case where if the user calls model.set_gradient_divide_factor(factor) and factor happens to equal the data parallel size, then a (I believe) unintended code path is taken: we end up dividing grads by factor and then averaging over the allreduce group, rather than summing. From the code below:

   pytorch/torch/distributed/fsdp/_fully_shard/_fsdp_collectives.py

   Lines 727 to 736 in dc00842

   	if not overflow_risk and not force_sum_reduction_for_comms:
   	if factor == data_parallel_size:
   	# Warning: NCCL ReduceOp.AVG may produce incorrect results with
   	# world size 1.
   	if data_parallel_size == 1:
   	return None, None, ReduceOp.SUM, ReduceOp.SUM
   	return None, None, ReduceOp.AVG, ReduceOp.AVG
   	else:
   	reduce_scatter_op = torch.distributed._make_nccl_premul_sum(1 / factor)
   	return None, None, reduce_scatter_op, ReduceOp.SUM

   if factor != data_parallel_size then we end up at the final line and return None, None, reduce_scatter_op, ReduceOp.SUM. But if we happen to set factor = data_parallel_size, then we enter the first if statement and return None, None, ReduceOp.AVG, ReduceOp.AVG, assuming data_parallel_size > 1. This is a change in semantics due to swapping the final ReduceOp.SUM return value for a ReduceOp.AVG in the latter case. I believe we should always want ReduceOp.SUM if a custom division factor is provided and changed the code to reflect that, but let me know if that is not desired.

2. I found the relevant _test_set_reduce_scatter_divide_factor test, and as written it wasn't sensitive enough to catch some of these issues. I updated this PR to increase the sensitivity of this test and cover more cases. This is how I caught the above edge case.

3. There's also the test _test_set_reduce_scatter_divide_factor_mixed_prevision, but this doesn't seem to actually test anything about custom division factors because we apply the division factor only to the outermost fully_shard wrapped module (the whole model) which doesn't seem to own any parameters itself, since it's a Sequential:
   

   pytorch/test/distributed/_composable/fsdp/test_fully_shard_comm.py

   Lines 474 to 482 in dc00842

   	model = nn.Sequential(*[MLP(16) for _ in range(3)])
   	ref_model = copy.deepcopy(model).to(device_type)
   	ref_model_bf16 = copy.deepcopy(ref_model).to(param_dtype)
   	ref_optim = torch.optim.AdamW(ref_model.parameters(), lr=1e-2)
   	for mlp in model:
   	fully_shard(mlp, mp_policy=mp_policy)
   	model = fully_shard(model, mp_policy=mp_policy)
   	optim = torch.optim.AdamW(model.parameters(), lr=1e-2)
   	model.set_reduce_scatter_divide_factor(divide_factor)

   Sound right? I didn't touch this test yet, but can if desired. Trying to keep the PR small.

CC @anshul-si @weifengpy @tianyu-l

Tested locally that everything in test/distributed/_composable/fsdp/test_fully_shard_comm.py passes.

@garrett361 thanks for the detailed unit tests

Thanks @weifengpy , do you want me to make any changes to _test_set_reduce_scatter_divide_factor_mixed_prevision as well?

Hi @weifengpy , checking in on this. Anything else you need from my end? Thanks!

Hi @weifengpy , checking in on this. Anything else you need from my end? Thanks!

@garrett361 just took another look and it's safe. we can land once CI passes. no need to cover _test_set_reduce_scatter_divide_factor_mixed_prevision

@pytorchmergebot merge

Merge started

Your change will be merged once all checks pass (ETA 0-4 Hours).

Learn more about merging in the wiki.

Questions? Feedback? Please reach out to the PyTorch DevX Team