Add NeMo waveform audio processor (data-side feature extractor) - #5570
Merged
Conversation
Introduces a self-contained audio encoder stack under megatron/core/models/audio for multimodal (audio-text) training: - nemo_transformer_encoder: vendored NeMo ASR TransformerEncoder with conv/stacking pre-encode, QK-norm, left-context windowing, activation checkpointing, and TE/SDPA/FlashAttention backends. - nemo_transformer_audio_model: NemoTransformerAudioModel + config (mel frames -> audio embeddings; dense and packed forward paths). - nemo_audio_preprocessing_standalone: vendored log-mel preprocessor with no NeMo/Lightning/Hydra runtime dependency. - nemo_audio_checkpoint: .nemo archive extraction/loading and checkpoint-local config persistence. - audio_projector / packed_audio: projection of audio embeddings to the language-model hidden size, with packed (THD) support. - audio_feature_config: NemoAudioFeatureConfig + NemoTransformerAudioTokenEstimator (mel hyperparameters and frame->token math describing the model frontend). The module depends only on megatron.core (vision MultimodalProjector, TransformerConfig/MLP specs) plus optional transformer_engine/flash_attn/ omegaconf; it has no data-loader or Megatron-Energon dependency. The concrete data-side waveform processor is injected separately. Unit tests cover the projector, the vendored preprocessor, the encoder forward/packed paths, token estimation, and .nemo checkpoint loading. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: Yongqiang Wang <yongqiang.seagull@gmail.com>
…ocessor - Set stack_factor default to 1 in AudioProjection and NemoTransformerAudioTokenEstimator. - Rename nemo_audio_preprocessing_standalone.py to nemo_audio_preprocessing.py and update references. - Add a module docstring explaining the pure-PyTorch mel feature extractor avoids a dependency on NVIDIA-NeMo/Speech; reference the upstream feature-parity PR. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: Yongqiang Wang <yongqiang.seagull@gmail.com>
yqwangustc
force-pushed
the
audio-model-processor
branch
from
July 16, 2026 01:45
ef52b47 to
77c4cb5
Compare
yqwangustc
marked this pull request as ready for review
July 16, 2026 01:49
Contributor
Author
|
/ok to test 77c4cb5 |
Adds NemoAudioProcessor under megatron/core/models/audio: the concrete waveform -> log-mel feature extractor and token-count estimator that the multimodal data pipeline uses to size audio placeholders for packing and to materialize encoder inputs. It composes the model-frontend descriptors NemoAudioFeatureConfig and NemoTransformerAudioTokenEstimator with the vendored standalone log-mel preprocessor. The audio reference it operates on is duck-typed (Any), so this module has no dependency on the data library's AudioRef type or on Megatron-Energon — the data pipeline consumes the processor through a small structural interface (compute_num_embeddings / compute_num_frames / materialize plus the cumulative-prefix num_*_from_num_samples primitives) and injects it. Tests cover the slice primitives (cumulative-prefix invariants), slice_range waveform cropping, and log-mel materialization. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: Yongqiang Wang <yongqiang.seagull@gmail.com>
Resolve audio-file conflicts (both-added from upstream PR NVIDIA#5565): take upstream's review-polished versions for all audio modules and tests, and preserve the local-only NemoAudioProcessor export in the package __init__. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: Yongqiang Wang <yongqiang.seagull@gmail.com>
yqwangustc
force-pushed
the
audio-model-processor
branch
from
July 16, 2026 18:24
77c4cb5 to
45056ff
Compare
desh2608
approved these changes
Jul 16, 2026
Contributor
Author
|
/ok to test 45056ff |
The codecov/patch gate failed at 53% (target 80%): the data-side NemoAudioProcessor's lazy AV-decoder decode chain and branch/error paths were unexercised. Add tests covering the decoder fakes (_decode_avdecoder / _resolve_lazy_media / _audio_clip_to_float32 / _decoder_sample_rate / _load_waveform_from_spec), sample-rate and tolerance resolution, the waveform-normalization and _infer_num_samples branches (mono averaging, num_samples pad/crop, dtype/shape/slice_range errors), and the public compute_num_frames / compute_num_embeddings methods and properties. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: Yongqiang Wang <yongqiang.seagull@gmail.com>
yqwangustc
enabled auto-merge
July 20, 2026 20:37
Contributor
Author
|
/ok to test 6a2defe |
Contributor
|
🔄 Merge queue validation started! You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/29782051025 |
terminator123
pushed a commit
to 021ai/Megatron-LM
that referenced
this pull request
Aug 3, 2026
…IA#5570) Signed-off-by: Yongqiang Wang <yongqiang.seagull@gmail.com> Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
svcnvidia-nemo-ci
pushed a commit
to dimapihtar/Megatron-LM
that referenced
this pull request
Aug 4, 2026
…IA#5570) Signed-off-by: Yongqiang Wang <yongqiang.seagull@gmail.com> Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: Dmytro Pykhtar <dpykhtar@nvidia.com>
This was referenced Aug 17, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What
Adds
NemoAudioProcessorundermegatron/core/models/audio/audio_processor.py— the concrete waveform → log-mel feature extractor + token-count estimator that the multimodal data pipeline uses to (a) size audio placeholders for packing and (b) materialize encoder inputs.It composes the model-frontend descriptors
NemoAudioFeatureConfig+NemoTransformerAudioTokenEstimator(from #5565) with the vendored standalone log-mel preprocessor.Design
The audio reference it operates on is duck-typed (
Any), so this module has no dependency on the data library'sAudioReftype or on Megatron-Energon. The data pipeline consumes the processor through a small structural interface —compute_num_embeddings/compute_num_frames/materializeplus the cumulative-prefixnum_*_from_num_samplesprimitives — and injects it at construction time. This keeps the dependency arrow pointing data → model and lets the energon-side packing code stay model-agnostic.Test plan
tests/unit_tests/models/test_nemo_audio_processor.py(no GPU required):slice_rangewaveform cropping and_infer_num_samplesLocal run of the full
models/audiosuite: 47 passed, 1 skipped (skip needs CUDA).🤖 Generated with Claude Code