(local_tensor_tutorial)=

LocalTensor Tutorial: Single-Process SPMD Debugging

This tutorial introduces LocalTensor, a powerful debugging tool for developing and testing distributed tensor operations without requiring multiple processes or GPUs.

{contents}

:local:
:depth: 2

What is LocalTensor?

LocalTensor is a torch.Tensor subclass that simulates distributed SPMD (Single Program, Multiple Data) computations on a single process. It internally maintains a mapping from rank IDs to their corresponding local tensor shards, allowing you to debug and test distributed code without infrastructure overhead.

Key Benefits

No Multi-Process Setup Required: Test distributed algorithms on a single CPU/GPU
Faster Debugging Cycles: Iterate quickly without launching multiple processes
Full Visibility: Inspect each rank's tensor state directly
CI-Friendly: Run distributed tests in single-process CI pipelines
DTensor Integration: Seamlessly test DTensor code locally

{note}

`LocalTensor` is intended for **debugging and testing** only, not production use.
The overhead of simulating multiple ranks locally is significant.

Installation and Setup

LocalTensor is part of PyTorch's distributed package. No additional installation is required beyond PyTorch itself.

Usage Examples

The following examples demonstrate core patterns for using LocalTensor. Each example's code is included directly from source files that are also tested to ensure correctness. The tests directly invoke these same functions.

Example 1: Basic LocalTensor Creation and Operations

Creating a LocalTensor from per-rank tensors:

{eval-rst}

.. literalinclude:: ../../../test/distributed/local_tensor_tutorial_examples/example_01_basic_operations.py
   :language: python
   :start-after: # [core_create_local_tensor]
   :end-before: # [end_core_create_local_tensor]
   :dedent: 0

Arithmetic operations (applied per-rank):

{eval-rst}

.. literalinclude:: ../../../test/distributed/local_tensor_tutorial_examples/example_01_basic_operations.py
   :language: python
   :start-after: # [core_arithmetic_operations]
   :end-before: # [end_core_arithmetic_operations]
   :dedent: 0

Extracting a tensor when all shards are identical:

{eval-rst}

.. literalinclude:: ../../../test/distributed/local_tensor_tutorial_examples/example_01_basic_operations.py
   :language: python
   :start-after: # [core_reconcile]
   :end-before: # [end_core_reconcile]
   :dedent: 0

Using LocalTensorMode for automatic LocalTensor creation:

{eval-rst}

.. literalinclude:: ../../../test/distributed/local_tensor_tutorial_examples/example_01_basic_operations.py
   :language: python
   :start-after: # [core_local_tensor_mode]
   :end-before: # [end_core_local_tensor_mode]
   :dedent: 0

Full source: example_01_basic_operations.py

Example 2: Simulating Collective Operations

Test collective operations like all_reduce, broadcast, and all_gather without multiple processes.

All-reduce with SUM:

{eval-rst}

.. literalinclude:: ../../../test/distributed/local_tensor_tutorial_examples/example_02_collective_operations.py
   :language: python
   :start-after: # [core_all_reduce]
   :end-before: # [end_core_all_reduce]
   :dedent: 0

Broadcast from a source rank:

{eval-rst}

.. literalinclude:: ../../../test/distributed/local_tensor_tutorial_examples/example_02_collective_operations.py
   :language: python
   :start-after: # [core_broadcast]
   :end-before: # [end_core_broadcast]
   :dedent: 0

All-gather to collect tensors from all ranks:

{eval-rst}

.. literalinclude:: ../../../test/distributed/local_tensor_tutorial_examples/example_02_collective_operations.py
   :language: python
   :start-after: # [core_all_gather]
   :end-before: # [end_core_all_gather]
   :dedent: 0

Full source: example_02_collective_operations.py

Example 3: Working with DTensor

LocalTensor integrates with DTensor for testing distributed tensor parallelism.

Distribute a tensor and verify reconstruction:

{eval-rst}

.. literalinclude:: ../../../test/distributed/local_tensor_tutorial_examples/example_03_dtensor_integration.py
   :language: python
   :start-after: # [core_dtensor_distribute]
   :end-before: # [end_core_dtensor_distribute]
   :dedent: 0

Distributed matrix multiplication:

{eval-rst}

.. literalinclude:: ../../../test/distributed/local_tensor_tutorial_examples/example_03_dtensor_integration.py
   :language: python
   :start-after: # [core_dtensor_matmul]
   :end-before: # [end_core_dtensor_matmul]
   :dedent: 0

Simulating a distributed linear layer:

{eval-rst}

.. literalinclude:: ../../../test/distributed/local_tensor_tutorial_examples/example_03_dtensor_integration.py
   :language: python
   :start-after: # [core_dtensor_nn_layer]
   :end-before: # [end_core_dtensor_nn_layer]
   :dedent: 0

Full source: example_03_dtensor_integration.py

Example 4: Handling Uneven Sharding

Real-world distributed systems often have uneven data distribution across ranks. LocalTensor handles this using LocalIntNode.

Creating LocalTensor with different sizes per rank:

{eval-rst}

.. literalinclude:: ../../../test/distributed/local_tensor_tutorial_examples/example_04_uneven_sharding.py
   :language: python
   :start-after: # [core_uneven_shards]
   :end-before: # [end_core_uneven_shards]
   :dedent: 0

LocalIntNode arithmetic operations:

{eval-rst}

.. literalinclude:: ../../../test/distributed/local_tensor_tutorial_examples/example_04_uneven_sharding.py
   :language: python
   :start-after: # [core_local_int_node]
   :end-before: # [end_core_local_int_node]
   :dedent: 0

DTensor with unevenly divisible dimensions:

{eval-rst}

.. literalinclude:: ../../../test/distributed/local_tensor_tutorial_examples/example_04_uneven_sharding.py
   :language: python
   :start-after: # [core_dtensor_uneven]
   :end-before: # [end_core_dtensor_uneven]
   :dedent: 0

Full source: example_04_uneven_sharding.py

Example 5: Rank-Specific Computations

Sometimes you need to perform different operations on different ranks.

Using rank_map() to create per-rank values:

{eval-rst}

.. literalinclude:: ../../../test/distributed/local_tensor_tutorial_examples/example_05_rank_specific.py
   :language: python
   :start-after: # [core_rank_map]
   :end-before: # [end_core_rank_map]
   :dedent: 0

Using tensor_map() to transform shards per-rank:

{eval-rst}

.. literalinclude:: ../../../test/distributed/local_tensor_tutorial_examples/example_05_rank_specific.py
   :language: python
   :start-after: # [core_tensor_map]
   :end-before: # [end_core_tensor_map]
   :dedent: 0

Temporarily exiting LocalTensorMode:

{eval-rst}

.. literalinclude:: ../../../test/distributed/local_tensor_tutorial_examples/example_05_rank_specific.py
   :language: python
   :start-after: # [core_disable_mode]
   :end-before: # [end_core_disable_mode]
   :dedent: 0

Full source: example_05_rank_specific.py

Example 6: Multi-Dimensional Meshes

Use 2D/3D device meshes for hybrid parallelism (e.g., data parallel + tensor parallel).

Creating a 2D mesh:

{eval-rst}

.. literalinclude:: ../../../test/distributed/local_tensor_tutorial_examples/example_06_multidim_mesh.py
   :language: python
   :start-after: # [core_2d_mesh]
   :end-before: # [end_core_2d_mesh]
   :dedent: 0

Hybrid parallelism (DP + TP):

{eval-rst}

.. literalinclude:: ../../../test/distributed/local_tensor_tutorial_examples/example_06_multidim_mesh.py
   :language: python
   :start-after: # [core_hybrid_parallel]
   :end-before: # [end_core_hybrid_parallel]
   :dedent: 0

3D mesh for DP + TP + PP:

{eval-rst}

.. literalinclude:: ../../../test/distributed/local_tensor_tutorial_examples/example_06_multidim_mesh.py
   :language: python
   :start-after: # [core_3d_mesh]
   :end-before: # [end_core_3d_mesh]
   :dedent: 0

Full source: example_06_multidim_mesh.py

Testing Tutorial Examples

All examples in this tutorial are tested to ensure correctness. The test suite directly invokes the same functions included above:

python

# From test_local_tensor_tutorial_examples.py
from example_01_basic_operations import create_local_tensor

def test_create_local_tensor(self):
    lt = create_local_tensor()
    self.assertIsInstance(lt, LocalTensor)
    self.assertEqual(lt.shape, torch.Size([2, 2]))

Test suite: test_local_tensor_tutorial_examples.py

API Reference

Core Classes

{eval-rst}

.. autoclass:: torch.distributed._local_tensor.LocalTensor
   :members: reconcile, is_contiguous, contiguous, tolist, numpy

.. autoclass:: torch.distributed._local_tensor.LocalTensorMode
   :members: disable, rank_map, tensor_map

.. autoclass:: torch.distributed._local_tensor.LocalIntNode

Utility Functions

{eval-rst}

.. autofunction:: torch.distributed._local_tensor.local_tensor_mode
.. autofunction:: torch.distributed._local_tensor.enabled_local_tensor_mode
.. autofunction:: torch.distributed._local_tensor.maybe_run_for_local_tensor
.. autofunction:: torch.distributed._local_tensor.maybe_disable_local_tensor_mode

Best Practices

Use for Testing Only: LocalTensor has significant overhead and should not be used in production code.
Initialize Process Groups: Even for local testing, you need to initialize a process group (use the "fake" backend).
Avoid requires_grad on Inner Tensors: LocalTensor expects inner tensors to not have requires_grad=True. Set gradients on the LocalTensor wrapper instead.
Reconcile for Assertions: Use reconcile() to extract a single tensor when all ranks should have identical values (e.g., after an all-reduce).
Debug with Direct Access: Access individual shards via tensor._local_tensors[rank] for debugging.

Common Pitfalls

Forgetting the Context Manager: Operations on LocalTensor outside LocalTensorMode still work but won't create new LocalTensors from factories.
Mismatched Ranks: Ensure all LocalTensors in an operation have compatible ranks.
Inner Tensor Gradients: Creating LocalTensor from tensors with requires_grad=True will raise an error.