Repository navigation
fix module and optimizer parameter mismatch before prepare_tp_ - #3845
Conversation
Add DTensor type class in get_named_parameters if DTensor exists before _prepare_tp
0411768 to
e7e949b
Compare
|
@SunMarc please take a look at this pull request addressing the failed test case from #3835 Based on the root cause section I wrote in the description, I don't think it is not caused by #3835 itself, but feel free to ping me for more discussion and modification. I'm located in the US, sorry if I respond late. Thank you so much!!! |
There was a problem hiding this comment.
Thanks, left a comment. Also, left a comment here, would you mind fixing it ? #3835 (comment)
| accessor_mapping[WeightWithDynamicFloat8CastTensor] = "_tensor" | ||
| # we know we're in FSDP2 so DTensor is available | ||
| if self.is_fsdp2: | ||
| if self.is_fsdp2 or self.initial_dtensor_count > 0: |
There was a problem hiding this comment.
the best fix here would be to just change the condition to one where we check if distributed is available + pytorch version.
There was a problem hiding this comment.
Do you mean change the condition to:
isinstance(tensor, DTensor)?
There was a problem hiding this comment.
this is why we do in transformers for example to check if dtensor is available:
_torch_distributed_available = torch.distributed.is_available()
_is_dtensor_available = _torch_distributed_available and is_torch_greater_or_equal("2.5")
if _is_dtensor_available:
from torch.distributed.tensor import DTensorThere was a problem hiding this comment.
Thanks for your advice, changed in this pull request.
3cdf5e3 to
c117685
Compare
Thanks for your help! I checked the prepare_tp won't modify the key value in the module.parameter. |
| def detect_module_and_optimizer_parameter_mismatch(self, *args, debug_str=""): | ||
| from torch.distributed.tensor import DTensor | ||
|
|
||
| param_reference_dict = {"module": [], "optimizer": [], "parameter_func": [], "dtensor": 0} | ||
|
|
||
| def _get_tensor_address(p): | ||
| if isinstance(p, DTensor): | ||
| return p._local_tensor.data_ptr() | ||
| return p.data_ptr() | ||
|
|
||
| for obj in args: |
There was a problem hiding this comment.
we don't need those then, can you can up a bit the pr ?
There was a problem hiding this comment.
sure! cleaned it up.
12587ad to
7a01ac9
Compare
|
The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update. |
What does this PR do?
Root Cause of Failing test_tp.py from #3835
Here what I believe, for torch.optim.Optimizer the only correct Tensor address is Dtensor._local_tensor.data_ptr() not the Dtensor abstract class reference. So we should use _local_tensor as long as DTensor exists in the module.parameter() even when self.is_fsdp2 = False.
Fixes # (issue) #3835
Before this pull request, test_working_of_tp failed due to address mismatch:
With this pull request, we can run test_working_of_tp successfully.
We can run test_working_of_tp_and_fsdp successfully.
Before submitting
Pull Request section?
to it if that's the case.
documentation guidelines, and
here are tips on formatting docstrings.
Who can review?
Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.