[FSDP] Do not clean FQNs even for use_orig_params=True

[awgu]

Stack from ghstack:

• [FSDP][RFC] Enforce rank r's current device is cuda:r #92035 [FSDP][RFC] Enforce rank r's current device is cuda:r
• [FSDP] Do not clean FQNs even for use_orig_params=True #91767 [FSDP] Do not clean FQNs even for use_orig_params=True
• [FSDP][BE] Improve device_id + CPU offload test #92031 [FSDP][BE] Improve device_id + CPU offload test
• [FSDP][BE] Rename prefixed_param_names -> fqns for consolidation #92028 [FSDP][BE] Rename prefixed_param_names -> fqns for consolidation
• [FSDP][BE] Better error msg for incorrect device for training #92027 [FSDP][BE] Better error msg for incorrect device for training

Cleaning FQN for FullyShardedDataParallel(use_orig_params=True) can cause some discrepancies with respect to the FQN compared to manually looping over named_modules() and named_parameters() together.

There is no requirement for the FQNs to be clean when using wrapper FSDP + use_orig_params=True. We can leave clean FQNs to fully_shard.

cc @mlazos @soumith @voznesenskym @yanboliang @penguinwu @anijain2305 @EikanWang @jgong5 @Guobing-Chen @chunyuan-w @XiaobingSuper @zhuhaozhe @blzheng @Xia-Weiwen @wenzhe-nrv @jiayisunx @desertfire

[pytorch-bot]

🔗 Helpful Links

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

• 📄 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 or our office hours

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

✅ No Failures

As of commit d50af79:
💚 Looks good so far! There are no failures yet. 💚

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

[pytorchmergebot]

@awgu your PR has been successfully reverted.

[awgu]

@wconstab

Context: This PR changes FSDP so that when use_orig_params=True, the parameter names returned from named_parameters() include the _fsdp_wrapped_module. prefixes. Before, FSDP overrode named_parameters() to clean the prefix (see the change in fully_sharded_data_parallel.py). I found that overriding it to clean the prefix creates unwanted discrepancy with the true module structure, which can lead to some issues. We can leave having clean parameter names to the currently-developed composable APIs (e.g. fully_shard).

This change here in _dynamo/testing.py to not call clean_tensor_name() allows the unit tests to pass. I still wanted to check: does Dynamo + FSDP rely on FSDP's named_parameters() returning the cleaned names (i.e. without _fsdp_wrapped_module.)?

[wconstab]

I don't think dynamo relies on the structure of the name either way. other than we do rely on finding the "is_fsdp_wrapped_module" boolean for enablement of fsdp handling.

question: what is the history of remove_optimized_module_prefix fn? at a glance, i have no memory of modifying this myself for fsdp but i see it had an fsdp clean param thing in it. did you put that there before and are taking it out now?

[awgu]

removed_optimized_module_prefix came from #89113

My understanding is as follows:

pytorch/torch/_dynamo/testing.py

Lines 63 to 65 in 46a81c8

	for name, param in model.named_parameters():
	if isinstance(model, eval_frame.OptimizedModule):
	name = remove_optimized_module_prefix(name)

pytorch/test/distributed/test_dynamo_distributed.py

Lines 368 to 370 in 46a81c8

	correct_results = collect_results(eager_model, correct_outputs.logits, correct_loss, inputs_flat)
	opt_results = collect_results(opt_model, opt_outputs.logits, opt_loss, inputs_flat)
	self.assertTrue(same(correct_results, opt_results))

The eager_model removed the prefixes but the opt_model did not, so for parity, we needed to manually remove the prefixes for opt_model. However, with this PR, eager_model no longer removes the prefixes, so similarly, we can remove the manual removal for opt_model.

In this case, we should be safe to land this.

[wconstab]

seems ok to me

[awgu]

@pytorchbot merge

[pytorchmergebot]

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

Advanced Debugging

Check the merge workflow status
here

[malfet]

@pytorchbot revert -m "Looks like it broke test_compatible_with_named_optimizer distribued tests, see https://hud.pytorch.org/pytorch/pytorch/commit/d6f3265e1add26abedb504910be93b393b9fb33c" -c nosignal

[pytorchmergebot]

@pytorchbot successfully started a revert job. Check the current status here.
Questions? Feedback? Please reach out to the PyTorch DevX Team

[pytorchmergebot]

@awgu your PR has been successfully reverted.

[awgu]

We had a land race. We will re-land this PR after a fix on the optimizer state side.