frozenleaves/pytorch
1
0
Fork 0
mirror of https://github.com/pytorch/pytorch.git synced 2025-10-20 12:54:11 +08:00
Code Issues Packages Projects Releases Wiki Activity

[pipelining] Update tutorials and documentation (#143045)

Browse Source 
Pull Request resolved: https://github.com/pytorch/pytorch/pull/143045
Approved by: https://github.com/wconstab, https://github.com/kwen2501
This commit is contained in:
Howard Huang
2024-12-11 13:15:44 -08:00
committed by PyTorch MergeBot
parent ee5bceaee6
commit b0c3d39e0d
2 changed files with 8 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files

9
docs/source/distributed.pipelining.rst
View File

@ -199,15 +199,8 @@ the model. For example:
stage_index,
num_stages,
device,
input_args=example_input_microbatch,
)
The ``PipelineStage`` requires an example argument ``input_args`` representing
the runtime input to the stage, which would be one microbatch worth of input
data. This argument is passed through the forward method of the stage module to
determine the input and output shapes required for communication.
When composing with other Data or Model parallelism techniques, ``output_args``
may also be required, if the output shape/dtype of the model chunk will be
affected.
@ -421,7 +414,7 @@ are subclasses of ``PipelineScheduleMulti``.
Logging
*******
You can turn on additional logging using the `TORCH_LOGS` environment variable from [`torch._logging`](https://pytorch.org/docs/main/logging.html#module-torch._logging):
You can turn on additional logging using the `TORCH_LOGS` environment variable from `torch._logging <https://pytorch.org/docs/main/logging.html#module-torch._logging>`_:
* `TORCH_LOGS=+pp` will display `logging.DEBUG` messages and all levels above it.
* `TORCH_LOGS=pp` will display `logging.INFO` messages and above.

7
torch/distributed/pipelining/_IR.py
View File

@ -1143,6 +1143,13 @@ class Pipe(torch.nn.Module):
class SplitPoint(Enum):
"""
Enum representing the points at which a split can occur in the execution of a submodule.
Attributes:
BEGINNING: Represents adding a split point *before* the execution of a certain submodule in the `forward` function.
END: Represents adding a split point *after* the execution of a certain submodule in the `forward` function.
"""
BEGINNING = 1
END = 2