Update on "[dtensor] add grad placements kwarg to to_local API"

wanchaol · wanchaol · commit 369be44f6d01 · 2023-10-05T11:53:53.000-07:00
When we convert to local tensor, dtensor can't track autograd or
gradient layout of the local tensor anymore, if user do sth not expected, there
needs to be a way for user to hint about the gradient layout of the
local tensor

[ghstack-poisoned]
diff --git a/torch/distributed/_tensor/api.py b/torch/distributed/_tensor/api.py
@@ -329,10 +329,10 @@ def to_local(
             grad_placements (List[:class:`Placement`], optional): the placements describes
                 the future layout of any gradient layout of the Tensor returned from this
                 function.
-                `to_local` converts DTensor to local tensor and after that the returned local
-                tensor might not be used as the original DTensor layout later, this argument
-                is the hint that user can give to autograd in case the gradient layout of
-                the returned tensor does not match the original DTensor layout.
+                `to_local` converts DTensor to local tensor and the returned local tensor
+                might not be used as the original DTensor layout later in the code. This
+                argument is the hint that user can give to autograd in case the gradient
+                layout of the returned tensor does not match the original DTensor layout.
                 If not specified, we will assume the gradient layout remains the same
                 as the original DTensor and use that for gradient computation.
 
