Merging AutoSP into DeepSpeed (#7860)

# AutoSP: Unlocking Long-Context LLM Training Via Compiler-Based
Sequence Parallelism

## Overview

AutoSP is a compiler optimization pass that shards inputs along the
sequence dimension and enables Ulysses styled sequence parallelism while
preventing graph breaks during `torch.compile()`. All the passes operate
at the Torch IR on the forward graph.

## API Design

### User-Facing Entry Point: `prepare_autosp_inputs()`

Users must explicitly call this function to prepare inputs for AutoSP
compilation:

```python
def prepare_autosp_inputs(
    input_id: torch.Tensor,
    label_id: torch.Tensor,
    position_id: torch.Tensor = None,
    attention_mask: torch.Tensor = None,
    seq_dim: int = 1
) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]
```

**Purpose**: Symbolize sequence dimension and annotate tensors for
identification.

**Operations**:
1. Mark sequence dimension as dynamic using
`torch._dynamo.decorators.mark_dynamic()`
2. Attach metadata tags for tensor identification for auto-sharding:
   - `input_id.tag = constants.INPUT_ID_KEY`
   - `label_id.tag = constants.LABEL_ID_KEY`
   - `position_id.tag = constants.POSITION_ID_KEY` (if provided)

**Rationale**: PyTorch's FX graph tracer requires explicit annotation of
data-dependent dimensions. Marking the sequence dimension as dynamic
prevents symbolic shape propagation from losing dimension information
through reshape/view operations.

## Compilation Passes

### Pass 1: `pass_shard_seq_dim()`

**Objective**: Propagate sharded sequence dimension to all consumers.

**Algorithm**:
1. Extract symbolic sequence dimension from `input_id` shape metadata
2. Locate the symbolic dimension node in the FX graph
3. Create a floor-divide node: `seq_dim / world_size`
4. Perform worklist-based graph traversal to find all direct and
indirect consumers of input node, label node and position id node.
5. Replace symbolic dimension references with sharded dimension in
consumer nodes

**Rationale**: Reshapes and views that consume the sequence dimension as
an argument do not get updated during propagation of symbolic shapes.
This pass explicitly rewires the computation graph to use sharded
dimensions, enabling proper shape inference downstream.

### Pass 2: `pass_shard_input_ids()` / `pass_shard_label_ids()` /
`pass_shard_position_ids()`

**Objective**: Insert slicing operations after input tensors.

**Implementation**: Call `shard_tensor_node()` utility which inserts
slice operations. Each rank retains only the portion of the tensor
corresponding to its sequence partition and drops the remaining buffer.

**Note on `attention_mask`**: Not sharded because it applies to the full
sequence length, not the partitioned dimension.

### Pass 3: `pass_insert_attention_all_to_all()`

**Objective**: Insert all-to-all collectives around attention (Ulysses
styled) to avoid graph breaks during compilation.

**Algorithm**:
1. Identify all SDPA (Scaled Dot-Product Attention) nodes in the graph
2. For each SDPA node with inputs Q, K, V, after each of Q, K, V: insert
A2A scatter heads (dim=1), gather sequence (dim=2)
3. Insert A2A after thre attention output O: scatter sequence (dim=2),
gather heads (dim=1)

**Graph Rewrite Example**:
```
Q [B, N, S/P, H] --A2A(scatter_heads,gather_seq)--> [B, N/P, S, H]
K [B, N, S/P, H] --A2A(scatter_heads,gather_seq)--> [B, N/P, S, H]
V [B, N, S/P, H] --A2A(scatter_heads,gather_seq)--> [B, N/P, S, H]
                     |
                    SDPA
                     |
O [B, N/P, S, H] --A2A(scatter_seq,gather_heads)--> [B, N, S/P, H]
```

**Current support**: Currently only supports
`torch.nn.functional.scaled_dot_product_attention()`. Composite
attention patterns require additional pattern matching logic.

### Pass 4: `pass_propagate_shapes()`

**Objective**: Compute static shapes for all nodes using fake tensor
execution.

**Implementation**:
1. Create `ShapeEnv` for symbolic dimension tracking
2. Construct `FakeTensorMode` with the shape environment
3. Execute `FakeTensorProp.propagate()` to compute shape metadata

### Pass 5: `pass_canonicalize()`

**Objective**: Finalize graph representation.

**Operations**:
1. `eliminate_dead_code()`: Remove unused operations
2. `lint()`: Validate graph structure
3. `recompile()`: Regenerate compiled representation

## Execution Order

```
prepare_autosp_inputs()
    ↓
pass_shard_seq_dim
    ↓
pass_shard_input_ids
    ↓
pass_shard_label_ids
    ↓
pass_shard_position_ids
    ↓
pass_insert_attention_all_to_all
    ↓
pass_propagate_shapes
    ↓
pass_canonicalize
    ↓
pass_selective_activation_checkpointing

```

## Memory savings
AutoSP adds some heuristics to torch.compile's partitioniner which
splits the joint graph into the forward and backward graph. Matmul and
related ops are not checkpointed since recomputing them is much cheaper
compared to the attention op, while reducing the peak active memory.

## Reducing gradients across ranks
AutoSP requires an all-reduce to reduce the gradients across ranks. This
is automatically called by DeepSpeed's engine
[here](https://github.com/deepspeedai/DeepSpeed/blob/93524c8931799a7631a2321d7ef4afaff6b6e54b/deepspeed/runtime/engine.py#L2433)

## Known Limitations

1. **Attention Pattern Matching**: Only
`torch.nn.functional.scaled_dot_product_attention()` is supported. Fused
attention implementations require pattern-specific handling.
2. **No Graph Break Requirement**: AutoSP will fail if there are graph
breaks because use-def chains are lost and it becomes tricky to
propagate auto-sharding information across graph modules.


## Example 

DeepSpeedExample PR:
deepspeedai/DeepSpeedExamples#999

---------

Signed-off-by: Neel Dani <neeldani98@gmail.com>
Signed-off-by: Ahan Gupta <ahangupta.96@gmail.com>
Signed-off-by: Masahiro Tanaka <mtanaka@anyscale.com>
Co-authored-by: Ahan Gupta <ahangupta.96@gmail.com>
Co-authored-by: Olatunji Ruwase <tunji.ruwase@snowflake.com>
Co-authored-by: Masahiro Tanaka <81312776+tohtana@users.noreply.github.com>
Co-authored-by: Zhipeng Wang <zhipeng.rainbowserie@gmail.com>

Author: 5 people

deepspeed/compile/config.py

@@ -3,8 +3,11 @@
 
 # DeepSpeed Team
 
+from typing import List, Optional, Literal
 from deepspeed.runtime.config_utils import DeepSpeedConfigModel
 
+PassName = Literal["z1", "z3", "autosp"]
+
 
 class CompileConfig(DeepSpeedConfigModel):
     """ Configure compile settings """
@@ -53,3 +56,6 @@ class CompileConfig(DeepSpeedConfigModel):
 
     keep_all_input_tensors: bool = False
     """ Keep real values for all input tensors in InputStorage instead of using dummy values """
+
+    passes: Optional[List[PassName]] = None
+    """ Composes different optimizations. """

deepspeed/compile/constants.py

@@ -0,0 +1,11 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+#########################################
+# AUTOSP
+#########################################
+AUTOSP_INPUT_ID_KEY = "input_id"
+AUTOSP_LABEL_ID_KEY = "label_id"
+AUTOSP_POSITION_ID_KEY = "position_id"

deepspeed/compile/custom_ops/__init__.py

@@ -0,0 +1,9 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+from .all_to_all import all_to_all
+from . import sp_dp_registry
+
+__all__ = ["all_to_all", "sp_dp_registry", "sp_compat"]

deepspeed/compile/custom_ops/all_to_all.py

@@ -0,0 +1,92 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import torch
+import deepspeed.comm as dist
+from torch.utils._sympy.functions import FloorDiv
+from .sp_dp_registry import get_group, is_setup, sp_size
+
+
+@torch.library.custom_op("autosp::all_to_all", mutates_args=())
+def all_to_all(
+    input: torch.Tensor,
+    scatter_idx: int,
+    gather_idx: int,
+    name: str,
+) -> torch.Tensor:
+    """
+    All-to-all collective for SDPA tensors [B, N, S, H].
+
+    For QKV (scatter_idx=1, gather_idx=2):
+        [B, N, S/P, H] -> [B, N/P, S, H]
+    For O (scatter_idx=2, gather_idx=1):
+        [B, N/P, S, H] -> [B, N, S/P, H]
+    """
+    assert is_setup(), 'Incorrect initialization of SP/DP mesh.'
+    B, dim1, dim2, H = input.shape
+    gid = dist.get_rank() // sp_size()
+    group = get_group(gid)
+
+    if scatter_idx == 1:
+        N, local_S = dim1, dim2
+        input_t = input.reshape(B, sp_size(), N // sp_size(), local_S, H)
+        input_t = input_t.permute(1, 0, 2, 3, 4).contiguous()
+
+        output = torch.empty_like(input_t)
+        dist.all_to_all_single(output, input_t, group=group)
+
+        output = output.permute(1, 2, 0, 3, 4).contiguous()
+        output = output.reshape(B, N // sp_size(), sp_size() * local_S, H)
+    else:
+        local_N, S = dim1, dim2
+        input_t = input.reshape(B, local_N, sp_size(), S // sp_size(), H)
+        input_t = input_t.permute(2, 0, 1, 3, 4).contiguous()
+
+        output = torch.empty_like(input_t)
+        dist.all_to_all_single(output, input_t, group=group)
+
+        output = output.permute(1, 0, 2, 3, 4).contiguous()
+        output = output.reshape(B, sp_size() * local_N, S // sp_size(), H)
+
+    return output
+
+
+@torch.library.register_fake("autosp::all_to_all")
+def all_to_all_fake(input: torch.Tensor, scatter_idx: int, gather_idx: int, name: str):
+
+    def maybe_restore_sharded_dim(dim: torch.SymInt, factor: int):
+        # Torch 2.9 may keep `P * (s // P)` distinct from the original `s` during
+        # fake shape propagation. When the local dim is exactly `FloorDiv(s, P)`,
+        # restore the original symbol so downstream ops see a consistent sequence dim.
+        node = getattr(dim, "node", None)
+        if node is None:
+            return dim * factor
+
+        expr = node.expr
+        if isinstance(expr, FloorDiv) and expr.args[1] == factor:
+            hint = node.hint * factor if node.has_hint() else None
+            return node.shape_env.create_symintnode(expr.args[0], hint=hint)
+
+        return dim * factor
+
+    B, dim1, dim2, H = input.shape
+    if scatter_idx == 1:
+        return input.new_empty(B, dim1 // sp_size(), maybe_restore_sharded_dim(dim2, sp_size()), H)
+    else:
+        return input.new_empty(B, dim1 * sp_size(), dim2 // sp_size(), H)
+
+
+def _all_to_all_backward_setup(ctx, inputs, output):
+    _, scatter_idx, gather_idx, name = inputs
+    ctx.scatter_idx = gather_idx
+    ctx.gather_idx = scatter_idx
+    ctx.name = name + "_grad"
+
+
+def _all_to_all_backward(ctx, grad):
+    return (all_to_all(grad, ctx.scatter_idx, ctx.gather_idx, ctx.name), None, None, None)
+
+
+torch.library.register_autograd("autosp::all_to_all", _all_to_all_backward, setup_context=_all_to_all_backward_setup)

deepspeed/compile/custom_ops/sp_compat.py

@@ -0,0 +1,24 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import torch
+from packaging.version import Version
+
+
+def _check_autosp_compatibility():
+    # Strip the local version segment (e.g. +cu128) so CUDA builds don't sort
+    # above the max bound when using packaging's local-version ordering rules.
+    torch_version = Version(torch.__version__.split("+")[0])
+    if torch_version < Version("2.9"):
+        raise RuntimeError("AutoSP requires PyTorch >= 2.9, found "
+                           f"{torch.__version__}.")
+
+    try:
+        import transformers
+        if Version(transformers.__version__) > Version("4.50.3"):
+            raise RuntimeError("AutoSP requires transformers <= 4.50.3, found "
+                               f"{transformers.__version__}.")
+    except ImportError:
+        pass  # transformers not installed; skip the check

deepspeed/compile/custom_ops/sp_dp_registry.py

@@ -0,0 +1,67 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import deepspeed.comm as dist
+
+GROUP_REGISTRY = {}  # int -> dist.ProcessGroup
+
+
+def register_groups(groups):
+    """groups: List[List[int]], e.g. [[0,1],[2,3]]"""
+    for gid, ranks in enumerate(groups):
+        if gid not in GROUP_REGISTRY:
+            GROUP_REGISTRY[gid] = dist.new_group(ranks)
+
+
+def get_group(gid: int):
+    return GROUP_REGISTRY[gid] if gid is not None else dist.get_world_group()
+
+
+def get_registry():
+    return GROUP_REGISTRY
+
+
+def is_setup():
+    return GROUP_REGISTRY['is_reg'] if 'is_reg' in GROUP_REGISTRY else False
+
+
+def extract_mesh_size(param_dict):
+    sp_size = param_dict.get('sequence_parallel_size', 1)
+    assert dist.get_world_size() % sp_size == 0, 'World mesh-size should be divisible by SP_SIZE'
+    dp_size = dist.get_world_size() // sp_size
+
+    return sp_size, dp_size
+
+
+def sp_size():
+    assert 'SP_SIZE' in GROUP_REGISTRY, 'SP_SIZE not init properly.'
+
+    return GROUP_REGISTRY['SP_SIZE']
+
+
+def dp_size():
+    assert 'DP_SIZE' in GROUP_REGISTRY, 'DP_SIZE not init properly'
+
+    return GROUP_REGISTRY['DP_SIZE']
+
+
+def populate_registry(SP_SIZE, DP_SIZE):
+    """ Populate rank to SP/DP mesh index.  """
+
+    if GROUP_REGISTRY.get('is_reg', False):
+        return
+
+    group_listing = []
+    offset = 0
+    for _ in range(DP_SIZE):
+        group_listing.append([i + offset for i in range(SP_SIZE)])
+        offset += SP_SIZE
+
+    register_groups(group_listing)
+
+    ## Extraneous metadata required for proper instatiation. ##
+    GROUP_REGISTRY['SP_SIZE'] = SP_SIZE
+    GROUP_REGISTRY['DP_SIZE'] = DP_SIZE
+    GROUP_REGISTRY['is_reg'] = True

deepspeed/compile/fx.py

@@ -3,11 +3,11 @@
 
 # DeepSpeed Team
 
-from typing import Callable, Any, List, Dict
+from typing import Callable, Any, List, Dict, Optional
 from collections import defaultdict
 
 import torch
-from torch.fx import Node, Graph
+from torch.fx import Node, Graph, GraphModule
 
 from .util import get_last_uses
 
@@ -138,3 +138,32 @@ def free_tensors(tensors: List[torch.Tensor]):
 
             # Python version for debugging
             # graph.create_node('call_function', free_tensors, args, {}, name=node_name)
+
+
+def find_node_by_name(gm: GraphModule, name: str) -> Optional[Node]:
+    for node in gm.graph.nodes:
+        if node.name == name:
+            return node
+    return None
+
+
+def get_node_shape_meta(node: Node) -> Optional[torch.Tensor]:
+    return node.meta.get("val") or 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:
+            input_id_node = node
+            break
+    return input_id_node
+
+
+def replace_node_users(node: Node, replacement: Node, exclude: Optional[List[Node]] = None):
+    exclude = exclude or []
+    to_replace = [u for u in node.users if u not in exclude]
+    for user in to_replace:
+        user.replace_input_with(node, replacement)

deepspeed/compile/init_sp.py

@@ -0,0 +1,23 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import torch
+from torch.fx import GraphModule
+from .passes.sp_compile import apply_autosp
+from .passes.long_context_checkpointing import register_long_context_checkpointing
+from .custom_ops.sp_dp_registry import extract_mesh_size
+from .custom_ops.sp_compat import _check_autosp_compatibility
+
+
+def init_autosp(config):
+    _check_autosp_compatibility()
+    sp_size, dp_size = extract_mesh_size(config._param_dict)
+    register_long_context_checkpointing()
+
+    def backend_fn(gm: GraphModule, real_inputs):
+        apply_autosp(gm, real_inputs, debug=False, sp_size=sp_size, dp_size=dp_size)
+        return torch._inductor.compile(gm, real_inputs)
+
+    return backend_fn