better guarding to handle torch compiled with USE_DISTRIBUTED=0 - #47619
Conversation
|
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. |
| for name in ( | ||
| "torch.distributed.tensor", | ||
| "torch.distributed.checkpoint", | ||
| "torch.distributed.fsdp", | ||
| "torch.distributed._composable", | ||
| ): | ||
| sys.modules[name] = None |
There was a problem hiding this comment.
These do exist in USE_DISTRIBUTED=0 build, it's the compiled torch._C._distributed_c10d that is missing.
sys.modules["torch._C._distributed_c10d"] = None would be more faithful to the actual case, the modules may still need evicting del sys.modules[name] though.
…ce/transformers into fix_distributed_guarding_audio
vasqu
left a comment
vasqu
left a comment
There was a problem hiding this comment.
Smaller comments, the one that is important to me is the test design to please use mocks and the sort so we dont play around real object -> less chance of any side effect
| if is_torch_available(): | ||
| import torch | ||
|
|
||
| if is_torch_distributed_available() and is_torch_greater_or_equal("2.5"): |
There was a problem hiding this comment.
| if is_torch_distributed_available() and is_torch_greater_or_equal("2.5"): | |
| if is_torch_distributed_available(): |
We removed support for 2.4.x 🫡
| if is_torch_available(): | ||
|
|
||
| @lru_cache | ||
| def is_torch_distributed_available() -> bool: |
There was a problem hiding this comment.
nit but why not have this live under import utils?
| materialize_device = param_device | ||
|
|
||
| if isinstance(empty_param, DTensor): | ||
| if is_torch_distributed_available() and isinstance(empty_param, DTensor): |
There was a problem hiding this comment.
Maybe small helper instead of all the time checking both manually
| # Forget transformers, so that importing it below actually re-runs its module-scope imports. | ||
| for name in list(sys.modules): | ||
| if name.startswith("transformers"): | ||
| del sys.modules[name] | ||
|
|
||
| # Emulate USE_DISTRIBUTED=0 by faking torch.distributed availability to False and deleting the distributed submodules in sys.modules. | ||
| torch.distributed.is_available = lambda: False | ||
| sys.modules["torch._C._distributed_c10d"] = None | ||
| for name in list(sys.modules): | ||
| if name.startswith( | ||
| ( | ||
| "torch.distributed.tensor", | ||
| "torch.distributed.checkpoint", | ||
| "torch.distributed.fsdp", | ||
| "torch.distributed._composable", | ||
| ) | ||
| ): | ||
| del sys.modules[name] |
There was a problem hiding this comment.
Could we rather temporarily monkey patch, not feeling too well about setting on "real" objects
…ce/transformers into fix_distributed_guarding_audio
…ce/transformers into fix_distributed_guarding_audio
…ingface#47619) * better guarding to handle torch compiled with USE_DISTRIBUTED=0 * better test * no private variable * remove 2.5 guarding * move is_torch_distributed_available() to utils.py * create an is_dtensor function * monkey patch test instead * linting
4 participants
This PR fixes #47603 when supported torch version has been compiled without distributed features which prevents import of transformers to work properly (due to non robust distributed guarding)