| /aosp_15_r20/external/pytorch/test/distributed/tensor/parallel/ |
| H A D | test_tp_style.py | 16 from torch.distributed.tensor.parallel import parallelize_module
51 colwise_mod = parallelize_module(deepcopy(model), mesh, default_col_parallel)
65 colwise_mod = parallelize_module(deepcopy(model), mesh, sharded_col_parallel)
92 colwise_mod = parallelize_module(deepcopy(model), mesh, default_col_parallel)
115 rowwise_mod = parallelize_module(deepcopy(model), mesh, default_row_parallel)
129 rowwise_mod = parallelize_module(deepcopy(model), mesh, sharded_row_parallel)
155 rowwise_mod = parallelize_module(
174 rowwise_mod = parallelize_module(deepcopy(model), mesh, sharded_row_parallel)
204 allgather_mod = parallelize_module(model, mesh, prepare_inp_style)
233 parallelize_module(test_mod.linear, mesh, ColwiseParallel())
[all …]
|
| H A D | test_parallelize_api.py | 7 from torch.distributed.tensor.parallel.api import parallelize_module
106 model_tp = parallelize_module(
126 model_tp = parallelize_module(
148 model_tp = parallelize_module(model_tp, device_mesh, rowwise)
166 model_tp = parallelize_module(model_tp, device_mesh, colwise)
174 parallelize_module(
189 parallelize_module(
210 model_tp = parallelize_module(
226 model_tp = parallelize_module(
242 model_tp = parallelize_module(
[all …]
|
| H A D | test_tp_examples.py | 26 parallelize_module,
106 model_tp = parallelize_module(model_tp, device_mesh, parallelize_plan)
176 model_tp = parallelize_module(model_tp, device_mesh, parallelize_plan)
469 parallelize_module(model, device_mesh, parallelize_plan)
|
| H A D | test_tp_random_state.py | 6 from torch.distributed.tensor.parallel.api import parallelize_module
56 model_tp = parallelize_module(
|
| H A D | test_micro_pipeline_tp.py | 27 parallelize_module,
397 model = parallelize_module(model, device_mesh, parallelize_plan)
|
| /aosp_15_r20/external/pytorch/test/distributed/_tensor/debug/ |
| H A D | test_comm_mode_features.py | 12 parallelize_module,
99 model = parallelize_module(model, device_mesh, parallelize_plan)
136 model = parallelize_module(model, device_mesh, parallelize_plan)
153 model2 = parallelize_module(model2, device_mesh, parallelize_plan)
190 model = parallelize_module(model, device_mesh, parallelize_plan)
243 model = parallelize_module(model, device_mesh, parallelize_plan)
|
| /aosp_15_r20/external/pytorch/test/distributed/checkpoint/ |
| H A D | test_tp_checkpoint.py | 14 parallelize_module,
56 model = parallelize_module(model, tp_mesh, parallelize_plan)
104 model = parallelize_module(model, tp_mesh, parallelize_plan=parallelize_plan)
114 model2 = parallelize_module(
|
| H A D | test_fsdp_tp_checkpoint_conversion.py | 11 parallelize_module,
57 tp_model = parallelize_module(model, device_mesh, parallelize_plan)
|
| /aosp_15_r20/external/pytorch/torch/distributed/tensor/parallel/ |
| H A D | api.py | 22 def parallelize_module( # type: ignore[return] function
107 parallelize_module(
112 parallelize_module(submodule, device_mesh, parallelize_style)
|
| H A D | __init__.py | 2 from torch.distributed.tensor.parallel.api import parallelize_module
|
| /aosp_15_r20/external/pytorch/test/distributed/_composable/test_composability/ |
| H A D | test_2d_composability.py | 34 parallelize_module,
362 model = parallelize_module(dummy_model, tp_mesh, parallelize_plan)
413 twod_model = parallelize_module(twod_model, mesh_2d["tp"], parallelize_plan)
501 model_2d = parallelize_module(
538 model_2d = parallelize_module(SimpleModel().cuda(), tp_mesh, parallelize_plan)
619 model_2d = parallelize_module(
657 model_2d = parallelize_module(simple_model().cuda(), tp_mesh, parallelize_plan)
705 model_2d = parallelize_module(simple_model().cuda(), tp_mesh, parallelize_plan)
766 model_2d = parallelize_module(
|
| /aosp_15_r20/external/pytorch/test/distributed/checkpoint/fsdp/ |
| H A D | test_fsdp_dsd.py | 21 parallelize_module,
269 fsdp2_tp_model = parallelize_module(
353 tp_model = parallelize_module(
403 fsdp2_tp_model = parallelize_module(
510 fsdp2_tp_model = parallelize_module(
564 tp_model = parallelize_module(
|
| /aosp_15_r20/external/pytorch/test/distributed/_tensor/ |
| H A D | test_dtensor_compile.py | 32 parallelize_module,
702 parallelize_module(
775 model = parallelize_module(
814 tp_model = parallelize_module(model, twod_mesh["tp"], parallelize_plan)
822 tp_model2 = parallelize_module(
862 tp_model = parallelize_module(model, mesh_2d["tp"], parallelize_plan)
871 tp_model2 = parallelize_module(model_copy, mesh_2d["tp"], parallelize_plan)
|
| H A D | test_attention.py | 18 from torch.distributed.tensor.parallel import parallelize_module
222 encoder_layer = parallelize_module(
277 model = parallelize_module(
|
| H A D | test_dtensor.py | 26 parallelize_module,
120 model_tp = parallelize_module(model, device_mesh, parallelize_plan)
125 model_regular_tp = parallelize_module(
|
| /aosp_15_r20/external/pytorch/torch/testing/_internal/distributed/_tensor/ |
| H A D | common_dtensor.py | 20 parallelize_module,
226 module_tp = parallelize_module(module, device_mesh, root_plan)
258 parallelize_module(layer, device_mesh, layer_parallelize_plan)
271 parallelize_module(module_tp.output, device_mesh, output_parallelize_plan)
|
| /aosp_15_r20/external/pytorch/test/distributed/_composable/fsdp/ |
| H A D | test_fully_shard_grad_scaler.py | 11 parallelize_module,
52 model = parallelize_module(
|
| H A D | test_fully_shard_state_dict.py | 16 parallelize_module,
124 model = parallelize_module(
|
| H A D | test_fully_shard_init.py | 33 parallelize_module,
81 parallelize_module(
113 parallelize_module(
405 parallelize_module(
636 parallelize_module(
|
| /aosp_15_r20/external/pytorch/docs/source/ |
| H A D | distributed.tensor.parallel.rst | 20 .. autofunction:: parallelize_module
39 the ``parallelize_plan`` when calling ``parallelize_module``:
|
| /aosp_15_r20/external/pytorch/test/distributed/_composable/ |
| H A D | test_replicate_with_compiler.py | 25 parallelize_module,
404 ref_model = parallelize_module(ref_model, tp_mesh, parallelize_plan)
406 compiled_replicate_model = parallelize_module(
|
| /aosp_15_r20/external/pytorch/test/distributed/fsdp/ |
| H A D | test_fsdp_tp_integration.py | 25 parallelize_module,
290 tp_fsdp_model = parallelize_module(
374 tp_model = parallelize_module(
|
| /aosp_15_r20/external/pytorch/torch/distributed/tensor/examples/ |
| H A D | comm_mode_features_example.py | 15 parallelize_module,
75 model = parallelize_module(model, device_mesh, parallelize_plan)
|
| /aosp_15_r20/external/pytorch/test/distributed/ |
| H A D | test_fake_pg.py | 14 parallelize_module,
191 my_module = parallelize_module(
|
| /aosp_15_r20/external/pytorch/test/distributed/checkpoint/e2e/ |
| H A D | test_e2e_save_and_load.py | 32 parallelize_module,
178 model = parallelize_module(dummy_model, tp_mesh, parallelize_plan)
|