mirror of
https://github.com/zebrajr/pytorch.git
synced 2025-12-06 12:20:52 +01:00
e74ba1b34a
See https://github.com/pytorch/pytorch/pull/129751#issue-2380881501. Most changes are auto-generated by linter. You can review these PRs via: ```bash git diff --ignore-all-space --ignore-blank-lines HEAD~1 ``` Pull Request resolved: https://github.com/pytorch/pytorch/pull/129767 Approved by: https://github.com/anijain2305
26 lines
594 B
Python
26 lines
594 B
Python
from typing import Optional
|
|
|
|
import torch.distributed as dist
|
|
|
|
from . import config
|
|
|
|
|
|
_COMPILE_PG: Optional[dist.ProcessGroup] = None
|
|
|
|
|
|
def get_compile_pg() -> Optional[dist.ProcessGroup]:
|
|
if (
|
|
config.enable_compiler_collectives
|
|
and dist.is_available()
|
|
and dist.is_initialized()
|
|
):
|
|
global _COMPILE_PG
|
|
if _COMPILE_PG is None:
|
|
# , timeout=datetime.timedelta(seconds=2)
|
|
_COMPILE_PG = dist.distributed_c10d._new_group_with_tag(
|
|
pg_tag="pt2_compile_pg"
|
|
)
|
|
return _COMPILE_PG
|
|
|
|
return None
|