chore: replace decord2 with torchvision/torchaudio for media loading (#3107)

kajalj22 · claude · chtruong814 · commit 15700a9d8adf · 2026-07-17T20:59:30.000-05:00
Signed-off-by: Kajal Jain &lt;kajalj@nvidia.com&gt;
Co-authored-by: Claude Sonnet 4.6 &lt;noreply@anthropic.com&gt;
Signed-off-by: Charlie Truong &lt;chtruong@nvidia.com&gt;
diff --git a/docs/guides/grpo-audio.md b/docs/guides/grpo-audio.md
@@ -1,5 +1,14 @@
 # Audio Post-training with NeMo RL
 
+> **Audio dependencies are not pre-installed in the NeMo-RL container.**
+> Run the following script once before training or evaluation:
+>
+> ```bash
+> bash tools/install_audio_deps.sh
+> ```
+>
+> The script is a no-op if they are already installed. Audio tests run it automatically.
+
 This guide explains how to use NeMo RL to train [Qwen2.5-Omni](https://huggingface.co/Qwen) (3B or 7B) and [Qwen3-Omni-30B-A3B-Instruct](https://huggingface.co/Qwen/Qwen3-Omni-30B-A3B-Instruct) with GRPO on audio question-answering data, convert the resulting Megatron checkpoint to Hugging Face format, and evaluate it on the [MMAU benchmark](https://huggingface.co/datasets/TwinkStart/MMAU).
 
 NeMo RL ships three recipes out of the box, but the pieces are independent and can be mixed:
diff --git a/nemo_rl/data/multimodal_utils.py b/nemo_rl/data/multimodal_utils.py
@@ -20,7 +20,6 @@
 from io import BytesIO
 from typing import Any, Optional, Union
 
-import decord
 import requests
 import torch
 from PIL import Image
@@ -357,17 +356,19 @@ def load_media_from_message(
                         load_audio(aud, **multimodal_load_kwargs["audio"])
                     )
                 except (RuntimeError, FileNotFoundError, OSError) as e:
-                    logger.warning("Audio loading failed. Fall back to decord.")
-                    # use decord
-                    loaded_audio = decord.AudioReader(
-                        aud,
-                        sample_rate=multimodal_load_kwargs["audio"]["sampling_rate"],
-                        mono=True,
-                    )
+                    logger.warning("Audio loading failed. Falling back to torchaudio.")
+                    import torchaudio
+
+                    waveform, sr = torchaudio.load(aud)
+                    target_sr = multimodal_load_kwargs["audio"]["sampling_rate"]
+                    if sr != target_sr:
+                        waveform = torchaudio.functional.resample(
+                            waveform, sr, target_sr
+                        )
+                    if waveform.shape[0] > 1:
+                        waveform = waveform.mean(0, keepdim=True)
                     loaded_media["audio"].append(
-                        loaded_audio[:].asnumpy()[
-                            get_dim_to_pack_along(processor, "audio")
-                        ]
+                        waveform.numpy()[get_dim_to_pack_along(processor, "audio")]
                     )
             else:
                 loaded_media["audio"].append(aud)
@@ -379,9 +380,8 @@ def load_media_from_message(
                     if "video" in multimodal_load_kwargs
                     else {}
                 )
-                # seems decord backend loads video faster with multithread ffmpeg and it is easier to install
                 loaded_media["video"].append(
-                    load_video(vid, backend="decord", **load_video_kwargs)[0]
+                    load_video(vid, backend="torchcodec", **load_video_kwargs)[0]
                 )
             else:
                 loaded_media["video"].append(vid)
diff --git a/nemo_rl/data_plane/codec.py b/nemo_rl/data_plane/codec.py
@@ -261,7 +261,7 @@ def materialize(
     The lazy ``BatchedDataDict`` import keeps
     ``import nemo_rl.data_plane`` cheap for unit tests that don't
     actually call this function (``BatchedDataDict`` transitively
-    pulls multimodal deps like decord / torchvision).
+    pulls multimodal deps like torchvision / torchaudio).
 
     Args:
         td: Wire TensorDict to materialize.
diff --git a/pyproject.toml b/pyproject.toml
@@ -50,6 +50,8 @@ dependencies = [
   "sympy>=1.14.0",
   "pillow>=12.1.1",
   "torchvision==0.26.0",
+  # torchaudio and torchcodec are NOT installed in the container image.
+  # Run tools/install_audio_deps.sh before using audio features or audio tests.
   # Keep the base range broad enough while allowing backend extras to tighten as needed.
   # Floor raised to 5.5.0 for Gemma 4 support.
   "transformers>=5.5.0,<5.9.0",
@@ -60,7 +62,6 @@ dependencies = [
   "nvidia-nvshmem-cu13>=3.6.5; sys_platform == 'linux' and (platform_machine == 'x86_64' or platform_machine == 'aarch64')", # for deep_ep build
   "swanlab",
   "pyzmq",
-  "decord2",
   "soundfile>=0.13.1",
   "nccl4py; sys_platform != 'darwin'",                                                                                       # for non-colocated refit
   "cuda-bindings; sys_platform != 'darwin'",                                                                                 # for non-colocated refit
@@ -255,7 +256,7 @@ nemo-automodel = { path = "3rdparty/Automodel-workspace/Automodel", editable = t
 megatron-bridge = { path = "3rdparty/Megatron-Bridge-workspace", editable = true }
 nemo_gym = { workspace = true }
 nemo_run = { git = "https://github.com/NVIDIA-NeMo/Run", rev = "414f0077c648fde2c71bb1186e97ccbf96d6844c" }
-# torch/torchaudio/torchvision/triton all come from the torch index in order to pick up aarch64 wheels
+# torch/torchvision/triton all come from the torch index in order to pick up aarch64 wheels
 torch = [
   { index = "pytorch-cu130", marker = "sys_platform != 'darwin'" },
   { index = "pypi", marker = "sys_platform == 'darwin'" },
@@ -264,10 +265,6 @@ torchvision = [
   { index = "pytorch-cu130", marker = "sys_platform != 'darwin'" },
   { index = "pypi", marker = "sys_platform == 'darwin'" },
 ]
-torchaudio = [
-  { index = "pytorch-cu130", marker = "sys_platform != 'darwin'" },
-  { index = "pypi", marker = "sys_platform == 'darwin'" },
-]
 triton = [
   { index = "pytorch-cu130", marker = "sys_platform != 'darwin'" },
   { index = "pypi", marker = "sys_platform == 'darwin'" },
@@ -349,7 +346,6 @@ override-dependencies = [
   "timm<=1.0.22",
   "nvidia-modelopt[torch]>=0.44.0a0",
   "torch==2.11.0",
-  "torchaudio==2.11.0",
   # sglang has conflicting llguidance versions than vllm, so enforcing vllm's version since it's newer
   "llguidance>=1.3.0,<1.4.0",
   # Override setuptools range in other dependencies to address CVE GHSA-58pv-8j8x-9vj2
@@ -378,6 +374,13 @@ override-dependencies = [
   # extra composes with mcore/automodel without version-mirroring TQ's
   # requirements.txt. Forward-compatible across TQ minor bumps.
   "numpy>=2.1.0",
+  # av (PyAV) carries CVE-bundled codec libs (libx264, libx265, libopenh264, libmp3lame).
+  # It is only needed by megatron-bridge's optional WAN diffusion path, which installs it
+  # at test time via scripts/install_diffusion_deps.sh. Exclude from the shipped container.
+  "av; sys_platform == 'never'",
+  # decord2 is pulled transitively by sglang → qwen-vl-utils[decord] on aarch64 Linux.
+  # NeMo-RL no longer calls any decord API (replaced by torchvision/torchaudio).
+  "decord2; sys_platform == 'never'",
 ]
 
 # CVE fixes
diff --git a/pyrefly.toml b/pyrefly.toml
@@ -10,6 +10,8 @@ replace-imports-with-any = [
   "math_verify.*",
   "sympy.*",
   "torchdata.*",
+  "torchaudio.*",
+  "torchcodec.*",
   "nemo.*",
   "megatron.*",
   "ray.*",
diff --git a/tests/functional/audio_grpo_megatron.sh b/tests/functional/audio_grpo_megatron.sh
@@ -7,6 +7,9 @@ git config --global --add safe.directory $PROJECT_ROOT
 
 set -eou pipefail
 
+# Audio deps (torchaudio/torchcodec) are not in the shipped container; install at test time.
+bash "$PROJECT_ROOT/tools/install_audio_deps.sh"
+
 EXP_NAME=$(basename $0 .sh)
 EXP_DIR=$SCRIPT_DIR/$EXP_NAME
 LOG_DIR=$EXP_DIR/logs
diff --git a/tests/functional/eval_audio.sh b/tests/functional/eval_audio.sh
@@ -7,6 +7,9 @@ git config --global --add safe.directory $PROJECT_ROOT
 
 set -eou pipefail
 
+# Audio deps (torchaudio/torchcodec) are not in the shipped container; install at test time.
+bash "$PROJECT_ROOT/tools/install_audio_deps.sh"
+
 EXP_NAME=$(basename $0 .sh)
 EXP_DIR=$SCRIPT_DIR/$EXP_NAME
 LOG_DIR=$EXP_DIR/logs
diff --git a/tests/functional/sft_avlm.sh b/tests/functional/sft_avlm.sh
@@ -10,6 +10,9 @@ git config --global --add safe.directory $PROJECT_ROOT
 
 set -eou pipefail
 
+# Audio/video deps (torchaudio/torchcodec/ffmpeg) are not in the shipped container.
+bash "$PROJECT_ROOT/tools/install_audio_deps.sh"
+
 EXP_NAME=$(basename $0 .sh)
 EXP_DIR=$SCRIPT_DIR/$EXP_NAME
 LOG_DIR=$EXP_DIR/logs
diff --git a/tests/test_suites/llm/common.env b/tests/test_suites/llm/common.env
@@ -9,6 +9,11 @@ SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd)
 # catch-22 to get the project root and mark it safe if you don't know the project root
 PROJECT_ROOT=$(realpath $SCRIPT_DIR/../../..)
 
+# Auto-install av deps for all VLM scripts and audio/omni/avqa scripts (no-op for LLM tests).
+if [[ "${BASH_SOURCE[1]:-}" =~ (omni|audio|avqa|vlm) ]]; then
+    bash "$PROJECT_ROOT/tools/install_audio_deps.sh"
+fi
+
 exit_if_max_steps_reached() {
   # Early stopping to save compute if max step has been reached
   STEPS_SO_FAR=$(jq 'to_entries | .[] | select(.key == "train/loss") | .value | keys | map(tonumber) | max' $JSON_METRICS || echo 0)
diff --git a/tests/unit/L0_Unit_Tests_Data.sh b/tests/unit/L0_Unit_Tests_Data.sh
@@ -17,4 +17,8 @@
 
 source "$(dirname "${BASH_SOURCE[0]}")/run_unit_shard_common.sh"
 
+# Audio/video deps (torchaudio/torchcodec/ffmpeg) are not in the shipped container.
+# Data unit tests (e.g. dailyomni) decode real audio/video via these deps.
+bash "$PROJECT_ROOT/tools/install_audio_deps.sh"
+
 uv run --no-sync bash -x ./tests/run_unit.sh "unit/data/" "${EXCLUDED_UNIT_TESTS[@]}" --cov=nemo_rl --cov-report=term-missing --cov-report=json --hf-gated
diff --git a/tools/install_audio_deps.sh b/tools/install_audio_deps.sh
@@ -0,0 +1,29 @@
+#!/bin/bash
+# Install audio dependencies that are NOT shipped in the NeMo-RL container.
+#
+# Run this script before using audio features or running audio tests:
+#
+#   bash tools/install_audio_deps.sh
+#
+# Safe to call multiple times — exits immediately if already installed.
+set -euo pipefail
+
+# Fast exit: if torchcodec imports cleanly it already has FFmpeg available.
+if python -c "import torchcodec" 2>/dev/null; then
+    echo "[audio-deps] Already installed and functional, skipping."
+    exit 0
+fi
+
+# Install system FFmpeg — torchcodec dlopens libavcodec.so.* at runtime.
+echo "[audio-deps] Installing system FFmpeg..."
+apt-get update && apt-get install -y --no-install-recommends ffmpeg
+
+# torchaudio 2.11+ routes torchaudio.load through torchcodec, so both are needed.
+echo "[audio-deps] Installing torchaudio==2.11.0 and torchcodec..."
+uv pip install \
+    --index-url https://download.pytorch.org/whl/cu130 \
+    --extra-index-url https://pypi.org/simple \
+    --reinstall-package torchaudio \
+    "torchaudio==2.11.0" \
+    "torchcodec>=0.3.0"
+echo "[audio-deps] Done."
diff --git a/uv.lock b/uv.lock
