Home
last modified time | relevance | path

Searched defs:GradBucket (Results 1 – 14 of 14) sorted by relevance

/aosp_15_r20/external/pytorch/torch/distributed/algorithms/ddp_comm_hooks/
H A Ddefault_hooks.py34 process_group: dist.ProcessGroup, bucket: dist.GradBucket
57 bucket: dist.GradBucket,
105 bucket: dist.GradBucket,
152 hook: Callable[[Any, dist.GradBucket], torch.futures.Future[torch.Tensor]]
169 hook_state, bucket: dist.GradBucket
190 hook: Callable[[Any, dist.GradBucket], torch.futures.Future[torch.Tensor]]
208 hook_state, bucket: dist.GradBucket
H A Dddp_zero_hook.py22 bucket: dist.GradBucket,
102 bucket: dist.GradBucket,
135 bucket: dist.GradBucket,
175 hook: Callable[[Any, dist.GradBucket], torch.futures.Future],
248 bucket: dist.GradBucket,
337 hook: Callable[[Any, dist.GradBucket], torch.futures.Future],
410 bucket: dist.GradBucket,
H A Doptimizer_overlap_hooks.py66 bucket: dist.GradBucket,
129 hook: Callable[[Any, dist.GradBucket], torch.futures.Future[torch.Tensor]],
139 hook_state, bucket: dist.GradBucket
H A Dquantization_hooks.py47 process_group: dist.ProcessGroup, bucket: dist.GradBucket
122 process_group: dist.ProcessGroup, bucket: dist.GradBucket, bucket_size=512
H A DpowerSGD_hook.py340 state: PowerSGDState, bucket: dist.GradBucket
651 state: PowerSGDState, bucket: dist.GradBucket
H A Ddebugging_hooks.py10 def noop_hook(_: Any, bucket: GradBucket) -> torch.futures.Future[torch.Tensor]:
H A Dmixed_precision_hooks.py26 hook_state: _AllreduceUpcastHookState, bucket: dist.GradBucket
H A Dpost_localSGD_hook.py70 state: PostLocalSGDState, bucket: dist.GradBucket
/aosp_15_r20/external/pytorch/torch/csrc/distributed/c10d/
H A Dcomm.hpp19 class TORCH_API GradBucket { class
21 explicit GradBucket( in GradBucket() function in c10d::GradBucket
/aosp_15_r20/external/pytorch/test/distributed/
H A Dtest_c10d_ucc.py895 def comm_hook(state: object, bucket: dist.GradBucket) -> int:
907 def comm_hook(state: object, bucket: dist.GradBucket):
963 state: object, bucket: dist.GradBucket
H A Dtest_c10d_gloo.py2067 def comm_hook(state: object, bucket: dist.GradBucket) -> int:
2079 def comm_hook(state: object, bucket: dist.GradBucket):
2133 state: object, bucket: dist.GradBucket
H A Dtest_c10d_nccl.py1794 state: object, bucket: dist.GradBucket
2020 state: object, bucket: dist.GradBucket
H A Dtest_c10d_common.py872 self, state: object, bucket: dist.GradBucket
/aosp_15_r20/external/pytorch/torch/testing/_internal/distributed/
H A Ddistributed_test.py5381 group_id: object, bucket: dist.GradBucket
5410 group_id: object, bucket: dist.GradBucket