Home
last modified time | relevance | path

Searched refs:FSDPTestMultiThread (Results 1 – 10 of 10) sorted by relevance

/aosp_15_r20/external/pytorch/test/distributed/_composable/fsdp/
H A Dtest_fully_shard_init.py38 from torch.testing._internal.common_fsdp import FSDPTestMultiThread, MLP
47 class TestFullyShardDeviceTensor(FSDPTestMultiThread):
65 class TestFullyShardDeviceDTensor(FSDPTestMultiThread):
127 class TestFullyShardMeshArg(FSDPTestMultiThread):
143 class TestFullyShardManagedModulesAndStates(FSDPTestMultiThread):
266 class TestFullyShardParamModuleInfos(FSDPTestMultiThread):
341 class TestFullyShardShardedParameterTensor(FSDPTestMultiThread):
387 class TestFullyShardShardedParameterDTensor(FSDPTestMultiThread):
432 class TestFullyShardLazyInit(FSDPTestMultiThread):
585 class TestFullyShardMetaDeviceInit(FSDPTestMultiThread):
[all …]
H A Dtest_fully_shard_state.py9 from torch.testing._internal.common_fsdp import FSDPTestMultiThread, MLP
13 class TestFullyShardState(FSDPTestMultiThread):
H A Dtest_fully_shard_state_dict.py21 from torch.testing._internal.common_fsdp import FSDPTest, FSDPTestMultiThread, MLP
206 class TestFullyShardStateDictMultiThread(FSDPTestMultiThread):
H A Dtest_fully_shard_extensions.py20 FSDPTestMultiThread,
143 TestFullyShardAllGatherExtensionsCommon, FSDPTestMultiThread
H A Dtest_fully_shard_autograd.py21 FSDPTestMultiThread,
243 class TestFullyShardPostAccGradHookMultiThread(FSDPTestMultiThread):
H A Dtest_fully_shard_training.py38 FSDPTestMultiThread,
62 class TestFullyShardForwardInputs(FSDPTestMultiThread):
95 class TestFullyShardRegisteredParams(FSDPTestMultiThread):
200 class TestFullyShardCastAfterInit(FSDPTestMultiThread):
H A Dtest_fully_shard_mixed_precision.py23 FSDPTestMultiThread,
324 class TestFullyShardMixedPrecisionCasts(FSDPTestMultiThread):
H A Dtest_fully_shard_comm.py44 FSDPTestMultiThread,
64 class TestFullyShardCollectiveOps(FSDPTestMultiThread):
1101 class TestFullyShardUnshardMultiThread(FSDPTestMultiThread):
/aosp_15_r20/external/pytorch/test/distributed/fsdp/
H A Dtest_fsdp_misc.py39 FSDPTestMultiThread,
628 class TestFSDPMiscMultiThread(FSDPTestMultiThread):
987 class TestFSDPMiscWorldSize1(FSDPTestMultiThread):
/aosp_15_r20/external/pytorch/torch/testing/_internal/
H A Dcommon_fsdp.py1091 class FSDPTestMultiThread(MultiThreadedTestCase): class