Searched refs:use_local_output (Results 1 – 5 of 5) sorted by relevance
| /external/pytorch/torch/distributed/tensor/parallel/ |
| D | style.py | 77 use_local_output: bool = True
86 self.use_local_output = use_local_output
121 def _prepare_output_fn(output_layouts, use_local_output, mod, outputs, device_mesh): argument
126 return outputs.to_local() if use_local_output else outputs
141 partial(self._prepare_output_fn, self.output_layouts, self.use_local_output),
183 use_local_output: bool = True
188 self.use_local_output = use_local_output
221 def _prepare_output_fn(output_layouts, use_local_output, mod, outputs, device_mesh): argument
228 return outputs.to_local() if use_local_output else outputs
247 partial(self._prepare_output_fn, self.output_layouts, self.use_local_output),
[all …]
|
| /external/pytorch/torch/testing/_internal/distributed/_tensor/ |
| D | common_dtensor.py | 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)
|
| /external/pytorch/docs/source/ |
| D | distributed.tensor.parallel.rst | 55 and use ``use_local_output=False`` to return DTensor after each ``ParallelStyle``, where
|
| /external/pytorch/torch/testing/_internal/ |
| D | common_fsdp.py | 902 "0.in_proj": ColwiseParallel(use_local_output=False),
903 "0.out_proj": RowwiseParallel(use_local_output=False),
904 "1.in_proj": ColwiseParallel(use_local_output=False),
905 "1.out_proj": RowwiseParallel(use_local_output=False),
906 "2.in_proj": ColwiseParallel(use_local_output=False),
|
| /external/pytorch/test/distributed/tensor/parallel/ |
| D | test_tp_style.py | 278 test_mod.linear, mesh, ColwiseParallel(use_local_output=False)
306 test_kwonly_mod.linear, mesh, ColwiseParallel(use_local_output=False)
|