.. DO NOT EDIT.
.. THIS FILE WAS AUTOMATICALLY GENERATED BY SPHINX-GALLERY.
.. TO MAKE CHANGES, EDIT THE SOURCE PYTHON FILE:
.. "generated/examples/ddp/spmd_ddp.py"
.. LINE NUMBERS ARE GIVEN BELOW.

.. only:: html

    .. note::
        :class: sphx-glr-download-link-note

        :ref:`Go to the end <sphx_glr_download_generated_examples_ddp_spmd_ddp.py>`
        to download the full example code

.. rst-class:: sphx-glr-example-title

.. _sphx_glr_generated_examples_ddp_spmd_ddp.py:


DDP using SPMDActor
============================

This example demonstrates how to run PyTorch's Distributed Data Parallel (DDP)
using Monarch's ``SPMDActor``. The actor configures torch elastic environment
variables and executes the training script, replicating torchrun behavior.

This example shows:

- How to spawn a process mesh on the local host
- How to use ``SPMDActor`` to run a DDP training script
- How ``SPMDActor`` configures RANK, LOCAL_RANK, WORLD_SIZE, etc.

Training Script
---------------

The training script (``train.py``) is a standard PyTorch DDP script::

    import os

    import torch
    import torch.distributed as dist
    import torch.nn as nn
    import torch.optim as optim
    from torch.nn.parallel import DistributedDataParallel as DDP


    def main():
        dist.init_process_group("nccl")
        rank = dist.get_rank()
        local_rank = int(os.environ["LOCAL_RANK"])
        torch.cuda.set_device(local_rank)

        model = nn.Linear(10, 1).cuda()
        ddp_model = DDP(model)

        optimizer = optim.SGD(ddp_model.parameters(), lr=0.01)

        for step in range(5):
            inputs = torch.randn(4, 10).cuda()
            outputs = ddp_model(inputs)
            loss = outputs.sum()
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
            print(f"[Rank {rank}] Step {step} loss={loss.item()}")

        dist.destroy_process_group()


    if __name__ == "__main__":
        main()

This script:

- Initializes NCCL process group (environment variables set by ``SPMDActor``)
- Creates a simple linear model wrapped in DDP
- Runs 5 training steps
- Cleans up the process group

.. GENERATED FROM PYTHON SOURCE LINES 70-73

Imports
-------
We import Monarch's actor API and SPMDActor.

.. GENERATED FROM PYTHON SOURCE LINES 73-87

.. code-block:: default


    import os

    from monarch.actor import this_host
    from monarch.spmd import SPMDActor


    GPUS_PER_HOST = 4

    # Get absolute path to train.py (in the same directory as this script)
    SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))
    TRAIN_SCRIPT = os.path.join(SCRIPT_DIR, "train.py")



.. GENERATED FROM PYTHON SOURCE LINES 88-91

Create Process Mesh
-------------------
Spawn a process mesh on the local host with 4 GPU processes.

.. GENERATED FROM PYTHON SOURCE LINES 91-95

.. code-block:: default


    local_proc_mesh = this_host().spawn_procs(per_host={"gpus": GPUS_PER_HOST})



.. GENERATED FROM PYTHON SOURCE LINES 96-101

Run DDP Training with SPMDActor
-------------------------------
Spawn ``SPMDActor`` on the process mesh. The actor configures torch elastic
environment variables (RANK, LOCAL_RANK, WORLD_SIZE, MASTER_ADDR, MASTER_PORT)
and executes the training script.

.. GENERATED FROM PYTHON SOURCE LINES 101-114

.. code-block:: default


    spmd_actors = local_proc_mesh.spawn("_SPMDActor", SPMDActor)

    # Get master address/port from first actor (all coordinates = 0)
    first_values = dict.fromkeys(local_proc_mesh._labels, 0)
    master_addr, master_port = (
        spmd_actors.slice(**first_values).get_host_port.call_one(None).get()
    )

    # Execute training script across the mesh
    spmd_actors.main.call(master_addr, master_port, [TRAIN_SCRIPT]).get()

    print("DDP example completed successfully!")


.. rst-class:: sphx-glr-timing

   **Total running time of the script:** (0 minutes 0.000 seconds)


.. _sphx_glr_download_generated_examples_ddp_spmd_ddp.py:

.. only:: html

  .. container:: sphx-glr-footer sphx-glr-footer-example




    .. container:: sphx-glr-download sphx-glr-download-python

      :download:`Download Python source code: spmd_ddp.py <spmd_ddp.py>`

    .. container:: sphx-glr-download sphx-glr-download-jupyter

      :download:`Download Jupyter notebook: spmd_ddp.ipynb <spmd_ddp.ipynb>`


.. only:: html

 .. rst-class:: sphx-glr-signature

    `Gallery generated by Sphinx-Gallery <https://sphinx-gallery.github.io>`_