chore: import upstream snapshot with attribution

docs/source/features/parallelisms.rst

@@ -0,0 +1,276 @@
+.. _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.