Files
2026-07-13 13:18:33 +08:00

625 lines
21 KiB
Python

# Copyright (c) Microsoft Corporation.
# SPDX-License-Identifier: Apache-2.0
# DeepSpeed Team
import functools
import operator
from typing import List, Tuple, Dict, Optional
from collections import defaultdict
import torch
from torch.fx import Node, Graph, GraphModule
from torch.fx.node import map_aggregate, Argument, map_arg
import torch.nn.functional as F
try:
from torch._subclasses.fake_tensor import unset_fake_temporarily
except ImportError:
# Unsupported torch version
pass
import deepspeed.comm as dist
from deepspeed.accelerator import get_accelerator
from deepspeed.utils.torch import required_torch_version
from deepspeed.ops.op_builder.dc import DeepCompileBuilder
from deepspeed.compile import constants
from .custom_ops import sp_dp_registry
def is_deepcompile_supported() -> bool:
return required_torch_version(min_version=2.6) and get_accelerator().device_name() == "cuda"
dc_handle = None
if is_deepcompile_supported():
sym_size_ops = {
operator.ge,
operator.le,
operator.eq,
operator.ne,
operator.gt,
operator.lt,
torch.ops.aten.sym_size.int,
operator.getitem,
}
def get_deepcompile_handle():
global dc_handle
if dc_handle is None:
dc_handle = DeepCompileBuilder().load()
return dc_handle
def is_backend_inductor(backend):
return backend == "inductor"
backward_started = False
pre_backward_hooks = []
post_backward_hooks = []
def add_pre_backward_hook(hook):
pre_backward_hooks.append(hook)
def add_post_backward_hook(hook):
post_backward_hooks.append(hook)
def deepcompile_backward_prologue(is_gradient_accumulation_boundary):
for hook in pre_backward_hooks:
hook(is_gradient_accumulation_boundary)
dc = get_deepcompile_handle()
dc.start_backward(is_gradient_accumulation_boundary)
def deepcompile_backward_epilogue():
for hook in post_backward_hooks:
hook()
def log_rank0(msg: str, enable: bool = False):
if dist.get_rank() == 0 and enable:
print(msg)
@functools.lru_cache
def get_no_copy_ops():
# Need to compile custom ops
get_deepcompile_handle()
no_copy_ops = {torch.ops.dc.wait_allgather.default}
# All operations whose return value aliases any of their inputs are included
# in the returned list to ensure that the last user of a node is computed
# correctly.
#
# This can be overly conservative if not all input tensors are aliased in
# the output. While we can determine exactly which tensors are aliased, a
# finer-grained algorithm is required in get_last_uses() and get_real_uses()
# to utilize that information. This is left as future work when real needs
# arise.
warned = False
for op_name in torch.ops.aten:
packet = getattr(torch.ops.aten, op_name)
for overload_name in packet:
op = getattr(packet, overload_name)
try:
for return_info in op._schema.returns:
if isinstance(return_info.type, torch.TensorType) and return_info.alias_info is not None:
no_copy_ops.add(op)
break
except AttributeError:
# In case no schema is available, conservatively assume the op
# may reuse tensor storage and print a one-time warning on its
# potential performance impact.
if not warned:
log_rank0(
f"WARNING: Schema is missing for some torch.ops.aten ops (e.g. {op_name}.{overload_name})."
"We assume those ops may reuse tensor storage. This may impact performance of compiled models.",
enable=True,
)
warned = True
no_copy_ops.add(op)
return no_copy_ops
def get_input_nodes(graph: Graph) -> List[Node]:
return [n for n in graph.nodes if n.op == "placeholder"]
def get_param_nodes(graph: Graph, index_to_ds_ids: List[Tuple[int, int]]) -> List[Node]:
all_input_nodes = get_input_nodes(graph)
return [all_input_nodes[i] for i, _, _ in index_to_ds_ids]
def is_comm_op(node: Node) -> bool:
return "comm" in node.meta and node.meta["comm"]
def is_cast_op(node: Node) -> Tuple[bool, Optional[torch.dtype]]:
if node.op == "call_function":
if node.target == torch.ops.prims.convert_element_type.default:
return (True, node.args[1])
elif node.target == torch.ops.aten._to_copy.default and set(node.kwargs.keys()) == {"dtype"}:
return (True, node.kwargs["dtype"])
return (False, None)
def exclude_from_act_offload(node: Node) -> bool:
return node.target in sym_size_ops
def dtype_to_elem_size(dtype: torch.dtype) -> int:
if dtype == torch.float32:
elem_size = 4
elif dtype == torch.float64:
elem_size = 8
elif dtype == torch.float16:
elem_size = 2
else:
raise ValueError(f"Unsupported dtype: {dtype}")
return elem_size
def tensor_meta_size(tensor_meta) -> int:
numel = 1 if len(tensor_meta.shape) == 0 else functools.reduce(operator.mul, tensor_meta.shape)
dtype = tensor_meta.dtype
if dtype == torch.float32:
elem_size = 4
elif dtype == torch.float64 or dtype == torch.int64:
elem_size = 8
elif dtype == torch.float16 or dtype == torch.bfloat16:
elem_size = 2
elif dtype == torch.bool:
elem_size = 1
else:
raise ValueError(f"Unsupported dtype: {dtype}")
return numel * elem_size
class NodeValueOffloadHelper:
def __init__(self, device):
self.device = device
self.env_values: Dict[str, Argument] = {}
self.original_device: Dict[torch.Tensor, torch.device] = {}
def _to_cpu(self, v):
if torch.is_tensor(v):
with unset_fake_temporarily():
device = v.device
offloaded = v.to('cpu').detach()
self.original_device[offloaded] = device
return offloaded
return v
def _from_cpu(self, v):
if torch.is_tensor(v) and v in self.original_device:
return v.to(self.original_device[v])
return v
def save(self, name: str, v: Argument, offload) -> None:
self.env_values[name] = map_aggregate(v, lambda x: self._to_cpu(x) if offload else x)
def load(self, name: str) -> Argument:
return map_aggregate(self.env_values[name], lambda x: self._from_cpu(x))
def get_offloaded_value(self, name: str) -> Argument:
return self.env_values[name]
def has_value(self, name: str) -> bool:
return name in self.env_values
def clear(self) -> None:
self.env_values.clear()
self.original_device.clear()
def materialize_fake(v, device=None):
from torch._subclasses.fake_tensor import is_fake
def convert(t):
if is_fake(t):
with unset_fake_temporarily():
if t.is_floating_point():
return torch.randn(t.shape,
dtype=t.dtype,
device=t.device if device is None else device,
layout=t.layout,
requires_grad=t.requires_grad,
pin_memory=t.is_pinned())
else:
return torch.zeros(t.shape,
dtype=t.dtype,
device=t.device if device is None else device,
requires_grad=t.requires_grad)
return t
return map_aggregate(v, lambda x: convert(x))
def get_last_uses(graph: Graph):
position = {node: i for i, node in enumerate(graph.nodes)}
node_to_last_use: Dict[Node, Node] = {}
user_to_last_uses: Dict[Node, List[Node]] = {}
no_copy_ops = get_no_copy_ops()
def register_last_uses(n: Node, user: Node):
update = False
known_last_use = None
if user.target in no_copy_ops and n in node_to_last_use:
last_user = node_to_last_use[user]
last_use_position = position[last_user]
known_last_use = node_to_last_use[n]
known_last_use_position = position[known_last_use]
update = last_use_position > known_last_use_position
if n not in node_to_last_use or update:
if user.target in no_copy_ops:
user = node_to_last_use[user]
node_to_last_use[n] = user
user_to_last_uses.setdefault(user, []).append(n)
if known_last_use:
user_to_last_uses[known_last_use].remove(n)
for node in reversed(graph.nodes):
map_arg(node.args, lambda n: register_last_uses(n, node))
map_arg(node.kwargs, lambda n: register_last_uses(n, node))
return node_to_last_use, user_to_last_uses
def get_real_uses(graph: Graph):
node_to_uses: Dict[Node, List[Node]] = defaultdict(list)
no_copy_ops = get_no_copy_ops()
def register_last_uses(n: Node, user: Node):
if user.target == "output":
return
if user.target in no_copy_ops:
users = node_to_uses[user]
node_to_uses[n].extend(users)
else:
node_to_uses[n].append(user)
for node in reversed(graph.nodes):
map_arg(node.args, lambda n: register_last_uses(n, node))
map_arg(node.kwargs, lambda n: register_last_uses(n, node))
return node_to_uses
def count_inflight_values(graph: Graph, file_path: str):
position = {node: i for i, node in enumerate(graph.nodes)}
node_to_last_use, user_to_last_uses = get_last_uses(graph)
max_inflight_size = 0
inflight_values = set()
# Output csv.
csv_filename = file_path
csv_data = []
header = [
'Node', 'tensor_size', 'inflight_size', 'inflight_size_in_output', 'args', 'users', 'node_to_last_use',
'lifetime', 'user_to_last_uses', 'inflight_values'
]
csv_data.append(header)
from .fx import get_output_node
output_node = get_output_node(graph)
values_in_output = set([n for n in output_node.args[0] if isinstance(n, Node)])
for node in graph.nodes:
inflight_values.add(node)
if node in user_to_last_uses:
for to_delete in user_to_last_uses[node]:
inflight_values.remove(to_delete)
assert "tensor_size" in node.meta, f"Node {node} does not have tensor_size"
inflight_size = sum(n.meta["tensor_size"] for n in inflight_values)
inflight_size_in_output = sum(n.meta["tensor_size"] for n in inflight_values if n in values_in_output)
lifetime = position[node_to_last_use[node]] - position[node] if node in node_to_last_use else 0
row = [
node.name, node.meta["tensor_size"], inflight_size, inflight_size_in_output,
[a.name for a in node.args if isinstance(a, Node)],
list(node.users.keys()), node_to_last_use[node] if node in node_to_last_use else 'NA', lifetime,
user_to_last_uses[node] if node in user_to_last_uses else 'NA',
list(inflight_values)
]
csv_data.append(row)
# print(
# f"Node: {node.name} users: {list(node.users.keys())} node_to_last_use: {node_to_last_use[node] if node in node_to_last_use else 'NA'} user_to_last_uses: {user_to_last_uses[node] if node in user_to_last_uses else 'NA'} inflight_values: {inflight_values} inflight_size: {inflight_size}"
# )
max_inflight_size = max(max_inflight_size, inflight_size)
import csv
with open(csv_filename, mode='w', newline='') as file:
writer = csv.writer(file)
writer.writerows(csv_data)
print(f"Max inflight size: {max_inflight_size}")
print(f"Data successfully written to {csv_filename}")
def get_activation_node_names(graph: Graph, param_nodes_bw: List[Node], fwd_output_names: List[str]):
input_nodes = get_input_nodes(graph)
param_node_names = set([n.name for n in param_nodes_bw])
activation_node_names = []
for in_node in input_nodes:
if in_node.name in fwd_output_names:
if in_node.name not in param_node_names:
activation_node_names.append(in_node.name)
return activation_node_names
class TensorOffloadHelper():
def __init__(self):
self.devices = {}
self.base_tensors = {}
self.views = {}
self.arg_list = []
self.offloaded = {}
self.non_tensor = {}
def offload(self, argument):
def is_base_tensor(tensor):
return torch.is_tensor(a) and not a._is_view() and not hasattr(tensor, "ds_id")
base_tensor_ids = set()
for a in argument:
if is_base_tensor(a):
base_tensor_ids.add(id(a))
for a in argument:
a_id = id(a)
if is_base_tensor(a):
# Base tensor
self.devices[a_id] = a.device
self.base_tensors[a_id] = a
# elif torch.is_tensor(a) and not hasattr(a, "ds_id") and id(a._base) in base_tensor_ids:
# # View
# self.views[a_id] = {
# "base_id": id(a._base),
# "size": a.size(),
# "stride": a.stride(),
# "offset": a.storage_offset(),
# }
else:
# other types or ds tensor
self.non_tensor[a_id] = a
self.arg_list.append(a_id)
for a in argument:
if is_base_tensor(a):
a.data = a.data.to("cpu")
def reload(self, in_place):
loaded_base_tensors = {}
for a_id in self.arg_list:
if a_id in self.base_tensors:
device = self.devices[a_id]
if in_place:
self.base_tensors[a_id].data = self.base_tensors[a_id].to(device)
loaded_base_tensors[a_id] = self.base_tensors[a_id]
else:
loaded_base_tensors[a_id] = self.base_tensors[a_id].to(device)
results = []
for a_id in self.arg_list:
if a_id in self.base_tensors:
results.append(loaded_base_tensors[a_id])
# elif a_id in self.views:
# view_info = self.views[a_id]
# # print(f"load_args loading view {a_id} base_id={view_info['base_id']} size={view_info['size']} stride={view_info['stride']} offset={view_info['offset']}")
# base_tensor = loaded_base_tensors[view_info["base_id"]]
# view_tensor = base_tensor.as_strided(
# view_info["size"], view_info["stride"], view_info["offset"]
# )
# results.append(view_tensor)
elif a_id in self.non_tensor:
results.append(self.non_tensor[a_id])
return results
def add_mem_profile_nodes(graph: Graph, prefix: str):
def show_memory(label: str):
if dist.get_rank() == 0:
print(
f"{prefix} {label} alloc_mem={get_accelerator().memory_allocated()} max_mem={get_accelerator().max_memory_allocated()}"
)
nodes = list(graph.nodes)
for node in nodes:
if node.op == "output":
continue
with graph.inserting_after(node):
msg = f"Mem {node.name}"
name = f"show_memory_{node.name}"
graph.create_node('call_function', show_memory, (msg, ), {}, name=name)
def is_release_node(n: Node) -> bool:
return n.target == torch.ops.dc.release_param.default
def get_index_by_graph_id(graph_order, target_graph_id):
for index, (graph_id, _) in enumerate(graph_order):
if graph_id == target_graph_id:
return index
return -1
def pad_tensors(specs: List[Tuple[torch.Tensor, int, int]]) -> List[torch.Tensor]:
"""
specs = [
(input_ids, 1, pad_token_id), # Example: Pad the right side with <pad>
(attention_mask, 1, 0), # Example: Pad the right side with 0
...
]
- Share the "maximum length of the dim dimension" across ranks for all specs
- Pad the right side for the missing parts and return
- Communication (`all_reduce`) happens only once
"""
assert len(specs) > 0, "specs is empty"
device = specs[0][0].device
# Vectorize local lengths
local_sizes = torch.tensor(
[tensor.size(dim) for tensor, dim, _ in specs],
dtype=torch.long,
device=device,
)
# Element-wise MAX across ranks
dist.all_reduce(local_sizes, op=dist.ReduceOp.MAX)
max_sizes = local_sizes.tolist()
# Pad each tensor as needed
padded: List[torch.Tensor] = []
# Don't use F.pad here:
# If you don't need to pad only on a certain rank, it will lead to different strides across ranks.
# This will cause recompilation on only some ranks and get the communication collective stuck.
for (tensor, dim, pad_val), max_len in zip(specs, max_sizes):
cur_len = tensor.size(dim)
# --- (1) Always allocate a new buffer with 'row-major, contiguous memory' -------------
out_shape = list(tensor.shape)
out_shape[dim] = max_len
out = torch.full(out_shape, pad_val, dtype=tensor.dtype, device=tensor.device)
# --- (2) Copy original data using slicing ------------------------------
slc = [slice(None)] * tensor.dim()
slc[dim] = slice(0, cur_len)
out[tuple(slc)] = tensor
# out is always row-major: for example, if shape is (..., 1, L), then
# stride = (..., L, 1)
padded.append(out)
return padded
def create_shard_offsets(gm: GraphModule, s0_node: Node) -> Tuple[Node, Node]:
sp_size: int = sp_dp_registry.sp_size()
sp_rank: int = dist.get_rank() % sp_dp_registry.sp_size()
with gm.graph.inserting_after(s0_node):
chunk_size_node = gm.graph.call_function(operator.floordiv, args=(s0_node, sp_size))
with gm.graph.inserting_after(chunk_size_node):
start_node = gm.graph.call_function(operator.mul, args=(sp_rank, chunk_size_node))
with gm.graph.inserting_after(start_node):
end_node = gm.graph.call_function(operator.add, args=(start_node, chunk_size_node))
return start_node, end_node
def get_sdpa_nodes(gm: GraphModule) -> List[Node]:
return list(gm.graph.find_nodes(
op="call_function",
target=F.scaled_dot_product_attention,
))
def get_input_id_node(gm: GraphModule) -> Node:
from .fx import find_node_by_tag
node = find_node_by_tag(gm, constants.AUTOSP_INPUT_ID_KEY)
if node is None:
raise RuntimeError("Failed to find a node for the input sequence.")
return node
def get_label_id_node(gm: GraphModule) -> Node:
from .fx import find_node_by_tag
node = find_node_by_tag(gm, constants.AUTOSP_LABEL_ID_KEY)
if node is None:
raise RuntimeError("Failed to find a node for the label.")
return node
def get_position_id_node(gm: GraphModule) -> Node:
from .fx import find_node_by_tag
node = find_node_by_tag(gm, constants.AUTOSP_POSITION_ID_KEY)
return node
def create_symbolic_slice_indices(
gm: GraphModule,
sym_seq_dim_node: Node,
) -> Tuple[Node, Node]:
start_node, end_node = create_shard_offsets(gm, sym_seq_dim_node)
with gm.graph.inserting_after(end_node):
slice_all = gm.graph.call_function(slice, args=(None, None, None))
with gm.graph.inserting_after(slice_all):
slice_range = gm.graph.call_function(slice, args=(start_node, end_node, None))
return slice_all, slice_range
def shard_tensor_node(gm: GraphModule, tensor_node: Node):
from .fx import find_node_by_name, get_node_shape_meta, replace_node_users
val = get_node_shape_meta(tensor_node)
assert val is not None, f"Node {tensor_node.name} has no shape metadata"
seq_len = val.shape[1]
assert isinstance(
seq_len,
torch.SymInt), (f"Expected sequence dimension to be {torch.SymInt!r} but instead found {type(seq_len)!r}")
symb_seq_int_node = find_node_by_name(gm, str(seq_len))
assert symb_seq_int_node, f"Unable to find symbolic placeholder for {seq_len}"
slice_all, slice_range = create_symbolic_slice_indices(gm, symb_seq_int_node)
indices = (slice_all, slice_range)
positions = {node: i for i, node in enumerate(gm.graph.nodes)}
# Insert after the later dependency so the new getitem does not appear
# before the symbolic slice nodes in graph order. Torch 2.9 bf16 can place
# the SymInt placeholder after the tensor placeholder.
anchor_node = slice_range if positions[slice_range] > positions[tensor_node] else tensor_node
with gm.graph.inserting_after(anchor_node):
sliced_node = gm.graph.call_function(
operator.getitem,
args=(tensor_node, indices),
)
replace_node_users(tensor_node, sliced_node, exclude=[sliced_node])