Home
last modified time | relevance | path

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

/aosp_15_r20/external/pytorch/test/distributed/fsdp/
H A Dtest_fsdp_unshard_params.py28 NestedWrappedModule,
125 local_model = NestedWrappedModule.init(
134 fsdp_model = NestedWrappedModule.init(
434 model = NestedWrappedModule.init(
445 NestedWrappedModule.init(
711 nested_wrapped_module = NestedWrappedModule.init(
724 nested_wrapped_module = NestedWrappedModule.init(
H A Dtest_fsdp_pure_fp16.py18 NestedWrappedModule,
61 NestedWrappedModule,
101 model = NestedWrappedModule.init(
H A Dtest_fsdp_misc.py41 NestedWrappedModule,
117 nested_wrapped_module = NestedWrappedModule.init(
126 nested_wrapped_module = NestedWrappedModule.init(
139 nested_wrapped_module = NestedWrappedModule.init(
555 nested_wrapped_module = NestedWrappedModule.init(
581 return NestedWrappedModule.init(
752 NestedWrappedModule.init(
949 model = NestedWrappedModule.init(
H A Dtest_fsdp_traversal.py12 NestedWrappedModule,
36 nested_wrapped_module = NestedWrappedModule.init(
H A Dtest_fsdp_apply.py14 NestedWrappedModule,
70 nested_wrapped_module = NestedWrappedModule.init(
H A Dtest_fsdp_core.py30 NestedWrappedModule,
116 NestedWrappedModule,
137 NestedWrappedModule,
H A Dtest_fsdp_clip_grad_norm.py24 NestedWrappedModule,
277 NestedWrappedModule.init(
H A Dtest_fsdp_sharded_grad_scaler.py28 NestedWrappedModule,
193 model_cls = NestedWrappedModule # type: ignore[assignment]
H A Dtest_fsdp_comm.py23 NestedWrappedModule,
63 model = NestedWrappedModule.init(
/aosp_15_r20/external/pytorch/torch/testing/_internal/
H A Dcommon_fsdp.py384 class NestedWrappedModule(FSDPTestModel): class
457 return NestedWrappedModule(
465 fsdp_model = NestedWrappedModule(
478 class AlwaysWrapNestedWrappedModule(NestedWrappedModule):
512 class NonUniformReqGradNWM(NestedWrappedModule):
521 super(NestedWrappedModule, self).__init__()
690 NestedWrappedModule,
710 class MixtureOfExperts(NestedWrappedModule):