Home
last modified time | relevance | path

Searched defs:FlatParamHandle (Results 1 – 10 of 10) sorted by relevance

/external/pytorch/torch/distributed/fsdp/
D_runtime_utils.py275 handle: FlatParamHandle,
308 handle: FlatParamHandle,
330 handle: Optional[FlatParamHandle],
337 handle: Optional[FlatParamHandle],
346 handle: Optional[FlatParamHandle],
409 handle: Optional[FlatParamHandle],
436 handle: Optional[FlatParamHandle],
484 handle: FlatParamHandle,
622 handle: FlatParamHandle,
692 handle: FlatParamHandle,
[all …]
D_exec_order_utils.py91 current_handle: FlatParamHandle,
112 current_handle: FlatParamHandle,
131 def record_post_forward(self, handle: Optional[FlatParamHandle]) -> None:
152 self, handle: Optional[FlatParamHandle], is_training: bool
173 def _check_order(self, handle: FlatParamHandle, is_training: bool) -> None:
309 handle: FlatParamHandle,
341 handle: FlatParamHandle,
D_common_utils.py368 state: _FSDPState, handle: "FlatParamHandle", logger: logging.Logger
380 state: _FSDPState, handle: "FlatParamHandle"
D_unshard_param_utils.py35 handle: FlatParamHandle,
Dfully_sharded_data_parallel.py2058 flat_param_handle: Optional[FlatParamHandle],
D_optim_utils.py704 handle: FlatParamHandle,
D_flat_param.py451 class FlatParamHandle: class
/external/pytorch/test/distributed/_composable/fully_shard/
Dtest_fully_shard_runtime.py190 handle: FlatParamHandle,
200 handle: FlatParamHandle,
276 composable_handle: FlatParamHandle,
277 wrapped_handle: FlatParamHandle,
/external/pytorch/test/distributed/fsdp/
Dtest_fsdp_flatten_params.py543 handle: FlatParamHandle,
Dtest_fsdp_core.py465 self: FlatParamHandle, as_params: bool