[AutoSP] Fix compiler correctness

deepspeed/compile/custom_ops/__init__.py

@@ -3,10 +3,11 @@
 
 # DeepSpeed Team
 
-from .all_to_all import all_to_all
+from .all_to_all import all_gather_sequence, all_to_all, aggregate_loss
 from .tp_collectives import copy_to_tp_region, gather_from_tp_region, reduce_from_tp_region
 from . import sp_dp_registry
 
 __all__ = [
-    "all_to_all", "copy_to_tp_region", "gather_from_tp_region", "reduce_from_tp_region", "sp_dp_registry", "sp_compat"
+    "all_gather_sequence", "all_to_all", "aggregate_loss", "copy_to_tp_region", "gather_from_tp_region",
+    "reduce_from_tp_region", "sp_dp_registry", "sp_compat"
 ]

deepspeed/compile/custom_ops/all_to_all.py

@@ -3,6 +3,8 @@
 
 # DeepSpeed Team
 
+from typing import Tuple
+
 import torch
 import deepspeed.comm as dist
 from torch.utils._sympy.functions import FloorDiv
@@ -31,6 +33,9 @@ def all_to_all(
 
     if scatter_idx == 1:
         N, local_S = dim1, dim2
+        if N % sp_size() != 0:
+            raise ValueError(f"AutoSP requires the Q/K/V head count ({N}) to be divisible by "
+                             f"sequence_parallel_size ({sp_size()})")
         input_t = input.reshape(B, sp_size(), N // sp_size(), local_S, H)
         input_t = input_t.permute(1, 0, 2, 3, 4).contiguous()
 
@@ -90,3 +95,58 @@ def _all_to_all_backward(ctx, grad):
 
 
 torch.library.register_autograd("autosp::all_to_all", _all_to_all_backward, setup_context=_all_to_all_backward_setup)
+
+
+@torch.library.custom_op("autosp::all_gather_sequence", mutates_args=())
+def all_gather_sequence(input: torch.Tensor, dim: int) -> torch.Tensor:
+    """Gather a local attention-mask dimension across the current SP group."""
+    assert is_setup(), 'Incorrect initialization of SP/DP mesh.'
+    gid = dist.get_rank() // sp_size()
+    group = get_group(gid)
+    outputs = [torch.empty_like(input) for _ in range(sp_size())]
+    dist.all_gather(outputs, input, group=group)
+    return torch.cat(outputs, dim=dim)
+
+
+@torch.library.register_fake("autosp::all_gather_sequence")
+def all_gather_sequence_fake(input: torch.Tensor, dim: int):
+    output_shape = list(input.shape)
+    output_shape[dim] *= sp_size()
+    return input.new_empty(output_shape)
+
+
+@torch.library.custom_op("autosp::aggregate_loss", mutates_args=())
+def aggregate_loss(loss: torch.Tensor, valid_tokens: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+    """Return the valid-token-weighted mean loss on every SP rank."""
+    assert is_setup(), 'Incorrect initialization of SP/DP mesh.'
+    gid = dist.get_rank() // sp_size()
+    group = get_group(gid)
+    finite_loss = torch.where(valid_tokens > 0, loss, torch.zeros_like(loss))
+    total = finite_loss * valid_tokens.to(loss.dtype)
+    total_tokens = valid_tokens.clone()
+    dist.all_reduce(total, group=group)
+    dist.all_reduce(total_tokens, group=group)
+    weight = valid_tokens.to(loss.dtype) / total_tokens.clamp_min(1).to(loss.dtype)
+    return total / total_tokens.clamp_min(1).to(loss.dtype), weight
+
+
+@torch.library.register_fake("autosp::aggregate_loss")
+def aggregate_loss_fake(loss: torch.Tensor, valid_tokens: torch.Tensor):
+    return torch.empty_like(loss), torch.empty_like(loss)
+
+
+def _aggregate_loss_backward_setup(ctx, inputs, output):
+    _, weight = output
+    ctx.mark_non_differentiable(weight)
+    ctx.save_for_backward(weight)
+
+
+def _aggregate_loss_backward(ctx, grad_loss, grad_weight):
+    del grad_weight
+    (weight, ) = ctx.saved_tensors
+    return grad_loss * weight, None
+
+
+torch.library.register_autograd("autosp::aggregate_loss",
+                                _aggregate_loss_backward,
+                                setup_context=_aggregate_loss_backward_setup)

deepspeed/compile/custom_ops/sp_dp_registry.py

@@ -50,7 +50,17 @@ def dp_size():
 def populate_registry(SP_SIZE, DP_SIZE):
     """ Populate rank to SP/DP mesh index.  """
 
+    world_size = dist.get_world_size()
+    if SP_SIZE * DP_SIZE != world_size:
+        raise ValueError(f"AutoSP mesh ({SP_SIZE} x {DP_SIZE}) must cover the distributed world size ({world_size})")
+
     if GROUP_REGISTRY.get('is_reg', False):
+        current_mesh = (GROUP_REGISTRY['SP_SIZE'], GROUP_REGISTRY['DP_SIZE'])
+        requested_mesh = (SP_SIZE, DP_SIZE)
+        if current_mesh != requested_mesh:
+            raise RuntimeError(f"AutoSP process groups are already initialized for mesh {current_mesh}, "
+                               f"but mesh {requested_mesh} was requested. Reinitialize the distributed "
+                               "process before changing sequence_parallel_size.")
         return
 
     group_listing = []

deepspeed/compile/fx.py

@@ -172,15 +172,17 @@ def find_node_by_name(gm: GraphModule, name: str) -> Optional[Node]:
 
 
 def get_node_shape_meta(node: Node) -> Optional[torch.Tensor]:
-    return node.meta.get("val") or node.meta.get("example_value")
+    value = node.meta.get("val")
+    return value if value is not None else node.meta.get("example_value")
 
 
 def find_node_by_tag(gm: GraphModule, tag: str) -> Optional[Node]:
     input_id_node = None
     for node in gm.graph.nodes:
         # https://github.com/pytorch/pytorch/blob/085b71eab05cbc7d474a173884269c62d2778f77/torch/_dynamo/utils.py#L5048
         tensor_dict = node.meta.get('tensor_dict')
-        if tensor_dict and tensor_dict.get('tag') == tag:
+        node_tag = tensor_dict.get('tag') if tensor_dict else None
+        if node_tag == tag or (isinstance(node_tag, tuple) and len(node_tag) == 2 and node_tag[0] == tag):
             input_id_node = node
             break
     return input_id_node

deepspeed/compile/passes/sp_compile.py

@@ -16,8 +16,9 @@
 from deepspeed.compile import constants
 
 from ..custom_ops import all_to_all, sp_dp_registry  # noqa: F401
-from ..fx import find_node_by_name, get_node_shape_meta
-from ..util import get_input_id_node, get_label_id_node, get_position_id_node, shard_tensor_node, get_sdpa_nodes
+from ..fx import get_node_shape_meta
+from ..util import (find_symbolic_shape_node, get_autosp_seq_dim, get_input_id_node, get_label_id_node,
+                    get_position_id_node, get_sdpa_nodes, shard_tensor_node)
 
 
 def prepare_autosp_inputs(input_id: torch.Tensor,
@@ -43,6 +44,8 @@ def prepare_autosp_inputs(input_id: torch.Tensor,
 
     if seq_dim < 0 or seq_dim >= input_id.ndim:
         raise ValueError(f"seq_dim {seq_dim} must be a valid index for input_id with shape {input_id.shape}")
+    if seq_dim >= label_id.ndim:
+        raise ValueError(f"seq_dim {seq_dim} is out of bounds for label_id with shape {label_id.shape}")
 
     if position_id is not None:
         if seq_dim >= position_id.ndim:
@@ -60,10 +63,10 @@ def prepare_autosp_inputs(input_id: torch.Tensor,
     if attention_mask is not None:
         torch._dynamo.decorators.mark_dynamic(attention_mask, seq_dim)
 
-    input_id.tag = constants.AUTOSP_INPUT_ID_KEY
-    label_id.tag = constants.AUTOSP_LABEL_ID_KEY
+    input_id.tag = (constants.AUTOSP_INPUT_ID_KEY, seq_dim)
+    label_id.tag = (constants.AUTOSP_LABEL_ID_KEY, seq_dim)
     if position_id is not None:
-        position_id.tag = constants.AUTOSP_POSITION_ID_KEY
+        position_id.tag = (constants.AUTOSP_POSITION_ID_KEY, seq_dim)
 
     return input_id, label_id, position_id, attention_mask
 
@@ -77,12 +80,13 @@ def pass_shard_seq_dim(gm: GraphModule, example_inputs):
 
     input_ids_node = get_input_id_node(gm)
     val = get_node_shape_meta(input_ids_node)
-    seq_symint = val.shape[1]
+    seq_dim = get_autosp_seq_dim(input_ids_node)
+    seq_symint = val.shape[seq_dim]
     assert isinstance(
         seq_symint,
         torch.SymInt), f"expected sequence dimension to be of type {torch.SymInt!r} but found {type(seq_symint)!r}"
 
-    sym_seq_dim_node = find_node_by_name(gm, str(seq_symint))
+    sym_seq_dim_node = find_symbolic_shape_node(gm, seq_symint)
     if sym_seq_dim_node is None:
         print(f"WARNING: Could not find the symbolic node for the sequence dimension")
         return
@@ -91,13 +95,10 @@ def pass_shard_seq_dim(gm: GraphModule, example_inputs):
         sharded_node = gm.graph.call_function(operator.floordiv, args=(sym_seq_dim_node, sp_size))
 
     sharded_input_nodes = set()
-    label_ids_node = get_label_id_node(gm)
     position_ids_node = get_position_id_node(gm)
 
     if input_ids_node is not None:
         sharded_input_nodes.add(input_ids_node)
-    if label_ids_node is not None:
-        sharded_input_nodes.add(label_ids_node)
     if position_ids_node is not None:
         sharded_input_nodes.add(position_ids_node)
 
@@ -133,7 +134,85 @@ def pass_shard_input_ids(gm: GraphModule, example_inputs):
 
 def pass_shard_label_ids(gm: GraphModule, example_inputs):
     label_ids_node = get_label_id_node(gm)
-    shard_tensor_node(gm, label_ids_node)
+    label_meta = get_node_shape_meta(label_ids_node)
+    seq_dim = get_autosp_seq_dim(label_ids_node)
+    seq_len = label_meta.shape[seq_dim]
+
+    def depends_on(node: Node, ancestor: Node) -> bool:
+        worklist = [node]
+        visited = set()
+        while worklist:
+            current = worklist.pop()
+            if current is ancestor:
+                return True
+            if current in visited:
+                continue
+            visited.add(current)
+            worklist.extend(current.all_input_nodes)
+        return False
+
+    loss_nodes = [node for node in gm.graph.nodes if node.target is torch.nn.functional.cross_entropy]
+    label_loss_nodes = [
+        node for node in loss_nodes
+        if len(node.args) > 1 and isinstance(node.args[1], Node) and depends_on(node.args[1], label_ids_node)
+    ]
+    if not label_loss_nodes:
+        shard_tensor_node(gm, label_ids_node, seq_dim)
+        return
+
+    causal_losses = []
+    for loss_node in label_loss_nodes:
+        weight = loss_node.kwargs.get("weight", loss_node.args[2] if len(loss_node.args) > 2 else None)
+        if weight is not None:
+            raise RuntimeError("AutoSP does not support class-weighted causal language-model cross entropy")
+        reduction = loss_node.kwargs.get("reduction", loss_node.args[6] if len(loss_node.args) > 6 else "mean")
+        if reduction != "mean":
+            raise RuntimeError(f"AutoSP only supports mean-reduced causal language-model loss, got {reduction!r}")
+        ignore_index = loss_node.kwargs.get("ignore_index", loss_node.args[4] if len(loss_node.args) > 4 else -100)
+
+        target_node = loss_node.args[1]
+        worklist = [target_node]
+        visited = set()
+        shifted_label_node = None
+        while worklist:
+            current = worklist.pop(0)
+            if current in visited or current is label_ids_node:
+                continue
+            visited.add(current)
+            current_meta = get_node_shape_meta(current)
+            if isinstance(current_meta, torch.Tensor) and current_meta.ndim == label_meta.ndim:
+                current_seq_len = current_meta.shape[seq_dim]
+                if str(current_seq_len) == str(seq_len):
+                    shifted_label_node = current
+                    break
+            worklist.extend(current.all_input_nodes)
+
+        if shifted_label_node is None:
+            shard_tensor_node(gm, label_ids_node, seq_dim)
+            return
+
+        causal_losses.append((loss_node, shifted_label_node, ignore_index))
+
+    sharded_candidates = {}
+    for loss_node, shifted_label_node, ignore_index in causal_losses:
+        if shifted_label_node not in sharded_candidates:
+            sharded_candidates[shifted_label_node] = shard_tensor_node(gm,
+                                                                       shifted_label_node,
+                                                                       seq_dim,
+                                                                       make_contiguous=True)
+        sharded_labels = sharded_candidates[shifted_label_node]
+
+        with gm.graph.inserting_after(sharded_labels):
+            valid_mask = gm.graph.call_function(operator.ne, args=(sharded_labels, ignore_index))
+        with gm.graph.inserting_after(valid_mask):
+            valid_tokens = gm.graph.call_function(torch.sum, args=(valid_mask, ))
+        with gm.graph.inserting_after(loss_node):
+            aggregated = gm.graph.call_function(torch.ops.autosp.aggregate_loss.default,
+                                                args=(loss_node, valid_tokens))
+        with gm.graph.inserting_after(aggregated):
+            global_loss = gm.graph.call_function(operator.getitem, args=(aggregated, 0))
+        loss_node.replace_all_uses_with(global_loss)
+        aggregated.update_arg(0, loss_node)
 
 
 def pass_shard_position_ids(gm: GraphModule, example_inputs):
@@ -168,6 +247,36 @@ def insert_a2a(node: Node, scatter_idx: int, gather_idx: int, name: str) -> Node
         q, k, v = attn_node.args[:3]
         suffix = f"_{idx}" if len(attention_nodes) > 1 else ""
 
+        for name, tensor in (("query", q), ("key", k), ("value", v)):
+            tensor_meta = get_node_shape_meta(tensor)
+            if tensor_meta is None or tensor_meta.ndim != 4:
+                raise RuntimeError(f"AutoSP expected a rank-4 {name} tensor for SDPA")
+            heads = tensor_meta.shape[1]
+            if isinstance(heads, int) and heads % sp_dp_registry.sp_size() != 0:
+                raise ValueError(f"AutoSP requires the {name} head count ({heads}) to be divisible by "
+                                 f"sequence_parallel_size ({sp_dp_registry.sp_size()})")
+
+        attn_mask = attn_node.kwargs.get("attn_mask")
+        mask_is_kwarg = attn_mask is not None
+        if attn_mask is None and len(attn_node.args) > 3:
+            attn_mask = attn_node.args[3]
+        if isinstance(attn_mask, Node):
+            mask_meta = get_node_shape_meta(attn_mask)
+            if mask_meta is not None and mask_meta.ndim >= 2:
+                q_meta = get_node_shape_meta(q)
+                k_meta = get_node_shape_meta(k)
+                if str(mask_meta.shape[-1]) == str(k_meta.shape[2]):
+                    raise RuntimeError("AutoSP cannot reconstruct an attention mask with a sharded key dimension. "
+                                       "Pass the full key padding mask to every SP rank.")
+                if str(mask_meta.shape[-2]) == str(q_meta.shape[2]):
+                    with gm.graph.inserting_after(attn_mask):
+                        global_mask = gm.graph.call_function(torch.ops.autosp.all_gather_sequence.default,
+                                                             args=(attn_mask, -2))
+                    if mask_is_kwarg:
+                        attn_node.update_kwarg("attn_mask", global_mask)
+                    else:
+                        attn_node.update_arg(3, global_mask)
+
         # QKV: [B, N, S/P, H] -> [B, N/P, S, H]
         insert_a2a(q, scatter_idx=1, gather_idx=2, name=f"q{suffix}")
         insert_a2a(k, scatter_idx=1, gather_idx=2, name=f"k{suffix}")
@@ -219,17 +328,27 @@ def pass_propagate_shapes(gm: torch.fx.GraphModule, real_inputs):
     saved_sdpa_masks = []
     for attn_node in get_sdpa_nodes(gm):
         attn_mask = attn_node.kwargs.get("attn_mask")
+        mask_location = "kwarg"
+        if attn_mask is None and len(attn_node.args) > 3:
+            attn_mask = attn_node.args[3]
+            mask_location = "arg"
         if attn_mask is not None:
-            saved_sdpa_masks.append((attn_node, attn_mask))
-            attn_node.update_kwarg("attn_mask", None)
+            saved_sdpa_masks.append((attn_node, mask_location, attn_mask))
+            if mask_location == "kwarg":
+                attn_node.update_kwarg("attn_mask", None)
+            else:
+                attn_node.update_arg(3, None)
 
     try:
         # fake_inputs are already created under fake_mode above, so run
         # propagation without reconverting them into a different fake mode.
         FakeTensorProp(gm, mode=fake_mode).propagate_dont_convert_inputs(*fake_inputs)
     finally:
-        for attn_node, attn_mask in saved_sdpa_masks:
-            attn_node.update_kwarg("attn_mask", attn_mask)
+        for attn_node, mask_location, attn_mask in saved_sdpa_masks:
+            if mask_location == "kwarg":
+                attn_node.update_kwarg("attn_mask", attn_mask)
+            else:
+                attn_node.update_arg(3, attn_mask)
 
 
 def apply_autosp(gm: GraphModule,
@@ -247,7 +366,7 @@ def apply_autosp(gm: GraphModule,
         debug: If True, print graph before/after each pass
         passes: Optional custom list of passes (default: DEFAULT_PASSES)
     """
-    assert sp_size * dp_size <= dist.get_world_size(), 'Insufficient device count for mesh size'
+    assert sp_size * dp_size == dist.get_world_size(), 'AutoSP mesh must cover the distributed world size'
 
     sp_dp_registry.populate_registry(sp_size, dp_size)
 
@@ -256,6 +375,7 @@ def apply_autosp(gm: GraphModule,
         pass_shard_input_ids,
         pass_shard_label_ids,
         pass_shard_position_ids,
+        pass_propagate_shapes,
         pass_insert_attention_all_to_all,
         pass_propagate_shapes,
         pass_canonicalize,