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:

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

Or in 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):

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:

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:

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:

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:

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:

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).

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

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.