784 lines
29 KiB
Python
784 lines
29 KiB
Python
import logging
|
|
import operator
|
|
from dataclasses import dataclass
|
|
from enum import auto, Enum
|
|
from functools import partial
|
|
from typing import Any, Callable, cast, Dict, List, Optional, Sequence, Tuple, Union
|
|
|
|
import torch
|
|
import torch.distributed._spmd.experimental_ops
|
|
import torch.fx as fx
|
|
|
|
from torch.distributed._spmd.comm_tensor import _get_tracer
|
|
from torch.distributed._spmd.graph_utils import OP
|
|
from torch.distributed._spmd.log_utils import get_logger
|
|
|
|
from torch.distributed._tensor import DeviceMesh, DTensor
|
|
from torch.distributed._tensor.op_schema import OpSchema
|
|
from torch.distributed._tensor.placement_types import (
|
|
_Partial,
|
|
DTensorSpec,
|
|
Placement,
|
|
Replicate,
|
|
Shard,
|
|
TensorMeta,
|
|
)
|
|
from torch.distributed._tensor.redistribute import redistribute_local_tensor
|
|
from torch.fx.experimental.proxy_tensor import make_fx, proxy_slot
|
|
from torch.utils import _pytree as pytree
|
|
from torch.utils._pytree import tree_flatten, tree_map, tree_map_only, tree_unflatten
|
|
|
|
|
|
logger: Optional[logging.Logger] = None
|
|
|
|
aten = torch.ops.aten
|
|
|
|
|
|
class TrainingPhase(Enum):
|
|
FORWARD = auto()
|
|
BACKWARD = auto()
|
|
|
|
|
|
@dataclass
|
|
class Schema:
|
|
mesh: DeviceMesh
|
|
placements: List[Placement]
|
|
|
|
|
|
@dataclass
|
|
class DSymInt:
|
|
"""DSymInt represents a value retrieved by a SymInt op from a DTensor.
|
|
|
|
DSymInt helps View and Factory ops to determine the placement and shape of the
|
|
output tensor, as those operators either do not have an input DTensor or
|
|
the input DTensor is insufficient to determine the output tensor's placement.
|
|
"""
|
|
|
|
global_value: int # value that the SymInt evaluates to
|
|
local_value: int # vaue that this SymInt evaluates to on the local shard
|
|
mesh: DeviceMesh # device mesh of the DTensor where this SymInt is retrieved from
|
|
|
|
def is_shard(self) -> bool:
|
|
return self.local_value != self.global_value
|
|
|
|
@classmethod
|
|
def from_node(cls, node: fx.Node, dtensor: DTensor) -> "DSymInt":
|
|
dim: int = 0
|
|
if node.target == aten.sym_size:
|
|
dim = cast(int, node.args[1])
|
|
return cls(
|
|
global_value=dtensor.size(dim),
|
|
local_value=dtensor.to_local().size(dim),
|
|
mesh=dtensor.device_mesh,
|
|
)
|
|
elif node.target == aten.sym_numel:
|
|
return cls(
|
|
global_value=dtensor.numel(),
|
|
local_value=dtensor.to_local().numel(),
|
|
mesh=dtensor.device_mesh,
|
|
)
|
|
elif node.target == aten.sym_stride:
|
|
dim = cast(int, node.args[1])
|
|
return cls(
|
|
global_value=dtensor.stride(dim),
|
|
local_value=dtensor.to_local().stride(dim),
|
|
mesh=dtensor.device_mesh,
|
|
)
|
|
else:
|
|
raise NotImplementedError(f"DSymInt does not support {node.target}")
|
|
|
|
|
|
def _is_partial_dtensor(obj: Any) -> bool:
|
|
"""Check if object is 1) DTensor and 2) with any placement of _Partial."""
|
|
if not isinstance(obj, DTensor):
|
|
return False
|
|
|
|
is_partial = False
|
|
for placement in obj.placements:
|
|
if isinstance(placement, _Partial):
|
|
is_partial = True
|
|
break
|
|
|
|
return is_partial
|
|
|
|
|
|
def _dispatch_with_local_tensors(
|
|
op: torch._ops.OpOverload,
|
|
local_args: Tuple[Any, ...],
|
|
kwargs: Optional[Dict[str, Any]] = None,
|
|
specs: Optional[
|
|
Dict[
|
|
torch.Tensor,
|
|
Tuple[torch.Size, DeviceMesh, Sequence[Placement], Sequence[Placement]],
|
|
]
|
|
] = None,
|
|
) -> Any:
|
|
if kwargs is None:
|
|
kwargs = {}
|
|
if specs is None:
|
|
specs = {}
|
|
|
|
def redistribute(arg: Any) -> Any:
|
|
tensor_shape, mesh, current_placement, target_placement = specs[arg]
|
|
tensor_meta = TensorMeta(
|
|
tensor_shape,
|
|
stride=arg.stride(),
|
|
dtype=arg.dtype,
|
|
)
|
|
current_spec = DTensorSpec(
|
|
mesh, tuple(current_placement), tensor_meta=tensor_meta
|
|
)
|
|
target_spec = DTensorSpec(
|
|
mesh, tuple(target_placement), tensor_meta=tensor_meta
|
|
)
|
|
|
|
return (
|
|
redistribute_local_tensor(arg, current_spec, target_spec) # type: ignore[index]
|
|
if isinstance(arg, torch.Tensor) and arg in specs # type: ignore[operator]
|
|
else arg
|
|
)
|
|
|
|
# TODO: this is broken because it won't redistributed potential tensors on the kwargs
|
|
return op(*tree_map(redistribute, local_args), **kwargs)
|
|
|
|
|
|
# Figure out how to specify a type spec for the return specs value
|
|
# without the entire structure.
|
|
# pyre-fixme
|
|
def _update_specs_for_redistribute(args, target_schema, redistribute):
|
|
# Code adapted from pack_args_kwargs_with_local_tensor
|
|
flatten_args, args_tree_spec = tree_flatten(args)
|
|
flatten_args_schema = pytree.tree_leaves(target_schema.args_schema)
|
|
|
|
specs: Dict[
|
|
torch.Tensor,
|
|
Tuple[
|
|
torch.Size,
|
|
DeviceMesh,
|
|
Sequence[Placement],
|
|
Sequence[Placement],
|
|
],
|
|
] = {}
|
|
for i, arg in enumerate(flatten_args):
|
|
if isinstance(arg, DTensor):
|
|
if redistribute:
|
|
specs[arg._local_tensor] = (
|
|
arg.size(),
|
|
flatten_args_schema[i].mesh,
|
|
arg.placements,
|
|
flatten_args_schema[i].placements,
|
|
)
|
|
flatten_args_schema[i] = arg._local_tensor
|
|
|
|
unflattened_args = tree_unflatten(flatten_args_schema, args_tree_spec)
|
|
return specs, unflattened_args
|
|
|
|
|
|
# When no tensor redistribution is required, we only need to update non-tensor args
|
|
# of the node according to op_schema and avoid building a GraphModule just for the
|
|
# node.
|
|
def _update_node_from_op_schema(node: torch.fx.Node, op_schema: OpSchema) -> None:
|
|
flat_args, args_tree_spec = tree_flatten(node.args)
|
|
flat_args_schema = pytree.tree_leaves(op_schema.args_schema)
|
|
|
|
def is_sym_int_or_int(arg: Union[int, torch.fx.Node]) -> bool:
|
|
if isinstance(arg, torch.fx.Node):
|
|
return arg.target in [
|
|
aten.sym_size,
|
|
aten.sym_numel,
|
|
aten.sym_stride,
|
|
]
|
|
return isinstance(arg, int)
|
|
|
|
assert len(flat_args) == len(flat_args_schema)
|
|
for i, (arg, arg_schema) in enumerate(zip(flat_args, flat_args_schema)):
|
|
if is_sym_int_or_int(arg) and isinstance(arg_schema, int):
|
|
flat_args[i] = arg_schema
|
|
|
|
args = tree_unflatten(flat_args, args_tree_spec)
|
|
for idx, arg in enumerate(args):
|
|
node.update_arg(idx, arg)
|
|
return None
|
|
|
|
|
|
def _remap_arg(node_to_obj: Dict[fx.Node, Any], arg: Any) -> Any:
|
|
if isinstance(arg, torch.fx.Node):
|
|
obj = node_to_obj[arg]
|
|
if _get_tracer():
|
|
# This is a shared arg, already has a tracer from previous
|
|
# tracing. Delete the tracer.
|
|
del cast(Dict[Any, Any], obj.__dict__)[proxy_slot]
|
|
return obj
|
|
else:
|
|
return arg
|
|
|
|
|
|
def unpack_sizes_and_dims(
|
|
sizes: List[Union[DSymInt, int]], mesh: DeviceMesh
|
|
) -> Tuple[List[int], List[Placement]]:
|
|
local_sizes: List[int] = [
|
|
s.local_value if isinstance(s, DSymInt) else s for s in sizes
|
|
]
|
|
placements: List[Placement] = [
|
|
Shard(i)
|
|
for i, a in enumerate(sizes)
|
|
if (isinstance(a, DSymInt) and a.is_shard())
|
|
] or [Replicate()]
|
|
|
|
assert len(placements) == mesh.ndim, (
|
|
f"The number of sharded dimensions ({len(placements)}) must "
|
|
f"match number of dimensions in device mesh ({mesh.ndim})."
|
|
)
|
|
|
|
return local_sizes, placements
|
|
|
|
|
|
def binop_sym_int_consumer_rule(node: fx.Node, args: Tuple[Any, ...]) -> DTensor:
|
|
assert len(args) == 2, f"Expect two args but got op {node.target} with args {args}"
|
|
assert isinstance(
|
|
args[0], DTensor
|
|
), f"Expect 1st argument to be DTensor but got {args[0]}"
|
|
assert isinstance(args[1], list), f"Expect 2nd argument as list but got {args[1]}"
|
|
|
|
# extract sharded dimensions in the size list, the output DTensor should
|
|
# follow these placements.
|
|
local_sizes, placements = unpack_sizes_and_dims(args[1], args[0].device_mesh)
|
|
|
|
# set node args to real int sizes.
|
|
node.args = (node.args[0], local_sizes)
|
|
op = cast(torch._ops.OpOverload, node.target)
|
|
return DTensor.from_local(
|
|
local_tensor=op(args[0]._local_tensor, local_sizes),
|
|
device_mesh=args[0].device_mesh,
|
|
placements=placements,
|
|
run_check=False,
|
|
)
|
|
|
|
|
|
def slice_backwad_sym_int_consumer_rule(
|
|
node: fx.Node, args: Tuple[Any, ...]
|
|
) -> DTensor:
|
|
grad_output, input_sizes, dim, start, end, step = args
|
|
|
|
local_sizes: List[int] = [
|
|
s.local_value if isinstance(s, DSymInt) else s for s in input_sizes
|
|
]
|
|
|
|
input_tensor = torch.zeros(
|
|
local_sizes, device=grad_output.device, dtype=grad_output.dtype
|
|
)
|
|
return DTensor.from_local(
|
|
local_tensor=torch.slice_scatter(
|
|
input_tensor, grad_output.to_local(), dim, start, end, step
|
|
),
|
|
device_mesh=grad_output.device_mesh,
|
|
placements=grad_output.placements,
|
|
run_check=False,
|
|
)
|
|
|
|
|
|
def factory_with_sizes_rule(
|
|
node: fx.Node,
|
|
args: Tuple[Any, ...],
|
|
kwargs: Dict[str, Any],
|
|
default_mesh: DeviceMesh,
|
|
) -> DTensor:
|
|
flat_args = pytree.arg_tree_leaves(*args)
|
|
assert not any(isinstance(a, DTensor) for a in flat_args), (
|
|
f"Not expect DTensor argument for factory op, but got {node.target} "
|
|
f"with arguments {args}."
|
|
)
|
|
assert isinstance(args[0], list), f"Expect 2nd argument as list but got {args[1]}"
|
|
|
|
local_sizes, placements = unpack_sizes_and_dims(args[0], default_mesh)
|
|
node.args = (local_sizes, *args[1:])
|
|
op = cast(torch._ops.OpOverload, node.target)
|
|
return DTensor.from_local(
|
|
local_tensor=op(*node.args, **kwargs),
|
|
device_mesh=default_mesh,
|
|
placements=placements,
|
|
run_check=False,
|
|
)
|
|
|
|
|
|
def factory_arange_rule(
|
|
node: fx.Node,
|
|
args: Tuple[Any, ...],
|
|
kwargs: Dict[str, Any],
|
|
default_mesh: DeviceMesh,
|
|
) -> DTensor:
|
|
node.args = tree_map(lambda a: a.local_value if isinstance(a, DSymInt) else a, args)
|
|
op = cast(torch._ops.OpOverload, node.target)
|
|
return DTensor.from_local(
|
|
local_tensor=op(*node.args, **kwargs),
|
|
device_mesh=default_mesh,
|
|
placements=[Replicate()],
|
|
run_check=False,
|
|
)
|
|
|
|
|
|
def default_factory_op_rule(
|
|
node: fx.Node,
|
|
args: Tuple[Any, ...],
|
|
kwargs: Dict[str, Any],
|
|
default_mesh: DeviceMesh,
|
|
) -> DTensor:
|
|
node.args, node.kwargs = args, kwargs
|
|
op = cast(torch._ops.OpOverload, node.target)
|
|
return DTensor.from_local(
|
|
local_tensor=op(*node.args, **node.kwargs),
|
|
device_mesh=default_mesh,
|
|
placements=[Replicate()],
|
|
run_check=False,
|
|
)
|
|
|
|
|
|
# Dispatch override for view and factory ops that consume SymInt arguments,
|
|
# where the output spec should follow dimension placement where the SymInt comes
|
|
# from.
|
|
VIEW_SYM_INT_CONSUMERS: Dict[torch._ops.OpOverload, Callable] = {
|
|
aten._unsafe_view.default: binop_sym_int_consumer_rule,
|
|
aten.expand.default: binop_sym_int_consumer_rule,
|
|
aten.slice_backward.default: slice_backwad_sym_int_consumer_rule,
|
|
aten.view.default: binop_sym_int_consumer_rule,
|
|
}
|
|
|
|
FACTORY_SYM_INT_CONSUMERS: Dict[torch._ops.OpOverload, Callable] = {
|
|
aten.full.default: factory_with_sizes_rule,
|
|
aten.arange.default: factory_arange_rule,
|
|
aten.arange.start: factory_arange_rule,
|
|
}
|
|
|
|
|
|
# Dispatch override for factory ops, as DTensor cannot propogate sharding spec
|
|
# without DTensor inputs.
|
|
FACTORY_OPS: Dict[torch._ops.OpOverload, Callable] = {
|
|
aten.scalar_tensor.default: default_factory_op_rule,
|
|
aten.arange.start: default_factory_op_rule,
|
|
aten.zeros.default: default_factory_op_rule,
|
|
}
|
|
|
|
|
|
def _get_dtensor_dispatch_graph(
|
|
node: fx.Node,
|
|
node_to_obj: Dict[fx.Node, Any],
|
|
*,
|
|
force_make_fx: bool = False,
|
|
default_mesh: Optional[DeviceMesh] = None,
|
|
) -> Optional[fx.GraphModule]:
|
|
with torch.no_grad():
|
|
# Args should be a list of objects post remapping.
|
|
args = tree_map(partial(_remap_arg, node_to_obj), node.args)
|
|
kwargs = tree_map(partial(_remap_arg, node_to_obj), node.kwargs)
|
|
|
|
op_overload = cast(torch._ops.OpOverload, node.target)
|
|
|
|
if any(
|
|
a.is_shard()
|
|
for a in pytree.arg_tree_leaves(*args)
|
|
if isinstance(a, DSymInt)
|
|
):
|
|
if op_overload in VIEW_SYM_INT_CONSUMERS:
|
|
assert len(kwargs) == 0, f"Expect empty kwargs, but got {kwargs}"
|
|
node_to_obj[node] = VIEW_SYM_INT_CONSUMERS[op_overload](node, args)
|
|
return None
|
|
elif op_overload in FACTORY_SYM_INT_CONSUMERS:
|
|
assert default_mesh is not None, "Requires default mesh for factory ops"
|
|
node_to_obj[node] = FACTORY_SYM_INT_CONSUMERS[op_overload](
|
|
node, args, kwargs, default_mesh
|
|
)
|
|
return None
|
|
else:
|
|
assert isinstance(logger, logging.Logger)
|
|
logger.warning(
|
|
"Assuming using local_value from SymInt for %s"
|
|
"is mathematically correct. Full args are %s.",
|
|
op_overload,
|
|
args,
|
|
)
|
|
|
|
if node.target == aten.view.default:
|
|
# HACK: this is a hack to get around with the fact that some
|
|
# view operations on a "global" tensor is invalid usage
|
|
# but somehow the view operation on the batch input might hit it
|
|
# so we convert the view op to reshape before calling DTensor
|
|
op_overload = aten.reshape.default
|
|
|
|
# DSymInt args are not sharded on any dimension, local value and global
|
|
# value should be the same
|
|
args = tree_map(lambda a: a.local_value if isinstance(a, DSymInt) else a, args)
|
|
kwargs = tree_map(
|
|
lambda a: a.local_value if isinstance(a, DSymInt) else a, kwargs
|
|
)
|
|
|
|
if op_overload in FACTORY_OPS:
|
|
# Don't pass factory ops to DTensor dispatch, as DTensor cannot
|
|
# propagate sharding spec without DTensor inputs.
|
|
node_to_obj[node] = FACTORY_OPS[op_overload](
|
|
node, args, kwargs, default_mesh
|
|
)
|
|
return None
|
|
|
|
dispatch = partial(
|
|
_dispatch_with_local_tensors,
|
|
op_overload,
|
|
kwargs=kwargs,
|
|
specs=args,
|
|
)
|
|
|
|
gm = make_fx(dispatch, _allow_non_fake_inputs=False)(args)
|
|
# FIXME(@wanchaol, @mrshenli): the above seems to accidentally captured
|
|
# DeviceMesh tensor ops when handling inplace operators? The ``_to_copy`` is
|
|
# not connected to graph output. So, using DCE to get rid of it, but this
|
|
# doesn't look correct.
|
|
#
|
|
# The following operators appear in the captured graph, where the dtype is
|
|
# torch.int64.
|
|
#
|
|
# get_attr _tensor_constant0 _tensor_constant0 ()
|
|
# call_function transpose aten.transpose.int (_tensor_constant0, -1, 0)
|
|
# call_function view aten.view.default (transpose, [-1, 2])
|
|
# call_function view_1 aten.view.default (view, [2])
|
|
# call_function _to_copy aten._to_copy.default (view_1,)
|
|
gm.graph.eliminate_dead_code()
|
|
|
|
return gm
|
|
|
|
|
|
def _build_dummy_add_graph(
|
|
dt: DTensor, node_to_obj: Dict[fx.Node, Any]
|
|
) -> Tuple[fx.GraphModule, Any]:
|
|
"""Create a graph for a dummy add function from a partial DTensor.
|
|
|
|
This dummy add is used for triggering all_reduce on a Partial DTensor
|
|
during the DTensor expansion of the traced graph.
|
|
Also returns the actual DTensor after resharding.
|
|
"""
|
|
|
|
def dummy_add(grad: torch.Tensor, zero: torch.Tensor) -> torch.Tensor:
|
|
return grad + zero
|
|
|
|
grad: torch.Tensor = dt._local_tensor
|
|
zero: torch.Tensor = torch.zeros_like(dt._local_tensor)
|
|
|
|
traced_add = make_fx(dummy_add)(grad, zero)
|
|
|
|
placeholders = [n for n in traced_add.graph.nodes if n.op == OP.PLACEHOLDER]
|
|
call_functions = [n for n in traced_add.graph.nodes if n.op == OP.CALL_FUNCTION]
|
|
assert len(placeholders) == 2
|
|
assert len(call_functions) == 1
|
|
node_to_obj[placeholders[0]] = dt
|
|
node_to_obj[placeholders[1]] = DTensor.from_local(
|
|
zero, dt.device_mesh, [Replicate()], run_check=False
|
|
)
|
|
|
|
traced_dispatch = _get_dtensor_dispatch_graph(
|
|
call_functions[0], node_to_obj, force_make_fx=True
|
|
)
|
|
assert traced_dispatch is not None
|
|
|
|
# TODO(anj): This depends on the call function node -> actual DTensor output
|
|
# mapping that we want to avoid for SPMD expansion
|
|
return traced_dispatch, node_to_obj[call_functions[0]]
|
|
|
|
|
|
def _convert_output(
|
|
gm: fx.GraphModule,
|
|
node: fx.Node,
|
|
node_to_obj: Dict[fx.Node, Any],
|
|
) -> fx.Node:
|
|
new_args = []
|
|
has_partial = False
|
|
for argument in node.args[0]: # type: ignore[union-attr]
|
|
if not isinstance(argument, fx.Node):
|
|
new_args.append(argument)
|
|
continue
|
|
|
|
obj = node_to_obj[argument]
|
|
|
|
if not _is_partial_dtensor(obj):
|
|
new_args.append(argument)
|
|
continue
|
|
|
|
has_partial = True
|
|
|
|
# we know it's a dtensor from is partial DT check...
|
|
dt = cast(DTensor, obj)
|
|
|
|
traced_dispatch, result_obj = _build_dummy_add_graph(dt, node_to_obj)
|
|
|
|
wait = [
|
|
n
|
|
for n in traced_dispatch.graph.nodes
|
|
if n.name == "wait_comm" or n.name == "wait_tensor"
|
|
]
|
|
add = [n for n in traced_dispatch.graph.nodes if n.name == "add"]
|
|
assert len(wait) == 1 and len(add) == 1
|
|
|
|
# remove add node and replace it with wait node
|
|
add[0].replace_all_uses_with(wait[0])
|
|
traced_dispatch.graph.eliminate_dead_code()
|
|
# also update the actual DTensor corresponding to the node
|
|
# TODO(anj): We require mapping of the final DTensor output to the wait
|
|
# comm node.
|
|
node_to_obj[wait[0]] = result_obj
|
|
|
|
value_remap: Dict[fx.Node, fx.Node] = {}
|
|
for dtn in traced_dispatch.graph.nodes:
|
|
if dtn.op == OP.PLACEHOLDER:
|
|
# do nothing, ignore placeholders, as it has
|
|
# already been prepared in value_remap
|
|
value_remap[dtn] = argument
|
|
elif dtn.op == OP.OUTPUT:
|
|
assert (
|
|
len(dtn.args) == 1 and len(dtn.args[0]) == 1
|
|
), f"Expecting single output, but got {dtn.args} {len(dtn.args)}"
|
|
new_args.append(value_remap[dtn.args[0][0]])
|
|
# the concrete DTensor value of output was added when creating the
|
|
# inner graph (in _build_dummy_add_graph). Just add it to the final
|
|
# output node so that we can report the final output specs correctly.
|
|
# TODO(anj): We are depending on the concrete DTensor output of the dummy add.
|
|
node_to_obj[value_remap[dtn.args[0][0]]] = node_to_obj[dtn.args[0][0]]
|
|
|
|
else:
|
|
if dtn.op == OP.GET_ATTR:
|
|
setattr(
|
|
gm,
|
|
dtn.target,
|
|
getattr(traced_dispatch, dtn.target),
|
|
)
|
|
with gm.graph.inserting_before(node):
|
|
value_remap[dtn] = gm.graph.node_copy(dtn, lambda n: value_remap[n])
|
|
if has_partial:
|
|
gm.graph.erase_node(node)
|
|
return gm.graph.output(new_args)
|
|
else:
|
|
return node
|
|
|
|
|
|
def _rebuild_graph(
|
|
gm: fx.GraphModule,
|
|
node_replacements: Dict[torch.fx.Node, torch.fx.GraphModule],
|
|
) -> None:
|
|
# replace nodes in local traced graph with DTensor's dispatch graph
|
|
for node in gm.graph.nodes:
|
|
if node not in node_replacements:
|
|
continue
|
|
|
|
traced_dispatch = node_replacements[node]
|
|
# Map DT's dispatch graph input placeholder nodes to the ones in
|
|
# local traced graph. It uses index-based accessing, which is
|
|
# brittle, just for testing purpose.
|
|
flatten_args = pytree.arg_tree_leaves(*node.args)
|
|
i, value_remap = 0, {}
|
|
for dtn in traced_dispatch.graph.nodes:
|
|
if dtn.op == OP.PLACEHOLDER:
|
|
value_remap[dtn] = flatten_args[i]
|
|
i += 1
|
|
|
|
# insert DT's dispatch graph to traced local graph.
|
|
with gm.graph.inserting_before(node):
|
|
for dtn in traced_dispatch.graph.nodes:
|
|
if dtn.op == OP.PLACEHOLDER:
|
|
# do nothing, ignore placeholders, as it has already
|
|
# been prepared in value_remap
|
|
pass
|
|
elif dtn.op == OP.OUTPUT:
|
|
assert (
|
|
len(dtn.args) == 1
|
|
), f"Expecting single output, but got {dtn.args} {len(dtn.args[0])}"
|
|
outputs = dtn.args[0]
|
|
# we currently support two very specific types of output
|
|
# 1. single output
|
|
# 2. multiple outputs resulting from getitem of all elements of tuple
|
|
if len(outputs) == 1:
|
|
# for single output, we replace the node with the single node
|
|
output = outputs[0]
|
|
else:
|
|
# for multiple outputs, we check that these outputs correspond
|
|
# to all elements of a tuple. In that case, we replace
|
|
# uses of the output directly with the original tuple
|
|
source = None
|
|
for i, out in enumerate(outputs):
|
|
# we allow None outputs for certain items in the tuple
|
|
if out is None:
|
|
continue
|
|
assert out.op == "call_function"
|
|
assert out.target.__module__ == "_operator"
|
|
assert out.target.__name__ == "getitem"
|
|
assert source is None or source == out.args[0]
|
|
source = out.args[0]
|
|
assert out.args[1] == i
|
|
assert source is not None
|
|
output = source
|
|
|
|
new_node = value_remap[output]
|
|
node.replace_all_uses_with(new_node)
|
|
else:
|
|
value_remap[dtn] = gm.graph.node_copy(dtn, lambda n: value_remap[n])
|
|
if all(
|
|
isinstance(n.target, torch._ops.OpOverload)
|
|
and n.target._schema.name.startswith(
|
|
("aten::_foreach", "aten::_fused_adam")
|
|
)
|
|
for n in [dtn, node]
|
|
):
|
|
# FIXME(@mrshenli): This is a temporary solution enable
|
|
# foreach ops. The problem is that foreach ops returns
|
|
# List[Tensor], but make_fx will flatten that before
|
|
# passing those tensors to output node, which will
|
|
# introduce additional getitem nodes. These redundant
|
|
# getitem nodes breaks graph correctness as we cannot do
|
|
# getitem(getitem(foreach_out, 0), 0). This temporary
|
|
# solution skips getitem nodes in DTensor expanded
|
|
# subgraphs.
|
|
node.replace_all_uses_with(value_remap[dtn])
|
|
break
|
|
# explicitly erase node instead of relying on DCE, as DCE does not
|
|
# remove inplace copy_ correctly.
|
|
gm.graph.erase_node(node)
|
|
|
|
gm.graph.eliminate_dead_code()
|
|
gm.recompile()
|
|
|
|
|
|
def _get_last_consumer_to_nodes(
|
|
graph: fx.Graph,
|
|
) -> Dict[fx.Node, List[fx.Node]]:
|
|
# Run through reverse nodes and record the first instance of a use
|
|
# of a given node. This represents the *last* use of the node in the
|
|
# execution order of the program, which we will use to free unused
|
|
# values
|
|
node_to_last_consumer: Dict[fx.Node, fx.Node] = {}
|
|
last_consumer_to_nodes: Dict[fx.Node, List[fx.Node]] = {}
|
|
|
|
def _register_final_consumer(arg_node: fx.Node, consumer: fx.Node) -> None:
|
|
if arg_node not in node_to_last_consumer:
|
|
node_to_last_consumer[arg_node] = consumer
|
|
last_consumer_to_nodes.setdefault(consumer, []).append(arg_node)
|
|
|
|
for node in reversed(graph.nodes):
|
|
fx.node.map_arg(
|
|
node.args, lambda arg_node: _register_final_consumer(arg_node, node)
|
|
)
|
|
fx.node.map_arg(
|
|
node.kwargs,
|
|
lambda kwarg_node: _register_final_consumer(kwarg_node, node),
|
|
)
|
|
|
|
return last_consumer_to_nodes
|
|
|
|
|
|
def _convert_to_distributed(
|
|
gm: fx.GraphModule,
|
|
inps: List[torch.Tensor],
|
|
schemas: List[Schema],
|
|
default_mesh: Optional[DeviceMesh] = None,
|
|
_allow_partial: bool = False,
|
|
) -> Tuple[fx.GraphModule, Dict[str, Schema]]:
|
|
"""Transform a graph module to a distributed graph module.
|
|
|
|
Returns:
|
|
- transformed graph module
|
|
- map from output name to DTensorSpec
|
|
|
|
"""
|
|
global logger
|
|
logger = get_logger("spmd_exp")
|
|
operators = {getattr(operator, name) for name in operator.__all__}
|
|
node_to_obj: Dict[fx.Node, Any] = {}
|
|
# map local op node in traced_f to its corresponding subgraph of
|
|
# DTensor ops.
|
|
node_replacements: Dict[torch.fx.Node, torch.fx.GraphModule] = {}
|
|
|
|
last_consumer_to_nodes = _get_last_consumer_to_nodes(gm.graph)
|
|
|
|
output_schemas: Dict[str, Schema] = {}
|
|
for i, node in enumerate(gm.graph.nodes):
|
|
assert logger is not None
|
|
logger.info("node%s: op=%s target=%s", i, node.op, node.target)
|
|
if node.op == OP.PLACEHOLDER:
|
|
assert i < len(
|
|
inps
|
|
), f"got more placeholder nodes ({i + 1}) than inputs ({len(inps)})"
|
|
|
|
# our example inputs are local shards. Create DTensors from them.
|
|
node_to_obj[node] = DTensor.from_local(
|
|
inps[i].clone(), # use clone to avoid modifications from inplace ops
|
|
schemas[i].mesh,
|
|
schemas[i].placements,
|
|
# prevent running this collective in backwards pass
|
|
run_check=False,
|
|
)
|
|
elif isinstance(node.target, torch._ops.OpOverloadPacket):
|
|
dtensor = cast(DTensor, node_to_obj[node.args[0]])
|
|
node_to_obj[node] = DSymInt.from_node(node, dtensor)
|
|
elif isinstance(node.target, torch._ops.OpOverload):
|
|
replacement = _get_dtensor_dispatch_graph(
|
|
node, node_to_obj, default_mesh=default_mesh
|
|
)
|
|
if replacement is not None:
|
|
node_replacements[node] = replacement
|
|
elif node.op == OP.OUTPUT:
|
|
if not _allow_partial:
|
|
# Returns an expanded dummy add node that ensures
|
|
# that the partial output tensor has been converted
|
|
# to a replicated tensor.
|
|
node = _convert_output(gm, node, node_to_obj)
|
|
|
|
# Save output sharding for the inputs to backward pass.
|
|
# TODO(anj): Pipe the output schema for the BW pass
|
|
# instead of requiring the full output DTensor to be
|
|
# materialized.
|
|
for inp_arg in node.args[0]:
|
|
if isinstance(inp_arg, fx.Node):
|
|
obj = node_to_obj[inp_arg]
|
|
if isinstance(obj, DTensor):
|
|
output_schemas[inp_arg.name] = Schema(
|
|
obj.device_mesh, obj.placements # type: ignore[arg-type]
|
|
)
|
|
elif node.op == OP.CALL_FUNCTION:
|
|
args = tree_map(partial(_remap_arg, node_to_obj), node.args)
|
|
kwargs = tree_map(partial(_remap_arg, node_to_obj), node.kwargs)
|
|
|
|
dsymints = list(
|
|
filter(lambda a: isinstance(a, DSymInt), args + tuple(kwargs.values()))
|
|
)
|
|
|
|
if node.target in operators and len(dsymints) > 0:
|
|
assert all(
|
|
dsymints[0].mesh == d.mesh for d in dsymints
|
|
), "all DSymInts must have the same mesh. "
|
|
|
|
local_args = tree_map_only(DSymInt, lambda a: a.local_value, args)
|
|
local_kwargs = tree_map_only(DSymInt, lambda a: a.local_value, kwargs)
|
|
|
|
global_args = tree_map_only(DSymInt, lambda a: a.global_value, args)
|
|
global_kwargs = tree_map_only(DSymInt, lambda a: a.global_value, kwargs)
|
|
|
|
node.args = local_args
|
|
node.kwargs = local_kwargs
|
|
|
|
node_to_obj[node] = DSymInt(
|
|
local_value=node.target(*local_args, **local_kwargs),
|
|
global_value=node.target(*global_args, **global_kwargs),
|
|
mesh=dsymints[0].mesh,
|
|
)
|
|
else:
|
|
assert len(dsymints) == 0, (
|
|
"SPMD expansion does not support SymInt in non-operator "
|
|
f"nodes, got {node.target}."
|
|
)
|
|
node_to_obj[node] = node.target(*args, **kwargs)
|
|
else:
|
|
raise ValueError(f"Unrecognized node.op type {node.op}")
|
|
|
|
if node in last_consumer_to_nodes:
|
|
# Save memory by deleting objs that wont be used anymore.
|
|
for arg_node in last_consumer_to_nodes[node]:
|
|
del node_to_obj[arg_node]
|
|
|
|
_rebuild_graph(gm, node_replacements)
|
|
|
|
return gm, output_schemas
|