[model]Add Qwen3‑Omni training support

docs/models/vlm/index.md

@@ -12,6 +12,7 @@ ministral3.md
 nemotron-3-omni.md
 nemotron-nano-v2-vl.md
 qwen2.5-vl.md
+qwen3-omni.md
 qwen3-vl.md
 qwen35-vl.md
 ```

docs/models/vlm/qwen3-omni.md

@@ -0,0 +1,43 @@
+# Qwen3-Omni
+
+Qwen3-Omni is a multimodal Qwen family model with text, image, video, and audio inputs. Megatron Bridge support for Qwen3-Omni reuses the existing Qwen3-VL language and vision path, and adds Qwen3-Omni-specific audio handling and checkpoint mappings.
+
+The current implementation focuses on checkpoint conversion, training-oriented multimodal forward paths, and smoke-level validation. It includes a full example workflow (HF -> Megatron -> HF export, single-rank inference) and a multi-node training recipe entrypoint.
+
+## Current Support
+
+- Hugging Face to Megatron Bridge checkpoint conversion for `Qwen/Qwen3-Omni-30B-A3B-Instruct`
+- Megatron Bridge to Hugging Face export for the same model family
+- Text, image, video, and audio multimodal forward paths
+- Qwen3-Omni-specific multimodal RoPE handling for Megatron Bridge runtime
+- Single-GPU smoke validation with a vertically trimmed checkpoint
+- Multi-node training recipe entrypoint (see Qwen3-Omni examples)
+- L0 conversion test coverage for Qwen3-Omni
+
+## Known Limitations
+
+- Megatron inference with `inference_params` is not implemented yet
+- `packed_seq_params` is not implemented yet
+- Automated validation coverage remains single-rank; multi-node training requires user execution
+- Functional smoke tests require user-provided local multimodal assets
+
+## Hugging Face Model Cards
+
+- Qwen3-Omni-30B-A3B-Instruct: `https://huggingface.co/Qwen/Qwen3-Omni-30B-A3B-Instruct`
+
+## Examples
+
+Qwen3-Omni examples are maintained here:
+
+- `examples/models/vlm/qwen3_omni/README.md`
+- `examples/models/vlm/qwen3_omni/conversion.sh`
+- `examples/models/vlm/qwen3_omni/inference.sh`
+- `examples/models/vlm/qwen3_omni/local_train_thinker_4node_tp2_ep8_sp.sh`
+
+## Related Docs
+
+- Related VLM: [Qwen3-VL](qwen3-vl.md)
+- Related VLM: [Qwen 3.5](qwen35-vl.md)
+- Recipe usage: [Recipe usage](../../recipe-usage.md)
+- Customizing the training recipe configuration: [Configuration overview](../../training/config-container-overview.md)
+- Training entry points: [Entry points](../../training/entry-points.md)

examples/conversion/hf_to_megatron_generate_omni_lm.py

@@ -266,6 +266,7 @@ def main(args) -> None:
         model_provider.expert_model_parallel_size = ep
         model_provider.expert_tensor_parallel_size = etp
         model_provider.pipeline_dtype = torch.bfloat16
+        model_provider.gradient_accumulation_fusion = False
         model_provider.finalize()
         model_provider.initialize_model_parallel(seed=0)
 
@@ -292,6 +293,7 @@ def main(args) -> None:
         model_provider.expert_model_parallel_size = ep
         model_provider.expert_tensor_parallel_size = etp
         model_provider.pipeline_dtype = torch.bfloat16
+        model_provider.gradient_accumulation_fusion = False
         model_provider.finalize()
         model_provider.initialize_model_parallel(seed=0)
         model = model_provider.provide_distributed_model(wrap_with_ddp=False)

examples/models/vlm/qwen3_omni/README.md

@@ -0,0 +1,144 @@
+# Qwen3-Omni Examples
+
+This directory contains example scripts for **Qwen3-Omni thinker-side support** in Megatron Bridge.
+
+For model introduction and implementation notes, see the [Qwen3-Omni documentation](../../../../docs/models/vlm/qwen3-omni.md).
+
+## Current Scope
+
+These examples cover:
+
+- reduced single-GPU smoke checkpoint creation
+- Hugging Face -> Megatron checkpoint import
+- Megatron -> Hugging Face checkpoint export
+- local inference with reduced HF, imported Megatron, and exported HF checkpoints
+
+These examples do **not** cover:
+
+- distributed parallel validation beyond single-rank smoke
+- talker / code2wav audio-output checkpoints
+- Megatron inference with `inference_params`
+
+## Workspace Configuration
+
+All scripts default to a repo-local cache workspace:
+
+```bash
+export WORKSPACE=$PWD/.cache/qwen3_omni_examples
+```
+
+You can override it if needed. The default directory structure is:
+
+- `${WORKSPACE}/hf/` - reduced local HF smoke checkpoints
+- `${WORKSPACE}/megatron/` - imported Megatron checkpoints
+- `${WORKSPACE}/export/` - exported HF checkpoints
+- `${WORKSPACE}/tmp/` - temporary files
+- `${WORKSPACE}/hf_home/` - Hugging Face cache used by the examples
+
+## Required Local Assets
+
+These examples assume the following local assets are available:
+
+```bash
+export SOURCE_HF_MODEL=/path/to/Qwen3-Omni-30B-A3B-Instruct
+export VIDEO_PATH=/path/to/example.mp4
+```
+
+The example smoke checkpoint keeps the original hidden dimensions intact and only trims layer counts, which keeps the HF config compatible while making single-GPU validation practical.
+
+## Data Preparation (Omni Bench)
+
+If you have the Omni Bench parquet data locally, you can convert it into JSONL plus extracted media assets:
+
+```bash
+python examples/models/vlm/qwen3_omni/convert_omni_bench_to_jsonl.py \
+  --input-root ./Omni_Bench_fix_simple \
+  --output-root ./omni_bench_fix_simple \
+  --splits train test
+```
+
+This produces:
+- `./omni_bench_fix_simple/train/train.jsonl`
+- `./omni_bench_fix_simple/test/test.jsonl`
+- extracted media under `./omni_bench_fix_simple/{split}/media`
+
+## Checkpoint Conversion
+
+Run the full local smoke conversion flow:
+
+```bash
+export SOURCE_HF_MODEL=/path/to/Qwen3-Omni-30B-A3B-Instruct
+bash examples/models/vlm/qwen3_omni/conversion.sh
+```
+
+This script will:
+
+1. create a reduced thinker-only HF smoke checkpoint
+2. import that checkpoint into Megatron format
+3. export the imported Megatron checkpoint back to HF format
+
+## Inference
+
+Run local inference across the reduced HF checkpoint, imported Megatron checkpoint, and exported HF checkpoint:
+
+```bash
+export VIDEO_PATH=/path/to/example.mp4
+bash examples/models/vlm/qwen3_omni/inference.sh
+```
+
+If you want to use the audio track from the video during understanding:
+
+```bash
+export VIDEO_PATH=/path/to/example.mp4
+export USE_AUDIO_IN_VIDEO=1
+bash examples/models/vlm/qwen3_omni/inference.sh
+```
+
+You can also use a remote video URL instead of a local file:
+
+```bash
+export VIDEO_URL=https://example.com/example.mp4
+bash examples/models/vlm/qwen3_omni/inference.sh
+```
+
+This script reuses [examples/conversion/hf_to_megatron_generate_omni_lm.py](../../../conversion/hf_to_megatron_generate_omni_lm.py) and runs:
+
+- inference with the reduced local HF smoke checkpoint
+- inference with the imported Megatron checkpoint
+- inference with the exported HF checkpoint
+
+Useful overrides:
+
+- `PROMPT` (default: `What is happening in this video?`)
+- `MAX_NEW_TOKENS` (default: `50`)
+- `NPROC_PER_NODE` / `TP` / `PP` / `EP` / `ETP` (default: `1`)
+- `DRY_RUN=1` to print commands without executing them
+
+The omni inference helper depends on `qwen-omni-utils[decord]` for video/audio preprocessing.
+
+## Training (local)
+
+The training recipe entrypoint is:
+
+```bash
+bash examples/models/vlm/qwen3_omni/local_train_thinker_full.sh
+```
+
+Required environment variables:
+
+```bash
+export HF_MODEL_PATH=/path/to/Qwen3-Omni-30B-A3B-Instruct
+export TRAIN_JSONL=/path/to/train.jsonl
+```
+
+Optional overrides:
+
+- `WORKSPACE` (default: `${PWD}/.cache/qwen3_omni_train`)
+- `RESULTS_DIR` / `LOG_DIR` (default: under `WORKSPACE`)
+
+To apply the 4-node TP2/PP2/EP8/SP preset used in our 32-GPU validation:
+
+```bash
+export PRESET=4node_tp2_ep8_sp
+bash examples/models/vlm/qwen3_omni/local_train_thinker_full.sh
+```

examples/models/vlm/qwen3_omni/conversion.sh

@@ -0,0 +1,64 @@
+#!/usr/bin/env bash
+# Copyright (c) 2026, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+set -euo pipefail
+
+SCRIPT_DIR=$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)
+ROOT_DIR=$(cd -- "${SCRIPT_DIR}/../../../.." && pwd)
+
+WORKSPACE=${WORKSPACE:-${ROOT_DIR}/.cache/qwen3_omni_examples}
+SOURCE_HF_MODEL=${SOURCE_HF_MODEL:-}
+SMOKE_MODEL_NAME=${SMOKE_MODEL_NAME:-Qwen3-Omni-30B-A3B-Instruct-smoke}
+TEXT_LAYERS=${TEXT_LAYERS:-2}
+VISION_DEPTH=${VISION_DEPTH:-2}
+AUDIO_LAYERS=${AUDIO_LAYERS:-2}
+TMPDIR=${TMPDIR:-${WORKSPACE}/tmp}
+HF_HOME=${HF_HOME:-${WORKSPACE}/hf_home}
+PYTHON=${PYTHON:-python}
+PYTHONPATH=${PYTHONPATH:-${ROOT_DIR}/src:${ROOT_DIR}/3rdparty/Megatron-LM}
+
+HF_SMOKE_PATH=${HF_SMOKE_PATH:-${WORKSPACE}/hf/${SMOKE_MODEL_NAME}}
+MEGATRON_PATH=${MEGATRON_PATH:-${WORKSPACE}/megatron/${SMOKE_MODEL_NAME}}
+HF_EXPORT_PATH=${HF_EXPORT_PATH:-${WORKSPACE}/export/${SMOKE_MODEL_NAME}}
+
+export TMPDIR HF_HOME PYTHONPATH
+mkdir -p "${WORKSPACE}/hf" "${WORKSPACE}/megatron" "${WORKSPACE}/export" "${TMPDIR}" "${HF_HOME}"
+cd "${ROOT_DIR}"
+
+if [[ -z "${SOURCE_HF_MODEL}" ]]; then
+  echo "SOURCE_HF_MODEL must point to a local HF Qwen3-Omni checkpoint or model id." >&2
+  exit 1
+fi
+
+"${PYTHON}" examples/models/vlm/qwen3_omni/create_smoke_checkpoint.py \
+  --source-model-path "${SOURCE_HF_MODEL}" \
+  --output-dir "${HF_SMOKE_PATH}" \
+  --text-layers "${TEXT_LAYERS}" \
+  --vision-depth "${VISION_DEPTH}" \
+  --audio-layers "${AUDIO_LAYERS}"
+
+"${PYTHON}" examples/conversion/convert_checkpoints.py import \
+  --hf-model "${HF_SMOKE_PATH}" \
+  --megatron-path "${MEGATRON_PATH}" \
+  --torch-dtype bfloat16
+
+"${PYTHON}" examples/conversion/convert_checkpoints.py export \
+  --hf-model "${HF_SMOKE_PATH}" \
+  --megatron-path "${MEGATRON_PATH}/iter_0000000" \
+  --hf-path "${HF_EXPORT_PATH}"
+
+echo "HF smoke checkpoint: ${HF_SMOKE_PATH}"
+echo "Megatron checkpoint: ${MEGATRON_PATH}/iter_0000000"
+echo "HF export checkpoint: ${HF_EXPORT_PATH}"