kolandr/gridbot
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Initial import: grid-bot — grid trading bot for BTC-USDT on Cifra Markets

Browse Source
This commit is contained in:
Kolp
2026-09-24 13:22:23 +07:00
commit 642cc11a9f
18968 changed files with 5683248 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
kronos-venv/lib/python3.12/site-packages/torch/distributed/pipelining/_IR.py
+1274
View File
File diff suppressed because it is too large Load Diff
kronos-venv/lib/python3.12/site-packages/torch/distributed/pipelining/__init__.py
+30
View File
@@ -0,0 +1,30 @@
# Copyright (c) Meta Platforms, Inc. and affiliates
from ._IR import Pipe, pipe_split, pipeline, SplitPoint
from .schedules import (
_ScheduleForwardOnly,
Schedule1F1B,
ScheduleDualPipeV,
ScheduleGPipe,
ScheduleInterleaved1F1B,
ScheduleInterleavedZeroBubble,
ScheduleLoopedBFS,
ScheduleZBVZeroBubble,
)
from .stage import build_stage, PipelineStage
__all__ = [
"Pipe",
"pipe_split",
"SplitPoint",
"pipeline",
"PipelineStage",
"build_stage",
"Schedule1F1B",
"ScheduleGPipe",
"ScheduleInterleaved1F1B",
"ScheduleLoopedBFS",
"ScheduleInterleavedZeroBubble",
"ScheduleZBVZeroBubble",
"ScheduleDualPipeV",
]
kronos-venv/lib/python3.12/site-packages/torch/distributed/pipelining/_backward.py
+460
View File
@@ -0,0 +1,460 @@
# mypy: allow-untyped-defs
# Copyright (c) Meta Platforms, Inc. and affiliates
import collections
import logging
from collections.abc import Iterator, Sequence
from typing import Any
import torch
from torch.autograd.graph import GradientEdge, Node
from torch.nn import Parameter
from ._debug import map_debug_info
logger = logging.getLogger(__name__)
def _get_grad_fn_or_grad_acc(t: torch.Tensor) -> Node | None:
"""
Get the grad function or grad accumulator for a tensor.
Accumulate grad nodes are lazily created, so we need to a
dummy view in order to trigger its creation.
"""
if t.requires_grad and t.grad_fn is None:
# if no grad function (leaf tensors) we use view
viewed_t = t.view_as(t)
grad_fn = viewed_t.grad_fn
if grad_fn is not None:
return grad_fn.next_functions[0][0]
else:
raise RuntimeError(
"Attempted to get grad_fn, but got None."
"Is this being created in a no-grad context?"
)
else:
return t.grad_fn
def reverse_closure(
roots: list[Node], target_nodes: set[Node], reverse_edges_dict
) -> tuple[set[Node], set[Node]]:
"""
This function returns the reverse closure of the given roots,
i.e. the set of nodes that can be reached from the roots by following the
reverse edges of the graph. The target_nodes are the nodes that we want to
include in the closure.
"""
# Recurse until we reach a target node
closure: set[Node] = set()
visited_target_nodes = set()
q: collections.deque[Node] = collections.deque()
for node in roots:
if node is not None and node not in closure:
closure.add(node)
q.append(node)
while q:
node = q.popleft()
reverse_edges = reverse_edges_dict[node]
for fn in reverse_edges:
if fn in closure or fn is None:
continue
if fn in target_nodes:
visited_target_nodes.add(fn)
continue
closure.add(fn)
q.append(fn)
return closure, visited_target_nodes
def construct_reverse_graph(roots: list[Node]) -> dict[Node, list[Node]]:
q: collections.deque[Node] = collections.deque()
root_seen: set[Node] = set()
reverse_edges_dict: dict[Node, list[Node]] = collections.defaultdict(list)
for node in roots:
if node is not None and node not in root_seen:
q.append(node)
root_seen.add(node)
while q:
node = q.popleft()
for fn, _ in node.next_functions:
if fn is not None:
if len(reverse_edges_dict[fn]) == 0:
q.append(fn)
reverse_edges_dict[fn].append(node)
return reverse_edges_dict
def get_param_groups(
inputs: list[Node], params: list[Node], reverse_edges_dict
) -> list[dict[str, Any]]:
"""
Given a list of inputs and a list of parameters, return a list of parameter
groups, where each group contains the parameters and the intermediates that
are connected to the parameters.
The returned list of parameter groups is a list of dictionaries, where each
dictionary contains the following keys:
- "params": a set of parameters
- "intermediates": a set of intermediates
The returned list of parameter groups is a list of dictionaries,
"""
# reverse graph that starts with inputs, and goes up to the dOutput or the loss,
# but omits weights and any subgraphs connecting weights to this closure
inputs_closure, _ = reverse_closure(inputs, set(), reverse_edges_dict)
param_groups: dict[Node, dict[str, set]] = dict() # keyed on intermediates
for param in params:
closure, intersected = reverse_closure(
[param], inputs_closure, reverse_edges_dict
)
param_group: dict[str, set] = {
"params": {param},
"intermediates": intersected,
}
for input_node in intersected:
existing = param_groups.get(input_node)
if existing is not None:
existing["params"] = existing["params"].union(param_group["params"])
existing["intermediates"] = existing["intermediates"].union(
param_group["intermediates"]
)
param_group = existing
else:
param_groups[input_node] = param_group
# Sanity check: union of all param_groups params should be equal to all params
union_params: set[Node] = set()
seen_ids: set[int] = set()
unique_param_groups = []
for param_group in param_groups.values():
if id(param_group) not in seen_ids:
seen_ids.add(id(param_group))
unique_param_groups.append(param_group)
union_params = union_params.union(param_group["params"])
# The assert will only be true if the input tensor requires gradients,
# otherwise the autograd graph will miss the first layer of inputs
# assert union_params == set(params)
return unique_param_groups
def _autograd_grad_for_inputs(
outputs: Sequence[torch.Tensor],
inputs: Sequence[torch.Tensor],
grad_outputs: Sequence[torch.Tensor | None] | None = None,
retain_graph: bool = False,
allow_unused: bool = False,
) -> tuple[torch.Tensor | None, ...]:
"""Compute input gradients, returning ``None`` for non-grad inputs."""
# Some inputs may not be used or may not require gradients, so we filter them out
# before calling autograd.grad and place None for those positions in the result.
grad_indices: list[int] = []
inputs_requiring_grad: list[torch.Tensor] = []
for i, inp in enumerate(inputs):
if isinstance(inp, torch.Tensor) and inp.requires_grad:
grad_indices.append(i)
inputs_requiring_grad.append(inp)
if not inputs_requiring_grad:
return tuple(None for _ in inputs)
grads = torch.autograd.grad(
outputs=outputs,
inputs=inputs_requiring_grad,
grad_outputs=grad_outputs,
retain_graph=retain_graph,
allow_unused=allow_unused,
)
result: list[torch.Tensor | None] = [None] * len(inputs)
for idx, g in zip(grad_indices, grads, strict=True):
result[idx] = g
return tuple(result)
def stage_backward_input(
stage_outputs_or_loss: list[torch.Tensor],
output_grads: list[torch.Tensor] | None,
input_values: list[torch.Tensor],
weights: Iterator[Parameter],
) -> tuple[tuple[torch.Tensor | None, ...], list[dict[str, Any]]]:
"""
Compute the gradients for only the stage inputs with
respect to the stage outputs (if non-last stage) or loss (if last stage)
After computing input gradients, we save the intermediate nodes in `param_groups`
for later use in stage_backward_weight. We don't need to save any other intermediate nodes
that aren't needed for dW because when we do dW calculation, we start from saved intermediates.
Detaching the stage_outputs_or_loss at the end of this function is important as
it frees up the memory that the autograd graph is anticipating to be used later (but doesn't actually need).
"""
stage_output_grad_fns: list[Node] = list(
filter(None, map(_get_grad_fn_or_grad_acc, stage_outputs_or_loss))
)
stage_input_grad_fns: list[Node] = list(
filter(None, map(_get_grad_fn_or_grad_acc, input_values))
)
weight_grad_fns: list[Node] = list(
filter(None, map(_get_grad_fn_or_grad_acc, weights))
)
reverse_edges_dict = construct_reverse_graph(stage_output_grad_fns)
param_groups = get_param_groups(
stage_input_grad_fns, weight_grad_fns, reverse_edges_dict
)
handles = []
for param_group in param_groups:
for i, intermediate in enumerate(param_group["intermediates"]):
def get_hook(param_group, i):
def hook(grad_inputs):
if param_group.get("grads", None) is None:
param_group["grads"] = [None] * len(
param_group["intermediates"]
)
param_group["grads"][i] = grad_inputs
return hook
# These are always "split" nodes that we need to recompute, so
# save their inputs.
handle = intermediate.register_prehook(get_hook(param_group, i))
handles.append(handle)
if output_grads is None:
# In case this is the loss and there are no output_grads, then we just use 1s
output_grads = [
torch.ones_like(stage_output) for stage_output in stage_outputs_or_loss
]
dinputs = _autograd_grad_for_inputs(
stage_outputs_or_loss,
input_values,
output_grads,
retain_graph=True,
)
# Accumulate into .grad
for inp, dinput in zip(input_values, dinputs):
if isinstance(inp, torch.Tensor) and dinput is not None:
if inp.grad is None:
inp.grad = dinput
else:
inp.grad += dinput
# stage_outputs_or_loss are not used in backwards after this point, so we can safely remove it from the autograd graph
# this allows autograd to clear up the graph dedicated for this tensor and free up significant memory
for t in stage_outputs_or_loss:
t.detach_()
# hooks are no longer necessary, clean up for consistency
for handle in handles:
handle.remove()
return dinputs, param_groups
def stage_backward_weight(
weights: Iterator[Parameter], param_groups: list[dict[str, Any]], retain_graph=False
) -> tuple[torch.Tensor | None, ...]:
# map weights to param_group_weights
grad_acc_to_weight = {}
weight_grads: list[torch.Tensor | None] = []
for index, weight in enumerate(weights):
grad_acc = _get_grad_fn_or_grad_acc(weight)
grad_acc_to_weight[grad_acc] = weight, index
weight_grads.append(weight.grad)
for param_group in param_groups:
valid_edges = []
valid_grad_outputs: list[torch.Tensor] = []
for grads_tuple, intermediate in zip(
param_group["grads"], param_group["intermediates"]
):
for i, grad in enumerate(grads_tuple):
if grad is not None:
valid_edges.append(GradientEdge(intermediate, i))
# pyrefly: ignore [bad-argument-type]
valid_grad_outputs.append(grad)
# Break a reference cycle caused inside stage_backward_input->get_hook->hook
# The summarized cycle is:
# `hook` -> cell -> param_group -> intermediates -> `hook`
# because we install the hook function onto each of the intermediate autograd nodes.
# We need to keep intermediates alive up until backward_weight, but we can free it now.
del param_group["intermediates"]
if valid_edges: # Only call autograd.grad if we have valid gradients
# [NEW!] Able to pass a GradientEdge to autograd.grad as output
weights_edges = tuple(GradientEdge(w, 0) for w in param_group["params"])
dweights = torch.autograd.grad(
valid_edges,
weights_edges,
grad_outputs=valid_grad_outputs,
retain_graph=retain_graph,
)
# release grad memory early after use
del param_group["grads"]
for grad_acc, dw in zip(param_group["params"], dweights):
weight, index = grad_acc_to_weight[grad_acc]
if weight.grad is None:
weight.grad = dw
else:
weight.grad += dw
# return grads in the original order weights were provided in
return tuple(weight_grads)
def stage_backward(
stage_output,
output_grads,
input_values,
outputs_with_grads_idxs: list[int] | None = None, # deprecated, not used
) -> tuple[torch.Tensor | None, ...]:
"""
This is a helper function to:
1. compute the gradients for the stage inputs, and
2. accumulate gradients for the stage module's parameters.
Given the input value(s) and the corresponding gradient for the output
value(s), compute and accumulate gradients for all parameter values (leaves
in the autograd trace) as well as return a list of the gradients for the
input values
"""
if outputs_with_grads_idxs is not None:
# Deprecated, not used in runtime calls, only exists in compiler
stage_output = [stage_output[i] for i in outputs_with_grads_idxs]
output_grads = [output_grads[i] for i in outputs_with_grads_idxs]
try:
# stage_output may be a composite datatype like dict. Extract all individual
# tensor values here
stage_output_tensors: list[torch.Tensor] = []
output_grad_tensors: list[torch.Tensor | None] = []
def extract_tensors_with_grads(
output_val,
grad_val,
# Don't delete me- see [Note: ref cycle]
extract_tensors_with_grads,
):
if isinstance(output_val, torch.Tensor):
if not output_val.requires_grad and output_val.grad_fn is None:
return
if not isinstance(grad_val, (torch.Tensor, type(None))):
raise AssertionError(
f"Expected Tensor or None gradient but got {type(grad_val)}"
)
stage_output_tensors.append(output_val)
output_grad_tensors.append(grad_val)
elif isinstance(output_val, (tuple, list)):
if grad_val is None:
return
if not isinstance(grad_val, (tuple, list)):
raise AssertionError(
f"grad_value expected to have type {type(output_val)} but got {type(grad_val)}"
)
if not len(output_val) == len(grad_val):
raise AssertionError(
f"Expected len(output_val) == len(grad_val), got {len(output_val)} != {len(grad_val)}"
)
for ov, gv in zip(output_val, grad_val):
extract_tensors_with_grads(
ov,
gv,
extract_tensors_with_grads,
)
elif isinstance(output_val, dict):
if grad_val is None:
return
if not isinstance(grad_val, dict):
raise AssertionError(f"Expected dict, got {type(grad_val)}")
if not set(output_val.keys()) == set(grad_val.keys()):
raise AssertionError(
f"Expected keys {set(output_val.keys())}, got {set(grad_val.keys())}"
)
for k in output_val:
extract_tensors_with_grads(
output_val[k], grad_val[k], extract_tensors_with_grads
)
else:
# Output is a non-tensor type; just ignore it
pass
# Note: ref cycle
# break a ref cycle that would keep tensors alive until GC runs
# 1. extract_tensors_with_grads refers to a cell that holds refs to any vars defined in stage_backward
# and used in extract_tensors_with_grads
# 2. extract_tensors_with_grads referred to both stage_output_tensors, output_grad_tensors,
# and to itself (extract_tensors_with_grads) since it makes a recursive call
# 3. stage_output_tensors was kept alive by the above refcycle, and it holds activation tensors, which is bad
# fix -> explicitly pass in the ref to the fn, so there is no gc cycle anymore
extract_tensors_with_grads(
stage_output, output_grads, extract_tensors_with_grads
)
torch.autograd.backward(
stage_output_tensors,
grad_tensors=output_grad_tensors, # type: ignore[arg-type]
)
# Extract gradients wrt the input values
grad_inputs: list[torch.Tensor | None] = []
for val in input_values:
if isinstance(val, torch.Tensor):
grad_inputs.append(val.grad)
# Since gradients that will pass back to previous stages do not require gradient accumulation,
# by decrementing the gradients' reference count at this point, the memory of gradients will be
# returned to the allocator as soon as the next micro batch's get_bwd_send_ops comes and current
# asynchronous send completes.
# This prevents the gradients from persisting in GPU memory for the entire duration of step_microbatches
# until clear_runtime_states() is called.
val.grad = None
else:
grad_inputs.append(None)
# Alternative impl: `torch.autograd.grad`.
# Note that `torch.autograd.grad` will not accumulate gradients into the
# model's parameters.
"""
inputs_with_grad = []
for val in input_values:
if isinstance(val, torch.Tensor) and val.requires_grad:
inputs_with_grad.append(val)
grad_inputs = torch.autograd.grad(
stage_output_tensors, inputs_with_grad, output_grad_tensors, # type: ignore[arg-type]
)
"""
except Exception as e:
exc_msg = f"""
Failed to run stage backward:
Stage output: {map_debug_info(stage_output)}
Output gradient: {map_debug_info(output_grads)}
Input: {map_debug_info(input_values)}
"""
raise RuntimeError(exc_msg) from e
return tuple(grad_inputs)
# TODO: handling requires_grad=False dynamically. Can we analyze this during initial
# IR emission?
def _null_coalesce_accumulate(lhs, rhs):
"""
Coalesce two values, even if one of them is null, returning the non-null
value.
"""
if lhs is None:
return rhs
elif rhs is None:
return lhs
else:
return torch.add(lhs, rhs)
kronos-venv/lib/python3.12/site-packages/torch/distributed/pipelining/_debug.py
+22
View File
@@ -0,0 +1,22 @@
# Copyright (c) Meta Platforms, Inc. and affiliates
import torch
from torch.fx.node import Argument
def friendly_debug_info(v: object) -> Argument:
"""
Helper function to print out debug info in a friendly way.
"""
if isinstance(v, torch.Tensor):
return f"Tensor({v.shape}, grad={v.requires_grad}, dtype={v.dtype})"
else:
return str(v)
def map_debug_info(a: Argument) -> Argument:
"""
Helper function to apply `friendly_debug_info` to items in `a`.
`a` may be a list, tuple, or dict.
"""
return torch.fx.node.map_aggregate(a, friendly_debug_info)
kronos-venv/lib/python3.12/site-packages/torch/distributed/pipelining/_schedule_visualizer.py
+447
View File
@@ -0,0 +1,447 @@
# Copyright (c) Meta Platforms, Inc. and affiliates
"""
This visualizer requires matplotlib to be installed.
Example usage:
ops = get_schedule_ops("InterleavedZeroBubble", 4, 8)
visualize_schedule(ops, "test.png")
"""
import collections
from typing import NamedTuple
from unittest import mock
from torch.distributed.pipelining.schedules import (
_Action,
_ComputationType,
_PipelineSchedule,
_PipelineScheduleRuntime,
get_schedule_class,
PipelineScheduleMulti,
PipelineScheduleSingle,
)
from torch.distributed.pipelining.stage import PipelineStage
class OpKey(NamedTuple):
stage_index: int
computation_type: _ComputationType
microbatch_index: int
def get_schedule_ops(
schedule: str | type[_PipelineSchedule],
pp_degree: int,
num_microbatches: int,
num_stages_per_rank: int | None = None,
add_spacing: bool = False,
with_comms: bool = False,
) -> list[list[_Action | None]]:
"""
Get all actions for a given schedule, pp_degree, and num_microbatches. The actions are returned in a list of lists
where each inner list represents a rank and each element in the inner list represents an action.
The schedule can be specified as a string which is passed into get_schedule_class() or a _PipelineSchedule instance.
"""
if add_spacing and with_comms:
raise ValueError("Cannot add spacing and view comms at the same time")
if isinstance(schedule, str):
schedule_class = get_schedule_class(schedule)
elif issubclass(schedule, _PipelineSchedule):
schedule_class = schedule
else:
raise ValueError(f"Invalid schedule: {schedule}")
# Create a mock of the PipelineStage class
mock_pipeline_stage = mock.create_autospec(PipelineStage, instance=True)
# Set the return values for group_rank and group_size methods
mock_pipeline_stage.group_rank = 0
mock_pipeline_stage.group_size = pp_degree
mock_pipeline_stage.submod = None
# Check num_stages_per_rank is valid
if issubclass(schedule_class, PipelineScheduleSingle):
if num_stages_per_rank is None:
num_stages_per_rank = 1
if not num_stages_per_rank == 1:
raise AssertionError(
f"Expected num_stages_per_rank to be 1, got {num_stages_per_rank}"
)
stages = mock_pipeline_stage
stages.num_stages = num_stages_per_rank * pp_degree
elif issubclass(schedule_class, PipelineScheduleMulti):
if num_stages_per_rank is None:
num_stages_per_rank = 2
if not num_stages_per_rank >= 2:
raise AssertionError(
f"Expected num_stages_per_rank >= 2, got {num_stages_per_rank}"
)
stages = [mock_pipeline_stage for _ in range(num_stages_per_rank)]
for stage in stages:
stage.num_stages = num_stages_per_rank * pp_degree
else:
raise ValueError(f"Invalid schedule: {schedule_class}")
# Instantiate the schedule class
# pyrefly: ignore [bad-argument-type]
schedule_instance = schedule_class(stages, num_microbatches)
if schedule_instance.pipeline_order is None:
raise AssertionError("Expected pipeline_order to not be None")
# Convert to List[List[_Action]]
all_actions: list[list[_Action | None]] = []
if with_comms:
runtime = _PipelineScheduleRuntime(stages, num_microbatches)
runtime._prepare_schedule_with_comms(schedule_instance.pipeline_order)
for rank in range(pp_degree):
all_actions.append(list(runtime.pipeline_order_with_comms[rank]))
else:
for rank in range(pp_degree):
all_actions.append(schedule_instance.pipeline_order[rank])
# Add spacing
if add_spacing:
# remove all Nones, then respace
# TODO: later we can change this at the schedule creation level to not use Nones
all_actions = [
[action for action in rank if action is not None] for rank in all_actions
]
all_actions = add_schedule_op_spacing(all_actions)
# Return the pipeline order
return all_actions
class _ComputationTypeVisual:
def __init__(
self,
color: str,
text: str = "",
width: int = 1,
):
self.color = color
self.width = width
self.text = text
# Update the mapping to use _ComputationTypeVisual instances
action_type_to_color_mapping = {
_ComputationType.FORWARD: _ComputationTypeVisual("blue", "Forward"),
_ComputationType.BACKWARD_INPUT: _ComputationTypeVisual("teal", "Backward Input"),
_ComputationType.BACKWARD_WEIGHT: _ComputationTypeVisual(
"green", "Backward Weight"
),
_ComputationType.FULL_BACKWARD: _ComputationTypeVisual(
"orange", "Full Backward", 2
),
_ComputationType.OVERLAP_F_B: _ComputationTypeVisual("purple", "Overlap F+B", 3),
}
def add_schedule_op_spacing(
schedule: list[list[_Action | None]],
) -> list[list[_Action | None]]:
"""
Add spacing to the schedule based on dependencies between ranks.
Before adding an operation to the list, this function checks if there are
dependencies from other ranks. If there are dependencies (other ranks have
not finished processing the required microbatch), it adds None instead.
For example, Forward microbatch 0 on rank 1 depends on rank 0 processing
Forward microbatch 0 first.
Args:
schedule: The original schedule as a list of lists where each inner list
represents a rank and each element represents an action.
Returns:
A new schedule with proper spacing based on dependencies.
"""
if not schedule:
return schedule
num_stages = (
max(
action.stage_index
for rank_actions in schedule
for action in rank_actions
if action is not None
)
+ 1
)
num_ranks = len(schedule)
spaced_schedule: list[list[_Action | None]] = [[] for _ in range(num_ranks)]
rank_ops = [collections.deque(ops) for ops in schedule]
# Track completion times: (stage_index, action_type, microbatch_index) -> completion_time
scheduled_ops: dict[OpKey, int] = {}
def is_dependency_ready(dependency_key: OpKey, timestep: int) -> bool:
"""Check if a dependency operation has completed by the given timestep."""
return (
dependency_key in scheduled_ops
and timestep >= scheduled_ops[dependency_key]
)
def get_dependencies(action: _Action) -> list[OpKey]:
"""Get the list of dependencies for an action."""
stage_idx = action.stage_index
comp_type = action.computation_type
mb_idx = action.microbatch_index
# Ensure mb_idx is not None for dependency tracking
if mb_idx is None:
raise AssertionError(f"Action {action} has None microbatch_index")
# First stage forward has no dependencies
if stage_idx == 0 and comp_type == _ComputationType.FORWARD:
return []
# Last stage backward depends on forward from previous stage
if stage_idx == num_stages - 1 and comp_type in (
_ComputationType.FULL_BACKWARD,
_ComputationType.BACKWARD_INPUT,
):
return [OpKey(stage_idx - 1, _ComputationType.FORWARD, mb_idx)]
# Forward depends on previous stage forward
if comp_type == _ComputationType.FORWARD:
return [OpKey(stage_idx - 1, _ComputationType.FORWARD, mb_idx)]
# Backward depends on next stage backward
if comp_type in (
_ComputationType.FULL_BACKWARD,
_ComputationType.BACKWARD_INPUT,
):
return [
OpKey(stage_idx + 1, _ComputationType.FULL_BACKWARD, mb_idx),
OpKey(stage_idx + 1, _ComputationType.BACKWARD_INPUT, mb_idx),
]
# Weight backward depends on input backward
if comp_type == _ComputationType.BACKWARD_WEIGHT:
return [OpKey(stage_idx, _ComputationType.BACKWARD_INPUT, mb_idx)]
raise RuntimeError(f"Unknown computation type: {comp_type}")
def is_action_ready(action: _Action, timestep: int) -> bool:
"""Check if an action is ready to be scheduled at the given timestep."""
# For OR dependencies (like backward), check if any dependency is satisfied
if action.computation_type in (
_ComputationType.FULL_BACKWARD,
_ComputationType.BACKWARD_INPUT,
_ComputationType.BACKWARD_WEIGHT,
):
dependencies = get_dependencies(action)
return any(is_dependency_ready(dep, timestep) for dep in dependencies)
# For AND dependencies, all must be satisfied
elif action.computation_type == _ComputationType.FORWARD:
dependencies = get_dependencies(action)
return all(is_dependency_ready(dep, timestep) for dep in dependencies)
elif action.computation_type == _ComputationType.OVERLAP_F_B:
if action.sub_actions is None:
raise AssertionError(
f"OVERLAP_F_B action {action} has None sub_actions"
)
dep_list: list[bool] = []
for sub_action in action.sub_actions:
dep_list.append(is_action_ready(sub_action, timestep))
return all(dep_list)
else:
raise RuntimeError(f"Unknown computation type: {action.computation_type}")
def schedule_action(action: _Action, rank: int, timestep: int) -> int:
"""Schedule an action and return completion time."""
spaced_schedule[rank].append(action)
comp_type = action.computation_type
comp_time = action_type_to_color_mapping[comp_type].width
completion_time = timestep + comp_time
if comp_type == _ComputationType.OVERLAP_F_B:
# For overlap actions, schedule each sub-action with cumulative timing
if action.sub_actions is None:
raise AssertionError(
f"OVERLAP_F_B action {action} has None sub_actions"
)
cumulative_time = 0
for sub_action in action.sub_actions:
if sub_action.microbatch_index is None:
raise AssertionError(
f"Sub-action {sub_action} has None microbatch_index"
)
sub_comp_time = action_type_to_color_mapping[
sub_action.computation_type
].width
cumulative_time += sub_comp_time
scheduled_ops[
OpKey(
sub_action.stage_index,
sub_action.computation_type,
sub_action.microbatch_index,
)
] = timestep + cumulative_time
else:
if action.microbatch_index is None:
raise AssertionError(f"Action {action} has None microbatch_index")
scheduled_ops[
OpKey(action.stage_index, comp_type, action.microbatch_index)
] = completion_time
return completion_time
# Main scheduling loop
current_timestep = 0
timesteps_without_progress = 0
rank_completion_times = dict.fromkeys(range(num_ranks), 0)
while rank_ops:
print(f"Current timestep: {current_timestep}")
# Process all operations during timestep until we run out of ready operations
for rank, op_queue in enumerate(rank_ops):
if not op_queue:
continue
op_queue = rank_ops[rank]
action = op_queue[0]
print(f"Rank: {rank}, {action=}")
if action is None:
spaced_schedule[rank].append(None)
op_queue.popleft()
timesteps_without_progress = 0
elif current_timestep >= rank_completion_times[rank] and is_action_ready(
action, current_timestep
):
rank_completion_times[rank] = schedule_action(
action, rank, current_timestep
)
op_queue.popleft()
timesteps_without_progress = 0
# Add None for ranks that are waiting
for rank in range(num_ranks):
if current_timestep >= rank_completion_times[rank]:
spaced_schedule[rank].append(None)
# Remove empty queues and advance timestep
rank_ops = [op_queue for op_queue in rank_ops if op_queue]
current_timestep += 1
timesteps_without_progress += 1
if timesteps_without_progress > max(
visual.width for visual in action_type_to_color_mapping.values()
):
raise RuntimeError("No progress made in scheduling - possible deadlock")
return spaced_schedule
def visualize_schedule(
schedule: list[list[_Action | None]],
filename: str | None = None,
) -> None:
"""
Visualize the schedule using matplotlib.
The schedule is a list of lists where each inner list represents a rank and each element in the inner list represents an action.
The actions are represented as rectangles with different colors based on their computation type.
The filename is optional and if provided, the plot will be saved to that file.
Args:
schedule: The schedule to visualize.
filename: The filename to save the plot to. If not provided, the plot will be displayed.
add_schedule_spacing: If True, add spacing to the schedule based on dependencies between ranks.
"""
import matplotlib.pyplot as plt
from matplotlib.patches import Rectangle
plt.rcParams["font.family"] = (
"DejaVu Sans" # or any other font available on your system
)
num_ranks = len(schedule)
max_actions = max(len(rank) for rank in schedule)
# Increase the figure size to provide more space for the legend
fig, ax = plt.subplots(figsize=(max_actions + 2, num_ranks + 2))
max_draw_position = -1
# Calculate dynamic font size based on figure size
font_size = min(max_actions, num_ranks) + 4
used_computation = set()
for rank_idx, actions in enumerate(schedule):
draw_position = 0 # Initialize drawing position for each rank
for action in actions:
if action is not None:
comp_type_color = action_type_to_color_mapping.get(
action.computation_type, _ComputationTypeVisual("black")
)
used_computation.add(action.computation_type)
color = comp_type_color.color
width = comp_type_color.width
# Check if action has sub_actions to determine styling
if action.sub_actions is not None:
linewidth = 2 # Thicker border for compound actions
text_weight = "normal" # Bold text for compound actions
else:
linewidth = 1 # Default linewidth for regular actions
text_weight = "normal" # Default text weight
# Draw the rectangle to represent the action duration
rect = Rectangle(
(draw_position, num_ranks - rank_idx - 1),
width,
1,
facecolor=color,
edgecolor="black",
linewidth=linewidth,
)
ax.add_patch(rect)
# Draw the text centered within the rectangle
ax.text(
draw_position + width / 2,
num_ranks - rank_idx - 1 + 0.5,
str(action),
ha="center",
va="center",
fontsize=font_size,
color="white",
weight=text_weight,
)
draw_position += width
else:
draw_position += 1 # Move to the next
max_draw_position = max(max_draw_position, draw_position)
ax.set_xlim(-0.5, max_draw_position + 1)
ax.set_ylim(-0.5, num_ranks + 0.5) # Add extra space at the top
# Set y-ticks to be in the middle of each rank's row
ax.set_yticks([num_ranks - rank_idx - 0.5 for rank_idx in range(num_ranks)])
ax.set_yticklabels([f"Rank {i}" for i in range(num_ranks)], fontsize=font_size)
ax.set_xticklabels([])
# Remove grid lines and ticks
ax.grid(False)
# Add legend with larger font size
legend_elements = [
Rectangle(
(0, 0),
1,
1,
facecolor=action_type_to_color_mapping[comp_type].color,
edgecolor="black",
label=action_type_to_color_mapping[comp_type].text,
)
for comp_type in used_computation
]
ax.legend(handles=legend_elements, loc="upper right", fontsize=font_size)
# Save to file if filename is provided, otherwise display the plot
if filename:
plt.savefig(filename, bbox_inches="tight")
else:
plt.show()
kronos-venv/lib/python3.12/site-packages/torch/distributed/pipelining/_unflatten.py
+30
View File
@@ -0,0 +1,30 @@
# Copyright (c) Meta Platforms, Inc. and affiliates
from collections import defaultdict
import torch
from torch.export.unflatten import _ModuleFrame, _SubmoduleEntry
def _outline_submodules(orig_graph: torch.fx.Graph) -> torch.fx.GraphModule:
# Create an empty GraphModule to hold the outlined modules
new_module = torch.fx.GraphModule(torch.nn.Module(), torch.fx.Graph())
seen_nodes: dict[str, torch.fx.Node] = {}
seen_modules: dict[int, list[_SubmoduleEntry]] = defaultdict(list)
seen_attrs: dict[str, set[str]] = defaultdict(set)
created_modules: dict[str, torch.nn.Module] = {}
_ModuleFrame(
orig_graph,
tuple(orig_graph.nodes),
seen_nodes,
seen_modules,
seen_attrs,
created_modules,
None,
[("", None, 0)],
"",
{},
module=new_module,
).run_outer()
new_module.graph.lint()
new_module.recompile()
return new_module
kronos-venv/lib/python3.12/site-packages/torch/distributed/pipelining/_utils.py
+938
View File
@@ -0,0 +1,938 @@
# mypy: allow-untyped-defs
# Copyright (c) Meta Platforms, Inc. and affiliates
from __future__ import annotations
import logging
import warnings
from dataclasses import dataclass, field
from enum import Enum
from typing import cast, Literal, overload, Protocol, TYPE_CHECKING, TypeAlias
import torch
from torch import fx
from torch.distributed._mesh_layout import _MeshLayout
from torch.distributed.tensor import DTensor
from torch.utils._pytree import tree_flatten, tree_unflatten
if TYPE_CHECKING:
from torch.distributed.device_mesh import DeviceMesh
from torch.distributed.tensor.placement_types import Placement
logger = logging.getLogger(__name__)
class GetMeshCallback(Protocol):
"""Callback to create/retrieve a DeviceMesh from its cache key components."""
def __call__(
self,
mesh_dim_names: tuple[str, ...],
mesh_layout: _MeshLayout | None,
) -> DeviceMesh: ...
# Key for mesh cache: (mesh_dim_names, mesh_layout)
# mesh_layout is the _MeshLayout object containing shape and stride (not actual ranks).
# This uniquely identifies a mesh within the same "universe" where all stages share
# the same rank tensor.
MeshCacheKey: TypeAlias = tuple[tuple[str, ...], _MeshLayout | None]
class PipeliningMetadataError(RuntimeError):
"""Raised on metadata mismatches during pipeline communication."""
@dataclass(frozen=True, slots=True)
class _TensorMeta:
"""Tensor metadata for recv buffer allocation and validation.
For plain tensors, these are the tensor's actual attributes.
For DTensors, these are LOCAL shard attributes; global attributes
are stored in :class:`_DTensorMeta`.
"""
shape: torch.Size
stride: tuple[int, ...]
dtype: torch.dtype
requires_grad: bool
@staticmethod
def from_tensor(tensor: torch.Tensor) -> _TensorMeta:
"""Create metadata from a plain tensor.
Args:
tensor: A plain ``torch.Tensor`` (not DTensor).
Returns:
Metadata capturing shape, stride, dtype, and requires_grad.
Raises:
TypeError: If ``tensor`` is a DTensor.
"""
if isinstance(tensor, DTensor):
raise PipeliningMetadataError(
"Expected plain tensor, got DTensor. Use _DTensorMeta.from_dtensor instead."
)
return _TensorMeta(
shape=tensor.shape,
stride=tensor.stride(),
dtype=tensor.dtype,
requires_grad=tensor.requires_grad,
)
def to_tensor(self, device: torch.device | str) -> torch.Tensor:
"""Reconstruct a tensor on ``device`` from this metadata.
Args:
device: Target device for the tensor.
Returns:
An empty strided tensor on ``device``.
"""
t = _make_tensor_from_meta(self, device)
t.requires_grad_(self.requires_grad)
return t
def get_diff(self, other: _TensorMeta) -> list[str]:
"""Return field-by-field differences with ``other``.
Args:
other: Metadata to compare against.
Returns:
List of human-readable difference strings (empty if equal).
"""
if self == other:
return []
diffs = []
if self.shape != other.shape:
diffs.append(f"shape mismatch: {self.shape} vs {other.shape}")
if self.stride != other.stride:
diffs.append(f"stride mismatch: {self.stride} vs {other.stride}")
if self.dtype != other.dtype:
diffs.append(f"dtype mismatch: {self.dtype} vs {other.dtype}")
# requires_grad is intentionally excluded: it is a runtime concern
# determined by has_backward and grad context, not a metadata invariant.
return diffs
@dataclass(frozen=True, slots=True)
class _DTensorMeta(_TensorMeta):
"""DTensor metadata extending :class:`_TensorMeta` with distribution info.
Inherited fields (shape, stride, etc.) are LOCAL shard attributes.
Additional fields capture global shape and placement information
needed to reconstruct a :class:`DTensor` via ``DTensor.from_local()``.
The :class:`DeviceMesh` is **not** stored (not serializable for P2P);
it is looked up from :class:`_MeshCache` using
``(mesh_dim_names, mesh_layout)`` as the key.
"""
# Global DTensor properties (for reconstruction)
global_shape: torch.Size = field(default_factory=lambda: torch.Size([]))
global_stride: tuple[int, ...] = field(default=())
# DTensor distribution properties
placements: tuple[Placement, ...] = field(
default=()
) # e.g., (Shard(0), Replicate())
# Mesh identification - used to look up the correct DeviceMesh from cache
mesh_dim_names: tuple[str, ...] = field(default=()) # e.g., ("tp",) or ("dp", "tp")
mesh_layout: _MeshLayout | None = field(
default=None
) # _MeshLayout with shape/stride - uniquely identifies mesh within the same universe
@staticmethod
def from_dtensor(dtensor: DTensor) -> _DTensorMeta:
"""Create metadata from a DTensor.
Args:
dtensor: The DTensor to extract metadata from.
Returns:
Metadata capturing both local and global attributes.
"""
device_mesh = dtensor.device_mesh
return _DTensorMeta(
# Local tensor attributes (for recv buffer allocation)
shape=dtensor._local_tensor.shape,
stride=dtensor._local_tensor.stride(),
dtype=dtensor.dtype,
requires_grad=dtensor.requires_grad,
# Global DTensor attributes (for reconstruction)
global_shape=dtensor.shape,
global_stride=dtensor.stride(),
# Distribution info
placements=dtensor._spec.placements,
mesh_dim_names=(
tuple(device_mesh.mesh_dim_names) if device_mesh.mesh_dim_names else ()
),
mesh_layout=device_mesh._layout,
)
@property
def mesh_cache_key(self) -> MeshCacheKey:
"""Cache key ``(mesh_dim_names, mesh_layout)`` for mesh lookup."""
return (self.mesh_dim_names, self.mesh_layout)
def to_dtensor(self, device: torch.device | str, mesh: DeviceMesh) -> DTensor:
"""Reconstruct a DTensor on ``device`` with placements.
Args:
device: Target device for the local tensor.
mesh: The ``DeviceMesh`` to attach.
Returns:
A DTensor on ``device``.
"""
local_tensor = _make_tensor_from_meta(self, device)
# Set requires_grad after from_local() so that the from_local
# operation itself is not recorded in the autograd graph.
return cast(
DTensor,
DTensor.from_local(
local_tensor,
device_mesh=mesh,
placements=self.placements,
shape=self.global_shape,
stride=self.global_stride,
run_check=False,
).requires_grad_(self.requires_grad),
)
def get_diff(self, other: _TensorMeta) -> list[str]:
"""Return field-by-field differences, including DTensor-specific fields.
Args:
other: Metadata to compare against.
Returns:
List of human-readable difference strings (empty if equal).
"""
if self == other:
return []
# Get base class differences (compares local shape/stride/dtype/requires_grad)
# NOTE: Use explicit class call instead of super() because
# @dataclass(slots=True) on both parent and child can break super().
diffs = _TensorMeta.get_diff(self, other)
# Add DTensor-specific comparisons if other is also _DTensorMeta
if isinstance(other, _DTensorMeta):
if self.global_shape != other.global_shape:
diffs.append(
f"global_shape mismatch: {self.global_shape} vs {other.global_shape}"
)
if self.global_stride != other.global_stride:
diffs.append(
f"global_stride mismatch: {self.global_stride} vs {other.global_stride}"
)
if self.placements != other.placements:
diffs.append(
f"placements mismatch: {self.placements} vs {other.placements}"
)
if self.mesh_dim_names != other.mesh_dim_names:
diffs.append(
f"mesh_dim_names mismatch: {self.mesh_dim_names} vs {other.mesh_dim_names}"
)
if self.mesh_layout != other.mesh_layout:
diffs.append(
f"mesh_layout mismatch: {self.mesh_layout} vs {other.mesh_layout}"
)
else:
diffs.append("type: _DTensorMeta vs _TensorMeta")
return diffs
# Type alias for union of tensor metadata types
TensorMeta: TypeAlias = _TensorMeta | _DTensorMeta
# Not frozen: fields are populated incrementally during forward and
# backward metadata inference or from user provided static metadata
@dataclass(slots=True)
class _StageMeta:
"""Consolidated tensor metadata for a pipeline stage's forward and backward passes."""
inputs: tuple[TensorMeta, ...] | None = None
outputs: tuple[TensorMeta, ...] | None = None
input_grads: tuple[TensorMeta | None, ...] | None = None
output_grads: tuple[TensorMeta | None, ...] | None = None
def has_any(self) -> bool:
"""Check if any metadata field is populated."""
return any(
v is not None
for v in [self.inputs, self.outputs, self.input_grads, self.output_grads]
)
def has_dtensors(self) -> bool:
"""Check if any input/output metadata is DTensor type."""
for metas in [self.inputs, self.outputs]:
if metas and any(isinstance(m, _DTensorMeta) for m in metas if m):
return True
return False
def is_complete_for_forward(self) -> bool:
"""Check if forward metadata is fully populated."""
return self.inputs is not None and self.outputs is not None
@dataclass(frozen=True, slots=True)
class _StageForwardMeta:
"""Forward metadata transmitted from stage *i* to stage *i+1* during inference."""
forward_metas: tuple[TensorMeta, ...] # Stage i's outputs → Stage i+1's inputs
@dataclass(frozen=True, slots=True)
class _StageBackwardMeta:
"""Backward metadata transmitted from stage *i* to stage *i-1* during inference.
Gradient placements may differ from forward activations
(e.g., ``Replicate`` → ``Partial``).
"""
backward_metas: tuple[
TensorMeta | None, ...
] # Stage i's input_grads → Stage i-1's output_grads
def _make_tensor_from_meta(
meta: _TensorMeta,
device: torch.device | str,
) -> torch.Tensor:
"""Create a tensor from metadata.
Args:
meta: Metadata with shape, stride, and dtype.
device: Target device for the tensor.
Returns:
Empty tensor preserving the exact memory layout.
"""
return torch.empty_strided(
size=meta.shape,
stride=meta.stride,
dtype=meta.dtype,
device=device,
)
def _derive_grad_metas(
tensor_metas: tuple[TensorMeta, ...],
) -> tuple[_TensorMeta | None, ...]:
"""Derive gradient metadata from tensor metadata.
Returns metadata with the same shape/stride/dtype but ``requires_grad=False``.
Entries where the source has ``requires_grad=False`` become ``None``.
"""
return tuple(
_TensorMeta(shape=m.shape, stride=m.stride, dtype=m.dtype, requires_grad=False)
if m.requires_grad
else None
for m in tensor_metas
)
class _MeshCache:
"""Cache for :class:`DeviceMesh` objects keyed by ``(mesh_dim_names, mesh_layout)``.
Assumes all pipeline stages share the same rank tensor (true for
TorchTitan-style frameworks where meshes derive from a common world).
"""
def __init__(self, get_mesh_cb: GetMeshCallback | None = None) -> None:
self._cache: dict[MeshCacheKey, DeviceMesh] = {}
self._get_mesh_cb = get_mesh_cb
def get_mesh(self, key: MeshCacheKey) -> DeviceMesh:
"""Return a cached mesh, or create one via the callback.
Args:
key: Cache key ``(mesh_dim_names, mesh_layout)``.
Returns:
The ``DeviceMesh``.
Raises:
PipeliningMetadataError: If not cached and no callback provided.
"""
if key in self._cache:
return self._cache[key]
mesh_dim_names, mesh_layout = key
if self._get_mesh_cb is None:
raise PipeliningMetadataError(
f"Mesh not found in cache for mesh_dim_names={mesh_dim_names}, "
f"mesh_layout={mesh_layout}, and no get_mesh callback provided. "
f"Provide a get_mesh callback or use DTensors in static mode."
)
mesh = self._get_mesh_cb(mesh_dim_names, mesh_layout)
if mesh is None:
raise PipeliningMetadataError(
f"Mesh lookup failed: callback returned None for "
f"mesh_dim_names={mesh_dim_names}, mesh_layout={mesh_layout}. "
f"Ensure all stages use meshes from the same universe."
)
self._cache[key] = mesh
return mesh
def put(self, key: MeshCacheKey, mesh: DeviceMesh) -> None:
"""Add a mesh to the cache."""
self._cache[key] = mesh
def update_from_tensors(self, tensors: tuple[torch.Tensor | None, ...]) -> None:
"""Extract and cache meshes from any :class:`DTensor` instances in *tensors*."""
for tensor in tensors:
if isinstance(tensor, DTensor):
mesh = tensor.device_mesh
dim_names = tuple(mesh.mesh_dim_names) if mesh.mesh_dim_names else ()
mesh_layout = mesh._layout
key = (dim_names, mesh_layout)
if key not in self._cache:
self._cache[key] = mesh
def __contains__(self, key: MeshCacheKey) -> bool:
return key in self._cache
def __len__(self) -> int:
return len(self._cache)
# ============================================================================
# Inference mode enum
# ============================================================================
class InferenceMode(Enum):
"""Pipeline-level metadata inference mode, determined collectively across all PP ranks.
The mode is set by the schedule (not individual stages) because
``has_backward`` is only known at schedule creation time and all
stages must agree to avoid P2P hangs.
.. attribute:: STATIC
All stages have sufficient metadata; runtime inference is skipped.
.. attribute:: DYNAMIC
At least one stage requires runtime metadata inference.
"""
STATIC = "static"
DYNAMIC = "dynamic"
@classmethod
def needs_dynamic(cls, meta: _StageMeta, stage_has_backward: bool) -> bool:
"""Determine whether dynamic metadata inference is needed for a stage.
Args:
meta: Stage metadata from user-provided args.
stage_has_backward: Whether a backward pass will be performed.
Returns:
``True`` if dynamic inference is needed.
"""
# Case 1: Forward metadata incomplete → needs DYNAMIC
if not meta.is_complete_for_forward():
return True
# Case 2: No DTensors → STATIC is fine (bwd metadata derivable from fwd metadata)
if not meta.has_dtensors():
return False
# Case 3: No backward needed → STATIC is fine (don't need grad metadata)
if not stage_has_backward:
return False
# Case 4: DTensors with backward but missing ANY grad metadata → needs DYNAMIC
# Both input_grads AND output_grads are required for static mode with DTensors
if meta.input_grads is None or meta.output_grads is None:
return True
# Case 5: DTensors with complete grads → STATIC is fine
return False
# ============================================================================
# Utility functions
# ============================================================================
def flatten_args(args, *, detach: bool = False):
"""Flatten ``args`` into a list, optionally detaching tensors.
Args:
args: Nested arguments to flatten.
detach: If ``True``, detach tensors while preserving ``requires_grad``.
Returns:
``(new_args, flat_detached_args)`` when ``detach=True``;
``flat_args`` list otherwise.
"""
flat_args, treespec = tree_flatten(args)
if detach:
flat_detached = [
a.detach().requires_grad_(a.requires_grad)
if isinstance(a, torch.Tensor)
else a
for a in flat_args
]
new_args = tree_unflatten(flat_detached, treespec)
return new_args, flat_detached
return flat_args
# Backward compatibility alias
def flatten_args_detach(args):
"""Flatten and detach. Deprecated: use ``flatten_args(args, detach=True)``."""
return flatten_args(args, detach=True)
def generate_stage_to_rank_mapping(
pp_size: int, num_stages: int, style: str = "loop"
) -> dict[int, int]:
"""
Compute the stage id to rank mapping for either a looped or V-style schedule.
Most commonly num_stages == pp_size * 2, but this function can be used to
compute the mapping for any number of stages per rank.
"""
mapping = {}
if style == "loop":
for stage_index in range(num_stages):
mapping[stage_index] = stage_index % pp_size
elif style == "v":
if num_stages % pp_size != 0:
raise ValueError(
f"num_stages {num_stages} must be evenly divisible by pp_size {pp_size} for V schedules"
)
rank_index = 0
for stage_index in range(num_stages):
mapping[stage_index] = rank_index
# dont change rank if we are on the border (to keep v shape)
if (stage_index + 1) % pp_size == 0:
continue
if (stage_index // pp_size) % 2 == 0:
rank_index += 1
else:
rank_index -= 1
else:
raise ValueError(f"Style {style} is not supported.")
return mapping
def generate_rank_to_stage_mapping(
pp_size: int, num_stages: int, style: str = "loop"
) -> dict[int, list[int]]:
"""
Compute the rank to stage id mapping for either a looped or V-style schedule.
This function inverts the stage_to_rank_mapping to get which stages are assigned to each rank.
Returns a dictionary mapping rank -> list of stage indices assigned to that rank.
"""
stage_to_rank = generate_stage_to_rank_mapping(pp_size, num_stages, style)
# Invert the mapping: rank -> list of stages
rank_to_stages: dict[int, list[int]] = {}
for stage_id, rank in stage_to_rank.items():
if rank not in rank_to_stages:
rank_to_stages[rank] = []
rank_to_stages[rank].append(stage_id)
# Sort the stage lists for each rank to ensure consistent ordering
for stages in rank_to_stages.values():
stages.sort()
return rank_to_stages
@dataclass(slots=True)
class PipeInfo:
"""
Captures information for a pipeline (`Pipe` object).
"""
graph: fx.Graph
num_stages: int
has_loss_and_backward: bool
# ============================================================================
# Metadata extraction helpers
# ============================================================================
def extract_tensor_meta(tensor: torch.Tensor) -> TensorMeta:
"""Extract metadata from a tensor.
Handles both plain Tensor and DTensor correctly: DTensors are
dispatched to ``_DTensorMeta.from_dtensor`` which captures local
shard attributes plus global shape/placement info, while plain
tensors use ``_TensorMeta.from_tensor``.
Args:
tensor: A plain tensor or DTensor.
Returns:
``_TensorMeta`` for plain tensors, ``_DTensorMeta`` for DTensors.
"""
if isinstance(tensor, DTensor):
return _DTensorMeta.from_dtensor(tensor)
else:
return _TensorMeta.from_tensor(tensor)
@overload
def extract_tensor_metas(
tensors: tuple[torch.Tensor, ...] | None,
*,
allow_none: Literal[False] = ...,
) -> tuple[TensorMeta, ...] | None: ...
@overload
def extract_tensor_metas(
tensors: tuple[torch.Tensor | None, ...] | None,
*,
allow_none: Literal[True],
) -> tuple[TensorMeta | None, ...] | None: ...
def extract_tensor_metas(
tensors: tuple[torch.Tensor | None, ...] | tuple[torch.Tensor, ...] | None,
*,
allow_none: bool = False,
) -> tuple[TensorMeta | None, ...] | None:
"""Extract metadata from a tuple of tensors.
Args:
tensors: Tuple of tensors (may include ``None`` when ``allow_none=True``).
allow_none: If ``True``, preserve ``None`` elements (for gradients).
Returns:
Tuple of ``TensorMeta``, or ``None`` if ``tensors`` is ``None``.
Raises:
PipeliningMetadataError: If ``None`` found and ``allow_none=False``.
"""
if tensors is None:
return None
metas_with_none: list[TensorMeta | None] = []
has_none = False
for t in tensors:
if isinstance(t, torch.Tensor):
metas_with_none.append(extract_tensor_meta(t))
else:
has_none = True
metas_with_none.append(None)
if not allow_none and has_none:
raise PipeliningMetadataError(
"None values are not allowed in tensor metadata tuples. "
"Use allow_none=True for optional values."
)
return tuple(metas_with_none)
def to_local_if_dtensor(tensor: torch.Tensor, detach: bool = False) -> torch.Tensor:
"""Convert a DTensor to its local shard, or return a plain tensor as-is.
When ``detach=True``, the tensor is detached before conversion —
this applies to both DTensors and plain tensors.
Args:
tensor: A tensor that may be a DTensor.
detach: If ``True``, detach before ``to_local()`` to avoid
redistribution during backward.
Returns:
The local tensor component.
"""
maybe_detached_tensor = tensor.detach() if detach else tensor
if isinstance(maybe_detached_tensor, DTensor):
return maybe_detached_tensor.to_local()
return maybe_detached_tensor
@overload
def validate_and_normalize_to_tuple(
args: torch.Tensor | tuple[torch.Tensor, ...] | list[torch.Tensor] | None,
allow_none: Literal[False] = ...,
) -> tuple[torch.Tensor, ...] | None: ...
@overload
def validate_and_normalize_to_tuple(
args: torch.Tensor
| tuple[torch.Tensor | None, ...]
| list[torch.Tensor | None]
| None,
allow_none: Literal[True] = ...,
) -> tuple[torch.Tensor | None, ...] | None: ...
def validate_and_normalize_to_tuple(
args: torch.Tensor
| tuple[torch.Tensor, ...]
| tuple[torch.Tensor | None, ...]
| list[torch.Tensor]
| list[torch.Tensor | None]
| None,
allow_none: bool = False,
) -> tuple[torch.Tensor | None, ...] | tuple[torch.Tensor, ...] | None:
"""Normalize ``args`` to a tuple and validate that all elements are tensors.
Args:
args: A single tensor, tuple/list of tensors, or ``None``.
allow_none: If ``True``, permit ``None`` elements (for gradients).
Returns:
Tuple of tensors, or ``None`` if ``args`` is ``None``.
Raises:
PipeliningMetadataError: On non-tensor values
(or ``None`` when ``allow_none=False``).
"""
if args is None:
return None
elif isinstance(args, torch.Tensor):
return (args,)
elif isinstance(args, (tuple, list)):
for i, arg in enumerate(args):
if arg is None:
if not allow_none:
raise PipeliningMetadataError(
f"Stage arg[{i}] is None. "
f"Stage args must be tensors. Use kwargs for optional values."
)
continue
if not isinstance(arg, torch.Tensor):
raise PipeliningMetadataError(
f"Stage arg[{i}] has type {type(arg).__name__}. "
f"All stage args must be tensors. Use kwargs for non-tensor inputs."
)
# Normalize list to tuple
return tuple(args) if isinstance(args, list) else args
else:
raise PipeliningMetadataError(
f"Stage args must be a tensor, tuple, or list of tensors, got {type(args).__name__}."
)
# ============================================================================
# Validation functions
# ============================================================================
def validate_metadata(
desc: str,
expected: TensorMeta,
actual: torch.Tensor | TensorMeta,
*,
raise_on_mismatch: bool = False,
warn_on_mismatch: bool = False,
) -> list[str]:
"""
Compare expected metadata against actual tensor or metadata.
This is the unified validation/comparison function that uses get_diff() from
metadata classes. Works with both plain tensors and DTensors.
For plain tensors: compares shape/stride/dtype/requires_grad.
For DTensors: compares all properties including global shape and placements.
Args:
desc: Description for error/warning messages.
expected: Expected tensor metadata (_TensorMeta or _DTensorMeta).
actual: Actual tensor or metadata to compare against.
raise_on_mismatch: If True, raise PipeliningMetadataError on mismatch.
warn_on_mismatch: If True, issue a warning on mismatch.
Returns:
List of differences (empty if metadata matches).
Raises:
PipeliningMetadataError: If raise_on_mismatch=True and differences exist.
"""
# Extract metadata if actual is a tensor
if isinstance(actual, torch.Tensor):
actual_meta = extract_tensor_meta(actual)
else:
actual_meta = actual
# Type check: ensure both are same type for meaningful comparison
if type(expected) is not type(actual_meta):
type_diff = [
f"type: expected {type(expected).__name__}, got {type(actual_meta).__name__}"
]
if raise_on_mismatch:
raise PipeliningMetadataError(f"{desc}: {type_diff[0]}")
if warn_on_mismatch:
warnings.warn(
f"{desc}: Metadata type mismatch. {type_diff[0]}. "
f"Using dynamically inferred metadata instead.",
UserWarning,
stacklevel=2,
)
return type_diff
# Use get_diff() from the metadata class
diffs = expected.get_diff(actual_meta)
if diffs:
if raise_on_mismatch:
raise PipeliningMetadataError(f"{desc}: {'; '.join(diffs)}")
if warn_on_mismatch:
warnings.warn(
f"{desc}: Metadata mismatch. {'; '.join(diffs)}. "
f"Using dynamically inferred metadata instead.",
UserWarning,
stacklevel=2,
)
return diffs
def validate_tensors_metadata(
desc: str,
expected: tuple[TensorMeta | None, ...],
actual: tuple[torch.Tensor | TensorMeta | None, ...],
*,
raise_on_mismatch: bool = True,
warn_on_mismatch: bool = False,
) -> list[str]:
"""Validate metadata for a tuple of tensors element-wise.
Args:
desc: Description prefix for error/warning messages.
expected: Tuple of expected metadata (may include ``None`` for grads).
actual: Tuple of actual tensors or metadata to compare against.
raise_on_mismatch: If ``True``, raise on the first mismatch.
warn_on_mismatch: If ``True``, issue warnings for mismatches.
Returns:
Aggregated list of difference strings.
Raises:
PipeliningMetadataError: If lengths differ or on mismatch.
"""
if len(expected) != len(actual):
msg = f"{desc}: expected {len(expected)} tensors, got {len(actual)}"
if raise_on_mismatch:
raise PipeliningMetadataError(msg)
if warn_on_mismatch:
warnings.warn(msg, UserWarning, stacklevel=2)
return [msg]
all_diffs: list[str] = []
for i, (exp, act) in enumerate(zip(expected, actual, strict=True)):
if exp is None and act is None:
continue
if exp is None or act is None:
msg = (
f"{desc}[{i}]: expected {'None' if exp is None else 'metadata'}, "
f"got {'None' if act is None else 'metadata'}"
)
if raise_on_mismatch:
raise PipeliningMetadataError(msg)
if warn_on_mismatch:
warnings.warn(msg, UserWarning, stacklevel=2)
all_diffs.append(msg)
continue
diffs = validate_metadata(
f"{desc}[{i}]",
exp,
act,
raise_on_mismatch=raise_on_mismatch,
warn_on_mismatch=warn_on_mismatch,
)
all_diffs.extend(diffs)
return all_diffs
def validate_static_arg_grad_correspondence(
stage_index: int,
args: tuple[torch.Tensor, ...],
grads: tuple[torch.Tensor | None, ...],
is_input: bool,
) -> None:
"""
Validate the args↔grads contract for static mode.
Enforces four rules for each (arg, grad) pair:
1. len(args) must equal len(grads).
2. If arg.requires_grad is False, grad must be None.
3. If arg.requires_grad is True and grad is None, emit a warning
(this is legal at pipeline boundaries but may indicate a bug).
4. If arg is a DTensor with requires_grad=True and grad is not None,
grad must also be a DTensor.
Args:
stage_index: The stage index for error messages.
args: Tuple of forward tensors.
grads: Tuple of gradient tensors (can include None).
is_input: True for input_args/input_grads, False for output_args/output_grads.
Raises:
PipeliningMetadataError: If any hard rule (1, 2, or 4) is violated.
"""
kind = "input" if is_input else "output"
args_name = f"{kind}_args"
grads_name = f"{kind}_grads"
# Rule 1: lengths must match
if len(args) != len(grads):
raise PipeliningMetadataError(
f"Stage {stage_index}: {grads_name} length ({len(grads)}) does not match "
f"{args_name} length ({len(args)}). Each forward tensor must have a "
f"corresponding gradient entry (use None for tensors that don't require grad)."
)
for i, (arg, grad) in enumerate(zip(args, grads, strict=True)):
# Rule 2: no grad for a non-differentiable arg
if not arg.requires_grad and grad is not None:
raise PipeliningMetadataError(
f"Stage {stage_index}: {args_name}[{i}] has requires_grad=False, "
f"but {grads_name}[{i}] is not None ({type(grad).__name__}). "
f"Non-differentiable tensors must have None as their gradient entry."
)
# Rule 3: missing grad for a differentiable arg (warn, don't raise)
if arg.requires_grad and grad is None:
warnings.warn(
f"Stage {stage_index}: {args_name}[{i}] has requires_grad=True, "
f"but {grads_name}[{i}] is None. This is legal at pipeline boundaries "
f"but may indicate a missing gradient.",
UserWarning,
stacklevel=2,
)
# Rule 4: DTensor arg must have DTensor grad
if (
isinstance(arg, DTensor)
and arg.requires_grad
and grad is not None
and not isinstance(grad, DTensor)
):
raise PipeliningMetadataError(
f"Stage {stage_index}: {args_name}[{i}] is a DTensor with requires_grad=True, "
f"but {grads_name}[{i}] is {type(grad).__name__}, expected DTensor or None. "
f"DTensor gradients may have different placements than forward tensors."
)
kronos-venv/lib/python3.12/site-packages/torch/distributed/pipelining/microbatch.py
+621
View File
@@ -0,0 +1,621 @@
# mypy: allow-untyped-defs
# Copyright (c) Meta Platforms, Inc. and affiliates
import logging
import operator
from collections.abc import Sequence
from typing import Any
import torch
from torch.distributed.tensor import DTensor
from torch.distributed.tensor.experimental import local_map
from torch.fx.node import map_aggregate
from torch.nn.attention.flex_attention import BlockMask
from torch.utils._pytree import tree_flatten, tree_map, tree_unflatten
__all__ = [
"TensorChunkSpec",
"split_args_kwargs_into_chunks",
"merge_chunks",
]
logger = logging.getLogger(__name__)
"""
_debug_mask_minibatches specifies to send masked versions of the mini-batch
through instead of micro-batch slices--this can be used for more stable
numerical testing (see [A Note About Correctness Testing])
"""
_debug_mask_minibatches = False
class _CustomReducer:
"""
Custom reducer class that can be used to specify a custom operation that
reduces losses of multiple microbatches into one value.
Example:
>>> # xdoctest: +SKIP
>>> sum_reducer = _CustomReducer(
>>> torch.tensor(0.0),
>>> lambda a, b: a + b
>>> )
"""
def __init__(self, init_value, reduce_fn):
self.init_value = init_value
self.reduce_fn = reduce_fn
class _LossReducer(_CustomReducer):
pass
sum_reducer = _LossReducer(torch.tensor(0.0), operator.add)
# Default chunking dimension is 0. This is used for the case where the user did
# not specify a chunking dimension.
DEFAULT_CHUNK_DIM = 0
class TensorChunkSpec:
"""
Class used to specify chunking of inputs
"""
def __init__(self, split_dim):
self.split_dim = split_dim
split_dim: int
def __repr__(self):
return (
f"{self.__class__.__module__}.{self.__class__.__name__}({self.split_dim})"
)
def __str__(self):
return f"TensorChunkSpec({self.split_dim})"
@staticmethod
def from_tuple(
chunk_dims: tuple[int, ...],
):
"""
A helper for creating a tuple of `TensorChunkSpec` from a tuple of chunk
dimensions (int's).
Example:
>>> # xdoctest: +SKIP
>>> # There are three positional arguments to the model, and
>>> # we are chunking them along dimension 0, 0 and 1, respectively
>>> args_chunk_spec = TensorChunkSpec.from_tuple((0, 0, 1))
"""
args_chunk_spec = map_aggregate(
chunk_dims,
lambda dim: TensorChunkSpec(dim), # type: ignore[arg-type,return-value]
)
return args_chunk_spec
@staticmethod
def from_dict(
chunk_dims: dict[str, int],
):
"""
A helper for creating a dictionary of `TensorChunkSpec` from a
dictionary of chunk dimensions (int's).
Example:
>>> # xdoctest: +SKIP
>>> # Chunk dimension 0 for the "id" argument, 1 for the "mask" argument
>>> kwargs_chunk_spec = TensorChunkSpec.from_dict({"id": 0, "mask": 1})
"""
kwargs_chunk_spec = map_aggregate(
chunk_dims,
lambda dim: TensorChunkSpec(dim), # type: ignore[arg-type,return-value]
)
return kwargs_chunk_spec
# Class used to specify replication of inputs
class _Replicate:
pass
def _split_block_mask(
block_mask: BlockMask,
num_chunks: int,
) -> list[BlockMask]:
"""Given a block mask, split the block mask along the batch dimension (dim0).
Args:
block_mask: Block mask to split
num_chunks: Number of chunks to split the block mask into
Returns:
chunk_block_masks: List of chunked block masks
"""
# BlockMask will broadcast if B is 1.
if block_mask.kv_num_blocks.size(0) == 1:
return [block_mask] * num_chunks
if not block_mask.kv_num_blocks.size(0) >= num_chunks:
raise AssertionError(
"Block mask has fewer batch size than the number of chunks. "
)
batch_dim = 0
kv_num_blocks_chunks = torch.tensor_split(
block_mask.kv_num_blocks, num_chunks, batch_dim
)
kv_indices_chunks = torch.tensor_split(block_mask.kv_indices, num_chunks, batch_dim)
full_kv_num_blocks_chunks = (
torch.tensor_split(block_mask.full_kv_num_blocks, num_chunks, batch_dim)
if block_mask.full_kv_num_blocks is not None
else [None] * num_chunks
)
full_kv_indices_chunks = (
torch.tensor_split(block_mask.full_kv_indices, num_chunks, batch_dim)
if block_mask.full_kv_indices is not None
else [None] * num_chunks
)
chunk_block_masks = []
batch_offset = 0
for chunk_idx in range(num_chunks):
def create_mask_mod(idx):
def batch_offset_mask_mod(b, h, q_idx, kv_idx):
b_offset = torch.full_like(b, idx)
return block_mask.mask_mod(b + b_offset, h, q_idx, kv_idx)
return batch_offset_mask_mod
chunk_block_masks.append(
BlockMask.from_kv_blocks(
kv_num_blocks=kv_num_blocks_chunks[chunk_idx],
kv_indices=kv_indices_chunks[chunk_idx],
full_kv_num_blocks=full_kv_num_blocks_chunks[chunk_idx],
full_kv_indices=full_kv_indices_chunks[chunk_idx],
BLOCK_SIZE=block_mask.BLOCK_SIZE,
mask_mod=create_mask_mod(batch_offset),
seq_lengths=block_mask.seq_lengths,
)
)
batch_offset += kv_num_blocks_chunks[chunk_idx].size(0)
return chunk_block_masks
def _split_tensor(
tensor: torch.Tensor,
spec: TensorChunkSpec,
num_chunks: int,
) -> Sequence[torch.Tensor]:
"""Given a tensor, and a chunking spec, split the tensor.
Args:
tensor: Tensor to split
spec: Chunking spec
num_chunks: Number of chunks to split the tensor into
Returns:
chunk_tensors: List of chunked tensors
"""
if not tensor.size(spec.split_dim) >= num_chunks:
raise AssertionError(
f"Tensor size {tensor.size(spec.split_dim)} is smaller than num_chunks"
)
_is_dtensor = isinstance(tensor, DTensor)
if _is_dtensor:
# Use local_map to split locally and preserve placements.
# Going through DTensor dispatch would convert Shard(split_dim) to
# Replicate() via an implicit all-gather, which is both wasteful and
# semantically wrong for PP microbatch splitting.
placements = tensor.placements
split_fn = local_map(
lambda t: torch.tensor_split(t, num_chunks, spec.split_dim),
out_placements=(placements,) * num_chunks,
in_placements=(placements,),
)
chunk_tensors: Sequence[torch.Tensor] = split_fn(tensor) # type: ignore[assignment]
else:
chunk_tensors = torch.tensor_split(tensor, num_chunks, spec.split_dim)
# tensor_split on a leaf tensor produces non-leaf views that won't
# accumulate .grad during torch.autograd.backward(). Call retain_grad()
# on those views so that stage_backward() can read .grad from them.
if tensor.requires_grad and tensor.is_leaf:
for chunk in chunk_tensors:
chunk.retain_grad()
if not _debug_mask_minibatches:
return chunk_tensors
def _expand_chunks(
orig: torch.Tensor, *chunks: torch.Tensor
) -> tuple[torch.Tensor, ...]:
expanded = []
idx = 0
for chunk in chunks:
new_val = torch.zeros_like(orig)
upper = idx + chunk.size(spec.split_dim)
slices: list[slice] = [slice(None)] * new_val.ndim
slices[spec.split_dim] = slice(idx, upper)
new_val[slices] = chunk
expanded.append(new_val)
idx += chunk.size(spec.split_dim)
return tuple(expanded)
if _is_dtensor:
placements = tensor.placements
n = len(chunk_tensors)
expand_fn = local_map(
_expand_chunks,
out_placements=(placements,) * n,
in_placements=(placements,) + (placements,) * n,
)
return list(expand_fn(tensor, *chunk_tensors)) # type: ignore[arg-type]
else:
return list(_expand_chunks(tensor, *chunk_tensors))
def _shard_dict_of_args(
args_dict,
args_chunk_spec,
num_chunks,
):
"""
Given a dictionary of args, and a dictionary of chunking specs, shard the
args according to the chunking specs.
Args:
args_dict: Dictionary of args
args_chunk_spec: Dictionary of chunking specs
num_chunks: Number of chunks to shard the args into
Returns:
args_split: List of sharded args
"""
if not args_dict:
return [{} for _ in range(num_chunks)]
if not len(args_dict) == len(args_chunk_spec):
raise AssertionError(
f"args_dict.keys() = {list(args_dict.keys())} "
f"args_chunk_spec.keys() = {list(args_chunk_spec.keys())}"
)
if args_chunk_spec is None:
raise AssertionError("args_chunk_spec should have been set by caller")
values, tree_spec = tree_flatten(
args_dict, is_leaf=lambda x: isinstance(x, BlockMask)
)
chunk_specs, _ = tree_flatten(
args_chunk_spec, is_leaf=lambda x: isinstance(x, BlockMask)
)
# First check and find the actual number of chunks
split_sizes = []
for v, spec in zip(values, chunk_specs, strict=True):
# The original logic is "spec is _Replicate". This doesn't seem to be
# correct. But we keep it for backward compatibility.
if spec is _Replicate or isinstance(spec, _Replicate):
split_sizes.append(num_chunks)
elif isinstance(v, torch.Tensor):
if not isinstance(spec, TensorChunkSpec):
raise AssertionError(f"Expected TensorChunkSpec, got {type(spec)}")
split_sizes.append(v.size(spec.split_dim))
elif isinstance(v, BlockMask):
if not isinstance(spec, TensorChunkSpec):
raise AssertionError(f"Expected TensorChunkSpec, got {type(spec)}")
if not spec.split_dim == 0:
raise AssertionError("BlockMask only supports split_dim=0")
# BlockMask will broadcast if B is 1.
if v.kv_num_blocks.size(0) == 1:
split_sizes.append(num_chunks)
else:
split_sizes.append(v.kv_num_blocks.size(0))
else:
raise ValueError(
f"Unsupported chunk spec: {spec} and value: {v} combination."
)
result_num_chunks = min(*split_sizes, num_chunks)
flat_split_results: list[Any] = [[] for _ in range(result_num_chunks)]
for v, spec in zip(values, chunk_specs, strict=True):
v_splits: Sequence[Any] = []
if spec is _Replicate or isinstance(spec, _Replicate):
v_splits = [v] * result_num_chunks
elif isinstance(v, torch.Tensor):
v_splits = _split_tensor(v, spec, result_num_chunks)
elif isinstance(v, BlockMask):
v_splits = _split_block_mask(v, result_num_chunks)
else:
raise ValueError(
f"Unsupported chunk spec: {spec} and value: {v} combination."
)
for _flat_split_result, _v_split in zip(
flat_split_results, v_splits, strict=True
):
_flat_split_result.append(_v_split)
return [
tree_unflatten(_flat_split_result, tree_spec)
for _flat_split_result in flat_split_results
]
def split_args_kwargs_into_chunks(
args: tuple[Any, ...],
kwargs: dict[str, Any] | None,
chunks: int,
args_chunk_spec: tuple[TensorChunkSpec, ...] | None = None,
kwargs_chunk_spec: dict[str, TensorChunkSpec] | None = None,
) -> tuple[list[tuple], list[dict]]:
"""
Given a sequence of args and kwargs, split them into a number of chunks
according to their respective chunking specs.
Args:
args: Tuple of args
kwargs: Dict of kwargs
chunks: Number of chunks to split the args and kwargs into
args_chunk_spec: chunking specs for args, in same shape as args
kwargs_chunk_spec: chunking specs for kwargs, in same shape as kwargs
Returns:
args_split: List of sharded args
kwargs_split: List of sharded kwargs
"""
# Given `args` and `kwargs`, we want to yield a set of `chunks` args and kwargs such that
# the constituent Tensor values have been sharded/replicated according to the `args_chunk_spec`
# and `kwargs_chunk_spec` specifications. The steps are as follows:
#
# 1. Use pytree.tree_flatten to flatten each arg and its spec into nto a 1d array of values.
# To use a running example: suppose our inputs look like
#
# args = ([A, [B, C]], D) args_spec = ([None, [None, TensorChunkSpec]], None)
# (kwargs not shown but it's a similar process)
#
# Then for this step we would end up with
#
# args = ([A, B, C], D) args_spec = ([None, None, TensorChunkSpec], None)
#
# 2. Shard or replicate the arguments subject to the policy in the spec. Suppose chunks = 2
#
# args = ([[A, A], [B, B], [C_1, C_2]], [D, D])
#
# 3. Rotate the nesting order such that chunks are the outer dimension
#
# args_chunks = [
# ([A, B, C_1], D),
# ([A, B, C_2], D),
# ]
#
# 4. Unflatten each chunk according to the spec
#
# args_chunks = [
# ([A, [B, C_1]], D),
# ([A, [B, C_2]], D),
# ]
# TODO: _debug_mask_minibatches
# Handle the case where kwargs is None
if kwargs is None:
kwargs = {}
# If user did not provide args_chunk_spec or kwargs_chunk_spec, we extend
# their format and use default chunking along dim 0
def default_spec(v):
if isinstance(v, torch.Tensor | BlockMask):
return TensorChunkSpec(DEFAULT_CHUNK_DIM)
else:
return _Replicate()
if args_chunk_spec is None:
args_chunk_spec = tree_map(
default_spec, args, is_leaf=lambda v: isinstance(v, BlockMask)
)
if kwargs_chunk_spec is None:
kwargs_chunk_spec = tree_map(
default_spec, kwargs, is_leaf=lambda v: isinstance(v, BlockMask)
)
args_split_dict = _shard_dict_of_args(
dict(enumerate(args)),
dict(enumerate(args_chunk_spec)),
chunks,
)
real_num_chunks = len(args_split_dict)
kwargs_split = _shard_dict_of_args(
kwargs,
kwargs_chunk_spec,
real_num_chunks,
)
if len(kwargs_split) < real_num_chunks:
# In case kwargs are sharded into less chunks
# e.g. when `args` has no tensor, just values
real_num_chunks = len(kwargs_split)
# Re-shard args
args_split_dict = _shard_dict_of_args(
dict(enumerate(args)),
dict(enumerate(args_chunk_spec)),
real_num_chunks,
)
if len(args_split_dict) != len(kwargs_split):
raise RuntimeError(
"args and kwargs are split into different number of chunks: "
f"{len(args_split_dict)}, {len(kwargs_split)}"
)
args_split = [
tuple(chunk_args[i] for i in range(len(chunk_args)))
for chunk_args in args_split_dict
]
return args_split, kwargs_split
def merge_chunks(
chunks: list[Any],
chunk_spec,
):
"""
Given a list of chunks, merge them into a single value according to
the chunk spec.
Args:
chunks: list of chunks
chunk_spec: Chunking spec for the chunks
Returns:
value: Merged value
"""
# This is essentially the inverse of `split_args_kwargs_into_chunks`, so the
# steps are similar to the steps in that function but in reverse. Given the
# input values:
#
# chunks = [
# ([A, [B, C_1]], D),
# ([A, [B, C_2]], D),
# ]
# args_spec = ([None, [None, TensorChunkSpec]], None)
#
# 1. Flatten the chunks according to the chunk_spec
#
# chunks_flat = [
# ([A, B, C_1], D),
# ([A, B, C_2], D),
# ]
#
# 2. Rotate the nesting order such that chunks are the inner dimension
#
# value_inner = ([A, B, [C_1, C_2]], D)
#
# 3. Concatenate sharded arguments
#
# value_combined = ([A, B, C], D)
#
# 4. Unflatten the combined args given the spec
#
# value = ([A, [B, C]], D)
# Preliminary: flatten the chunk spec
if chunk_spec is not None:
spec_flattened, flatten_spec = tree_flatten(chunk_spec)
else:
# If chunk_spec is not provided, we will merge chunks along the default dimension (0), for all output fields
# We obtain the output structure by flattening chunk 0 and generate the chunk_spec
chunk0_flat, flatten_spec = tree_flatten(chunks[0])
spec_flattened = [TensorChunkSpec(DEFAULT_CHUNK_DIM)] * len(chunk0_flat)
# Stage 1: flatten chunks
# chunks_flattened : [num chunks, num args]
chunks_flattened = []
for chunk in chunks:
chunk_flattened, _ = tree_flatten(chunk)
if len(chunk_flattened) != len(spec_flattened):
raise ValueError(f"Chunk {chunk} did not match chunk spec {chunk_spec}")
chunks_flattened.append(chunk_flattened)
# Stage 2 and 3: Rotate nesting order s.t. chunks are inner dimension and
# concatenate sharded operands
# args_flattened : [num args]
args_flattened = []
for arg_idx, arg in enumerate(spec_flattened):
if isinstance(arg, TensorChunkSpec):
partial_values = [
chunks_flattened[chunk_idx][arg_idx]
for chunk_idx in range(len(chunks_flattened))
]
if _debug_mask_minibatches:
# Infer size of individual chunks by running `tensor_split` again
overall_shape = partial_values[0].shape
for val in partial_values[1:]:
if not val.shape == overall_shape:
raise AssertionError(
f"Expected shape {overall_shape}, got {val.shape}"
)
meta_chunks = torch.tensor_split(
torch.empty(*overall_shape, device="meta"),
sections=len(partial_values),
dim=arg.split_dim,
)
values_to_cat = []
chunk_start_idx = 0
if not len(partial_values) == len(meta_chunks):
raise AssertionError(
f"Expected len(partial_values) == len(meta_chunks), got {len(partial_values)} != {len(meta_chunks)}"
)
for partial_value, meta_chunk in zip(
partial_values, meta_chunks, strict=True
):
chunk_end_idx = chunk_start_idx + meta_chunk.size(arg.split_dim)
slice_indices = [slice(None, None, None)] * partial_value.ndim
slice_indices[arg.split_dim] = slice(chunk_start_idx, chunk_end_idx)
sliced = partial_value[slice_indices]
values_to_cat.append(sliced)
chunk_start_idx = chunk_end_idx
else:
values_to_cat = partial_values
# Validate DTensor consistency: either all values are DTensors
# or none are. A mix indicates a bug in the pipeline stage.
dtensor_flags = [isinstance(v, DTensor) for v in values_to_cat]
if any(dtensor_flags):
if not all(dtensor_flags):
raise AssertionError(
"merge_chunks: expected all values to be DTensors or "
"none to be DTensors, got a mix"
)
# All DTensors must have matching placements.
placements = values_to_cat[0].placements
for i, v in enumerate(values_to_cat[1:], 1):
if v.placements != placements:
raise AssertionError(
f"merge_chunks: placement mismatch at chunk {i}: "
f"expected {placements}, got {v.placements}"
)
cat_fn = local_map(
lambda *chunks: torch.cat(chunks, dim=arg.split_dim),
out_placements=(placements,),
in_placements=tuple(placements for _ in range(len(values_to_cat))),
)
args_flattened.append(cat_fn(*values_to_cat))
else:
args_flattened.append(torch.cat(values_to_cat, dim=arg.split_dim))
elif isinstance(arg, _CustomReducer):
reduced_val = arg.init_value
for chunk_idx in range(len(chunks_flattened)):
reduced_val = arg.reduce_fn(
reduced_val, chunks_flattened[chunk_idx][arg_idx]
)
args_flattened.append(reduced_val)
else:
value = chunks_flattened[0][arg_idx]
for chunk_idx in range(1, len(chunks_flattened)):
if not chunks_flattened[chunk_idx][arg_idx] == value:
raise AssertionError(
f"Expected {value}, got {chunks_flattened[chunk_idx][arg_idx]}"
)
args_flattened.append(value)
# Stage 4: Unflatten combined args
return tree_unflatten(args_flattened, flatten_spec)
kronos-venv/lib/python3.12/site-packages/torch/distributed/pipelining/schedules.py
+3582
View File
File diff suppressed because it is too large Load Diff
kronos-venv/lib/python3.12/site-packages/torch/distributed/pipelining/stage.py
+2249
View File
File diff suppressed because it is too large Load Diff