chore: import upstream snapshot with attribution

docs/source/features/mixed_precision.rst

@@ -0,0 +1,160 @@
+.. _mix_precision:
+
+Mixed Precision Training
+========================
+
+Mixed precision training enhances computational efficiency by conducting operations in low-precision
+format while selectively maintaining critical data in single-precision. NeMo supports FP16 and BF16
+precision via PyTorch Lightning, in mixed, true, and flash half-precision modes.
+
+Precision Modes
+---------------
+
+PyTorch Lightning provides two categories of half-precision training:
+
+**Mixed Precision** (``"bf16-mixed"`` / ``"16-mixed"``):
+    Operations run in half-precision where safe, but model weights are kept in FP32.
+    Gradients are computed in half-precision and accumulated in FP32. This is the safest
+    option and generally a good default for ASR and TTS training.
+
+**True Half Precision** (``"bf16-true"`` / ``"fp16-true"``):
+    The entire model -- weights, activations, and gradients -- runs in half-precision.
+    This uses less memory than mixed precision (no FP32 weight copy) and is faster,
+    but requires the model to be numerically stable in half-precision.
+    SpeechLM2 models use ``"bf16-true"`` by default for training.
+
+**Flash Precision** (``"bf16-flash"`` / ``"fp16-flash"``):
+    The model also runs in half-precision, but NeMo avoids Lightning's global
+    default-dtype override and autocast context. This mode is intended for use
+    with FlashOptim, a library of drop-in optimizers that reduces training
+    memory by shrinking optimizer states, master weights, and gradients. In
+    practice, this may be a better fit than AMP / mixed precision when
+    optimizer-state memory or checkpoint size is the bottleneck, and may lead to
+    improved convergence compared to Lightning's true half-precision as it keeps
+    track of the residual between half and full precision weights.
+
+Configuration
+-------------
+
+Set precision through the PyTorch Lightning trainer's ``precision`` argument.
+
+In YAML (with Hydra):
+
+.. code-block:: yaml
+
+    trainer:
+        precision: "bf16-mixed"    # BF16 mixed precision
+        # precision: "16-mixed"    # FP16 mixed precision
+        # precision: "bf16-true"   # True BF16 half precision
+        # precision: "fp16-true"   # True FP16 half precision
+        # precision: "bf16-flash"  # BF16 flash precision
+        # precision: "fp16-flash"  # FP16 flash precision
+
+In Python:
+
+.. code-block:: python
+
+    import lightning.pytorch as pl
+
+    trainer = pl.Trainer(
+        precision="bf16-mixed",
+        devices=2,
+        accelerator="gpu",
+    )
+
+Choosing a Precision Format
+----------------------------
+
+- **BF16** has the same dynamic range as FP32, which makes it more numerically stable and generally
+  easier to use. It is the recommended choice for most Speech AI training workloads.
+- **FP16** offers slightly higher throughput on some hardware but has a reduced dynamic range.
+  In mixed precision mode, PyTorch Lightning handles loss scaling automatically.
+
+HalfPrecisionForAudio
+----------------------
+
+Audio waveform tensors are sensitive to precision loss -- downcasting raw audio samples to half-precision
+can degrade signal quality and hurt model accuracy. NeMo provides the ``HalfPrecisionForAudio`` plugin
+(in ``nemo.utils.trainer_utils``) that extends Lightning's ``HalfPrecision`` plugin to preserve
+full-precision for audio tensors while still casting all other inputs to half-precision.
+
+Specifically, when the training mini-batch is a dictionary, any tensor whose key contains
+the substring ``"audio"`` is kept in its original precision (typically FP32). All other floating-point
+tensors are cast to the target half-precision dtype.
+
+This plugin is used automatically when you launch training with NeMo's ``resolve_trainer_cfg``
+utility (used by many NeMo example training scripts). When the trainer config specifies
+``precision: "bf16-true"`` or ``precision: "fp16-true"``, ``resolve_trainer_cfg`` replaces
+the precision setting with the ``HalfPrecisionForAudio`` plugin:
+
+.. code-block:: python
+
+    from nemo.utils.trainer_utils import resolve_trainer_cfg
+
+    # In YAML: trainer.precision = "bf16-true"
+    # resolve_trainer_cfg automatically installs HalfPrecisionForAudio
+    trainer = pl.Trainer(**resolve_trainer_cfg(cfg.trainer))
+
+If you construct the trainer manually, you can install the plugin directly:
+
+.. code-block:: python
+
+    from nemo.utils.trainer_utils import HalfPrecisionForAudio
+
+    trainer = pl.Trainer(
+        plugins=[HalfPrecisionForAudio("bf16-true")],
+        devices=2,
+        accelerator="gpu",
+    )
+
+FlashPrecision
+---------------
+
+NeMo provides the ``FlashPrecision`` plugin (in
+``nemo.utils.trainer_utils``) primarily for FlashOptim-backed training.
+According to the official FlashOptim README, FlashOptim provides drop-in
+optimizer replacements that reduce training memory by compressing optimizer
+states, master weights, and gradients while preserving the standard PyTorch
+optimizer API.
+
+FlashOptim generally expects the model parameters to already be in bf16/fp16,
+while the optimizer manages reduced-precision state and master-weight
+correction internally. ``FlashPrecision`` fits that model: it preserves the
+same audio-aware input casting behavior as ``HalfPrecisionForAudio``, but does
+not enter autocast and does not change PyTorch's global default dtype. This
+avoids layering Lightning's global precision policy on top of FlashOptim's own
+reduced-precision optimizer behavior.
+
+When the trainer config specifies ``precision: "bf16-flash"`` or
+``precision: "fp16-flash"``, ``resolve_trainer_cfg`` replaces the precision
+setting with the ``FlashPrecision`` plugin:
+
+.. code-block:: python
+
+    from nemo.utils.trainer_utils import resolve_trainer_cfg
+
+    # In YAML: trainer.precision = "bf16-flash"
+    trainer = pl.Trainer(**resolve_trainer_cfg(cfg.trainer))
+
+If you construct the trainer manually, you can install the plugin directly:
+
+.. code-block:: python
+
+    from nemo.utils.trainer_utils import FlashPrecision
+
+    trainer = pl.Trainer(
+        plugins=[FlashPrecision("bf16-flash")],
+        devices=2,
+        accelerator="gpu",
+    )
+
+If you're going to use ``FlashPrecision``, make sure to set up ``flashoptim`` optimizer, e.g.:
+
+.. code-block:: yaml
+
+    optimizer:
+      _target_: flashoptim.FlashAdamW
+      lr: 1e-4
+      betas: [0.9, 0.999]
+      weight_decay: 5e-2
+

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.