Skip to content
Merged
Show file tree
Hide file tree
Changes from 24 commits
Commits
Show all changes
27 commits
Select commit Hold shift + click to select a range
0d90010
chore: replace decord2 with torchvision/torchaudio for media loading
kajalj22 Jul 7, 2026
aff511d
chore: update uv.lock to remove decord2
kajalj22 Jul 7, 2026
25b19e2
chore: merge main into replace-decord-torchvision
kajalj22 Jul 7, 2026
2445e54
Merge remote-tracking branch 'origin/main' into replace-decord-torchv…
kajalj22 Jul 9, 2026
029b9ee
fix: replace decord in datasets/utils.py and add torchaudio to pyrefly
kajalj22 Jul 9, 2026
7bc9016
chore: exclude av (PyAV) from container via uv override
kajalj22 Jul 9, 2026
a66a82e
fix: replace remaining decord backend in processors.py; exclude decor…
kajalj22 Jul 9, 2026
3e21f62
Merge remote-tracking branch 'origin/main' into replace-decord-torchv…
kajalj22 Jul 9, 2026
e52aac7
fix: add torchaudio to main project dependencies
kajalj22 Jul 10, 2026
cd54863
fix: add torchcodec to main deps for torchaudio video→audio decoding
kajalj22 Jul 12, 2026
8dba7d2
feat: exclude torchaudio FFmpeg build from container; add audio insta…
kajalj22 Jul 13, 2026
3e22196
docs: remove internal compliance language from comments and docs
kajalj22 Jul 13, 2026
a49e732
refactor: consolidate audio-deps install into common.env name check
kajalj22 Jul 13, 2026
3204c3c
chore: merge main into replace-decord-torchvision
kajalj22 Jul 13, 2026
82caa36
fix: remove torchcodec from audio install script
kajalj22 Jul 15, 2026
c2c2384
fix: install ffmpeg at test time; switch video backend to torchcodec
kajalj22 Jul 15, 2026
18361c7
chore: merge main into replace-decord-torchvision
kajalj22 Jul 15, 2026
b755507
fix: add apt-get update before ffmpeg install; add audio deps to sft_…
kajalj22 Jul 15, 2026
55fc91b
fix: expand av-deps hook to cover all vlm test scripts
kajalj22 Jul 15, 2026
436db97
Merge remote-tracking branch 'origin/main' into replace-decord-torchv…
kajalj22 Jul 15, 2026
7dd4251
fix: install audio/video deps before data unit tests in GitHub CI
kajalj22 Jul 16, 2026
2fe5480
Merge remote-tracking branch 'origin/main' into replace-decord-torchv…
kajalj22 Jul 16, 2026
6204d56
docs: add audio/video dependency install section to grpo-audio-visual…
kajalj22 Jul 16, 2026
bf18ead
Merge remote-tracking branch 'origin/main' into replace-decord-torchv…
kajalj22 Jul 16, 2026
17e1154
refactor: move install_audio_deps.sh from scripts/ to tools/
kajalj22 Jul 16, 2026
3a0d641
Merge branch 'main' into replace-decord-torchvision
kajalj22 Jul 16, 2026
0af12fe
Merge branch 'main' into replace-decord-torchvision
chtruong814 Jul 17, 2026
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
12 changes: 12 additions & 0 deletions docs/guides/grpo-audio-visual.md
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,18 @@ This guide explains how to use NeMo RL to train [Qwen2.5-Omni-7B](https://huggin

Each training sample feeds the Qwen2.5-Omni processor both the video stream (8 frames) and the audio track decoded from the same file at 16 kHz mono. Audio and video flow as two **independent multimodal items** per prompt: the dataset emits `{type: video}` + `{type: audio}` content items, the Qwen2.5-Omni chat template renders both `<|VIDEO|>` and `<|AUDIO|>` placeholders, and vLLM rollouts populate `multi_modal_data["video"]` and `multi_modal_data["audio"]` from the same sample.

## 0. Install Audio/Video Dependencies

The NeMo RL container does not ship `torchaudio`, `torchcodec`, or system FFmpeg. Run the helper script once before training or evaluation:

```bash
bash scripts/install_audio_deps.sh
```

This installs system FFmpeg (required by `torchcodec` at runtime) and pins `torchaudio==2.11.0` + `torchcodec>=0.3.0` into the container venv. The script is idempotent — re-running it on a machine where the deps are already present exits immediately.

> **Note:** `decord` and `av` (PyAV) are **not** used. Audio and video are decoded via `torchcodec` (backed by system FFmpeg) and `torchaudio` respectively.

## 1. Train the Model

Run GRPO training with the provided config:
Expand Down
9 changes: 9 additions & 0 deletions docs/guides/grpo-audio.md
Original file line number Diff line number Diff line change
@@ -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 scripts/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:
Expand Down
22 changes: 9 additions & 13 deletions nemo_rl/data/datasets/utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -37,19 +37,15 @@


def load_audio_from_file(path: str, sampling_rate: int = 16000) -> np.ndarray:
"""Decode an audio file (or the audio track of a video) as a 1-D float32 array.

Uses decord's ``AudioReader`` (already a project dependency for video
decoding) to produce a mono waveform at the requested sampling rate.
"""
import decord

reader = decord.AudioReader(path, sample_rate=sampling_rate, mono=True)
# Shape: (channels, T). With mono=True channels=1; squeeze to (T,).
audio = reader[:].asnumpy()
if audio.ndim > 1:
audio = audio[0]
return audio.astype(np.float32)
"""Decode an audio file (or the audio track of a video) as a 1-D float32 array."""
import torchaudio

waveform, sr = torchaudio.load(path)
if sr != sampling_rate:
waveform = torchaudio.functional.resample(waveform, sr, sampling_rate)
if waveform.shape[0] > 1:
waveform = waveform.mean(0, keepdim=True)
return waveform.squeeze(0).numpy().astype(np.float32)


def assert_no_double_bos(token_ids: torch.Tensor, tokenizer: TokenizerType) -> None:
Expand Down
26 changes: 13 additions & 13 deletions nemo_rl/data/multimodal_utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -20,7 +20,6 @@
from io import BytesIO
from typing import Any, Optional, Union

import decord
import requests
import torch
from PIL import Image
Expand Down Expand Up @@ -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)
Expand All @@ -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]
Comment thread
kajalj22 marked this conversation as resolved.
)
else:
loaded_media["video"].append(vid)
Expand Down
2 changes: 1 addition & 1 deletion nemo_rl/data/processors.py
Original file line number Diff line number Diff line change
Expand Up @@ -538,7 +538,7 @@ def vlm_hf_data_processor(
video_value = content["video"]
if isinstance(video_value, str):
video_value = load_video(
video_value, backend="decord", **load_video_kwargs
video_value, backend="torchcodec", **load_video_kwargs
)[0]
# Replace path with loaded frames so apply_chat_template can consume it
user_message["content"].append({"type": "video", "video": video_value})
Expand Down
2 changes: 1 addition & 1 deletion nemo_rl/data_plane/codec.py
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down
15 changes: 9 additions & 6 deletions pyproject.toml
Original file line number Diff line number Diff line change
Expand Up @@ -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 scripts/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",
Expand All @@ -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
Expand Down Expand Up @@ -264,10 +265,6 @@ torchvision = [
{ index = "pytorch-cu130", marker = "sys_platform != 'darwin'" },
{ index = "pypi", marker = "sys_platform == 'darwin'" },
]
torchaudio = [
Comment thread
kajalj22 marked this conversation as resolved.
{ 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'" },
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down
2 changes: 2 additions & 0 deletions pyrefly.toml
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,8 @@ replace-imports-with-any = [
"math_verify.*",
"sympy.*",
"torchdata.*",
"torchaudio.*",
"torchcodec.*",
"nemo.*",
"megatron.*",
"ray.*",
Expand Down
29 changes: 29 additions & 0 deletions scripts/install_audio_deps.sh
Original file line number Diff line number Diff line change
@@ -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 scripts/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."
3 changes: 3 additions & 0 deletions tests/functional/audio_grpo_megatron.sh
Original file line number Diff line number Diff line change
Expand Up @@ -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/scripts/install_audio_deps.sh"

EXP_NAME=$(basename $0 .sh)
EXP_DIR=$SCRIPT_DIR/$EXP_NAME
LOG_DIR=$EXP_DIR/logs
Expand Down
3 changes: 3 additions & 0 deletions tests/functional/eval_audio.sh
Original file line number Diff line number Diff line change
Expand Up @@ -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/scripts/install_audio_deps.sh"

EXP_NAME=$(basename $0 .sh)
EXP_DIR=$SCRIPT_DIR/$EXP_NAME
LOG_DIR=$EXP_DIR/logs
Expand Down
3 changes: 3 additions & 0 deletions tests/functional/eval_daily_omni.sh
Original file line number Diff line number Diff line change
Expand Up @@ -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/scripts/install_audio_deps.sh"

EXP_NAME=$(basename $0 .sh)
EXP_DIR=$SCRIPT_DIR/$EXP_NAME
LOG_DIR=$EXP_DIR/logs
Expand Down
3 changes: 3 additions & 0 deletions tests/functional/sft_avlm.sh
Original file line number Diff line number Diff line change
Expand Up @@ -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/scripts/install_audio_deps.sh"

EXP_NAME=$(basename $0 .sh)
EXP_DIR=$SCRIPT_DIR/$EXP_NAME
LOG_DIR=$EXP_DIR/logs
Expand Down
5 changes: 5 additions & 0 deletions tests/test_suites/llm/common.env
Original file line number Diff line number Diff line change
Expand Up @@ -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/scripts/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)
Expand Down
4 changes: 4 additions & 0 deletions tests/unit/L0_Unit_Tests_Data.sh
Original file line number Diff line number Diff line change
Expand Up @@ -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/scripts/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
Loading
Loading