| /aosp_15_r20/external/pytorch/test/distributed/tensor/parallel/ |
| H A D | test_parallelize_api.py | 9 ColwiseParallel,
110 "net1": ColwiseParallel(output_layouts=Replicate()),
111 "net2": ColwiseParallel(output_layouts=Replicate()),
130 "dummy_encoder.net1": ColwiseParallel(output_layouts=Replicate()),
131 "dummy_encoder.net2": ColwiseParallel(output_layouts=Replicate()),
158 colwise = ColwiseParallel(output_layouts=Replicate())
214 "net*": ColwiseParallel(output_layouts=Replicate()),
230 "net?": ColwiseParallel(output_layouts=Replicate()),
246 "net[1-2]": ColwiseParallel(output_layouts=Replicate()),
262 "layers.*.net[1]": ColwiseParallel(),
|
| H A D | test_tp_style.py | 18 ColwiseParallel,
50 default_col_parallel = ColwiseParallel()
64 sharded_col_parallel = ColwiseParallel(input_layouts=Shard(0))
91 default_col_parallel = ColwiseParallel()
233 parallelize_module(test_mod.linear, mesh, ColwiseParallel())
250 parallelize_module(test_mod.linear, mesh, ColwiseParallel())
278 test_mod.linear, mesh, ColwiseParallel(use_local_output=False)
306 test_kwonly_mod.linear, mesh, ColwiseParallel(use_local_output=False)
|
| H A D | test_tp_random_state.py | 7 from torch.distributed.tensor.parallel.style import ColwiseParallel
60 "net1": ColwiseParallel(output_layouts=Replicate()),
61 "net2": ColwiseParallel(output_layouts=Replicate()),
|
| H A D | test_tp_examples.py | 24 ColwiseParallel,
96 ColwiseParallel(input_layouts=Shard(0))
98 else ColwiseParallel()
173 "net1": ColwiseParallel(),
465 "embedding": ColwiseParallel(),
|
| H A D | test_micro_pipeline_tp.py | 26 ColwiseParallel,
394 "net1": ColwiseParallel(input_layouts=Shard(shard_dim)),
|
| /aosp_15_r20/external/pytorch/test/distributed/checkpoint/fsdp/ |
| H A D | test_fsdp_dsd.py | 20 ColwiseParallel,
273 "0.in_proj": ColwiseParallel(),
275 "1.in_proj": ColwiseParallel(),
277 "2.in_proj": ColwiseParallel(),
329 "0.in_proj": ColwiseParallel(),
331 "1.in_proj": ColwiseParallel(),
333 "2.in_proj": ColwiseParallel(),
469 "0.in_proj": ColwiseParallel(),
471 "1.in_proj": ColwiseParallel(),
473 "2.in_proj": ColwiseParallel(),
[all …]
|
| /aosp_15_r20/external/pytorch/test/distributed/_tensor/debug/ |
| H A D | test_comm_mode_features.py | 11 ColwiseParallel,
95 "net1": ColwiseParallel(),
132 "net1": ColwiseParallel(),
147 "MLPStacked.layers.0.net1": ColwiseParallel(),
149 "MLPStacked.layers.1.net1": ColwiseParallel(),
186 "net1": ColwiseParallel(),
239 "net1": ColwiseParallel(),
|
| /aosp_15_r20/external/pytorch/test/distributed/_tensor/experimental/ |
| H A D | test_tp_transform.py | 10 ColwiseParallel,
84 {"fc": ColwiseParallel},
105 "mlps.0.0": ColwiseParallel,
107 "mlps.1.0": ColwiseParallel,
143 "mlps.0.0": ColwiseParallel,
|
| /aosp_15_r20/external/pytorch/test/distributed/_composable/test_composability/ |
| H A D | test_2d_composability.py | 33 ColwiseParallel,
358 "net1": ColwiseParallel(),
360 "net3": ColwiseParallel(),
410 "net1": ColwiseParallel(),
498 "net1": ColwiseParallel(),
535 "net1": ColwiseParallel(),
615 "net1": ColwiseParallel(),
617 "net3": ColwiseParallel(),
654 "net1": ColwiseParallel(),
702 "net1": ColwiseParallel(),
[all …]
|
| /aosp_15_r20/external/pytorch/torch/testing/_internal/distributed/_tensor/ |
| H A D | common_dtensor.py | 19 ColwiseParallel,
238 … layer_parallelize_plan["attention.wq"] = ColwiseParallel(use_local_output=local_output_for_attn)
239 … layer_parallelize_plan["attention.wk"] = ColwiseParallel(use_local_output=local_output_for_attn)
240 … layer_parallelize_plan["attention.wv"] = ColwiseParallel(use_local_output=local_output_for_attn)
248 ColwiseParallel(input_layouts=Shard(1))
250 else ColwiseParallel()
264 ColwiseParallel(
269 else ColwiseParallel(output_layouts=Replicate())
|
| /aosp_15_r20/external/pytorch/test/distributed/_composable/fsdp/ |
| H A D | test_fully_shard_grad_scaler.py | 10 ColwiseParallel,
45 "0.in_proj": ColwiseParallel(),
47 "1.in_proj": ColwiseParallel(),
49 "2.in_proj": ColwiseParallel(),
|
| H A D | test_fully_shard_state_dict.py | 15 ColwiseParallel,
128 "0.in_proj": ColwiseParallel(),
130 "1.in_proj": ColwiseParallel(),
132 "2.in_proj": ColwiseParallel(),
|
| H A D | test_fully_shard_init.py | 32 ColwiseParallel,
84 {"in_proj": ColwiseParallel(), "out_proj": RowwiseParallel()},
116 {"in_proj": ColwiseParallel(), "out_proj": RowwiseParallel()},
408 {"in_proj": ColwiseParallel(), "out_proj": RowwiseParallel()},
639 {"in_proj": ColwiseParallel(), "out_proj": RowwiseParallel()},
|
| /aosp_15_r20/external/pytorch/test/distributed/checkpoint/ |
| H A D | test_tp_checkpoint.py | 13 ColwiseParallel,
53 "net1": ColwiseParallel(),
100 "net1": ColwiseParallel(),
102 "net3": ColwiseParallel(),
|
| H A D | test_fsdp_tp_checkpoint_conversion.py | 10 ColwiseParallel,
54 "net1": ColwiseParallel(),
|
| /aosp_15_r20/external/pytorch/test/distributed/_tensor/ |
| H A D | test_dtensor_compile.py | 31 ColwiseParallel,
696 "attn.wq": ColwiseParallel(),
697 "attn.wk": ColwiseParallel(),
698 "attn.wv": ColwiseParallel(),
739 ColwiseParallel(input_layouts=Shard(0))
741 else ColwiseParallel()
761 "mlp_0.net1": ColwiseParallel(),
809 "mlp_0.net1": ColwiseParallel(),
811 "mlp_1.net1": ColwiseParallel(),
857 "mlp_0.net1": ColwiseParallel(),
[all …]
|
| /aosp_15_r20/external/pytorch/torch/distributed/tensor/examples/ |
| H A D | comm_mode_features_example.py | 14 ColwiseParallel,
62 "net1": ColwiseParallel(),
139 "MLPStacked.layers.0.net1": ColwiseParallel(),
141 "MLPStacked.layers.1.net1": ColwiseParallel(),
|
| /aosp_15_r20/external/pytorch/test/distributed/ |
| H A D | test_fake_pg.py | 13 ColwiseParallel,
183 "net1": ColwiseParallel(input_layouts=Shard(0)),
187 "net1": ColwiseParallel(),
|
| /aosp_15_r20/external/pytorch/docs/source/ |
| H A D | distributed.tensor.parallel.rst | 24 .. autoclass:: torch.distributed.tensor.parallel.ColwiseParallel
58 For models like Transformer, we recommend users to use ``ColwiseParallel``
|
| /aosp_15_r20/external/pytorch/test/distributed/_composable/ |
| H A D | test_replicate_with_compiler.py | 24 ColwiseParallel,
399 "fc1": ColwiseParallel(),
401 "fc3": ColwiseParallel(),
|
| /aosp_15_r20/external/pytorch/test/distributed/fsdp/ |
| H A D | test_fsdp_tp_integration.py | 24 ColwiseParallel,
287 "net1": ColwiseParallel(input_layouts=Shard(0)),
378 "mlp.net1": ColwiseParallel(input_layouts=Shard(0)),
|
| /aosp_15_r20/external/pytorch/torch/distributed/tensor/experimental/ |
| H A D | _tp_transform.py | 17 from torch.distributed.tensor.parallel.style import ColwiseParallel, ParallelStyle
120 Shard(0) if parallel_style == ColwiseParallel else Shard(1)
124 Shard(0) if parallel_style == ColwiseParallel else Replicate()
|
| /aosp_15_r20/external/pytorch/torch/distributed/tensor/parallel/ |
| H A D | __init__.py | 5 ColwiseParallel,
|
| /aosp_15_r20/external/pytorch/test/distributed/checkpoint/e2e/ |
| H A D | test_e2e_save_and_load.py | 31 ColwiseParallel,
175 "net1": ColwiseParallel(),
|
| /aosp_15_r20/external/pytorch/torch/testing/_internal/ |
| H A D | common_fsdp.py | 50 ColwiseParallel,
902 "0.in_proj": ColwiseParallel(use_local_output=False),
904 "1.in_proj": ColwiseParallel(use_local_output=False),
906 "2.in_proj": ColwiseParallel(use_local_output=False),
|