spmd_types 0.2.1__tar.gz → 0.2.3__tar.gz

{spmd_types-0.2.1/spmd_types.egg-info → spmd_types-0.2.3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: spmd_types
-Version: 0.2.1
+Version: 0.2.3
 Summary: A type system for distributed (SPMD) tensor computations in PyTorch
 Author: Meta Platforms, Inc.
 License: BSD 3-Clause License
@@ -47,17 +47,31 @@ Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
 Requires-Python: >=3.10
 Description-Content-Type: text/markdown
 License-File: LICENSE
-Requires-Dist: torch>=2.10.0
 Dynamic: license-file
 
-# spmd_types
+# `spmd_types`
 
-A type system for distributed (SPMD) tensor computations in PyTorch.
+A type system for distributed (SPMD) tensor computations in PyTorch.  This
+package provides two type systems:
 
-spmd_types tracks per-mesh-axis types on tensors -- Replicate (R), Invariant (I),
-Varying (V), Partial (P), and Shard (S) -- and enforces type-correct transitions
-through collective operations and local rewrites. It catches distributed
-programming errors at development time without requiring a GPU cluster.
+* **Local SPMD types**, which allow you to use Megatron-style differentiable
+  collectives in a safe way by tracking whether or not your backward gradients
+  are pending reduction or not.
+
+* **Global SPMD types**, a DTensor-like abstraction for writing code that has
+  the same semantics whether run on a single device or in a distributed
+  fashion, but with explicit communication operations so you are never
+  guessing when a redistribute occurs.
+
+In both cases, the SPMD types makes it possible for you to check that your
+code computes correct gradients (local SPMD) or gives equivalent results
+across different parallelizations (global SPMD), without having to actually
+run a full E2E distributed training run to check for loss matching.
+
+The goal of this package is to provide a flexible type system that can
+typecheck realistic training code.  We have used local SPMD types to typecheck
+a realistic pretraining codebase, and global SPMD types is actively under
+construction!
 
 ## Installation
 
@@ -71,6 +85,7 @@ pip install spmd_types
 import torch
 import torch.distributed as dist
 import spmd_types as spmd
+import spmd_types.checker
 from torch.distributed.device_mesh import init_device_mesh
 
 # Set up a fake process group (no GPUs needed)
@@ -79,12 +94,12 @@ mesh = init_device_mesh("cpu", (2, 4), mesh_dim_names=("dp", "tp"))
 dp = mesh.get_group("dp")
 tp = mesh.get_group("tp")
 
-with spmd.set_current_mesh(mesh), spmd.typecheck():
+with spmd.set_current_mesh(mesh), spmd.checker.typecheck():
     x = torch.randn(4)
-    spmd.assert_type(x, {dp: spmd.R, tp: spmd.P})  # R on dp, partial on tp
+    spmd.assert_type(x, {dp: spmd.R, tp: spmd.P})       # R on dp, partial on tp
     y = spmd.all_reduce(x, tp, src=spmd.P, dst=spmd.R)  # sum across tp ranks
-    spmd.assert_type(y, {dp: spmd.R, tp: spmd.R})   # now replicated everywhere
-    z = torch.mul(y, y)                              # type inference: R * R -> R
+    spmd.assert_type(y, {dp: spmd.R, tp: spmd.R})       # now replicated everywhere
+    z = torch.mul(y, y)                                 # type inference: R * R -> R
     spmd.assert_type(z, {dp: spmd.R, tp: spmd.R})
 
 dist.destroy_process_group()
@@ -92,12 +107,14 @@ dist.destroy_process_group()
 
 ## Documentation
 
-See [MEGATRON_QUICKSTART.md](MEGATRON_QUICKSTART.md) for a guide on porting
-Megatron-derived training frameworks to use spmd_types.
+See [Local SPMD types](docs/local_spmd_types.md) for a hands-on guide on
+porting Megatron-derived training frameworks, including the Megatron-to-spmd_types
+function mapping table and advice on Invariant vs Replicate.
 
-See [DESIGN.md](DESIGN.md) for the full type system specification, including
-type inference rules, collective signatures, and forward-backward pairs.
+See [Design](docs/design.md) for the full type system specification, including
+local vs global SPMD modes, collective signatures with diagrams, forward-backward
+pairs, expert mode, cross-mesh compatibility, and partition spec redistribute.
 
 ## License
 
-BSD 3-Clause License. See [LICENSE](LICENSE) for details.
+BSD 3-Clause License. See [CONTRIBUTING.md](CONTRIBUTING.md) for how to contribute.

spmd_types-0.2.3/README.md

@@ -0,0 +1,69 @@
+# `spmd_types`
+
+A type system for distributed (SPMD) tensor computations in PyTorch.  This
+package provides two type systems:
+
+* **Local SPMD types**, which allow you to use Megatron-style differentiable
+  collectives in a safe way by tracking whether or not your backward gradients
+  are pending reduction or not.
+
+* **Global SPMD types**, a DTensor-like abstraction for writing code that has
+  the same semantics whether run on a single device or in a distributed
+  fashion, but with explicit communication operations so you are never
+  guessing when a redistribute occurs.
+
+In both cases, the SPMD types makes it possible for you to check that your
+code computes correct gradients (local SPMD) or gives equivalent results
+across different parallelizations (global SPMD), without having to actually
+run a full E2E distributed training run to check for loss matching.
+
+The goal of this package is to provide a flexible type system that can
+typecheck realistic training code.  We have used local SPMD types to typecheck
+a realistic pretraining codebase, and global SPMD types is actively under
+construction!
+
+## Installation
+
+```bash
+pip install spmd_types
+```
+
+## Quick start
+
+```python
+import torch
+import torch.distributed as dist
+import spmd_types as spmd
+import spmd_types.checker
+from torch.distributed.device_mesh import init_device_mesh
+
+# Set up a fake process group (no GPUs needed)
+dist.init_process_group(backend="fake", rank=0, world_size=8)
+mesh = init_device_mesh("cpu", (2, 4), mesh_dim_names=("dp", "tp"))
+dp = mesh.get_group("dp")
+tp = mesh.get_group("tp")
+
+with spmd.set_current_mesh(mesh), spmd.checker.typecheck():
+    x = torch.randn(4)
+    spmd.assert_type(x, {dp: spmd.R, tp: spmd.P})       # R on dp, partial on tp
+    y = spmd.all_reduce(x, tp, src=spmd.P, dst=spmd.R)  # sum across tp ranks
+    spmd.assert_type(y, {dp: spmd.R, tp: spmd.R})       # now replicated everywhere
+    z = torch.mul(y, y)                                 # type inference: R * R -> R
+    spmd.assert_type(z, {dp: spmd.R, tp: spmd.R})
+
+dist.destroy_process_group()
+```
+
+## Documentation
+
+See [Local SPMD types](docs/local_spmd_types.md) for a hands-on guide on
+porting Megatron-derived training frameworks, including the Megatron-to-spmd_types
+function mapping table and advice on Invariant vs Replicate.
+
+See [Design](docs/design.md) for the full type system specification, including
+local vs global SPMD modes, collective signatures with diagrams, forward-backward
+pairs, expert mode, cross-mesh compatibility, and partition spec redistribute.
+
+## License
+
+BSD 3-Clause License. See [CONTRIBUTING.md](CONTRIBUTING.md) for how to contribute.

{spmd_types-0.2.1 → spmd_types-0.2.3}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "spmd_types"
-version = "0.2.1"
+version = "0.2.3"
 description = "A type system for distributed (SPMD) tensor computations in PyTorch"
 readme = "README.md"
 license = {file = "LICENSE"}
@@ -24,7 +24,6 @@ classifiers = [
     "Topic :: Scientific/Engineering :: Artificial Intelligence",
 ]
 dependencies = [
-    "torch>=2.10.0",
 ]
 
 [project.urls]

{spmd_types-0.2.1 → spmd_types-0.2.3}/spmd_types/__init__.py

@@ -41,7 +41,6 @@ from spmd_types._state import (  # noqa: F401
     current_mesh_all_names,
     current_mesh_names,
     is_type_checking,
-    no_typecheck,
 )
 from spmd_types._traceback import traceback_filtering  # noqa: F401
 from spmd_types._type_attr import (  # noqa: F401
@@ -53,12 +52,16 @@ from spmd_types._type_attr import (  # noqa: F401
 # Collectives and operations -- runtime API (no _checker dependency)
 from spmd_types.runtime import (  # noqa: F401
     assert_local_type,
+    assert_local_type_like,
     assert_type,
     assert_type_like,
+    get_partition_spec,
+    has_local_type,
     Infer,
     local,
     local_map,
     mutate_type,
+    no_typecheck,
     register_autograd_function,
     register_decomposition,
     register_local_autograd_function,
@@ -84,6 +87,7 @@ from spmd_types.types import (  # noqa: F401
     Replicate,
     S,
     Shard,
+    SpmdType,
     SpmdTypeError,
     TensorSharding,
     V,

{spmd_types-0.2.1 → spmd_types-0.2.3}/spmd_types/_checker/__init__.py

@@ -68,14 +68,13 @@ from spmd_types.runtime import (  # noqa: F401
     get_partition_spec,
     has_local_type,
     mutate_type,
+    no_typecheck,
     register_autograd_function,
     register_decomposition,
     register_local_autograd_function,
     trace,
 )
 from spmd_types.types import (
-    _canonicalize_shard,
-    _check_orthogonality,
     DeviceMeshAxis,
     DTensorPropagationError,
     format_axis,
@@ -885,13 +884,15 @@ class _OpSpec:
         tensor_kwargs: Kwarg names that are tensor inputs.
         tensor_varargs_from: If set, all positional args from this index onward
             are tensor inputs (for ops with *args like einsum).
+        nonlinear_args: Positional arg indices whose operand enters the op
+            nonlinearly and therefore must not be P.
     """
 
     linearity: OpLinearity
     tensor_args: tuple[int, ...] = ()
     tensor_kwargs: tuple[str, ...] = ()
     tensor_varargs_from: int | None = None
-    fixed_args: tuple[int, ...] = ()
+    nonlinear_args: tuple[int, ...] = ()
 
 
 _OP_REGISTRY: dict[Callable, _OpSpec] = {
@@ -941,6 +942,8 @@ _OP_REGISTRY: dict[Callable, _OpSpec] = {
     torch.Tensor.mean: _OpSpec(OpLinearity.LINEAR, (0,)),
     torch.nansum: _OpSpec(OpLinearity.LINEAR, (0,)),
     torch.nanmean: _OpSpec(OpLinearity.LINEAR, (0,)),
+    torch.cumsum: _OpSpec(OpLinearity.LINEAR, (0,)),
+    torch.Tensor.cumsum: _OpSpec(OpLinearity.LINEAR, (0,)),
     # =================================================================
     # LINEAR -- concat / stack (tensor list at pos 0)
     # =================================================================
@@ -1028,31 +1031,68 @@ _OP_REGISTRY: dict[Callable, _OpSpec] = {
     torch.Tensor.repeat_interleave: _OpSpec(OpLinearity.LINEAR, (0,)),
     torch.Tensor.unfold: _OpSpec(OpLinearity.LINEAR, (0,)),
     # =================================================================
-    # LINEAR with fixed_args -- division (linear in numerator only)
+    # MULTILINEAR with nonlinear_args -- division (linear in numerator only)
+    # Division is multiplicative, not additive: a lone P numerator scaled by an
+    # R/scalar divisor stays P, while LINEAR would reject this as affine.
+    # nonlinear_args additionally forbids P in the divisor because 1/y is
+    # nonlinear in y.
     # =================================================================
-    torch.div: _OpSpec(OpLinearity.LINEAR, (0, 1), ("input", "other"), fixed_args=(1,)),
+    torch.div: _OpSpec(
+        OpLinearity.MULTILINEAR,
+        (0, 1),
+        ("input", "other"),
+        nonlinear_args=(1,),
+    ),
     torch.divide: _OpSpec(
-        OpLinearity.LINEAR, (0, 1), ("input", "other"), fixed_args=(1,)
+        OpLinearity.MULTILINEAR,
+        (0, 1),
+        ("input", "other"),
+        nonlinear_args=(1,),
     ),
     torch.true_divide: _OpSpec(
-        OpLinearity.LINEAR, (0, 1), ("input", "other"), fixed_args=(1,)
+        OpLinearity.MULTILINEAR,
+        (0, 1),
+        ("input", "other"),
+        nonlinear_args=(1,),
+    ),
+    torch.Tensor.div: _OpSpec(
+        OpLinearity.MULTILINEAR,
+        (0, 1),
+        ("other",),
+        nonlinear_args=(1,),
     ),
-    torch.Tensor.div: _OpSpec(OpLinearity.LINEAR, (0, 1), ("other",), fixed_args=(1,)),
     torch.Tensor.divide: _OpSpec(
-        OpLinearity.LINEAR, (0, 1), ("other",), fixed_args=(1,)
+        OpLinearity.MULTILINEAR,
+        (0, 1),
+        ("other",),
+        nonlinear_args=(1,),
     ),
     torch.Tensor.true_divide: _OpSpec(
-        OpLinearity.LINEAR, (0, 1), ("other",), fixed_args=(1,)
+        OpLinearity.MULTILINEAR,
+        (0, 1),
+        ("other",),
+        nonlinear_args=(1,),
+    ),
+    torch.Tensor.div_: _OpSpec(
+        OpLinearity.MULTILINEAR,
+        (0, 1),
+        ("other",),
+        nonlinear_args=(1,),
     ),
-    torch.Tensor.div_: _OpSpec(OpLinearity.LINEAR, (0, 1), ("other",), fixed_args=(1,)),
     torch.Tensor.__truediv__: _OpSpec(
-        OpLinearity.LINEAR, (0, 1), ("other",), fixed_args=(1,)
+        OpLinearity.MULTILINEAR,
+        (0, 1),
+        ("other",),
+        nonlinear_args=(1,),
     ),
     torch.Tensor.__rtruediv__: _OpSpec(
-        OpLinearity.LINEAR, (0, 1), ("other",), fixed_args=(0,)
+        OpLinearity.MULTILINEAR, (0, 1), ("other",), nonlinear_args=(0,)
     ),
     torch.Tensor.__itruediv__: _OpSpec(
-        OpLinearity.LINEAR, (0, 1), ("other",), fixed_args=(1,)
+        OpLinearity.MULTILINEAR,
+        (0, 1),
+        ("other",),
+        nonlinear_args=(1,),
     ),
     # =================================================================
     # MULTILINEAR -- multiplication
@@ -1539,80 +1579,40 @@ def _set_result_type(result: object, output_type: LocalSpmdType) -> None:
             _apply(item)
 
 
-def _apply_fixed_args(  # noqa: C901
+def _validate_nonlinear_args(
     func: Callable,
     args: tuple,
-    kwargs: dict,
     spec: _OpSpec,
-    input_types_list: list[LocalSpmdType],
-) -> list[LocalSpmdType]:
-    """Filter fixed_args for LINEAR ops when Partial is present.
-
-    ``fixed_args`` lists positional arg indices that must be held fixed (not P)
-    for the op to be linear in the remaining args.  For example, ``div(a, b)``
-    is linear in ``a`` when ``b`` is fixed, so ``fixed_args=(1,)``.
-
-    When P is present among the non-fixed tensor args:
-    1. Validate that tensor args at fixed_args positions don't have P on any axis.
-    2. Exclude their types from the returned list so LINEAR inference sees only
-       the "free" args (and doesn't reject P + R mixing).
-
-    When no P is present in the free args, return the original list unchanged
-    so normal inference rules apply (e.g., R + V -> V).
+) -> None:
+    """Validate that operands entering the op nonlinearly are not Partial.
 
-    Args:
-        func: The torch function being called.
-        args: Positional arguments to the torch operation.
-        kwargs: Keyword arguments to the torch operation.
-        spec: The op specification (must have non-empty fixed_args).
-        input_types_list: The already-collected input types list.
+    Division is the current user: ``x / y`` is linear (scaling) in ``x`` only,
+    since ``1/y`` is nonlinear in ``y``. Nonlinear operands may be untyped
+    Python scalars or typed non-Partial values, but they cannot be P.
     """
-    # Identify which input_types_list entries came from fixed_args positions.
-    # We re-walk args to figure out which types are "fixed" vs "free".
-    fixed_positions = set(spec.fixed_args)
-    free_types: list[LocalSpmdType] = []
-    fixed_types: list[LocalSpmdType] = []
 
-    for i in spec.tensor_args:
-        if i < len(args):
-            val = args[i]
-            if isinstance(val, torch.Tensor):
-                if i in fixed_positions:
-                    fixed_types.append(get_local_type(val))
-                else:
-                    free_types.append(get_local_type(val))
-            elif isinstance(val, (list, tuple)):
-                for item in val:
-                    if isinstance(item, torch.Tensor):
-                        if i in fixed_positions:
-                            fixed_types.append(get_local_type(item))
-                        else:
-                            free_types.append(get_local_type(item))
-
-    # Check if P is present in any free arg on any axis
-    has_p_in_free = any(P in typ.values() for typ in free_types)
-    # Also check default V (missing from dict means V, not P) -- V != P, so fine.
-
-    if not has_p_in_free:
-        # No P in free args -- normal inference, include all types.
-        return input_types_list
-
-    # P is present in free args -- validate fixed args don't have P
-    for fixed_type in fixed_types:
-        for axis, typ in fixed_type.items():
-            if typ is P:
+    def _check_type(typ: LocalSpmdType) -> None:
+        for axis, axis_type in typ.items():
+            if axis_type is P:
                 raise SpmdTypeError(
-                    f"{func.__name__}: Partial type in fixed argument "
-                    f"(denominator/divisor) on axis {format_axis(axis)} is not allowed. "
-                    f"Division is only linear in the numerator."
+                    f"{func.__name__}: Partial type in nonlinear argument "
+                    f"on axis {format_axis(axis)} is not allowed. Reduce or "
+                    f"convert the Partial operand before using it in a "
+                    f"nonlinear position."
                 )
 
-    # Exclude fixed_types from input_types_list: return only free_types
-    # plus any scalar types that were appended by _collect_scalar_types.
-    # scalar_types count = len(input_types_list) - len(free_types) - len(fixed_types)
-    n_tensor_types = len(free_types) + len(fixed_types)
-    scalar_types = input_types_list[n_tensor_types:]
-    return free_types + scalar_types
+    def _check_value(val: object) -> None:
+        if isinstance(val, torch.Tensor):
+            _check_type(get_local_type(val))
+        elif isinstance(val, Scalar):
+            _check_type(get_local_type(val))
+        elif isinstance(val, (list, tuple)):
+            for item in val:
+                _check_value(item)
+
+    for i in spec.nonlinear_args:
+        if i < len(args):
+            _check_value(args[i])
 
 
 # Every function in this registry must accept (x: Tensor, axis, *, src=..., dst=...)
@@ -1680,7 +1680,7 @@ def _set_result_partition_spec(
             _set_partition_spec(item, spec)
 
 
-def _collect_shard_axes(
+def _collect_shard_axes(  # noqa: C901
     partition_specs: list[PartitionSpec | None],
     global_axes: set[MeshAxis],
 ) -> tuple[list[MeshAxis], dict[MeshAxis, set[MeshAxis]]]:
@@ -2086,11 +2086,16 @@ class _ShardPropagator:
         """
         mesh = self._get_mesh(axis.size())
 
+        from torch.utils._python_dispatch import _disable_current_modes
+
+        # Keep checker-internal fake DTensor propagation invisible to active
+        # TorchDispatchModes such as selective activation checkpointing's
+        # caching mode.
         # Disable LocalTensorMode during meta DTensor propagation.
         # DTensor.from_local calls mesh.get_coordinate() which, under
         # LocalTensorMode(N), tries to map all N simulated ranks into
         # the propagator's fake 1-D mesh.
-        with maybe_disable_local_tensor_mode():
+        with maybe_disable_local_tensor_mode(), _disable_current_modes():
             meta_args = torch.utils._pytree.tree_map(
                 lambda x: self._to_meta_dtensor(x, axis, mesh), args
             )
@@ -2527,6 +2532,8 @@ def _get_autograd_function_class(func) -> type | None:
 def _validate_partition_spec_for_global_spmd(
     local_type: LocalSpmdType,
     spec: PartitionSpec | None,
+    *,
+    is_global_axis: Callable[[MeshAxis], bool],
 ) -> None:
     """Validate PartitionSpec <-> local type consistency.
 
@@ -2539,10 +2546,12 @@ def _validate_partition_spec_for_global_spmd(
     Args:
         local_type: The tensor's local SPMD type dict.
         spec: The tensor's PartitionSpec (None if absent).
+        is_global_axis: Returns whether an axis requires global PartitionSpec
+            validation in the current mesh context.
     """
     spec_axes = spec.axes_with_partition_spec() if spec is not None else set()
     for axis, typ in local_type.items():
-        if typ is V and axis not in spec_axes:
+        if is_global_axis(axis) and typ is V and axis not in spec_axes:
             ax = format_axis(axis)
             raise SpmdTypeError(
                 f"Tensor has Varying type on axis {ax} but no PartitionSpec entry. "
@@ -2551,7 +2560,7 @@ def _validate_partition_spec_for_global_spmd(
             )
     for axis in spec_axes:
         typ = local_type.get(axis)
-        if typ is not V:
+        if is_global_axis(axis) and typ is not V:
             ax = format_axis(axis)
             raise SpmdTypeError(
                 f"Tensor has PartitionSpec on axis {ax} but local type is "
@@ -2619,6 +2628,23 @@ class _SpmdTypeMode(torch.overrides.TorchFunctionMode):
 
         self._disabled = False
 
+    def _is_global_axis(self, axis: MeshAxis) -> bool:
+        if self._local:
+            return False
+        local_axes = _state._current_local_axes()
+        if not local_axes:
+            return True
+        mesh = _state.current_mesh()
+        if mesh is None or axis not in mesh:
+            raise SpmdTypeError(
+                f"Axis {format_axis(axis)} is not in the current mesh, so its "
+                "local/global typechecking policy is undefined"
+            )
+        return axis not in local_axes
+
+    def _global_axes_for(self, axes: Iterable[MeshAxis]) -> set[MeshAxis]:
+        return {axis for axis in axes if self._is_global_axis(axis)}
+
     def __torch_function__(self, func, types, args=(), kwargs=None):  # noqa: C901
         kwargs = kwargs or {}
         # Paused via no_typecheck(): run without type checking.
@@ -2675,7 +2701,11 @@ class _SpmdTypeMode(torch.overrides.TorchFunctionMode):
             # Global SPMD validation.
             if not self._local:
                 for local_type, spec in zip(info.tensor_types, info.partition_specs):
-                    _validate_partition_spec_for_global_spmd(local_type, spec)
+                    _validate_partition_spec_for_global_spmd(
+                        local_type,
+                        spec,
+                        is_global_axis=self._is_global_axis,
+                    )
 
             # Typecheck-registered autograd function: dispatch to typecheck_forward
             # INSTEAD of executing func.  The mode is popped off the
@@ -2740,6 +2770,29 @@ class _SpmdTypeMode(torch.overrides.TorchFunctionMode):
                 # Resolve axis to a list of matching axes on the tensor.
                 # Direct match: [axis]. Flattened decomposition: [sub1, sub2, ...].
                 resolved_axes = _resolve_collective_axes(axis, local_type)
+                collective_axes = resolved_axes or [axis]
+                global_collective_axes = [
+                    ax for ax in collective_axes if self._is_global_axis(ax)
+                ]
+                if global_collective_axes and len(global_collective_axes) != len(
+                    collective_axes
+                ):
+                    global_axis_set = set(global_collective_axes)
+                    local_collective_axes = [
+                        ax for ax in collective_axes if ax not in global_axis_set
+                    ]
+                    local_names = ", ".join(
+                        format_axis(ax) for ax in local_collective_axes
+                    )
+                    global_names = ", ".join(
+                        format_axis(ax) for ax in global_collective_axes
+                    )
+                    raise SpmdTypeError(
+                        f"Collective axis {format_axis(axis)} spans mesh axes with "
+                        f"mixed typechecking policies: local axes ({local_names}) "
+                        f"and global axes ({global_names}). Pass a collective axis "
+                        f"whose component axes are all local or all global."
+                    )
 
                 if resolved_axes is not None:
                     # Validate input types on all resolved axes
@@ -2770,7 +2823,7 @@ class _SpmdTypeMode(torch.overrides.TorchFunctionMode):
                         output_type[ax] = local_dst
                 _set_local_type(result, output_type)
 
-                if _state._is_global():
+                if global_collective_axes:
                     input_spec = get_partition_spec(x)
                     # Find S(dim) for this axis in the input spec, validating
                     # it is innermost in its multi-axis group.
@@ -2902,10 +2955,8 @@ class _SpmdTypeMode(torch.overrides.TorchFunctionMode):
                     linearity = (
                         spec.linearity if spec is not None else OpLinearity.NONLINEAR
                     )
-                    if spec is not None and spec.fixed_args:
-                        input_types_list = _apply_fixed_args(
-                            func, args, kwargs, spec, input_types_list
-                        )
+                    if spec is not None and spec.nonlinear_args:
+                        _validate_nonlinear_args(func, original_args, spec)
                     # NB: output_type may be further updated because global
                     # shard propagation below may promote local type V to P.
                     output_type = infer_output_type(
@@ -2914,15 +2965,16 @@ class _SpmdTypeMode(torch.overrides.TorchFunctionMode):
                     )
 
                 # Stage 3: global shard propagation.
-                if self._local:
-                    global_shard_axes, shard_edges = [], {}
-                else:
-                    all_axes: set[DeviceMeshAxis] = set()
-                    for typ in input_types_list:
-                        all_axes.update(typ.keys())
+                all_axes: set[MeshAxis] = set()
+                for typ in input_types_list:
+                    all_axes.update(typ.keys())
+                global_axes = self._global_axes_for(all_axes)
+                if global_axes:
                     global_shard_axes, shard_edges = _collect_shard_axes(
-                        partition_specs, all_axes
+                        partition_specs, global_axes
                     )
+                else:
+                    global_shard_axes, shard_edges = [], {}
 
                 # Validate mutation safety for in-place/out operations.
                 mutated = _get_mutated_tensors(func, args, kwargs, result)
@@ -3084,9 +3136,6 @@ def typecheck(
             _remove_autograd_apply_patch()
 
 
-from spmd_types._state import no_typecheck  # noqa: F401
-
-
 class _SpmdTypeBackwardCompatibleMode:
     """Backwards-compatible wrapper around ``typecheck()``.