nvidia-nemo--speech/docs/source/features/parallelisms.rst

.. _parallelisms:

Parallelisms
============

NeMo uses native PyTorch parallelism primitives for distributed training, enabling efficient multi-GPU and multi-node
model training for Speech AI workloads.

DDP (all collections)
---------------------

Distributed Data Parallelism (DDP) is the default strategy for all NeMo collections (ASR, TTS, Audio, SpeechLM2).
It replicates the entire model on every GPU, runs each GPU on a different data shard, and synchronizes
parameter gradients via all-reduce after each backward pass.

**When to use:** DDP works well when the full model fits in a single GPU's memory.
This covers the vast majority of ASR, TTS, and Audio training workloads.

DDP is enabled by default in NeMo. You can configure it explicitly in YAML:

.. code-block:: yaml

    trainer:
        strategy:
            _target_: lightning.pytorch.strategies.DDPStrategy
            gradient_as_bucket_view: true
            find_unused_parameters: true

Or in Python:

.. code-block:: python

    from lightning.pytorch.strategies import DDPStrategy

    trainer = pl.Trainer(
        strategy=DDPStrategy(gradient_as_bucket_view=True, find_unused_parameters=True),
        devices=8,
        accelerator="gpu",
    )

AutomodelParallelStrategy (SpeechLM2)
-------------------------------------

For SpeechLM2 models that use NeMo Automodel (for example ``SALMAutomodel``), the backbone LLM can be
too large for a single GPU. NeMo provides
``nemo.collections.speechlm2.parts.parallel.AutomodelParallelStrategy``, a Lightning strategy that
delegates device mesh creation to NeMo Automodel and supports FSDP2, Tensor Parallelism (TP),
Sequence Parallelism (SP), Context Parallelism (CP), Expert Parallelism (EP) for MoE models, and
Hybrid Sharded Data Parallelism (HSDP).

**When to use:** When training or fine-tuning SpeechLM2 models whose LLM backbone does not fit
in a single GPU's memory, or when you want to scale training to many GPUs more efficiently
than DDP allows. Use ``AutomodelParallelStrategy`` for ``SALMAutomodel`` and MoE LLM backbones such
as NVIDIA Nemotron Nano V3.

**Requirements:** Each model must implement a ``configure_model()`` method that defines how its
layers are sharded and parallelized. ``SALMAutomodel`` already implements this and receives the
Automodel device mesh during ``configure_model()``. You cannot simply switch an arbitrary model
from DDP to ``AutomodelParallelStrategy`` without providing this implementation.

Concepts
^^^^^^^^

**FSDP2 (Fully Sharded Data Parallelism):**
    Shards model parameters, gradients, and optimizer states across GPUs in the data-parallel
    dimension. Dramatically reduces per-GPU memory -- enabling training of models that would not
    fit with DDP. Controlled via ``dp_size``; when ``dp_size`` is ``null``, NeMo Automodel infers
    it from the world size and the other parallelism dimensions.

**Tensor Parallelism (TP):**
    Splits individual weight matrices across GPUs. For example, a large linear layer's weight
    is partitioned column-wise or row-wise so each GPU holds only a slice. Controlled via
    ``tp_size``. The model must define a TP sharding plan (which layers are split and how).
    Automodel-backed SpeechLM2 models use the Automodel plan for the backbone LLM.

**Sequence Parallelism (SP):**
    Distributes activation memory along the sequence dimension across the TP group.
    SP is typically enabled alongside TP and reduces activation memory further. Enable it with
    ``distributed_config.sequence_parallel: true``.

**Context Parallelism (CP):**
    Splits long-context sequence processing across GPUs in the context-parallel group. Controlled
    via ``cp_size``. For SpeechLM2 models, CP is intended for packed-sequence training where each
    utterance is handled as its own attention segment.

**Expert Parallelism (EP):**
    Routes MoE experts across GPUs for MoE LLM backbones. Controlled via ``ep_size``. EP reuses
    the FSDP2 data-parallel axis: dense layers are sharded via FSDP2, while MoE expert layers use
    all-to-all expert routing on the same ranks.

**Hybrid Sharded Data Parallelism (HSDP):**
    Adds replication groups around FSDP2 sharding. Controlled via ``dp_replicate_size``.

Configuration
^^^^^^^^^^^^^

To enable ``AutomodelParallelStrategy`` for Automodel-backed SpeechLM2 models, replace the DDP
strategy block in the trainer config. The configured sizes must be compatible with the total
number of GPUs (``devices * num_nodes``). Leave ``dp_size: null`` to let NeMo Automodel infer the
data-parallel size from the remaining dimensions. ``ep_size`` controls MoE expert routing on the
data-parallel axis rather than adding a separate data-parallel dimension.

In YAML (with Hydra):

.. code-block:: yaml

    trainer:
        devices: 8
        num_nodes: 1
        accelerator: gpu
        precision: bf16-true
        strategy:
            _target_: nemo.collections.speechlm2.parts.parallel.AutomodelParallelStrategy
            dp_size: null          # inferred from world_size / other dimensions
            dp_replicate_size: 1   # HSDP replication group size
            tp_size: 1             # tensor parallel size
            cp_size: 1             # context parallel size
            ep_size: 8             # expert parallel size for MoE models

            distributed_config:
                sequence_parallel: false

            activation_checkpointing_llm: false
            activation_checkpointing_perception: false

In Python:

.. code-block:: python

    from nemo.collections.speechlm2.parts.parallel import AutomodelParallelStrategy

    trainer = pl.Trainer(
        strategy=AutomodelParallelStrategy(
            dp_size=None,
            dp_replicate_size=1,
            tp_size=1,
            cp_size=1,
            ep_size=8,
        ),
        devices=8,
        accelerator="gpu",
        precision="bf16-true",
        use_distributed_sampler=False,
    )

.. note::

    When using ``AutomodelParallelStrategy``, set ``use_distributed_sampler=False`` in the trainer.
    NeMo's data modules handle distributed sampling internally.

Activation Checkpointing
^^^^^^^^^^^^^^^^^^^^^^^^

``AutomodelParallelStrategy`` exposes two activation-checkpointing knobs that can be enabled
independently:

* ``activation_checkpointing_llm`` checkpoints LLM transformer blocks. This single switch covers
  both the standard FSDP2 path and the EP/MoE parallelizer path, so use it for MoE LLM backbones
  whether ``ep_size`` is 1 or larger.
* ``activation_checkpointing_perception`` checkpoints the speech perception encoder layers before
  FSDP2 sharding.

Both options default to ``false``. Enable them to reduce activation memory at the cost of extra
recomputation during backward:

.. code-block:: yaml

    trainer:
        strategy:
            _target_: nemo.collections.speechlm2.parts.parallel.AutomodelParallelStrategy
            activation_checkpointing_llm: true
            activation_checkpointing_perception: true

Example: SALMAutomodel with FSDP2 only
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^

The simplest ``AutomodelParallelStrategy`` setup uses FSDP2 alone. This works when individual
layers fit in GPU memory:

.. code-block:: yaml

    trainer:
        devices: 8
        strategy:
            _target_: nemo.collections.speechlm2.parts.parallel.AutomodelParallelStrategy
            dp_size: 8
            tp_size: 1
            ep_size: 1

Example: SALMAutomodel with MoE Expert Parallelism
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^

For MoE LLM backbones such as NVIDIA Nemotron Nano V3, use EP to distribute experts across GPUs.
Here, the dense layers use FSDP2 and MoE layers use 8-way expert routing:

.. code-block:: yaml

    trainer:
        devices: 8
        strategy:
            _target_: nemo.collections.speechlm2.parts.parallel.AutomodelParallelStrategy
            dp_size: null
            tp_size: 1
            ep_size: 8

Example: SALMAutomodel with TP + FSDP2
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^

For larger dense LLM backbones, combine TP with FSDP2. Here, 2-way TP splits each layer across
2 GPUs and NeMo Automodel infers the FSDP2 data-parallel size from the remaining ranks:

.. code-block:: yaml

    trainer:
        devices: 8
        strategy:
            _target_: nemo.collections.speechlm2.parts.parallel.AutomodelParallelStrategy
            dp_size: null
            tp_size: 2
            ep_size: 1

ModelParallelStrategy (SALM and Duplex)
---------------------------------------

The original SpeechLM2 ``SALM`` and Duplex model configs use PyTorch Lightning's
``ModelParallelStrategy`` directly. This path is separate from ``SALMAutomodel`` and supports
FSDP2, TP, and SP using PyTorch-native DTensor.

**When to use:** Use ``ModelParallelStrategy`` for non-Automodel SpeechLM2 models, such as
``SALM`` and Duplex models. Use ``AutomodelParallelStrategy`` only for Automodel-backed models such
as ``SALMAutomodel``.

**Requirements:** As with ``AutomodelParallelStrategy``, the model must implement
``configure_model()`` to define how layers are sharded and parallelized. The SpeechLM2 SALM and
Duplex models already implement this.

ModelParallelStrategy Configuration
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^

The product of ``data_parallel_size`` and ``tensor_parallel_size`` must equal the total number of
GPUs (``devices * num_nodes``).

.. code-block:: yaml

    trainer:
        devices: 8
        num_nodes: 1
        accelerator: gpu
        precision: bf16-true
        strategy:
            _target_: lightning.pytorch.strategies.ModelParallelStrategy
            data_parallel_size: 4   # FSDP2: shard across 4 GPUs
            tensor_parallel_size: 2  # TP: split layers across 2 GPUs

.. code-block:: python

    from lightning.pytorch.strategies import ModelParallelStrategy

    trainer = pl.Trainer(
        strategy=ModelParallelStrategy(
            data_parallel_size=4,
            tensor_parallel_size=2,
        ),
        devices=8,
        accelerator="gpu",
        precision="bf16-true",
        use_distributed_sampler=False,
    )

.. note::

    When using ``ModelParallelStrategy``, set ``use_distributed_sampler=False`` in the trainer.
    NeMo's data modules handle distributed sampling internally.

See the SpeechLM2 example configs in ``examples/speechlm2/conf/`` for complete training
configurations including data and optimizer settings.