From 991eac7f2141525ae45be85ef27837beeebdb8bc Mon Sep 17 00:00:00 2001 From: niuyazhe Date: Mon, 13 Apr 2026 11:57:40 +0800 Subject: [PATCH 01/11] polish(nyz): simplify r1-aqa script --- .gitignore | 3 +- examples/r1_aqa/audio_dataset.py | 17 ++++++---- .../r1_aqa/run_grpo_r1_aqa_qwen2_audio_7b.sh | 33 ++++++++----------- lightrft/strategy/strategy_base.py | 2 +- 4 files changed, 27 insertions(+), 28 deletions(-) diff --git a/.gitignore b/.gitignore index d3a81dae..33d356e2 100644 --- a/.gitignore +++ b/.gitignore @@ -1219,4 +1219,5 @@ wandb* examples/demo_grpo/results* build/* examples/math_benchmarks/eval_results/ -.llmconfig.yaml \ No newline at end of file +.llmconfig.yaml +tb/* diff --git a/examples/r1_aqa/audio_dataset.py b/examples/r1_aqa/audio_dataset.py index 888e4a50..44a0d450 100644 --- a/examples/r1_aqa/audio_dataset.py +++ b/examples/r1_aqa/audio_dataset.py @@ -493,6 +493,7 @@ def build_audio_multimodal_inputs( images_num: Optional[List[int]], all_videos: Optional[List] = None, videos_num: Optional[List[int]] = None, + all_prompt_token_ids: Optional[List[List[int]]] = None, ) -> List[Dict[str, Any]]: """ Replacement for ``strategy._build_multimodal_inputs`` that maps @@ -530,10 +531,10 @@ def build_audio_multimodal_inputs( if not multi_modal_data: # Remove audio placeholder tokens if no audio data - prompt = re.sub( + cleaned_prompt = re.sub( r'<\|audio_bos\|>.*?<\|audio_eos\|>', '', prompt ) - inputs.append({"prompt": prompt}) + inputs.append({"prompt": cleaned_prompt}) else: inputs.append({ "prompt": prompt, @@ -554,9 +555,15 @@ def patch_strategy_for_audio(strategy): original_build = strategy._build_multimodal_inputs def patched_build(all_prompts, all_images=None, images_num=None, - all_videos=None, videos_num=None): + all_videos=None, videos_num=None, all_prompt_token_ids=None): return build_audio_multimodal_inputs( - strategy, all_prompts, all_images, images_num, all_videos, videos_num + strategy, + all_prompts, + all_images, + images_num, + all_videos, + videos_num, + all_prompt_token_ids, ) strategy._build_multimodal_inputs = patched_build @@ -594,5 +601,3 @@ def patch_experience_maker_for_audio(exp_maker, processor, tokenizer, prompt_max print("[PATCH] FastExperienceMaker patched for audio") return exp_maker - - diff --git a/examples/r1_aqa/run_grpo_r1_aqa_qwen2_audio_7b.sh b/examples/r1_aqa/run_grpo_r1_aqa_qwen2_audio_7b.sh index 171234e7..7f7a42b6 100644 --- a/examples/r1_aqa/run_grpo_r1_aqa_qwen2_audio_7b.sh +++ b/examples/r1_aqa/run_grpo_r1_aqa_qwen2_audio_7b.sh @@ -71,23 +71,16 @@ SAVE_STEPS=50 # Save checkpoint every N steps # Part 3: Distributed Training Setup # ################################################################################ -# --- Single-Node Setup --- -export MLP_WORKER_NUM=1 -export MLP_WORKER_GPU=1 # Number of GPUs per node -export MLP_ROLE_INDEX=0 -export MLP_WORKER_0_HOST="localhost" -export MLP_WORKER_0_PORT=20092 - # --- PyTorch Distributed --- -export MASTER_ADDR=$MLP_WORKER_0_HOST -export MASTER_PORT=$MLP_WORKER_0_PORT -export NNODES=$MLP_WORKER_NUM -export NODE_RANK=$MLP_ROLE_INDEX -export GPUS_PER_NODE=$MLP_WORKER_GPU +export MASTER_ADDR="localhost" +export MASTER_PORT=20092 +export NNODES=1 +export NODE_RANK=0 +export GPUS_PER_NODE=8 # --- Inference Engine --- -ENGINE_TP=1 # Tensor parallelism for inference engine -ENGINE_MEM_UTIL=0.3 # Memory utilization for inference engine (reduced from 0.6) +ENGINE_TP=2 # Tensor parallelism for inference engine +ENGINE_MEM_UTIL=0.6 # Memory utilization for inference engine (reduced from 0.6) # --- Checkpoint local path (use if NFS causes OSError) --- CKPT_PATH_LOCAL="" @@ -159,13 +152,13 @@ torchrun \ --gradient_checkpointing \ --save_steps ${SAVE_STEPS} \ --max_ckpt_num 3 \ - --engine_type vllm \ + --engine_type sglang \ --engine_mem_util ${ENGINE_MEM_UTIL} \ --engine_tp_size $ENGINE_TP \ --enable_engine_sleep \ --l2 1.0e-2 \ --adam_offload \ - --use_wandb "${WANDB_API_KEY}" \ + --use_tensorboard "tb/r1-aqa-baseline" \ --wandb_project "${WANDB_PROJECT}" \ --wandb_run_name "${WANDB_RUN_NAME}" \ 2>&1 | tee "rft_logs/${EXPERIMENT_NAME}/node${NODE_RANK}_${current_time}.log" @@ -186,7 +179,7 @@ torchrun \ # Edit "Part 1: User Configuration" above: # # - Set PATH_TO_YOUR_BASE_MODEL (Qwen2-Audio-7B-Instruct) # # - Set PATH_TO_YOUR_AVQA_DATASET # -# - Set GPU count in MLP_WORKER_GPU # +# - Set GPU count in GPU_PER_NODE # # # # Step 3: Run Training # # bash examples/r1_aqa/run_grpo_r1_aqa_qwen2_audio_7b.sh # @@ -201,8 +194,8 @@ torchrun \ # python /path/to/mmau/evaluation.py --input results/res_mmau_mini.json # # # # Notes: # -# - This uses the AUDIO pipeline (not VL). Audio is processed via Qwen2-Audio. # +# - This uses the AUDIO pipeline. Audio is processed via Qwen2-Audio. # # - TBS must >= RBS * N_SAMPLES for GRPO constraint. # -# - For dry-run: set EPISODE=1, RBS=4, TBS=32, N_SAMPLES=4. # -# - For 1-GPU: set MLP_WORKER_GPU=1, ENGINE_TP=1, ENGINE_MEM_UTIL=0.5. # +# - For dry-run: set EPISODE=1, RBS=4, TBS=32, N_SAMPLES=4. # +# - For 1-GPU: set GPU_PER_NODE=1, ENGINE_TP=1, ENGINE_MEM_UTIL=0.5. # ################################################################################ diff --git a/lightrft/strategy/strategy_base.py b/lightrft/strategy/strategy_base.py index f6c20706..e3d55cdf 100644 --- a/lightrft/strategy/strategy_base.py +++ b/lightrft/strategy/strategy_base.py @@ -979,7 +979,7 @@ def gather_and_generate( inputs = gather_inputs_object_for_inference(input_data=inputs, group=self.engine_mp_group) - self.print(f"Start VLM gather_and_generate ..., total prompts: {len(inputs)}") + self.print(f"Start MLLM gather_and_generate ..., total prompts: {len(inputs)}") all_outputs = self.engine_generate_local( sampling_params=sampling_params, From 006b2e26289328ff8d3e1c089286226c0baf4d26 Mon Sep 17 00:00:00 2001 From: niuyazhe Date: Tue, 14 Apr 2026 15:22:51 +0800 Subject: [PATCH 02/11] fix(nyz): clean data and add sglang audio data input --- examples/r1_aqa/README.md | 51 +++- examples/r1_aqa/audio_dataset.py | 243 +++++++++++++++--- .../data_preprocess/clean_audio_dataset.py | 115 +++++++++ examples/r1_aqa/eval.py | 6 +- .../strategy/sglang_utils/sglang_engine.py | 6 + lightrft/strategy/strategy_base.py | 2 +- lightrft/trainer/spmd_ppo_trainer.py | 4 +- 7 files changed, 386 insertions(+), 41 deletions(-) create mode 100644 examples/r1_aqa/data_preprocess/clean_audio_dataset.py diff --git a/examples/r1_aqa/README.md b/examples/r1_aqa/README.md index b5deb4e6..4df60f3d 100644 --- a/examples/r1_aqa/README.md +++ b/examples/r1_aqa/README.md @@ -11,7 +11,8 @@ R1-AQA applies Group Relative Policy Optimization (GRPO) to Qwen2-Audio-7B-Instr ``` examples/r1_aqa/ ├── data_preprocess/ -│ └── avqa.py # Convert R1-AQA JSONL → LightRFT parquet +│ ├── avqa.py # Convert R1-AQA JSONL → LightRFT parquet +│ └── clean_audio_dataset.py # Drop rows whose audio files are missing/unreadable ├── audio_dataset.py # Audio multimodal pipeline extensions and patches ├── reward_models_utils.py # Rule-based reward (accuracy + format) ├── train_colocate.py # GRPO training entry point @@ -67,13 +68,44 @@ python examples/r1_aqa/data_preprocess/avqa.py \\ --local_save_dir ./avqa_lightrft ``` -### Step 2: Configure and Run Training +### Step 2: Clean Missing / Broken Audio Rows + +Before training, strongly recommend cleaning the parquet once. In distributed GRPO training, +rows whose prompt still contains audio placeholders but whose `audio_path` points to a missing +file can make one rank fall into a text-only branch while other ranks still process audio, +which often shows up later as a hang in actor forward. + +Run: +```bash +python examples/r1_aqa/data_preprocess/clean_audio_dataset.py \\ + --input_dataset ./avqa_lightrft \\ + --output_dir ./avqa_lightrft_clean +``` + +Optional stricter validation: +```bash +python examples/r1_aqa/data_preprocess/clean_audio_dataset.py \\ + --input_dataset ./avqa_lightrft \\ + --output_dir ./avqa_lightrft_clean \\ + --verify_decode +``` + +What this script writes: +- `train.parquet`: cleaned split with only valid audio rows +- `train.dropped.jsonl`: dropped rows with original dataset index, `audio_path`, and reason + +Recommended workflow: +1. Run `avqa.py` once to build the parquet dataset. +2. Run `clean_audio_dataset.py` once on that parquet directory. +3. Point training to the cleaned output directory, not the raw parquet directory. + +### Step 3: Configure and Run Training Edit the shell script to set your paths: ```bash # In run_grpo_r1_aqa_qwen2_audio_7b.sh: PATH_TO_YOUR_BASE_MODEL="Qwen/Qwen2-Audio-7B-Instruct" -PATH_TO_YOUR_AVQA_DATASET="/path/to/your/avqa_lightrft" +PATH_TO_YOUR_AVQA_DATASET="/path/to/your/avqa_lightrft_clean" ``` Run training: @@ -81,7 +113,7 @@ Run training: bash examples/r1_aqa/run_grpo_r1_aqa_qwen2_audio_7b.sh ``` -### Step 3: Evaluate on MMAU / MMAR +### Step 4: Evaluate on MMAU / MMAR ```bash # MMAU (test-mini) @@ -128,6 +160,16 @@ For R1-AQA defaults (n_samples=8): ### 1. Audio Path Not Found Ensure `audio_dir` in the preprocessing script points to the directory containing `.wav` files. Audio paths in the JSONL can be relative or absolute. +If training logs show per-rank audio counts becoming inconsistent, for example one rank logs fewer +`<|AUDIO|>` prompts or fewer loaded audios than other ranks, clean the parquet first and train on +the cleaned directory: +```bash +python examples/r1_aqa/data_preprocess/clean_audio_dataset.py \\ + --input_dataset /path/to/avqa_lightrft \\ + --output_dir /path/to/avqa_lightrft_clean +``` +Then update `PATH_TO_YOUR_AVQA_DATASET` to the cleaned output. + ### 2. VRAM / OOM - Reduce `MICRO_TRAIN` and `MICRO_ROLLOUT` (e.g., 1) - Reduce `N_SAMPLES` (e.g., 4 instead of 8) @@ -169,4 +211,3 @@ Qwen2-Audio uses `Qwen2AudioForConditionalGeneration` (not `AutoModelForVision2S ### 4. Chat Template R1-AQA embeds audio URLs in the chat message content as `{"type": "audio", "audio_url": path}`. We preserve this format and use the Qwen2-Audio processor's `apply_chat_template` to convert it to the correct token format with audio placeholders. - diff --git a/examples/r1_aqa/audio_dataset.py b/examples/r1_aqa/audio_dataset.py index 44a0d450..b5b99cc8 100644 --- a/examples/r1_aqa/audio_dataset.py +++ b/examples/r1_aqa/audio_dataset.py @@ -26,7 +26,9 @@ from __future__ import annotations +import copy import inspect +import io import os import re from typing import Any, Dict, List, Optional, Tuple, Union @@ -34,6 +36,7 @@ import librosa from easydict import EasyDict import numpy as np +import soundfile as sf import torch import torch.nn as nn from torch.utils.data import Dataset @@ -48,6 +51,117 @@ def load_audio(audio_path: str, sr: int = 16000) -> Tuple[np.ndarray, int]: return librosa.load(audio_path, sr=sr) +def sanitize_qwen2_audio_messages(messages) -> List[Dict[str, Any]]: + """ + Remove misleading keys from text segments before applying Qwen2-Audio's chat template. + + The upstream template treats the mere presence of ``audio_url`` as an audio segment, even if the + value is ``None``. Our parquet stores text items as ``{"audio_url": None, "text": ..., "type": "text"}``, + which causes every text segment to be rendered as another audio placeholder. + """ + sanitized = copy.deepcopy(messages) + for message in sanitized: + content = message.get("content") + if not isinstance(content, list): + continue + for segment in content: + if not isinstance(segment, dict): + continue + segment_type = segment.get("type") + if segment_type == "text": + segment.pop("audio_url", None) + elif segment_type == "audio": + segment.pop("text", None) + return sanitized + + +def extract_audio_array(audio_item: Any, default_sr: int = 16000) -> Tuple[np.ndarray, int]: + """Normalize supported audio payloads to ``(waveform, sampling_rate)``.""" + if isinstance(audio_item, tuple) and len(audio_item) == 2: + audio_array, sr = audio_item + return np.asarray(audio_array, dtype=np.float32), int(sr) + if isinstance(audio_item, list) and len(audio_item) == 2 and isinstance(audio_item[0], np.ndarray): + audio_array, sr = audio_item + return np.asarray(audio_array, dtype=np.float32), int(sr) + if isinstance(audio_item, np.ndarray): + return np.asarray(audio_item, dtype=np.float32), default_sr + raise TypeError(f"Unsupported audio payload type: {type(audio_item).__name__}") + + +def serialize_audio_for_sglang(audio_item: Any, default_sr: int = 16000) -> Union[str, bytes, Dict[str, Any]]: + """ + Convert local audio payloads into a SGLang-compatible audio input. + + SGLang accepts file paths / URLs / bytes, but not ``(waveform, sr)`` tuples directly. + """ + if audio_item is None: + return None + if isinstance(audio_item, dict): + return audio_item + if isinstance(audio_item, (str, bytes)): + return audio_item + + audio_array, sr = extract_audio_array(audio_item, default_sr=default_sr) + buffer = io.BytesIO() + sf.write(buffer, audio_array, sr, format="WAV") + return buffer.getvalue() + + +def normalize_qwen2_audio_features( + input_features: torch.Tensor, + feature_attention_mask: Optional[torch.Tensor], + expected_mel_len: int = 3000, +) -> Tuple[torch.Tensor, Optional[torch.Tensor], Tuple[int, ...], Tuple[int, ...]]: + """ + Normalize Qwen2-Audio features to ``(B, mel_bins, expected_mel_len)``. + + Returns the normalized tensor, normalized mask, original shape, and normalized shape. + """ + if input_features is None: + raise RuntimeError("input_features must not be None for audio batches") + + if not isinstance(input_features, torch.Tensor): + input_features = torch.as_tensor(input_features) + if feature_attention_mask is not None and not isinstance(feature_attention_mask, torch.Tensor): + feature_attention_mask = torch.as_tensor(feature_attention_mask) + + original_shape = tuple(input_features.shape) + if input_features.dim() != 3: + raise RuntimeError( + f"Expected 3D audio features, but got shape {original_shape}. " + "Qwen2-Audio should return (batch, mel_bins, time)." + ) + + # Most processors return (B, mel_bins, T). Keep that if mel bins look valid. + if input_features.shape[1] in (80, 128): + normalized = input_features + # Some variants may return (B, T, mel_bins). + elif input_features.shape[-1] in (80, 128): + normalized = input_features.transpose(1, 2).contiguous() + else: + raise RuntimeError( + f"Unexpected Qwen2-Audio feature shape {original_shape}: unable to identify mel-bin dimension. " + "This usually means the processor received malformed audio inputs." + ) + + current_len = normalized.shape[-1] + if current_len < expected_mel_len: + normalized = torch.nn.functional.pad(normalized, (0, expected_mel_len - current_len), value=0.0) + elif current_len > expected_mel_len: + normalized = normalized[..., :expected_mel_len] + + if feature_attention_mask is not None: + current_mask_len = feature_attention_mask.shape[-1] + if current_mask_len < expected_mel_len: + feature_attention_mask = torch.nn.functional.pad( + feature_attention_mask, (0, expected_mel_len - current_mask_len), value=0, + ) + elif current_mask_len > expected_mel_len: + feature_attention_mask = feature_attention_mask[..., :expected_mel_len] + + return normalized, feature_attention_mask, original_shape, tuple(normalized.shape) + + # ============================================================================ # Audio Prompt Dataset # ============================================================================ @@ -125,6 +239,8 @@ def __getitem__(self, idx: int) -> Tuple[str, Any, str, str]: prompt_messages = [{"role": "user", "content": prompt_messages}] # ---- 2. Render via processor's chat template ---- + prompt_messages = sanitize_qwen2_audio_messages(prompt_messages) + try: prompt_text = self.processor.apply_chat_template( prompt_messages, @@ -224,6 +340,32 @@ def process_multimodal_batch( if all_images is None: all_images = [None] * L + audio_token_str = getattr(self.processor, "audio_token", "<|AUDIO|>") + audio_bos_str = getattr(self.processor, "audio_bos_token", "<|audio_bos|>") + audio_eos_str = getattr(self.processor, "audio_eos_token", "<|audio_eos|>") + rank = ( + torch.distributed.get_rank() + if torch.distributed.is_available() and torch.distributed.is_initialized() + else 0 + ) + + # In the audio-only AQA pipeline, a prompt that still contains audio placeholders but + # has no loaded audio payload is a hard data inconsistency. Let it fail here instead of + # silently falling into the text-only branch and hanging later in distributed forward. + # Clean the parquet first with examples/r1_aqa/data_preprocess/clean_audio_dataset.py. + missing_audio_indices = [ + idx + for idx, (prompt, audio) in enumerate(zip(all_prompts, all_images)) + if audio is None and ( + audio_token_str in prompt or audio_bos_str in prompt or audio_eos_str in prompt + ) + ] + if missing_audio_indices: + raise RuntimeError( + f"[AudioPipeline][rank={rank}] Found {len(missing_audio_indices)} prompts with audio " + f"placeholders but no audio payload before processing. sample_indices={missing_audio_indices[:8]}" + ) + # ===== Stage 1: Separation (text-only vs audio) ===== all_prompts_text, all_prompts_audio = [], [] all_audios_valid = [] @@ -265,11 +407,10 @@ def process_multimodal_batch( # Extract audio arrays for the processor flat_audios = [] for audio_tuple in all_audios_valid: - if isinstance(audio_tuple, tuple) and len(audio_tuple) == 2: - flat_audios.append(audio_tuple[0]) # numpy array - elif isinstance(audio_tuple, np.ndarray): - flat_audios.append(audio_tuple) - else: + try: + audio_array, _ = extract_audio_array(audio_tuple, default_sr=16000) + flat_audios.append(audio_array) + except TypeError: # Fallback: create silence flat_audios.append(np.zeros(16000, dtype=np.float32)) @@ -278,9 +419,6 @@ def process_multimodal_batch( # Each audio prompt must contain exactly one <|AUDIO|> token. # Mismatches can occur when the chat template or data serialization # inserts duplicate tokens (e.g. template + processor both add one). - audio_token_str = getattr(self.processor, "audio_token", "<|AUDIO|>") - audio_bos_str = getattr(self.processor, "audio_bos_token", "<|audio_bos|>") - audio_eos_str = getattr(self.processor, "audio_eos_token", "<|audio_eos|>") sanitized_prompts = [] n_fixed = 0 for prompt_str in all_prompts_audio: @@ -345,12 +483,17 @@ def process_multimodal_batch( "truncation": True, "padding": True, "return_tensors": "pt", + "sampling_rate": getattr(self.processor.feature_extractor, "sampling_rate", 16000), } print(f"[AudioPipeline] Processor type: {type(self.processor).__name__}") print(f"[AudioPipeline] Processing {len(flat_audios)} audio samples") total_audio_tokens = sum(p.count(audio_token_str) for p in all_prompts_audio) - print(f"[AudioPipeline] Total <|AUDIO|> tokens in text: {total_audio_tokens}, audios: {len(flat_audios)}") + print( + f"[AudioPipeline][rank={rank}] Total <|AUDIO|> tokens in text: {total_audio_tokens}, " + f"audios: {len(flat_audios)}, text_only_prompts: {len(all_prompts_text)}, " + f"audio_prompts: {len(all_prompts_audio)}" + ) inputs_audio = self.processor(**processor_kwargs) print(f"[AudioPipeline] Processor output keys: {list(inputs_audio.keys())}") @@ -361,26 +504,18 @@ def process_multimodal_batch( "feature_attention_mask", None ) - # ------------------------------------------------------------------ - # Qwen2Audio's Whisper encoder requires mel features of exactly 3000 - # frames. When ``padding=True`` is forwarded to the feature - # extractor it pads to the batch-max instead of 3000. Fix here. - # ------------------------------------------------------------------ if all_input_features is not None: - EXPECTED_MEL_LEN = 3000 - actual_len = all_input_features.shape[-1] - if actual_len < EXPECTED_MEL_LEN: - pad_len = EXPECTED_MEL_LEN - actual_len - all_input_features = torch.nn.functional.pad( - all_input_features, (0, pad_len), value=0.0, + all_input_features, all_feature_attention_mask, input_shape_before, input_shape_after = ( + normalize_qwen2_audio_features( + all_input_features, + all_feature_attention_mask, + expected_mel_len=3000, ) - if all_feature_attention_mask is not None: - all_feature_attention_mask = torch.nn.functional.pad( - all_feature_attention_mask, (0, pad_len), value=0, - ) + ) + if input_shape_before != input_shape_after: print( - f"[AudioPipeline] Padded input_features from " - f"{actual_len} → {EXPECTED_MEL_LEN} frames" + f"[AudioPipeline] Normalized input_features shape " + f"{input_shape_before} -> {input_shape_after}" ) if all_input_features is None: @@ -446,6 +581,7 @@ def process_multimodal_batch( all_images_grid_thw=all_images_grid_thw, all_videos_grid_thw=all_videos_grid_thw, all_references=all_references, + all_feature_attention_mask=all_feature_attention_mask, # Audio-specific: store feature mask separately _audio_feature_attention_mask=all_feature_attention_mask, ) @@ -519,11 +655,10 @@ def build_audio_multimodal_inputs( idx = audio_start_idx + j if idx < len(all_images) and all_images[idx] is not None: audio_item = all_images[idx] - if isinstance(audio_item, tuple): - audio_list.append(audio_item) # (array, sr) - elif isinstance(audio_item, np.ndarray): - audio_list.append((audio_item, 16000)) - # else skip + try: + audio_list.append(serialize_audio_for_sglang(audio_item)) + except TypeError: + continue multi_modal_data = {} if audio_list: @@ -553,6 +688,7 @@ def patch_strategy_for_audio(strategy): Replaces ``_build_multimodal_inputs`` with audio-aware version. """ original_build = strategy._build_multimodal_inputs + original_engine_generate_local = strategy.engine_generate_local def patched_build(all_prompts, all_images=None, images_num=None, all_videos=None, videos_num=None, all_prompt_token_ids=None): @@ -566,8 +702,52 @@ def patched_build(all_prompts, all_images=None, images_num=None, all_prompt_token_ids, ) + def patched_engine_generate_local( + sampling_params, + prompt_token_ids=None, + multi_modal_inputs=None, + ): + if multi_modal_inputs is None or strategy.inference_engine_type != "sglang": + return original_engine_generate_local( + sampling_params=sampling_params, + prompt_token_ids=prompt_token_ids, + multi_modal_inputs=multi_modal_inputs, + ) + + has_audio = any( + "audio" in prompt.get("multi_modal_data", {}) + for prompt in multi_modal_inputs + ) + has_image = any( + "image" in prompt.get("multi_modal_data", {}) + for prompt in multi_modal_inputs + ) + if not has_audio or has_image: + return original_engine_generate_local( + sampling_params=sampling_params, + prompt_token_ids=prompt_token_ids, + multi_modal_inputs=multi_modal_inputs, + ) + + prompt = [p["prompt"] for p in multi_modal_inputs] + audio_data = [p.get("multi_modal_data", {}).get("audio") for p in multi_modal_inputs] + + sglang_outputs = strategy.inference_engine.generate( + sampling_params=sampling_params, + prompt=prompt, + audio_data=audio_data, + ) + return [ + EasyDict( + prompt_token_ids=None, + output_token_ids=sglang_outputs[i]["output_ids"], + ) for i in range(len(sglang_outputs)) + ] + strategy._build_multimodal_inputs = patched_build + strategy.engine_generate_local = patched_engine_generate_local strategy.print("[PATCH] Strategy._build_multimodal_inputs patched for audio") + strategy.print("[PATCH] Strategy.engine_generate_local patched for SGLang audio_data") return strategy @@ -600,4 +780,3 @@ def patch_experience_maker_for_audio(exp_maker, processor, tokenizer, prompt_max print("[PATCH] FastExperienceMaker patched for audio") return exp_maker - diff --git a/examples/r1_aqa/data_preprocess/clean_audio_dataset.py b/examples/r1_aqa/data_preprocess/clean_audio_dataset.py new file mode 100644 index 00000000..324b79e0 --- /dev/null +++ b/examples/r1_aqa/data_preprocess/clean_audio_dataset.py @@ -0,0 +1,115 @@ +#!/usr/bin/env python3 +""" +Clean an audio dataset by dropping rows whose audio files are missing or unreadable. + +This is intended for LightRFT R1-AQA parquet datasets, but it also works for any +dataset split that contains an ``audio_path`` column. + +Recommended workflow: +1. Build parquet with ``examples/r1_aqa/data_preprocess/avqa.py``. +2. Run this script once on the parquet directory. +3. Train on the cleaned output directory, not on the raw parquet directory. + +Why this exists: +- In audio GRPO training, the prompt text and the loaded audio payload must stay aligned. +- If a row still contains audio placeholders but its ``audio_path`` no longer exists on disk, + one distributed rank can silently treat it as text-only while others still process audio. +- That mismatch often surfaces later as a hang during actor forward or replay/PPO processing. + +Outputs: +- ``.parquet``: cleaned split with only valid rows +- ``.dropped.jsonl``: dropped rows with original index, ``audio_path``, and reason +""" + +from __future__ import annotations + +import argparse +import json +from pathlib import Path +from typing import Any + +from datasets import load_dataset + + +def can_decode_audio(audio_path: str) -> tuple[bool, str | None]: + try: + import soundfile as sf + + with sf.SoundFile(audio_path): + return True, None + except Exception as exc: # pragma: no cover - best-effort validation helper + return False, str(exc) + + +def clean_split(dataset, verify_decode: bool) -> tuple[list[int], list[dict[str, Any]]]: + """Return kept row indices and dropped-row metadata for one dataset split.""" + keep_indices: list[int] = [] + dropped_rows: list[dict[str, Any]] = [] + + for idx, row in enumerate(dataset): + audio_path = row.get("audio_path") + if not audio_path: + dropped_rows.append({"index": idx, "audio_path": audio_path, "reason": "missing_audio_path"}) + continue + + path = Path(audio_path) + if not path.exists(): + dropped_rows.append({"index": idx, "audio_path": audio_path, "reason": "missing_file"}) + continue + + if verify_decode: + ok, error = can_decode_audio(audio_path) + if not ok: + dropped_rows.append( + {"index": idx, "audio_path": audio_path, "reason": "decode_error", "error": error} + ) + continue + + keep_indices.append(idx) + + return keep_indices, dropped_rows + + +def main() -> None: + """CLI entrypoint for dataset cleaning.""" + parser = argparse.ArgumentParser(description="Clean LightRFT audio parquet dataset by removing bad audio rows.") + parser.add_argument("--input_dataset", required=True, help="Path to the input dataset directory or parquet file") + parser.add_argument("--output_dir", required=True, help="Directory to write cleaned parquet split files") + parser.add_argument( + "--verify_decode", + action="store_true", + help="Also verify each existing audio file can be opened by soundfile", + ) + args = parser.parse_args() + + input_path = Path(args.input_dataset) + output_dir = Path(args.output_dir) + output_dir.mkdir(parents=True, exist_ok=True) + + dataset_dict = load_dataset(str(input_path)) + summary: dict[str, Any] = {} + + for split_name, split_dataset in dataset_dict.items(): + keep_indices, dropped_rows = clean_split(split_dataset, verify_decode=args.verify_decode) + cleaned_dataset = split_dataset.select(keep_indices) + split_out = output_dir / f"{split_name}.parquet" + report_out = output_dir / f"{split_name}.dropped.jsonl" + + cleaned_dataset.to_parquet(str(split_out)) + with report_out.open("w", encoding="utf-8") as fout: + for row in dropped_rows: + fout.write(json.dumps(row, ensure_ascii=False) + "\n") + + summary[split_name] = { + "total": len(split_dataset), + "kept": len(keep_indices), + "dropped": len(dropped_rows), + "parquet": str(split_out), + "report": str(report_out), + } + + print(json.dumps(summary, indent=2, ensure_ascii=False)) + + +if __name__ == "__main__": + main() diff --git a/examples/r1_aqa/eval.py b/examples/r1_aqa/eval.py index 98f774f2..55081335 100644 --- a/examples/r1_aqa/eval.py +++ b/examples/r1_aqa/eval.py @@ -183,7 +183,11 @@ def get_audio_path(sample: Dict) -> str: try: audio, _ = librosa.load(audio_path, sr=sr) inputs = processor( - text=text, audios=[audio], return_tensors="pt", padding=True + text=text, + audios=[audio], + sampling_rate=sr, + return_tensors="pt", + padding=True, ) except Exception as e: print(f"[WARNING] Audio load failed for {audio_path}: {e}") diff --git a/lightrft/strategy/sglang_utils/sglang_engine.py b/lightrft/strategy/sglang_utils/sglang_engine.py index 3378f246..faea487b 100644 --- a/lightrft/strategy/sglang_utils/sglang_engine.py +++ b/lightrft/strategy/sglang_utils/sglang_engine.py @@ -14,6 +14,7 @@ such as batch processing, custom sampling parameters, and LoRA fine-tuning. """ +import inspect import os from typing import Dict, List, Optional, Union @@ -121,6 +122,7 @@ def generate( # The image input. It can be a file name, a url, or base64 encoded string. # See also python/sglang/srt/utils.py:load_image. image_data: Optional[Union[List[str], str]] = None, + audio_data: Optional[Union[List[str], str, bytes, List[bytes]]] = None, return_logprob: Optional[Union[List[bool], bool]] = False, logprob_start_len: Optional[Union[List[int], int]] = None, top_logprobs_num: Optional[Union[List[int], int]] = None, @@ -146,6 +148,8 @@ def generate( :type input_ids: Optional[Union[List[List[int]], List[int]]] :param image_data: Image input as file name, URL, or base64 encoded string :type image_data: Optional[Union[List[str], str]] + :param audio_data: Audio input as file name, URL, WAV bytes, or a batch of them + :type audio_data: Optional[Union[List[str], str, bytes, List[bytes]]] :param return_logprob: Whether to return log probabilities for generated tokens :type return_logprob: Optional[Union[List[bool], bool]] :param logprob_start_len: Start position for log probability calculation @@ -184,6 +188,8 @@ def generate( input_ids = gather_inputs_object_for_inference(input_ids, group=self.tp_group_cpu) if image_data is not None: image_data = gather_inputs_object_for_inference(image_data, group=self.tp_group_cpu) + if audio_data is not None: + audio_data = gather_inputs_object_for_inference(audio_data, group=self.tp_group_cpu) if self._tp_rank == 0: output = self._engine.generate( diff --git a/lightrft/strategy/strategy_base.py b/lightrft/strategy/strategy_base.py index e3d55cdf..1b36fe17 100644 --- a/lightrft/strategy/strategy_base.py +++ b/lightrft/strategy/strategy_base.py @@ -728,7 +728,7 @@ def wakeup_inference_engine(self): raise ValueError(f"Unsupported engine type: {self.inference_engine_type}") # torch.cuda.reset_max_memory_allocated() self.report_memory("after ppo training, after wakeup inference engine") - self.print(f"Finished {self.inference_engine_type} wakeup, TIMECOST {time.time() - wkup_t0}") + self.print(f"Finished {self.inference_engine_type} wakeup, TIMECOST {time.time() - wkup_t0:.4f}s") self.inference_engine_status = EngineStatus.WAKEUP diff --git a/lightrft/trainer/spmd_ppo_trainer.py b/lightrft/trainer/spmd_ppo_trainer.py index d79a7458..fd4a91ac 100644 --- a/lightrft/trainer/spmd_ppo_trainer.py +++ b/lightrft/trainer/spmd_ppo_trainer.py @@ -387,7 +387,7 @@ def ppo_train(self, global_steps=0): # Currently using this rewritten ppo_train model_tensor = torch.tensor(all_model_rewards, dtype=torch.float32, device=device) if model_tensor.abs().sum() > 0: # Only log if model rewards are non-zero status_mean["model_reward_mean"] = model_tensor.mean().item() - self.strategy.print(f" model_reward_mean: {status_mean['model_reward_mean']}") + self.strategy.print(f"model_reward_mean: {status_mean['model_reward_mean']}") if all_rule_rewards: # [TENSOR-FIX] Handle both tensor lists and scalar lists @@ -467,7 +467,7 @@ def ppo_train(self, global_steps=0): # Currently using this rewritten ppo_train self.strategy.maybe_offload_optimizer(self.actor_optim) torch.cuda.synchronize() torch.cuda.empty_cache() - self.strategy.print(f"PPO Train TIMECOST {time.time() - train_begin}") + self.strategy.print(f"PPO Train TIMECOST {time.time() - train_begin:.4f}s") self.strategy.report_memory("after train, opt offloaded, before update weights") self.strategy.print(torch.cuda.memory_summary()) self.strategy.update_engine_weights(self.actor) From 999a144dea6ac5c19dfd7e522963bf32284a6f0f Mon Sep 17 00:00:00 2001 From: niuyazhe Date: Tue, 14 Apr 2026 18:14:12 +0800 Subject: [PATCH 03/11] refactor(nyz): audio language rl pipeline --- examples/r1_aqa/README.md | 13 +- examples/r1_aqa/audio_dataset.py | 706 +---------- examples/r1_aqa/train_colocate.py | 17 +- lightrft/models/actor_al.py | 74 +- lightrft/models/actor_modality.py | 2 +- lightrft/models/utils.py | 46 + lightrft/strategy/fake_strategy.py | 6 +- .../strategy/sglang_utils/sglang_engine.py | 1 + lightrft/strategy/strategy_base.py | 97 +- lightrft/trainer/experience_maker_vl.py | 12 + lightrft/trainer/fast_exp_maker.py | 229 +++- lightrft/trainer/ppo_trainer_vl.py | 1084 ++++++++--------- lightrft/trainer/replay_buffer_utils.py | 15 + lightrft/trainer/utils.py | 6 + 14 files changed, 981 insertions(+), 1327 deletions(-) diff --git a/examples/r1_aqa/README.md b/examples/r1_aqa/README.md index 4df60f3d..b08e1252 100644 --- a/examples/r1_aqa/README.md +++ b/examples/r1_aqa/README.md @@ -197,14 +197,11 @@ The reward function automatically handles both modes. When `enable_think=True`, ### 1. Reward Summation (not Weighting) R1-AQA sums accuracy and format rewards (max=2.0) while GSM8K/Geo3K in LightRFT uses weighted combination (0.9×accuracy + 0.1×format, max=1.0). We keep R1-AQA's summation to ensure identical reward signal. The GRPO normalization handles the scale difference. -### 2. Audio Pipeline via Image Slot -LightRFT's VL pipeline is built for images/videos. We repurpose the image data slots to carry audio data: -- `pixel_values` → `input_features` (audio features) -- `image_grid_thw` → `feature_attention_mask` -- `raw_images` → raw audio tuples `(np.array, sr)` -- `multi_modal_data["image"]` → `multi_modal_data["audio"]` - -This is done via targeted monkey patches in `audio_dataset.py` rather than modifying core LightRFT code. +### 2. Native Audio Rollout Path +Audio RL now uses a dedicated rollout path in core LightRFT code: +- raw audio payloads stay on the generation side and are passed to SGLang as `audio_data` +- processed mel features are stored explicitly as `audio_values` +- Qwen2-Audio feature masking is stored explicitly as `feature_attention_mask` ### 3. ActorAL (Audio Language Actor) Qwen2-Audio uses `Qwen2AudioForConditionalGeneration` (not `AutoModelForVision2Seq`), and its forward pass expects `audio_values` instead of `pixel_values` + `image_grid_thw`. We use `ActorAL` from `lightrft.models.actor_al`, which natively supports Qwen2-Audio's parameter interface. diff --git a/examples/r1_aqa/audio_dataset.py b/examples/r1_aqa/audio_dataset.py index b5b99cc8..4e008300 100644 --- a/examples/r1_aqa/audio_dataset.py +++ b/examples/r1_aqa/audio_dataset.py @@ -1,44 +1,27 @@ """ -Audio Pipeline Extensions for LightRFT +Audio dataset helpers for the R1-AQA LightRFT example. -This module provides audio-specific adaptations for running Qwen2-Audio -within LightRFT's VL (Vision-Language) training pipeline. Since LightRFT's -pipeline is built around image/video processing, we repurpose the image -data slots to carry audio data through the pipeline. +Historically this example carried audio through the old VL pipeline with example-side +patches. The rollout path is now native in core LightRFT, so this module is reduced +to the example-specific data layer: Architecture: - 1. AudioPromptDataset: Returns (prompt_text, audio_data, reference, label) - where audio_data flows through the 'images' slot. - 2. AudioMultimodalProcessor: Replaces the image processor to handle audio. - - Calls processor(text=..., audios=...) instead of processor(text=..., images=...) - - Stores audio features in pixel_values (VL slot) and audio_values (actor API). - 3. Monkey patches: Adapt normalize/count/build functions for audio data. + 1. AudioPromptDataset returns ``(prompt_text, audio_data, reference, label)``. + 2. ``prompt_text`` is rendered with the Qwen2-Audio chat template. + 3. ``audio_data`` stays as raw waveform + sampling rate for core rollout code. -The Actor model is provided by lightrft.models.actor_al.ActorAL, which -expects audio via the audio_values parameter (forward and generate). - -Key Mapping (image pipeline → audio pipeline): - pixel_values → VL slot (same tensor as audio_values) - audio_values → passed to ActorAL.forward / ActorAL.generate - raw_images (PIL) → raw_audios (numpy_array, sr) tuples - multi_modal_data["image"] → multi_modal_data["audio"] +The actor is still ``lightrft.models.actor_al.ActorAL``, which expects audio through +the explicit audio-language interface in the trainer/model stack. """ from __future__ import annotations import copy -import inspect -import io +import json import os -import re -from typing import Any, Dict, List, Optional, Tuple, Union +from typing import Any, Dict, List, Optional, Tuple import librosa -from easydict import EasyDict -import numpy as np -import soundfile as sf -import torch -import torch.nn as nn from torch.utils.data import Dataset @@ -46,18 +29,17 @@ # Audio Loading # ============================================================================ -def load_audio(audio_path: str, sr: int = 16000) -> Tuple[np.ndarray, int]: - """Load audio file with librosa; returns (waveform, sample_rate).""" +def load_audio(audio_path: str, sr: int = 16000) -> Tuple[Any, int]: + """Load an audio file as ``(waveform, sampling_rate)``.""" return librosa.load(audio_path, sr=sr) def sanitize_qwen2_audio_messages(messages) -> List[Dict[str, Any]]: """ - Remove misleading keys from text segments before applying Qwen2-Audio's chat template. + Remove misleading keys before applying the Qwen2-Audio chat template. - The upstream template treats the mere presence of ``audio_url`` as an audio segment, even if the - value is ``None``. Our parquet stores text items as ``{"audio_url": None, "text": ..., "type": "text"}``, - which causes every text segment to be rendered as another audio placeholder. + Some parquet rows keep ``audio_url=None`` on text segments, and the upstream + template interprets the presence of that key as an audio placeholder. """ sanitized = copy.deepcopy(messages) for message in sanitized: @@ -75,113 +57,14 @@ def sanitize_qwen2_audio_messages(messages) -> List[Dict[str, Any]]: return sanitized -def extract_audio_array(audio_item: Any, default_sr: int = 16000) -> Tuple[np.ndarray, int]: - """Normalize supported audio payloads to ``(waveform, sampling_rate)``.""" - if isinstance(audio_item, tuple) and len(audio_item) == 2: - audio_array, sr = audio_item - return np.asarray(audio_array, dtype=np.float32), int(sr) - if isinstance(audio_item, list) and len(audio_item) == 2 and isinstance(audio_item[0], np.ndarray): - audio_array, sr = audio_item - return np.asarray(audio_array, dtype=np.float32), int(sr) - if isinstance(audio_item, np.ndarray): - return np.asarray(audio_item, dtype=np.float32), default_sr - raise TypeError(f"Unsupported audio payload type: {type(audio_item).__name__}") - - -def serialize_audio_for_sglang(audio_item: Any, default_sr: int = 16000) -> Union[str, bytes, Dict[str, Any]]: - """ - Convert local audio payloads into a SGLang-compatible audio input. - - SGLang accepts file paths / URLs / bytes, but not ``(waveform, sr)`` tuples directly. - """ - if audio_item is None: - return None - if isinstance(audio_item, dict): - return audio_item - if isinstance(audio_item, (str, bytes)): - return audio_item - - audio_array, sr = extract_audio_array(audio_item, default_sr=default_sr) - buffer = io.BytesIO() - sf.write(buffer, audio_array, sr, format="WAV") - return buffer.getvalue() - - -def normalize_qwen2_audio_features( - input_features: torch.Tensor, - feature_attention_mask: Optional[torch.Tensor], - expected_mel_len: int = 3000, -) -> Tuple[torch.Tensor, Optional[torch.Tensor], Tuple[int, ...], Tuple[int, ...]]: - """ - Normalize Qwen2-Audio features to ``(B, mel_bins, expected_mel_len)``. - - Returns the normalized tensor, normalized mask, original shape, and normalized shape. - """ - if input_features is None: - raise RuntimeError("input_features must not be None for audio batches") - - if not isinstance(input_features, torch.Tensor): - input_features = torch.as_tensor(input_features) - if feature_attention_mask is not None and not isinstance(feature_attention_mask, torch.Tensor): - feature_attention_mask = torch.as_tensor(feature_attention_mask) - - original_shape = tuple(input_features.shape) - if input_features.dim() != 3: - raise RuntimeError( - f"Expected 3D audio features, but got shape {original_shape}. " - "Qwen2-Audio should return (batch, mel_bins, time)." - ) - - # Most processors return (B, mel_bins, T). Keep that if mel bins look valid. - if input_features.shape[1] in (80, 128): - normalized = input_features - # Some variants may return (B, T, mel_bins). - elif input_features.shape[-1] in (80, 128): - normalized = input_features.transpose(1, 2).contiguous() - else: - raise RuntimeError( - f"Unexpected Qwen2-Audio feature shape {original_shape}: unable to identify mel-bin dimension. " - "This usually means the processor received malformed audio inputs." - ) - - current_len = normalized.shape[-1] - if current_len < expected_mel_len: - normalized = torch.nn.functional.pad(normalized, (0, expected_mel_len - current_len), value=0.0) - elif current_len > expected_mel_len: - normalized = normalized[..., :expected_mel_len] - - if feature_attention_mask is not None: - current_mask_len = feature_attention_mask.shape[-1] - if current_mask_len < expected_mel_len: - feature_attention_mask = torch.nn.functional.pad( - feature_attention_mask, (0, expected_mel_len - current_mask_len), value=0, - ) - elif current_mask_len > expected_mel_len: - feature_attention_mask = feature_attention_mask[..., :expected_mel_len] - - return normalized, feature_attention_mask, original_shape, tuple(normalized.shape) - - -# ============================================================================ -# Audio Prompt Dataset -# ============================================================================ - class AudioPromptDataset(Dataset): """ - PyTorch Dataset for audio question-answering prompts. - - This dataset reads R1-AQA formatted data (with audio content type in - the prompt) and returns (prompt_text, audio_data, reference, label). + PyTorch dataset for the R1-AQA audio prompt format. - The audio data flows through LightRFT's 'images' slot. The prompt - is rendered using the Qwen2-Audio processor's chat template. - - :param dataset: Underlying HuggingFace dataset. - :param tokenizer: HuggingFace tokenizer. - :param processor: Qwen2-Audio processor (for chat template + feature extraction). - :param max_length: Maximum prompt length. - :param strategy: LightRFT strategy object. - :param input_template: Optional template for formatting input text. + Each item returns ``(prompt_text, audio_payload, reference, label)`` where: + - ``prompt_text`` is rendered through the Qwen2-Audio chat template + - ``audio_payload`` is kept as raw waveform + sampling rate for rollout-side processing + - ``reference`` and ``label`` are passed through to reward computation """ def __init__( @@ -201,38 +84,26 @@ def __init__( self.strategy = strategy self.input_template = input_template - # Field keys from strategy args + # Field keys from strategy args. self.prompt_key = getattr(strategy.args, "input_key", "prompt") self.reference_key = getattr(strategy.args, "reference_key", "reference") self.label_key = getattr(strategy.args, "label_key", "label") self.audio_path_key = "audio_path" - # Audio loading configuration + # Audio loading configuration. self.target_sr = 16000 if hasattr(processor, "feature_extractor") and processor.feature_extractor is not None: - self.target_sr = getattr( - processor.feature_extractor, "sampling_rate", 16000 - ) + self.target_sr = getattr(processor.feature_extractor, "sampling_rate", 16000) def __len__(self) -> int: return len(self.dataset) def __getitem__(self, idx: int) -> Tuple[str, Any, str, str]: - """ - Return (prompt_text, audio_data, reference, label). - - - prompt_text: Rendered chat template string with audio tokens. - - audio_data: Tuple (audio_array, sample_rate) for the inference engine. - - reference: Ground truth answer string. - - label: Reward recipe label (e.g., "avqa_rule"). - """ data = self.dataset[idx] # ---- 1. Extract prompt (chat messages with audio content) ---- prompt_messages = data.get(self.prompt_key, []) if isinstance(prompt_messages, str): - # If stored as string, try to parse as JSON - import json try: prompt_messages = json.loads(prompt_messages) except (json.JSONDecodeError, TypeError): @@ -240,18 +111,15 @@ def __getitem__(self, idx: int) -> Tuple[str, Any, str, str]: # ---- 2. Render via processor's chat template ---- prompt_messages = sanitize_qwen2_audio_messages(prompt_messages) - try: prompt_text = self.processor.apply_chat_template( prompt_messages, tokenize=False, add_generation_prompt=True, ) - except Exception as e: - # Fallback: extract text and format manually - self.strategy.print(f"[WARNING] Chat template failed for idx {idx}: {e}") - user_text = self._extract_text_from_messages(prompt_messages) - prompt_text = user_text + except Exception as exc: + self.strategy.print(f"[WARNING] Chat template failed for idx {idx}: {exc}") + prompt_text = self._extract_text_from_messages(prompt_messages) # ---- 3. Load audio ---- audio_path = data.get(self.audio_path_key, "") @@ -259,8 +127,8 @@ def __getitem__(self, idx: int) -> Tuple[str, Any, str, str]: if audio_path and os.path.exists(audio_path): try: audio_data = load_audio(audio_path, sr=self.target_sr) - except Exception as e: - self.strategy.print(f"[WARNING] Failed to load audio {audio_path}: {e}") + except Exception as exc: + self.strategy.print(f"[WARNING] Failed to load audio {audio_path}: {exc}") audio_data = None # ---- 4. Reference and label (defaults if missing) ---- @@ -269,514 +137,24 @@ def __getitem__(self, idx: int) -> Tuple[str, Any, str, str]: return prompt_text, audio_data, reference, label def collate_fn(self, batch: List[Tuple]) -> Tuple[List, List, List, List]: - """Collate a batch of (prompt, audio, reference, label) tuples.""" + """Keep prompts/audios/references/labels as plain Python lists for the rollout stack.""" prompts, audios, refs, labels = zip(*batch) return list(prompts), list(audios), list(refs), list(labels) @staticmethod def _extract_text_from_messages(messages) -> str: - """Extract text content from chat messages.""" + """Fallback text extraction used when the upstream chat template fails.""" texts = [] for msg in messages: - if isinstance(msg, dict): - content = msg.get("content", "") - if isinstance(content, str): - texts.append(content) - elif isinstance(content, list): - for seg in content: - if isinstance(seg, dict) and seg.get("type") == "text": - texts.append(seg.get("text", "")) - return " ".join(texts) - - -# ============================================================================ -# Audio Multimodal Processor -# ============================================================================ - -class AudioMultimodalProcessor: - """ - Multimodal data processor adapted for audio (Qwen2-Audio). - - Replaces LightRFT's ``MultimodalDataProcessor`` for audio models. - Instead of processing images, it processes audio data through - the Qwen2-Audio processor. - - Key differences from image processor: - - Calls ``processor(text=..., audios=...)`` instead of ``processor(text=..., images=...)`` - - Returns ``input_features`` (stored as pixel_values in pipeline) - - Returns ``feature_attention_mask`` (stored as image_grid_thw in pipeline) - """ - - def __init__(self, tokenizer, processor, prompt_max_len: int): - self.tokenizer = tokenizer - self.processor = processor - self.prompt_max_len = prompt_max_len - - def process_multimodal_batch( - self, - all_prompts: List[str], - all_images: List, # Actually audio data: List[Optional[Tuple[np.ndarray, int]]] - all_references: Optional[List[str]], - images_num: List[int], - n_samples_per_prompt: int, - all_videos: Optional[List] = None, - videos_num: Optional[List[int]] = None, - ) -> EasyDict: - """ - Process a batch of audio-multimodal data. - - :param all_prompts: List of prompt strings (with audio tokens). - :param all_images: List of audio data tuples (repurposed images slot). - :param all_references: Reference answers. - :param images_num: Audio count per sample (always [1, 1, ...]). - :param n_samples_per_prompt: Number of GRPO samples per prompt. - :param all_videos: Unused (None for audio). - :param videos_num: Unused (None for audio). - :return: EasyDict with processed data matching LightRFT's expected format. - """ - N = n_samples_per_prompt - L = len(all_prompts) - - if all_images is None: - all_images = [None] * L - - audio_token_str = getattr(self.processor, "audio_token", "<|AUDIO|>") - audio_bos_str = getattr(self.processor, "audio_bos_token", "<|audio_bos|>") - audio_eos_str = getattr(self.processor, "audio_eos_token", "<|audio_eos|>") - rank = ( - torch.distributed.get_rank() - if torch.distributed.is_available() and torch.distributed.is_initialized() - else 0 - ) - - # In the audio-only AQA pipeline, a prompt that still contains audio placeholders but - # has no loaded audio payload is a hard data inconsistency. Let it fail here instead of - # silently falling into the text-only branch and hanging later in distributed forward. - # Clean the parquet first with examples/r1_aqa/data_preprocess/clean_audio_dataset.py. - missing_audio_indices = [ - idx - for idx, (prompt, audio) in enumerate(zip(all_prompts, all_images)) - if audio is None and ( - audio_token_str in prompt or audio_bos_str in prompt or audio_eos_str in prompt - ) - ] - if missing_audio_indices: - raise RuntimeError( - f"[AudioPipeline][rank={rank}] Found {len(missing_audio_indices)} prompts with audio " - f"placeholders but no audio payload before processing. sample_indices={missing_audio_indices[:8]}" - ) - - # ===== Stage 1: Separation (text-only vs audio) ===== - all_prompts_text, all_prompts_audio = [], [] - all_audios_valid = [] - text_idx = [] - - for idx, (prompt, audio) in enumerate(zip(all_prompts, all_images)): - if audio is None: - all_prompts_text.append(prompt) - text_idx.append(idx) - else: - all_prompts_audio.append(prompt) - all_audios_valid.append(audio) - - # ===== Stage 2: Expansion for n_samples_per_prompt ===== - all_prompts_text = sum([[p] * N for p in all_prompts_text], []) - all_prompts_audio = sum([[p] * N for p in all_prompts_audio], []) - all_audios_valid = [a for a in all_audios_valid for _ in range(N)] - all_images_num = sum([[num] * N for num in images_num], []) if images_num else [0] * (L * N) - all_videos_num = [0] * (L * N) - - # ===== Stage 3-A: Text-only processing ===== - if len(all_prompts_text) > 0: - inputs_text = self.tokenizer( - all_prompts_text, - max_length=self.prompt_max_len, - truncation=True, - add_special_tokens=False, - ) - all_prompt_token_ids_text = inputs_text["input_ids"] - else: - all_prompt_token_ids_text = [] - - # ===== Stage 3-B: Audio processing ===== - all_prompt_token_ids_audio = [] - all_input_features = None - all_feature_attention_mask = None - - if len(all_prompts_audio) > 0: - # Extract audio arrays for the processor - flat_audios = [] - for audio_tuple in all_audios_valid: - try: - audio_array, _ = extract_audio_array(audio_tuple, default_sr=16000) - flat_audios.append(audio_array) - except TypeError: - # Fallback: create silence - flat_audios.append(np.zeros(16000, dtype=np.float32)) - - # Process through Qwen2-Audio processor - # --- Sanitize <|AUDIO|> token count --- - # Each audio prompt must contain exactly one <|AUDIO|> token. - # Mismatches can occur when the chat template or data serialization - # inserts duplicate tokens (e.g. template + processor both add one). - sanitized_prompts = [] - n_fixed = 0 - for prompt_str in all_prompts_audio: - count = prompt_str.count(audio_token_str) - if count == 1: - sanitized_prompts.append(prompt_str) - elif count == 0: - # No audio token found – insert one at the beginning of the - # user turn (after system preamble) so the processor can - # expand it later. - insert_marker = f"{audio_bos_str}{audio_token_str}{audio_eos_str}" - # Try to place it after <|im_start|>user\n - user_tag = "<|im_start|>user\n" - idx_user = prompt_str.find(user_tag) - if idx_user >= 0: - ins = idx_user + len(user_tag) - prompt_str = prompt_str[:ins] + insert_marker + "\n" + prompt_str[ins:] - else: - prompt_str = insert_marker + "\n" + prompt_str - sanitized_prompts.append(prompt_str) - n_fixed += 1 - else: - # More than one <|AUDIO|> token – keep only the first one. - # Remove the full <|audio_bos|><|AUDIO|><|audio_eos|> group - # for duplicates, or bare <|AUDIO|> tokens. - full_group = f"{audio_bos_str}{audio_token_str}{audio_eos_str}" - if full_group in prompt_str: - # Keep first occurrence of the full group, remove the rest - first_end = prompt_str.find(full_group) + len(full_group) - before = prompt_str[:first_end] - after = prompt_str[first_end:] - after = after.replace(full_group, "") - # Also remove any bare <|AUDIO|> that remain - after = after.replace(audio_token_str, "") - prompt_str = before + after - else: - # No full group – just keep first bare <|AUDIO|> - first_end = prompt_str.find(audio_token_str) + len(audio_token_str) - before = prompt_str[:first_end] - after = prompt_str[first_end:] - after = after.replace(audio_token_str, "") - prompt_str = before + after - sanitized_prompts.append(prompt_str) - n_fixed += 1 - - if n_fixed > 0: - print( - f"[AudioPipeline] Fixed <|AUDIO|> token count in {n_fixed}/{len(all_prompts_audio)} prompts" - ) - all_prompts_audio = sanitized_prompts - - # Determine the correct kwarg name for the processor (older - # transformers versions use "audios", newer use "audio"). - proc_sig = inspect.signature(self.processor.__call__) - audio_kwarg = "audio" if "audio" in proc_sig.parameters else "audios" - - processor_kwargs = { - "text": all_prompts_audio, - audio_kwarg: flat_audios, - "add_special_tokens": False, - "max_length": self.prompt_max_len, - "truncation": True, - "padding": True, - "return_tensors": "pt", - "sampling_rate": getattr(self.processor.feature_extractor, "sampling_rate", 16000), - } - - print(f"[AudioPipeline] Processor type: {type(self.processor).__name__}") - print(f"[AudioPipeline] Processing {len(flat_audios)} audio samples") - total_audio_tokens = sum(p.count(audio_token_str) for p in all_prompts_audio) - print( - f"[AudioPipeline][rank={rank}] Total <|AUDIO|> tokens in text: {total_audio_tokens}, " - f"audios: {len(flat_audios)}, text_only_prompts: {len(all_prompts_text)}, " - f"audio_prompts: {len(all_prompts_audio)}" - ) - - inputs_audio = self.processor(**processor_kwargs) - print(f"[AudioPipeline] Processor output keys: {list(inputs_audio.keys())}") - - all_prompt_token_ids_audio = inputs_audio["input_ids"].tolist() - all_input_features = inputs_audio.get("input_features", None) - all_feature_attention_mask = inputs_audio.get( - "feature_attention_mask", None - ) - - if all_input_features is not None: - all_input_features, all_feature_attention_mask, input_shape_before, input_shape_after = ( - normalize_qwen2_audio_features( - all_input_features, - all_feature_attention_mask, - expected_mel_len=3000, - ) - ) - if input_shape_before != input_shape_after: - print( - f"[AudioPipeline] Normalized input_features shape " - f"{input_shape_before} -> {input_shape_after}" - ) - - if all_input_features is None: - raise RuntimeError( - f"Processor {type(self.processor).__name__} returned no " - f"'input_features'. Available keys: {list(inputs_audio.keys())}. " - "Ensure the processor is Qwen2AudioProcessor (not a generic text processor)." - ) - - # ===== Stage 4: Merge back in original order ===== - total_samples = L * N - all_prompts_out = [None] * total_samples - all_images_out = [None] * total_samples # Raw audio data for engine - all_prompt_token_ids_out = [None] * total_samples - - # 4-A: Fill text-only slots - text_ptr = 0 - for orig_idx in text_idx: - for n in range(N): - gid = orig_idx * N + n - all_prompts_out[gid] = all_prompts_text[text_ptr] - all_prompt_token_ids_out[gid] = all_prompt_token_ids_text[text_ptr] - text_ptr += 1 - - # 4-B: Fill audio slots - audio_ptr = 0 - for orig_idx in range(L): - if orig_idx in text_idx: + if not isinstance(msg, dict): continue - for n in range(N): - gid = orig_idx * N + n - all_prompts_out[gid] = all_prompts_audio[audio_ptr] - all_images_out[gid] = all_audios_valid[audio_ptr] # Raw audio for engine - all_prompt_token_ids_out[gid] = all_prompt_token_ids_audio[audio_ptr] - audio_ptr += 1 - - # Expand references - if all_references is not None: - all_references = sum([[ref] * N for ref in all_references], []) - - # Build grid_thw entries for audio. - # Each audio sample needs a (1, 1, 1) grid entry so that the VL pipeline's - # per-sample slicing of pixel_values (input_features) works correctly: - # num_patch = 1*1*1 = 1 → slices exactly 1 row from input_features per audio. - total_audio_count = sum(all_images_num) - if total_audio_count > 0: - all_images_grid_thw = torch.ones((total_audio_count, 3), dtype=torch.long) - else: - all_images_grid_thw = torch.empty((0, 3), dtype=torch.long) - all_videos_grid_thw = torch.empty((0, 3), dtype=torch.long) - - # Store audio as both pixel_values (VL slot) and audio_values (actor API) - return EasyDict( - all_prompt_token_ids=all_prompt_token_ids_out, - all_prompts=all_prompts_out, - all_images=all_images_out, # Raw audio data for engine - all_videos=[None] * total_samples, - all_images_num=all_images_num, - all_videos_num=all_videos_num, - all_images_pixel_values=all_input_features, - all_audio_values=all_input_features, # ActorAL expects audio_values - all_videos_pixel_values=None, - all_images_grid_thw=all_images_grid_thw, - all_videos_grid_thw=all_videos_grid_thw, - all_references=all_references, - all_feature_attention_mask=all_feature_attention_mask, - # Audio-specific: store feature mask separately - _audio_feature_attention_mask=all_feature_attention_mask, - ) - - -# ============================================================================ -# Audio-aware image utilities (monkey-patch replacements) -# ============================================================================ - -def normalize_audios(raw_items: List) -> List: - """ - Replacement for ``normalize_images`` that handles audio tuples. - - Audio data is already in the correct format (numpy array, sr) so - we just pass it through unchanged. - """ - return raw_items - - -def get_audios_num(all_items: Optional[List]) -> Optional[List[int]]: - """ - Replacement for ``get_images_num`` that handles audio tuples. - - Returns 1 for each non-None audio item, 0 for None. - """ - if all_items is None: - return None - counts = [] - for item in all_items: - if item is None: - counts.append(0) - else: - counts.append(1) - return counts - - -# ============================================================================ -# Strategy patching for audio -# ============================================================================ - -def build_audio_multimodal_inputs( - strategy, - all_prompts: List[str], - all_images: List, # Actually raw audio data - images_num: Optional[List[int]], - all_videos: Optional[List] = None, - videos_num: Optional[List[int]] = None, - all_prompt_token_ids: Optional[List[List[int]]] = None, -) -> List[Dict[str, Any]]: - """ - Replacement for ``strategy._build_multimodal_inputs`` that maps - audio data to ``multi_modal_data["audio"]`` instead of ``["image"]``. - - :param all_prompts: List of prompt strings. - :param all_images: List of raw audio data (tuples or numpy arrays). - :param images_num: Audio count per sample. - :return: List of dicts with 'prompt' and optional 'multi_modal_data'. - """ - inputs = [] - audio_start_idx = 0 - - if images_num is None: - images_num = [0] * len(all_prompts) - - for i, prompt in enumerate(all_prompts): - audio_num = images_num[i] if i < len(images_num) else 0 - - audio_list = [] - if audio_num > 0 and all_images is not None: - for j in range(audio_num): - idx = audio_start_idx + j - if idx < len(all_images) and all_images[idx] is not None: - audio_item = all_images[idx] - try: - audio_list.append(serialize_audio_for_sglang(audio_item)) - except TypeError: - continue - - multi_modal_data = {} - if audio_list: - multi_modal_data["audio"] = audio_list - - if not multi_modal_data: - # Remove audio placeholder tokens if no audio data - cleaned_prompt = re.sub( - r'<\|audio_bos\|>.*?<\|audio_eos\|>', '', prompt - ) - inputs.append({"prompt": cleaned_prompt}) - else: - inputs.append({ - "prompt": prompt, - "multi_modal_data": multi_modal_data, - }) - - audio_start_idx += audio_num - - return inputs - - -def patch_strategy_for_audio(strategy): - """ - Monkey-patch the strategy object for audio multimodal support. - - Replaces ``_build_multimodal_inputs`` with audio-aware version. - """ - original_build = strategy._build_multimodal_inputs - original_engine_generate_local = strategy.engine_generate_local - - def patched_build(all_prompts, all_images=None, images_num=None, - all_videos=None, videos_num=None, all_prompt_token_ids=None): - return build_audio_multimodal_inputs( - strategy, - all_prompts, - all_images, - images_num, - all_videos, - videos_num, - all_prompt_token_ids, - ) - - def patched_engine_generate_local( - sampling_params, - prompt_token_ids=None, - multi_modal_inputs=None, - ): - if multi_modal_inputs is None or strategy.inference_engine_type != "sglang": - return original_engine_generate_local( - sampling_params=sampling_params, - prompt_token_ids=prompt_token_ids, - multi_modal_inputs=multi_modal_inputs, - ) - - has_audio = any( - "audio" in prompt.get("multi_modal_data", {}) - for prompt in multi_modal_inputs - ) - has_image = any( - "image" in prompt.get("multi_modal_data", {}) - for prompt in multi_modal_inputs - ) - if not has_audio or has_image: - return original_engine_generate_local( - sampling_params=sampling_params, - prompt_token_ids=prompt_token_ids, - multi_modal_inputs=multi_modal_inputs, - ) - - prompt = [p["prompt"] for p in multi_modal_inputs] - audio_data = [p.get("multi_modal_data", {}).get("audio") for p in multi_modal_inputs] - - sglang_outputs = strategy.inference_engine.generate( - sampling_params=sampling_params, - prompt=prompt, - audio_data=audio_data, - ) - return [ - EasyDict( - prompt_token_ids=None, - output_token_ids=sglang_outputs[i]["output_ids"], - ) for i in range(len(sglang_outputs)) - ] - - strategy._build_multimodal_inputs = patched_build - strategy.engine_generate_local = patched_engine_generate_local - strategy.print("[PATCH] Strategy._build_multimodal_inputs patched for audio") - strategy.print("[PATCH] Strategy.engine_generate_local patched for SGLang audio_data") - return strategy - - -# ============================================================================ -# Experience Maker patching for audio -# ============================================================================ - -def patch_experience_maker_for_audio(exp_maker, processor, tokenizer, prompt_max_len): - """ - Monkey-patch the FastExperienceMaker for audio support. - - Replaces the multimodal_processor with AudioMultimodalProcessor - and patches normalize/count functions. - """ - # Replace multimodal processor - exp_maker.multimodal_processor = AudioMultimodalProcessor( - tokenizer=tokenizer, - processor=processor, - prompt_max_len=prompt_max_len, - ) - - # Store original functions for potential restoration - from lightrft.trainer.image_utils import normalize_images as _orig_normalize - from lightrft.trainer.image_utils import get_images_num as _orig_get_num - import lightrft.trainer.fast_exp_maker as fem_module - - # Patch module-level functions used by make_experience_list - fem_module.normalize_images = normalize_audios - fem_module.get_images_num = get_audios_num - - print("[PATCH] FastExperienceMaker patched for audio") - return exp_maker + content = msg.get("content", "") + if isinstance(content, str): + texts.append(content) + continue + if not isinstance(content, list): + continue + for segment in content: + if isinstance(segment, dict) and segment.get("type") == "text": + texts.append(segment.get("text", "")) + return " ".join(texts) diff --git a/examples/r1_aqa/train_colocate.py b/examples/r1_aqa/train_colocate.py index fee2fb23..c7e053d5 100644 --- a/examples/r1_aqa/train_colocate.py +++ b/examples/r1_aqa/train_colocate.py @@ -34,11 +34,7 @@ # Local imports sys.path.append(os.path.dirname(os.path.abspath(__file__))) from reward_models_utils import reward_fn -from audio_dataset import ( - AudioPromptDataset, - patch_strategy_for_audio, - patch_experience_maker_for_audio, -) +from audio_dataset import AudioPromptDataset def train(args): @@ -52,7 +48,7 @@ def train(args): 4. Setup audio prompt dataloader 5. Configure optimizers and schedulers 6. Setup inference engine (vLLM or SGLang) - 7. Apply audio pipeline patches + 7. Route audio rollout through the native core trainer/strategy audio path 8. Run training loop via SPMDPPOTrainerVL 9. Save final model """ @@ -277,10 +273,6 @@ def train(args): strategy.setup_inference_engine(args, engine_type=args.engine_type, actor=actor) strategy.report_memory("after setup_inference_engine") - # ==================== Apply Audio Patches ==================== - # Patch strategy for audio multimodal inputs - patch_strategy_for_audio(strategy) - # ==================== Trainer ==================== trainer = SPMDPPOTrainerVL( strategy, @@ -337,11 +329,6 @@ def train(args): print_replay_buffer_stats=args.print_replay_buffer_stats, ) - # Patch the experience maker for audio processing - patch_experience_maker_for_audio( - trainer.experience_maker, processor, tokenizer, args.prompt_max_len - ) - # ==================== Training ==================== trainer.fit( args, diff --git a/lightrft/models/actor_al.py b/lightrft/models/actor_al.py index 4c7e076a..10f7faae 100644 --- a/lightrft/models/actor_al.py +++ b/lightrft/models/actor_al.py @@ -7,6 +7,7 @@ """ +import os from typing import Optional, Tuple, Union import torch @@ -16,7 +17,12 @@ from transformers.integrations.deepspeed import HfDeepSpeedConfig from .actor_modality import ActorModality -from .utils import apply_lora_configuration, log_probs_from_logits, reset_position_ids +from .utils import ( + apply_lora_configuration, + canonicalize_left_padded_inputs, + log_probs_from_logits, + reset_position_ids, +) class _AudioEmbedPositions(nn.Module): @@ -291,6 +297,7 @@ def forward( return_output=False, packed_seq_lens: Optional[list[int]] = None, audio_values: Optional[torch.Tensor] = None, + feature_attention_mask: Optional[torch.Tensor] = None, ) -> torch.Tensor: """ Forward pass to compute action log probabilities for reinforcement learning. @@ -344,6 +351,14 @@ def forward( ) """ if not self.packing_samples: + pad_token_id = getattr(self.model.config, "pad_token_id", 0) + if pad_token_id is None: + pad_token_id = 0 + sequences, attention_mask = canonicalize_left_padded_inputs( + sequences=sequences, + attention_mask=attention_mask, + pad_token_id=pad_token_id, + ) position_ids = attention_mask.long().cumsum(-1) - 1 position_ids.masked_fill_(attention_mask == 0, 1) else: @@ -383,7 +398,22 @@ def forward( has_audio_placeholder = (audio_token_id is not None and (sequences == audio_token_id).any().item()) if has_audio_placeholder: - input_features, feature_attention_mask = self._prepare_audio_features(input_features) + input_features, feature_attention_mask = self._prepare_audio_features( + input_features, + feature_attention_mask=feature_attention_mask, + sequences=sequences, + audio_token_id=audio_token_id, + ) + if os.environ.get("LIGHTRFT_AUDIO_DEBUG", "0") == "1": + rank = dist.get_rank() if dist.is_initialized() else 0 + print( + f"[ActorAL][rank={rank}] sequences={tuple(sequences.shape)} " + f"audio_values={tuple(input_features.shape)} " + f"feature_attention_mask={tuple(feature_attention_mask.shape)} " + f"audio_token_count={(sequences == audio_token_id).sum(dim=1).tolist()} " + f"feature_len={feature_attention_mask.sum(dim=1).tolist()}", + flush=True, + ) model_kwargs["input_features"] = input_features model_kwargs["feature_attention_mask"] = feature_attention_mask # else: audio_token_id absent → text-only forward (see comment above) @@ -417,6 +447,9 @@ def forward( def _prepare_audio_features( input_features: torch.Tensor, expected_mel_len: int = 3000, + feature_attention_mask: Optional[torch.Tensor] = None, + sequences: Optional[torch.Tensor] = None, + audio_token_id: Optional[int] = None, ) -> Tuple[torch.Tensor, torch.Tensor]: """ Normalize audio features to the expected mel length and build a feature mask. @@ -433,13 +466,36 @@ def _prepare_audio_features( input_features = input_features[..., :expected_mel_len] actual_len = expected_mel_len - feature_attention_mask = torch.zeros( - input_features.shape[0], - expected_mel_len, - dtype=torch.long, - device=input_features.device, - ) - feature_attention_mask[:, :actual_len] = 1 + if feature_attention_mask is None: + inferred_lengths = None + if sequences is not None and audio_token_id is not None: + # Qwen2AudioProcessor expands one audio placeholder to N consecutive audio_token_ids, where: + # N = floor(floor((mel_len + 1) / 2) / 2) + # So the original mel length lies in [4N - 1, 4N]. We choose 4N and clamp to 3000. + audio_token_counts = (sequences == audio_token_id).sum(dim=1) + inferred_lengths = torch.clamp(audio_token_counts * 4, min=1, max=expected_mel_len) + + feature_attention_mask = torch.zeros( + input_features.shape[0], + expected_mel_len, + dtype=torch.long, + device=input_features.device, + ) + + if inferred_lengths is None: + feature_attention_mask[:, :actual_len] = 1 + else: + for row_idx, inferred_len in enumerate(inferred_lengths.tolist()): + feature_attention_mask[row_idx, :inferred_len] = 1 + else: + feature_attention_mask = feature_attention_mask.to(device=input_features.device, dtype=torch.long) + if feature_attention_mask.shape[-1] < expected_mel_len: + feature_attention_mask = torch.nn.functional.pad( + feature_attention_mask, (0, expected_mel_len - feature_attention_mask.shape[-1]), value=0 + ) + elif feature_attention_mask.shape[-1] > expected_mel_len: + feature_attention_mask = feature_attention_mask[..., :expected_mel_len] + return input_features, feature_attention_mask def gradient_checkpointing_enable(self, gradient_checkpointing_kwargs={"use_reentrant": False}): diff --git a/lightrft/models/actor_modality.py b/lightrft/models/actor_modality.py index 5d386cda..724a2026 100644 --- a/lightrft/models/actor_modality.py +++ b/lightrft/models/actor_modality.py @@ -33,7 +33,7 @@ class ActorModality(Enum): }, ActorModality.AUDIO_LANGUAGE: { "audio_values", - "image_grid_thw", # Audio pipeline stores dummy grid entries for compatibility + "feature_attention_mask", }, ActorModality.OMNI: { "pixel_values", diff --git a/lightrft/models/utils.py b/lightrft/models/utils.py index 2107d542..617f2e33 100644 --- a/lightrft/models/utils.py +++ b/lightrft/models/utils.py @@ -311,6 +311,52 @@ def reset_position_ids(attention_mask: torch.Tensor) -> torch.Tensor: return position_ids +def canonicalize_left_padded_inputs( + sequences: torch.Tensor, + attention_mask: Optional[torch.Tensor], + pad_token_id: int, +) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: + """ + Convert dual-sided padding into left-only padding while preserving token order. + + Qwen2-Audio rejects batches whose ``attention_mask`` contains zeros on both the + left and right side because it cannot infer a consistent padding direction. + RL rollouts in LightRFT naturally create that layout (left-padded prompts plus + right-padded responses). This helper compacts each sample's active span and + re-pads it on the left so downstream models see a valid, one-sided mask. + + :param sequences: Batched token ids of shape ``(batch_size, seq_len)``. + :type sequences: torch.Tensor + :param attention_mask: Optional binary mask aligned with ``sequences``. + :type attention_mask: Optional[torch.Tensor] + :param pad_token_id: Token id used to fill padding positions. + :type pad_token_id: int + + :return: Possibly rewritten ``(sequences, attention_mask)`` pair. + :rtype: Tuple[torch.Tensor, Optional[torch.Tensor]] + """ + if attention_mask is None or attention_mask.ndim != 2 or attention_mask.size(0) == 0: + return sequences, attention_mask + + has_left_padding = torch.any(attention_mask[:, 0] == 0).item() + has_right_padding = torch.any(attention_mask[:, -1] == 0).item() + if not (has_left_padding and has_right_padding): + return sequences, attention_mask + + normalized_sequences = torch.full_like(sequences, pad_token_id) + normalized_attention_mask = torch.zeros_like(attention_mask) + active_lengths = attention_mask.long().sum(dim=-1) + + for row_idx, active_len in enumerate(active_lengths.tolist()): + if active_len <= 0: + continue + active_tokens = sequences[row_idx, attention_mask[row_idx].bool()] + normalized_sequences[row_idx, -active_len:] = active_tokens + normalized_attention_mask[row_idx, -active_len:] = 1 + + return normalized_sequences, normalized_attention_mask + + def apply_lora_configuration( model: "nn.Module", lora_rank: int, diff --git a/lightrft/strategy/fake_strategy.py b/lightrft/strategy/fake_strategy.py index e1ed71ed..7568efa8 100644 --- a/lightrft/strategy/fake_strategy.py +++ b/lightrft/strategy/fake_strategy.py @@ -332,7 +332,11 @@ def gather_and_generate( all_prompts=None, all_images=None, sleep_engine=True, - images_num=None + images_num=None, + all_videos=None, + videos_num=None, + all_audios=None, + audios_num=None, ): """ Fake gather and generate - returns empty results. diff --git a/lightrft/strategy/sglang_utils/sglang_engine.py b/lightrft/strategy/sglang_utils/sglang_engine.py index faea487b..353df918 100644 --- a/lightrft/strategy/sglang_utils/sglang_engine.py +++ b/lightrft/strategy/sglang_utils/sglang_engine.py @@ -197,6 +197,7 @@ def generate( sampling_params=sampling_params, input_ids=input_ids, image_data=image_data, + audio_data=audio_data, return_logprob=return_logprob, logprob_start_len=logprob_start_len, top_logprobs_num=top_logprobs_num, diff --git a/lightrft/strategy/strategy_base.py b/lightrft/strategy/strategy_base.py index 1b36fe17..caecb098 100644 --- a/lightrft/strategy/strategy_base.py +++ b/lightrft/strategy/strategy_base.py @@ -10,6 +10,7 @@ import re import random import time +import io from loguru import logger from abc import ABC, abstractmethod from collections import defaultdict @@ -20,6 +21,7 @@ import deepspeed import numpy as np +import soundfile as sf import torch from easydict import EasyDict from torch import distributed as dist @@ -45,6 +47,36 @@ ModelOrModelOptimPair = Union[nn.Module, ModelOptimPair] +def _extract_audio_array(audio_item: Any, default_sr: int = 16000) -> Tuple[np.ndarray, int]: + """Normalize supported raw audio payloads to ``(waveform, sampling_rate)``.""" + if isinstance(audio_item, tuple) and len(audio_item) == 2: + audio_array, sr = audio_item + return np.asarray(audio_array, dtype=np.float32), int(sr) + if isinstance(audio_item, list) and len(audio_item) == 2 and isinstance(audio_item[0], np.ndarray): + audio_array, sr = audio_item + return np.asarray(audio_array, dtype=np.float32), int(sr) + if isinstance(audio_item, np.ndarray): + return np.asarray(audio_item, dtype=np.float32), default_sr + raise TypeError(f"Unsupported audio payload type: {type(audio_item).__name__}") + + +def _serialize_audio_for_sglang(audio_item: Any, default_sr: int = 16000): + """ + Convert local audio payloads into the form accepted by SGLang. + + SGLang accepts file paths / URLs / bytes, but not ``(waveform, sr)`` tuples directly. + """ + if audio_item is None: + return None + if isinstance(audio_item, (str, bytes, dict)): + return audio_item + + audio_array, sr = _extract_audio_array(audio_item, default_sr=default_sr) + buffer = io.BytesIO() + sf.write(buffer, audio_array, sr, format="WAV") + return buffer.getvalue() + + class EngineStatus(Enum): """ Enum class for inference engine status. @@ -771,6 +803,11 @@ def engine_generate_local( # - If `prompt_token_ids` is provided, it indicates a pure LLM (text-only) generation. # - If `prompts` (i.e., `multi_modal_inputs`) is provided, it indicates a VLM (multimodal) generation. if multi_modal_inputs is not None: + # Audio generation is only wired through SGLang for now, so fail early instead of + # silently passing an unsupported payload into vLLM. + has_audio = any("audio" in payload.get("multi_modal_data", {}) for payload in multi_modal_inputs) + if has_audio: + raise ValueError("Audio-language rollout currently requires the SGLang inference backend.") prompt = multi_modal_inputs elif prompt_token_ids is not None: prompt = prompt_token_ids @@ -793,6 +830,25 @@ def engine_generate_local( if multi_modal_inputs is not None: # VLM case logger.debug(f"rank {dist.get_rank()} VLM branch") prompt = [p["prompt"] for p in multi_modal_inputs] + has_audio = any("audio" in p.get("multi_modal_data", {}) for p in multi_modal_inputs) + has_image = any("image" in p.get("multi_modal_data", {}) for p in multi_modal_inputs) + has_video = any("video" in p.get("multi_modal_data", {}) for p in multi_modal_inputs) + + if has_audio and not has_image and not has_video: + # Pure audio-language generation uses SGLang's dedicated ``audio_data`` input + # instead of overloading the image branch. + audio = [p.get("multi_modal_data", {}).get("audio") for p in multi_modal_inputs] + sglang_outputs = self.inference_engine.generate( + sampling_params=sampling_params, + prompt=prompt, + audio_data=audio, + ) + return [ + EasyDict( + prompt_token_ids=None, + output_token_ids=sglang_outputs[i]["output_ids"], + ) for i in range(len(sglang_outputs)) + ] # Handle cases where some prompts might not have images # Flatten nested list format if needed: [[PIL.Image]] -> [PIL.Image] @@ -832,7 +888,16 @@ def engine_generate_local( raise ValueError(f"Unsupported engine type: {self.inference_engine_type}") @classmethod - def _build_multimodal_inputs(cls, all_prompts, all_images, images_num, all_videos, videos_num): + def _build_multimodal_inputs( + cls, + all_prompts, + all_images, + images_num, + all_videos, + videos_num, + all_audios=None, + audios_num=None, + ): """ Build multimodal inputs for inference engine (vLLM/SGLang). @@ -859,9 +924,11 @@ def _build_multimodal_inputs(cls, all_prompts, all_images, images_num, all_video inputs = [] img_start_idx = 0 vid_start_idx = 0 + audio_start_idx = 0 for i, prompt in enumerate(all_prompts): img_num = images_num[i] if images_num is not None else 0 vid_num = videos_num[i] if videos_num is not None else 0 + audio_num = audios_num[i] if audios_num is not None else 0 # Support two input formats: # 1. Nested list: all_images[i] is already a list of images for this prompt @@ -883,17 +950,35 @@ def _build_multimodal_inputs(cls, all_prompts, all_images, images_num, all_video else: vid_list = [] + if all_audios is not None: + if i < len(all_audios) and isinstance(all_audios[i], list) and len(all_audios[i]) == audio_num: + raw_audio_list = all_audios[i] + else: + raw_audio_list = all_audios[audio_start_idx:audio_start_idx + audio_num] + # Serialize in one place so the rest of the rollout stack can keep audio payloads + # in their native Python forms. + audio_list = [ + _serialize_audio_for_sglang(audio_item) for audio_item in raw_audio_list if audio_item is not None + ] + else: + audio_list = [] + multi_modal_data = {} if len(img_list) > 0 and img_list[0] is not None: multi_modal_data["image"] = img_list if len(vid_list) > 0 and vid_list[0] is not None: multi_modal_data["video"] = vid_list + if len(audio_list) > 0 and audio_list[0] is not None: + multi_modal_data["audio"] = audio_list if not multi_modal_data: # remove the vision start and end tokens for data after apply chat template. # Use regex to handle multiple <|image_pad|> tokens (e.g., for high-res images) prompt = re.sub(r'<\|vision_start\|>(<\|image_pad\|>)+<\|vision_end\|>', '', prompt) prompt = re.sub(r'<\|vision_start\|>(<\|video_pad\|>)+<\|vision_end\|>', '', prompt) + # Audio prompts should also degrade cleanly to text-only when their payload is absent. + prompt = re.sub(r'<\|audio_bos\|>.*?<\|audio_eos\|>', '', prompt, flags=re.DOTALL) + prompt = prompt.replace("<|AUDIO|>", "") inputs.append({ "prompt": prompt, }) @@ -904,6 +989,7 @@ def _build_multimodal_inputs(cls, all_prompts, all_images, images_num, all_video }) img_start_idx += img_num vid_start_idx += vid_num + audio_start_idx += audio_num return inputs def gather_and_generate( @@ -916,6 +1002,8 @@ def gather_and_generate( images_num=None, all_videos=None, videos_num=None, + all_audios=None, + audios_num=None, ): """ Gather prompts across distributed ranks and perform text/multimodal generation. @@ -961,9 +1049,10 @@ def gather_and_generate( # is_multimodal = all_images is not None # NOTE: not only check if all_images is None, but also check if it contains non-None elements # If all_images is [None, None, ...], any(img is not None for img in all_images) will return False - # Same logic applies to all_videos + # Same logic applies to all_videos and all_audios. is_multimodal = (((all_images is not None) and any(img is not None for img in all_images)) - or ((all_videos is not None) and any(vid is not None for vid in all_videos))) + or ((all_videos is not None) and any(vid is not None for vid in all_videos)) + or ((all_audios is not None) and any(audio is not None for audio in all_audios))) if is_multimodal: inputs = self._build_multimodal_inputs( @@ -972,6 +1061,8 @@ def gather_and_generate( images_num=images_num, all_videos=all_videos, videos_num=videos_num, + all_audios=all_audios, + audios_num=audios_num, ) else: inputs = all_prompt_token_ids diff --git a/lightrft/trainer/experience_maker_vl.py b/lightrft/trainer/experience_maker_vl.py index bf6a205f..a6052aff 100644 --- a/lightrft/trainer/experience_maker_vl.py +++ b/lightrft/trainer/experience_maker_vl.py @@ -110,6 +110,7 @@ class ExperienceVL: sequences: torch.Tensor # Image processing related pixel_values: Optional[torch.Tensor] = None + audio_values: Optional[torch.Tensor] = None image_grid_thws: Optional[torch.Tensor] = None raw_images: Optional[List[Image.Image]] = None @@ -129,6 +130,7 @@ class ExperienceVL: action_entropy: Optional[torch.Tensor] = None # Entropy for high-entropy token filtering labels: Optional[List[str]] = None # data source labels (if available, e.g., "gsm8k_rule") references: Optional[List[str]] = None # ground truth references (if available, e.g., correct answers) + feature_attention_mask: Optional[torch.Tensor] = None # audio feature mask for audio-language models @torch.no_grad() def to_device(self, device: torch.device): @@ -147,12 +149,16 @@ def to_device(self, device: torch.device): self.advantages = to(self.advantages, device) if self.pixel_values is not None: self.pixel_values = to(self.pixel_values, device) + if self.audio_values is not None: + self.audio_values = to(self.audio_values, device) if self.image_grid_thws is not None: self.image_grid_thws = to(self.image_grid_thws, device) if self.pixel_values_videos is not None: self.pixel_values_videos = to(self.pixel_values_videos, device) if self.video_grid_thws is not None: self.video_grid_thws = to(self.video_grid_thws, device) + if self.feature_attention_mask is not None: + self.feature_attention_mask = to(self.feature_attention_mask, device) self.values = to(self.values, device) self.attention_mask = to(self.attention_mask, device) self.action_mask = to(self.action_mask, device) @@ -176,12 +182,16 @@ def pin_memory(self): self.advantages = pin_memory(self.advantages) if self.pixel_values is not None: self.pixel_values = pin_memory(self.pixel_values) + if self.audio_values is not None: + self.audio_values = pin_memory(self.audio_values) if self.image_grid_thws is not None: self.image_grid_thws = pin_memory(self.image_grid_thws) if self.pixel_values_videos is not None: self.pixel_values_videos = pin_memory(self.pixel_values_videos) if self.video_grid_thws is not None: self.video_grid_thws = pin_memory(self.video_grid_thws) + if self.feature_attention_mask is not None: + self.feature_attention_mask = pin_memory(self.feature_attention_mask) self.values = pin_memory(self.values) self.attention_mask = pin_memory(self.attention_mask) self.action_mask = pin_memory(self.action_mask) @@ -260,6 +270,7 @@ class SamplesVL: action_mask: Optional[torch.BoolTensor] = None pixel_values: Optional[torch.Tensor] = None + audio_values: Optional[torch.Tensor] = None image_grid_thws: Optional[torch.Tensor] = None raw_images: Optional[List[Image.Image]] = None image_num: Optional[List[int]] = None @@ -278,6 +289,7 @@ class SamplesVL: prompts: list[str] = None output_texts: list[str] = None + feature_attention_mask: Optional[torch.Tensor] = None class NaiveExperienceMakerVL(ABC): diff --git a/lightrft/trainer/fast_exp_maker.py b/lightrft/trainer/fast_exp_maker.py index 91850606..e3ec6ec5 100644 --- a/lightrft/trainer/fast_exp_maker.py +++ b/lightrft/trainer/fast_exp_maker.py @@ -55,7 +55,9 @@ from lightrft.utils import Timer, get_current_device from .utils import RunningMoments, compute_clip_fraction, get_cpgd_advantages_returns, fire_sampling, vllm_ge_0130 from .advantage_calculator import get_advantage_calculator, normalize_advantages_cross_batch +from .audio_utils import AudioDataProcessor, get_audios_num, normalize_audios from .image_utils import normalize_images, get_images_num +from .modality_utils import build_supported_model_kwargs from .video_utils import normalize_videos, get_videos_num # ============================================================================ @@ -114,6 +116,7 @@ class _SamplesOutput: # Vision-Language Model fields pixel_values: Optional[torch.Tensor] = None + audio_values: Optional[torch.Tensor] = None image_grid_thw: Optional[torch.Tensor] = None pixel_values_videos: Optional[torch.Tensor] = None video_grid_thw: Optional[torch.Tensor] = None @@ -121,6 +124,7 @@ class _SamplesOutput: references: Optional[list] = None image_num: Optional[List[int]] = None video_num: Optional[List[int]] = None + feature_attention_mask: Optional[torch.Tensor] = None # Model outputs action_log_probs: Optional[torch.Tensor] = None @@ -914,14 +918,25 @@ def __init__(self, *args, packing_samples: bool = False, processor=None, **kwarg self.advantage_calculator = get_advantage_calculator(advantage_estimator, self.strategy.config) # Initialize helper modules + actor_modality = self.actor.modality if self.processor is not None: - self.multimodal_processor = MultimodalDataProcessor( - tokenizer=self.tokenizer, - processor=self.processor, - prompt_max_len=self.prompt_max_len, - ) + if actor_modality == ActorModality.AUDIO_LANGUAGE: + self.audio_processor = AudioDataProcessor( + tokenizer=self.tokenizer, + processor=self.processor, + prompt_max_len=self.prompt_max_len, + ) + self.multimodal_processor = None + else: + self.multimodal_processor = MultimodalDataProcessor( + tokenizer=self.tokenizer, + processor=self.processor, + prompt_max_len=self.prompt_max_len, + ) + self.audio_processor = None else: self.multimodal_processor = None + self.audio_processor = None self.reward_engine = RewardComputationEngine( reward_model=self.reward_model, @@ -937,7 +952,6 @@ def __init__(self, *args, packing_samples: bool = False, processor=None, **kwarg # Cache actor's supported parameters based on its modality # Default to VISION_LANGUAGE for backward compatibility with models without modality attribute - actor_modality = self.actor.modality self._actor_supported_params = get_supported_parameters(actor_modality) # ======================================================================== @@ -950,6 +964,7 @@ def make_experience_list( all_prompts: List[str], all_images: Optional[List] = None, all_videos: Optional[List] = None, + all_audios: Optional[List] = None, all_references: Optional[List[str]] = None, all_labels: Optional[List] = None, **generate_kwargs, @@ -977,7 +992,7 @@ def make_experience_list( """ config = self.strategy.config - # Normalize images if provided + # Normalize vision inputs if provided if all_images is not None: if self.multimodal_processor is None: raise ValueError( @@ -995,12 +1010,21 @@ def make_experience_list( ) all_videos = normalize_videos(all_videos) + if all_audios is not None: + if self.audio_processor is None: + raise ValueError( + "Audio data provided but the experience maker was not initialized for audio-language rollout." + ) + all_audios = normalize_audios(all_audios) + # Get image counts images_num = (get_images_num(all_images) if self.multimodal_processor and all_images is not None else None) # Get video counts videos_num = (get_videos_num(all_videos) if self.multimodal_processor and all_videos is not None else None) + audio_num = (get_audios_num(all_audios) if self.audio_processor and all_audios is not None else None) + # ========== Stage 1: Sample Generation ========== Timer.start(' generate_samples') samples_list = self.generate_samples( @@ -1009,6 +1033,8 @@ def make_experience_list( images_num=images_num, all_videos=all_videos, videos_num=videos_num, + all_audios=all_audios, + audios_num=audio_num, all_references=all_references, all_labels=all_labels, **generate_kwargs, @@ -1057,8 +1083,10 @@ def generate_samples( all_prompts: List[str], all_images: Optional[List] = None, all_videos: Optional[List] = None, + all_audios: Optional[List] = None, images_num: Optional[List[int]] = None, videos_num: Optional[List[int]] = None, + audios_num: Optional[List[int]] = None, all_references: Optional[List[str]] = None, all_labels: Optional[List] = None, **generate_kwargs, @@ -1097,7 +1125,8 @@ def generate_samples( start_time = time.time() config = self.strategy.config - is_multimodal = all_images is not None or all_videos is not None + is_audio_batch = all_audios is not None + is_multimodal = all_images is not None or all_videos is not None or is_audio_batch n_samples = config.n_samples_per_prompt # Initialize multimodal-specific variables to None @@ -1107,6 +1136,9 @@ def generate_samples( all_videos_pixel_values = None all_images_grid_thw = None all_videos_grid_thw = None + all_feature_attention_mask = None + all_audio_num = None + all_audio_values = None # ========== Configure Sampling Parameters ========== if config.engine_type == "vllm": @@ -1158,26 +1190,45 @@ def generate_samples( # ========== Process Multimodal Data ========== if is_multimodal: - processed_data = self.multimodal_processor.process_multimodal_batch( - all_prompts=all_prompts, - all_images=all_images, - all_references=all_references, - images_num=images_num, - n_samples_per_prompt=n_samples, - all_videos=all_videos, - videos_num=videos_num, - ) - all_prompt_token_ids = processed_data["all_prompt_token_ids"] - all_prompts = processed_data["all_prompts"] - all_images = processed_data["all_images"] - all_videos = processed_data["all_videos"] - all_images_num = processed_data["all_images_num"] - all_videos_num = processed_data["all_videos_num"] - all_images_grid_thw = processed_data["all_images_grid_thw"] - all_videos_grid_thw = processed_data["all_videos_grid_thw"] - all_images_pixel_values = processed_data["all_images_pixel_values"] - all_videos_pixel_values = processed_data["all_videos_pixel_values"] - all_references = processed_data.get("all_references", None) + if is_audio_batch: + processed_data = self.audio_processor.process_audio_batch( + all_prompts=all_prompts, + all_audios=all_audios, + all_references=all_references, + n_samples_per_prompt=n_samples, + ) + all_prompt_token_ids = processed_data["all_prompt_token_ids"] + all_prompts = processed_data["all_prompts"] + all_audios = processed_data["all_audios"] + all_audio_num = processed_data["all_audio_num"] + all_audio_values = processed_data["all_audio_values"] + all_feature_attention_mask = processed_data["all_feature_attention_mask"] + all_references = processed_data.get("all_references", None) + else: + processed_data = self.multimodal_processor.process_multimodal_batch( + all_prompts=all_prompts, + all_images=all_images, + all_references=all_references, + images_num=images_num, + n_samples_per_prompt=n_samples, + all_videos=all_videos, + videos_num=videos_num, + ) + all_prompt_token_ids = processed_data["all_prompt_token_ids"] + all_prompts = processed_data["all_prompts"] + all_images = processed_data["all_images"] + all_videos = processed_data["all_videos"] + all_images_num = processed_data["all_images_num"] + all_videos_num = processed_data["all_videos_num"] + all_images_grid_thw = processed_data["all_images_grid_thw"] + all_videos_grid_thw = processed_data["all_videos_grid_thw"] + all_images_pixel_values = processed_data["all_images_pixel_values"] + all_videos_pixel_values = processed_data["all_videos_pixel_values"] + all_feature_attention_mask = processed_data.get( + "all_feature_attention_mask", + processed_data.get("_audio_feature_attention_mask", None), + ) + all_references = processed_data.get("all_references", None) else: # Text-only processing tokenized = self.tokenize_fn(all_prompts, self.prompt_max_len, padding=False) @@ -1199,8 +1250,10 @@ def generate_fn( all_prompts=None, all_images=None, all_videos=None, + all_audios=None, images_num=None, videos_num=None, + audios_num=None, ): return self.strategy.gather_and_generate( sampling_params=sampling_params, @@ -1208,8 +1261,10 @@ def generate_fn( all_prompts=all_prompts, all_images=all_images, all_videos=all_videos, + all_audios=all_audios, images_num=images_num, videos_num=videos_num, + audios_num=audios_num, sleep_engine=sleep_engine, ) @@ -1228,8 +1283,10 @@ def generate_fn( all_prompts=all_prompts, all_images=all_images, all_videos=all_videos, + all_audios=all_audios, all_images_num=all_images_num, all_videos_num=all_videos_num, + all_audio_num=all_audio_num, sampling_params=sampling_params, tokenizer=self.tokenizer, ) @@ -1243,8 +1300,10 @@ def generate_fn( sleep_engine=self.strategy.args.enable_engine_sleep, all_images=all_images if is_multimodal else None, all_videos=all_videos if is_multimodal else None, + all_audios=all_audios if is_multimodal else None, images_num=all_images_num if is_multimodal else None, videos_num=all_videos_num if is_multimodal else None, + audios_num=all_audio_num if is_multimodal else None, ) except ValueError as e: if "prompt" in str(e) and "too long" in str(e): @@ -1263,13 +1322,32 @@ def generate_fn( for i in range(0, len(all_outputs), config.micro_rollout_batch_size): micro_batch_outputs = all_outputs[i:i + config.micro_rollout_batch_size] micro_batch_prompts = all_prompts[i:i + config.micro_rollout_batch_size] + micro_batch_references = (all_references[i:i + config.micro_rollout_batch_size] if all_references else None) + micro_batch_labels = (all_labels[i:i + config.micro_rollout_batch_size] if all_labels else None) # Extract micro-batch data micro_batch_grid_thw = None micro_batch_video_grid_thw = None micro_batch_raw_images = None - if is_multimodal: + if is_audio_batch: + micro_batch_audio_values = ( + all_audio_values[i:i + config.micro_rollout_batch_size] if all_audio_values is not None else None + ) + micro_batch_feature_attention_mask = ( + all_feature_attention_mask[i:i + config.micro_rollout_batch_size] + if all_feature_attention_mask is not None else None + ) + sample = self._build_unpacked_audio_sample( + outputs=micro_batch_outputs, + prompts=micro_batch_prompts, + labels=micro_batch_labels, + references=micro_batch_references, + audio_values=micro_batch_audio_values, + feature_attention_mask=micro_batch_feature_attention_mask, + ) + samples_list.append(sample) + elif is_multimodal: rollout_image_count = sum(all_images_num[i:i + config.micro_rollout_batch_size]) micro_batch_grid_thw = all_images_grid_thw[image_start_idx:image_start_idx + rollout_image_count] micro_batch_raw_images = all_images[i:i + config.micro_rollout_batch_size] @@ -1279,10 +1357,10 @@ def generate_fn( micro_batch_video_grid_thw = all_videos_grid_thw[video_start_idx:video_start_idx + rollout_video_count] video_start_idx += rollout_video_count - micro_batch_references = (all_references[i:i + config.micro_rollout_batch_size] if all_references else None) - micro_batch_labels = (all_labels[i:i + config.micro_rollout_batch_size] if all_labels else None) - # Build samples + if is_audio_batch: + continue + if not self.packing_samples: sample, updated_patch_idx, updated_video_patch_idx = self._build_unpacked_sample( outputs=micro_batch_outputs, @@ -1295,6 +1373,7 @@ def generate_fn( raw_images=micro_batch_raw_images, pixel_values=all_images_pixel_values if is_multimodal else None, pixel_values_videos=all_videos_pixel_values if is_multimodal else None, + feature_attention_mask=all_feature_attention_mask if is_multimodal else None, images_num=all_images_num[i:i + config.micro_rollout_batch_size] if is_multimodal else None, videos_num=all_videos_num[i:i + config.micro_rollout_batch_size] if is_multimodal else None, image_patch_idx=image_patch_idx, @@ -1699,23 +1778,7 @@ def _preprocess_sample( # Only include parameters that the actor's modality supports extra_kwargs = {} if vlm: - # Candidate parameters to pass - candidate_params = { - "pixel_values": sample.pixel_values, - "image_grid_thw": sample.image_grid_thws, - "pixel_values_videos": sample.pixel_values_videos, - "video_grid_thw": sample.video_grid_thws, - } - # Audio-language actors expect audio_values; pipeline stores them in pixel_values slot - if "audio_values" in self._actor_supported_params: - candidate_params["audio_values"] = candidate_params.get("pixel_values") - - # Filter to only include supported parameters - extra_kwargs = { - key: value - for key, value in candidate_params.items() - if key in self._actor_supported_params - } + extra_kwargs = build_supported_model_kwargs(sample, self._actor_supported_params) # Fix Qwen-VL image token count bug self._fix_qwen_vl_image_tokens(sequences, sample, vlm) @@ -1731,9 +1794,11 @@ def _preprocess_sample( prompts=prompts, labels=labels, pixel_values=getattr(sample, "pixel_values", None), + audio_values=getattr(sample, "audio_values", None), image_grid_thw=getattr(sample, "image_grid_thws", None), pixel_values_videos=getattr(sample, "pixel_values_videos", None), video_grid_thw=getattr(sample, "video_grid_thws", None), + feature_attention_mask=getattr(sample, "feature_attention_mask", None), raw_images=getattr(sample, "raw_images", None), image_num=getattr(sample, "image_num", None), video_num=getattr(sample, "video_num", None), @@ -1844,10 +1909,12 @@ def _pack_experience( return ExperienceVL( sequences=output.sequences, pixel_values=output.pixel_values, + audio_values=output.audio_values, image_grid_thws=output.image_grid_thw, raw_images=output.raw_images, pixel_values_videos=output.pixel_values_videos, video_grid_thws=output.video_grid_thw, + feature_attention_mask=output.feature_attention_mask, action_log_probs=output.action_log_probs, base_action_log_probs=output.base_action_log_probs, values=output.value, @@ -1921,6 +1988,7 @@ def _build_unpacked_sample( pixel_values = [] image_grid_thw_list = [] all_img_num = [] + feature_attention_mask_list = [] pixel_values_videos = [] video_grid_thw_list = [] @@ -1929,6 +1997,7 @@ def _build_unpacked_sample( grid_thw = kwargs["grid_thw"] raw_images = kwargs["raw_images"] pixel_values_tensor = kwargs["pixel_values"] + feature_attention_mask_tensor = kwargs.get("feature_attention_mask") images_num = kwargs["images_num"] image_patch_idx = kwargs["image_patch_idx"] @@ -1965,6 +2034,10 @@ def _build_unpacked_sample( if num_patch > 0: pixel_slice = pixel_values_tensor[image_patch_idx:image_patch_idx + num_patch] pixel_values.append(pixel_slice.clone()) + if feature_attention_mask_tensor is not None: + mask_slice = feature_attention_mask_tensor[local_grid_idx + img_idx:local_grid_idx + + img_idx + 1] + feature_attention_mask_list.append(mask_slice.clone()) image_patch_idx += num_patch local_grid_idx += image_num @@ -2031,6 +2104,9 @@ def _build_unpacked_sample( raw_images=raw_images, pixel_values=pixel_values, pixel_values_videos=pixel_values_videos, + feature_attention_mask=( + torch.cat(feature_attention_mask_list, dim=0).to("cuda") if feature_attention_mask_list else None + ), num_actions=action_mask.size(1), packed_seq_lens=None, response_length=action_mask.float().sum(dim=-1), @@ -2043,6 +2119,61 @@ def _build_unpacked_sample( video_num=all_vid_num, ), image_patch_idx, video_patch_idx + def _build_unpacked_audio_sample( + self, + outputs: List, + prompts: List[str], + labels: Optional[List], + references: Optional[List], + audio_values: Optional[torch.Tensor], + feature_attention_mask: Optional[torch.Tensor], + ) -> SamplesVL: + """ + Build an unpacked audio-language sample batch. + """ + max_input_len = max(len(out.prompt_token_ids) for out in outputs) + max_output_len = max(len(out.output_token_ids) for out in outputs) + + pad_token_id = self.tokenizer.pad_token_id + eos_token_id = self.tokenizer.eos_token_id + + sequences = [] + all_output_ids = [] + for output in outputs: + input_len = len(output.prompt_token_ids) + input_ids = [pad_token_id] * (max_input_len - input_len) + list(output.prompt_token_ids) + + output_len = len(output.output_token_ids) + output_ids = list(output.output_token_ids) + [pad_token_id] * (max_output_len - output_len) + all_output_ids.append(output.output_token_ids) + sequences.append(input_ids + output_ids) + + output_texts = self.tokenizer.batch_decode(all_output_ids) + + sequences = torch.tensor(sequences) + sequences, attention_mask, action_mask = self.actor.process_sequences( + sequences, max_input_len, eos_token_id, pad_token_id + ) + sequences = sequences.to("cuda") + attention_mask = attention_mask.to("cuda") + action_mask = action_mask.to("cuda") + + return SamplesVL( + sequences=sequences, + attention_mask=attention_mask, + action_mask=action_mask, + audio_values=(audio_values.to("cuda") if audio_values is not None else None), + feature_attention_mask=(feature_attention_mask.to("cuda") if feature_attention_mask is not None else None), + num_actions=action_mask.size(1), + packed_seq_lens=None, + response_length=action_mask.float().sum(dim=-1), + total_length=attention_mask.float().sum(dim=-1), + references=references, + labels=labels, + prompts=prompts, + output_texts=output_texts, + ) + def _build_packed_sample( self, outputs: List, diff --git a/lightrft/trainer/ppo_trainer_vl.py b/lightrft/trainer/ppo_trainer_vl.py index e3b44091..d34b412e 100644 --- a/lightrft/trainer/ppo_trainer_vl.py +++ b/lightrft/trainer/ppo_trainer_vl.py @@ -1,11 +1,11 @@ +import math import os -import sys import os.path +import sys from abc import ABC from typing import Any, Callable, Dict, List, Optional import torch -import math import torch.nn as nn from torch.optim import Optimizer from torch.utils.data import DataLoader @@ -13,86 +13,35 @@ from lightrft.models import ActorVL, GPTLMLoss, PolicyLoss, ValueLoss from lightrft.models.actor_modality import ActorModality, get_supported_parameters -from lightrft.models.utils import masked_mean, unpacking_samples, compute_approx_kl -from lightrft.utils.distributed_sampler import DistributedSampler +from lightrft.models.utils import compute_approx_kl, masked_mean, unpacking_samples +from lightrft.trainer import ( + AdaptiveKLController, + ExperienceVL, + FixedKLController, + NaiveExperienceMakerVL, + NaiveReplayBufferVL, +) +from lightrft.trainer.modality_utils import build_supported_model_kwargs from lightrft.utils import rotate_ckpt_dirs -from lightrft.trainer import AdaptiveKLController, ExperienceVL, FixedKLController, NaiveExperienceMakerVL, NaiveReplayBufferVL # noqa +from lightrft.utils.distributed_sampler import DistributedSampler class PPOTrainerVL(ABC): """ - Trainer for Proximal Policy Optimization (PPO) algorithm for Vision-Language Models. - - :param strategy: The training strategy to use. - :type strategy: Strategy - :param actor: The actor model in the PPO algorithm. - :type actor: ActorVL - :param critic: The critic model in the PPO algorithm. - :type critic: nn.Module - :param reward_model: The reward model for calculating rewards in the RLHF setup. - :type reward_model: nn.Module - :param initial_model: The initial model for reference logits to limit actor updates in RLHF. - :type initial_model: ActorVL - :param ema_model: The exponential moving average model for stable training. - :type ema_model: ActorVL - :param actor_optim: The optimizer for the actor model. - :type actor_optim: Optimizer - :param critic_optim: The optimizer for the critic model. - :type critic_optim: Optimizer - :param actor_scheduler: The learning rate scheduler for the actor. - :type actor_scheduler: Scheduler - :param critic_scheduler: The learning rate scheduler for the critic. - :type critic_scheduler: Scheduler - :param ema_beta: EMA decay rate for model stability, defaults to 0.992. - :type ema_beta: float - :param init_kl_coef: Initial coefficient for KL divergence, defaults to 0.001. - :type init_kl_coef: float - :param kl_target: Target value for KL divergence, defaults to None. - :type kl_target: float, optional - :param kl_horizon: Horizon for KL annealing, defaults to 10000. - :type kl_horizon: int - :param ptx_coef: Coefficient for supervised loss from pre-trained data, defaults to 0. - :type ptx_coef: float - :param micro_train_batch_size: Micro-batch size for actor training, defaults to 8. - :type micro_train_batch_size: int - :param buffer_limit: Maximum size of the replay buffer, defaults to 0. - :type buffer_limit: int - :param buffer_cpu_offload: If True, offloads replay buffer to CPU, defaults to True. - :type buffer_cpu_offload: bool - :param eps_clip: Clipping coefficient for policy loss, defaults to 0.2. - :type eps_clip: float - :param value_clip: Clipping coefficient for value function loss, defaults to 0.2. - :type value_clip: float - :param micro_rollout_batch_size: Micro-batch size for generating rollouts, defaults to 8. - :type micro_rollout_batch_size: int - :param gradient_checkpointing: If True, enables gradient checkpointing, defaults to False. - :type gradient_checkpointing: bool - :param max_epochs: Number of epochs to train, defaults to 1. - :type max_epochs: int - :param max_norm: Maximum gradient norm for gradient clipping, defaults to 1.0. - :type max_norm: float - :param tokenizer: Tokenizer for input data, defaults to None. - :type tokenizer: Callable, optional - :param processor: Processor for multimodal input data, defaults to None. - :type processor: Callable, optional - :param prompt_max_len: Maximum length for prompts, defaults to 128. - :type prompt_max_len: int - :param dataloader_pin_memory: If True, pins memory in the data loader, defaults to True. - :type dataloader_pin_memory: bool - :param remote_rm_url: URL for remote reward model API, defaults to None. - :type remote_rm_url: str, optional - :param reward_fn: Custom reward function for computing rewards, defaults to None. - :type reward_fn: Callable, optional - :param reward_fn_label_map: Label mapping for reward function, defaults to None. - :type reward_fn_label_map: dict, optional - :param reward_recipe: Recipe configuration for reward computation, defaults to None. - :type reward_recipe: dict, optional - :param save_hf_ckpt: Whether to save huggingface-format model weight, defaults to False. - :type save_hf_ckpt: bool - :param disable_ds_ckpt: Whether not to save deepspeed-format model weight (used for training recovery). - :type disable_ds_ckpt: bool - :param generate_kwargs: Additional arguments for model generation. - :type generate_kwargs: dict + Trainer for Proximal Policy Optimization (PPO) over multimodal actors. + + The trainer keeps the original PPO/VL training structure intact: rollout collection, + replay buffering, PPO optimization, logging, evaluation, and checkpoint rotation all + remain in the same layer. The refactor in this file is limited to modality-aware data + plumbing so audio-language models can follow an explicit audio path while vision-language + models continue to use the existing image/video path. + + In practice this means prompt batches are normalized once, replay items preserve the + modality-specific tensors they carry, and actor/critic forwards only receive the kwargs + they actually support. Audio-language models therefore use dedicated fields such as + ``audio_values`` and ``feature_attention_mask`` instead of overloading the vision path. + + Parameter details are documented on :meth:`__init__`. """ def __init__( self, @@ -132,11 +81,88 @@ def __init__( disable_ds_ckpt: bool = False, **generate_kwargs, ) -> None: + """ + Initialize the PPO trainer for multimodal RL fine-tuning. + + :param strategy: Distributed or single-process training strategy wrapper. + :type strategy: Strategy + :param actor: Actor model optimized by PPO. + :type actor: ActorVL + :param critic: Critic model used for value prediction. + :type critic: nn.Module + :param reward_model: Reward model used in RLHF/RLAIF reward computation. + :type reward_model: nn.Module + :param initial_model: Frozen reference model used for KL regularization. + :type initial_model: ActorVL + :param ema_model: Exponential moving average copy of the actor, if enabled. + :type ema_model: ActorVL + :param actor_optim: Optimizer for actor updates. + :type actor_optim: Optimizer + :param critic_optim: Optimizer for critic updates. + :type critic_optim: Optimizer + :param actor_scheduler: Learning-rate scheduler for the actor optimizer. + :type actor_scheduler: Scheduler + :param critic_scheduler: Learning-rate scheduler for the critic optimizer. + :type critic_scheduler: Scheduler + :param ema_beta: EMA decay used when updating ``ema_model``. + :type ema_beta: float + :param init_kl_coef: Initial KL penalty coefficient. + :type init_kl_coef: float + :param kl_target: Target KL value for adaptive control. If ``None``, a fixed controller is used. + :type kl_target: float, optional + :param kl_horizon: Horizon used by the adaptive KL controller. + :type kl_horizon: int + :param ptx_coef: Coefficient applied to the optional PTX loss. + :type ptx_coef: float + :param micro_train_batch_size: Micro-batch size used by the replay buffer and PPO dataloader. + :type micro_train_batch_size: int + :param buffer_limit: Maximum replay-buffer capacity. ``0`` keeps the default behavior. + :type buffer_limit: int + :param buffer_cpu_offload: Whether replay items may be offloaded to CPU memory. + :type buffer_cpu_offload: bool + :param eps_clip: PPO policy clipping coefficient. + :type eps_clip: float + :param value_clip: Value-function clipping coefficient. + :type value_clip: float + :param micro_rollout_batch_size: Micro-batch size used during rollout generation. + :type micro_rollout_batch_size: int + :param gradient_checkpointing: Whether actor/critic gradient checkpointing is enabled upstream. + :type gradient_checkpointing: bool + :param max_epochs: Number of PPO epochs run over each replay-buffer snapshot. + :type max_epochs: int + :param max_norm: Gradient clipping threshold. + :type max_norm: float + :param tokenizer: Tokenizer used for text decode/encode helpers. + :type tokenizer: Callable, optional + :param processor: Multimodal processor used by rollout collection. It may be a vision or audio processor + depending on the actor modality. + :type processor: Callable, optional + :param prompt_max_len: Maximum prompt length used by the experience maker. + :type prompt_max_len: int + :param dataloader_pin_memory: Whether PPO dataloaders should pin host memory. + :type dataloader_pin_memory: bool + :param remote_rm_url: Optional remote reward-model endpoint. + :type remote_rm_url: str, optional + :param reward_fn: Optional custom reward function applied on rollout outputs. + :type reward_fn: Callable, optional + :param reward_fn_label_map: Optional label mapping passed to reward helpers. + :type reward_fn_label_map: dict, optional + :param reward_recipe: Optional structured reward configuration. + :type reward_recipe: dict, optional + :param save_hf_ckpt: Whether to additionally export Hugging Face checkpoints/adapters. + :type save_hf_ckpt: bool + :param disable_ds_ckpt: Whether to disable DeepSpeed-format checkpoints. + :type disable_ds_ckpt: bool + :param generate_kwargs: Extra generation kwargs forwarded to rollout/eval collection. + Modality-specific fields are filtered later according to the actor modality. + :type generate_kwargs: dict + """ assert ( not isinstance(reward_model, List) or len(reward_model) == 1 or reward_fn is not None ), "reward_fn must be specified if using multiple reward models" ABC.__init__(self) + self.strategy = strategy self.args = strategy.args self.save_hf_ckpt = save_hf_ckpt @@ -175,18 +201,17 @@ def __init__( self.actor_scheduler = actor_scheduler self.critic_scheduler = critic_scheduler - # Cache actor's supported parameters based on its modality - # Default to VISION_LANGUAGE for backward compatibility with models without modality attribute + # Cache actor modality once so rollout/training can route inputs without branching on model types. + # This plays the same role as the old supported-parameter cache, but now audio has its own path too. actor_modality = self.actor.modality self._actor_supported_params = get_supported_parameters(actor_modality) + self._is_audio_actor = actor_modality == ActorModality.AUDIO_LANGUAGE self.actor_loss_fn = PolicyLoss(eps_clip, use_cpg_loss=self.args.use_cpg_loss) - self.critic_loss_fn = ValueLoss(value_clip) self.ptx_loss_fn = GPTLMLoss() self.freezing_actor_steps = getattr(self.args, "freezing_actor_steps", -1) - self.aux_loss = self.args.aux_loss_coef > 1e-8 if self.kl_target: @@ -212,11 +237,12 @@ def __init__( micro_train_batch_size, buffer_limit, buffer_cpu_offload, packing_samples ) - # Initialize wandb/tensorboard for logging self._wandb = None self._tensorboard = None - self.eval_step_counter = 0 # Independent counter for eval X-axis - self.wandb_log_counter = 0 # Global counter for unique wandb system steps + # Independent counters keep eval plots monotonic and avoid wandb step collisions. + # This preserves the old intent where eval metrics were not forced onto sparse training steps. + self.eval_step_counter = 0 + self.wandb_log_counter = 0 if self.strategy.args.use_wandb and self.strategy.is_rank_0(): import wandb @@ -232,21 +258,16 @@ def __init__( config=strategy.args.__dict__, reinit=True, ) - - # Define custom metrics to allow different X-axes - # rollout/* and train/* use the main training step + # Define custom metrics to allow different X-axes: + # rollout/* and train/* use the main training step, + # while eval/* uses its own counter. wandb.define_metric("rollout/global_step") wandb.define_metric("rollout/*", step_metric="rollout/global_step") - wandb.define_metric("train/global_step") wandb.define_metric("train/*", step_metric="train/global_step") - - # eval/* uses its own counter, allowing it to be plotted sequentially - # even if evaluations happen rarely wandb.define_metric("eval/global_step") wandb.define_metric("eval/*", step_metric="eval/global_step") - # Initialize TensorBoard writer if wandb is not available if self.strategy.args.use_tensorboard and self._wandb is None and self.strategy.is_rank_0(): from torch.utils.tensorboard import SummaryWriter @@ -254,6 +275,97 @@ def __init__( log_dir = os.path.join(self.strategy.args.use_tensorboard, strategy.args.wandb_run_name) self._tensorboard = SummaryWriter(log_dir=log_dir) + @staticmethod + def _ensure_device_and_contiguous(value, device): + """ + Move tensors to the target GPU and make them contiguous for downstream kernels. + + :param value: Tensor or nested tensor list to normalize. + :type value: torch.Tensor or list or Any + :param device: CUDA device index expected by the current rank. + :type device: int + :return: Value moved to the requested device with contiguous layout preserved recursively. + :rtype: Any + """ + if value is None: + return None + if isinstance(value, list): + return [PPOTrainerVL._ensure_device_and_contiguous(v, device) for v in value] + if not isinstance(value, torch.Tensor): + return value + if value.device.type != "cuda" or value.device.index != device: + value = value.to(device) + if not value.is_contiguous(): + value = value.contiguous() + return value + + def _build_model_kwargs(self, source, device: Optional[int] = None) -> Dict[str, Any]: + """ + Select and optionally relocate only the multimodal kwargs supported by the current actor modality. + + :param source: Replay item or mapping containing candidate multimodal tensors. + :type source: Any + :param device: Optional CUDA device index used to normalize tensor placement. + :type device: int, optional + :return: Filtered kwargs that can be passed directly into actor/critic forward. + :rtype: Dict[str, Any] + """ + kwargs = build_supported_model_kwargs(source, self._actor_supported_params) + if device is not None: + kwargs = {key: self._ensure_device_and_contiguous(value, device) for key, value in kwargs.items()} + return kwargs + + def _unpack_prompt_batch(self, batch): + """ + Normalize prompt-dataloader outputs across text, vision, video, and audio variants. + + Audio example datasets still produce a 4-field batch, but the second field now maps to + ``audios`` instead of overloading the image slot. + + :param batch: Raw batch emitted by the prompt dataloader. + :type batch: tuple or list + :return: Tuple of ``(prompts, images, videos, audios, references, labels)`` used by rollout code. + :rtype: tuple + """ + if len(batch) == 5: + prompts, images, videos, references, labels = batch + return prompts, images, videos, None, references, labels + if len(batch) == 4: + prompts, modality_inputs, references, labels = batch + if self._is_audio_actor: + return prompts, None, None, modality_inputs, references, labels + return prompts, modality_inputs, None, None, references, labels + raise ValueError(f"Unsupported prompt batch format with {len(batch)} fields.") + + def _make_experience_list(self, prompts, images, videos, audios, references, labels): + """ + Shared rollout helper used by both training and evaluation. + + :param prompts: Prompt strings for the current batch. + :type prompts: list + :param images: Optional image inputs. + :type images: Any + :param videos: Optional video inputs. + :type videos: Any + :param audios: Optional audio inputs. + :type audios: Any + :param references: Optional references used by reward functions. + :type references: Any + :param labels: Optional labels used by reward functions. + :type labels: Any + :return: List of rollout experiences produced from the prompt batch. + :rtype: list + """ + return self.experience_maker.make_experience_list( + prompts, + all_images=images, + all_videos=videos, + all_audios=audios, + all_references=references, + all_labels=labels, + **self.generate_kwargs, + ) + def fit( self, args, @@ -264,32 +376,34 @@ def fit( num_update_steps_per_episodes=1, ) -> None: """ - Main training loop for PPO. + Main PPO loop: rollout, aggregate replay items, optimize, log, evaluate, and checkpoint. - :param args: Training arguments. + :param args: Runtime training arguments. :type args: Namespace - :param prompts_dataloader: DataLoader for prompt data. + :param prompts_dataloader: Prompt dataloader. Batches may be text-only, image/video multimodal, + or audio multimodal, and are normalized by :meth:`_unpack_prompt_batch`. :type prompts_dataloader: DataLoader - :param pretrain_dataloader: DataLoader for pre-training data. + :param pretrain_dataloader: Optional PTX dataloader consumed during actor updates. :type pretrain_dataloader: DataLoader - :param eval_dataloader: DataLoader for evaluation data, defaults to None. + :param eval_dataloader: Optional evaluation dataloader using the same rollout path. :type eval_dataloader: DataLoader, optional - :param consumed_samples: Number of samples already consumed, defaults to 0. + :param consumed_samples: Number of rollout samples already consumed when resuming training. :type consumed_samples: int - :param num_update_steps_per_episodes: Number of update steps per episode, defaults to 1. + :param num_update_steps_per_episodes: Planned PPO update steps per episode. :type num_update_steps_per_episodes: int + :return: ``None``. + :rtype: None """ - - # Calculate samples per rollout and per training iteration samples_per_rollout = args.rollout_batch_size * args.n_samples_per_prompt samples_per_train = args.train_batch_size * args.n_samples_per_prompt - # Print training mode information + # Report whether each rollout leads to multiple updates or vice versa. if args.train_batch_size < args.rollout_batch_size: updates_per_rollout = samples_per_rollout / samples_per_train self.strategy.print( f"\n{'=' * 80}\n" - f"HIGH FREQUENCY UPDATE MODE: train_batch_size ({args.train_batch_size}) < rollout_batch_size ({args.rollout_batch_size})\n" # noqa + f"HIGH FREQUENCY UPDATE MODE: train_batch_size ({args.train_batch_size}) < " + f"rollout_batch_size ({args.rollout_batch_size})\n" f"{'=' * 80}\n" f"Behavior:\n" f" - Each rollout generates {samples_per_rollout} samples.\n" @@ -300,7 +414,8 @@ def fit( elif args.train_batch_size > args.rollout_batch_size: self.strategy.print( f"\n{'=' * 80}\n" - f"ACCUMULATION MODE: train_batch_size ({args.train_batch_size}) > rollout_batch_size ({args.rollout_batch_size})\n" # noqa + f"ACCUMULATION MODE: train_batch_size ({args.train_batch_size}) > " + f"rollout_batch_size ({args.rollout_batch_size})\n" f"{'=' * 80}\n" f"Behavior:\n" f" - Multiple rollouts needed for one update.\n" @@ -308,41 +423,32 @@ def fit( ) # Calculate number of rollouts per episode. - # Regardless of TBS and RBS relationship, rollout count should be determined by "total data / rollout size". - # Numerator (num_update_steps * train_batch_size) equals "total samples planned for this episode". - # Denominator (rollout_batch_size * n_samples) equals "samples produced per rollout". - # This calculation ensures data collection volume is constant. - # When TBS=64, num_update_steps is naturally twice as large as when TBS=128. - # Substituting into formula: (2N * 0.5T) / R = (N * T) / R. - # Conclusion: Rollout count unchanged, but internal update loop count doubles due to smaller TBS. - + # Regardless of the TBS/RBS relationship, rollout count should depend on total sample volume, + # not on how those samples are internally split across optimizer steps. num_rollouts_per_episodes = ( num_update_steps_per_episodes * args.train_batch_size // args.max_epochs // args.rollout_batch_size // args.n_samples_per_prompt ) - - # Safeguard to prevent num_rollouts_per_episodes from being 0 if num_rollouts_per_episodes == 0: - # Try recalculating with ceil to prevent fractional values from being discarded by integer division - val = (num_update_steps_per_episodes * - args.train_batch_size) / (args.max_epochs * args.rollout_batch_size * args.n_samples_per_prompt) - num_rollouts_per_episodes = math.ceil(val) - + # Use ceil as a safeguard when integer division would otherwise drop a fractional rollout. + num_rollouts_per_episodes = math.ceil( + (num_update_steps_per_episodes * args.train_batch_size) / + (args.max_epochs * args.rollout_batch_size * args.n_samples_per_prompt) + ) if num_rollouts_per_episodes == 0: self.strategy.print("[WARNING] Calculated num_rollouts_per_episodes is 0. Forcing to 1.") num_rollouts_per_episodes = 1 - # Get eval and save steps if args.eval_steps == -1: - args.eval_steps = num_rollouts_per_episodes # Evaluate once per epoch + args.eval_steps = num_rollouts_per_episodes if args.save_steps == -1: - args.save_steps = float("inf") # Do not save checkpoint + args.save_steps = float("inf") self.prompts_dataloader = prompts_dataloader self.pretrain_dataloader = pretrain_dataloader - self.eval_dataloader = eval_dataloader # Save for evaluation + self.eval_dataloader = eval_dataloader - # Restore step and start_episode + # Recover where the previous run left off when resuming from checkpoints. steps = consumed_samples // args.rollout_batch_size + 1 start_episode = consumed_samples // args.rollout_batch_size // num_rollouts_per_episodes consumed_samples = consumed_samples % (num_rollouts_per_episodes * args.rollout_batch_size) @@ -350,8 +456,10 @@ def fit( for episode in range(start_episode, args.num_episodes): if isinstance(self.prompts_dataloader.sampler, DistributedSampler): self.prompts_dataloader.sampler.set_epoch( - episode, consumed_samples=0 if episode > start_episode else consumed_samples + episode, + consumed_samples=0 if episode > start_episode else consumed_samples, ) + pbar = tqdm( range(self.prompts_dataloader.__len__()), desc=f"Episode [{episode + 1}/{args.num_episodes}]", @@ -359,47 +467,20 @@ def fit( ) for batch in self.prompts_dataloader: - # Compatible with both image-only (4 args) and video (5 args) dataloaders - if len(batch) == 5: - rand_prompts, rand_images, rand_videos, rand_references, rand_labels = batch - else: - rand_prompts, rand_images, rand_references, rand_labels = batch - rand_videos = None - - # TODO: Remove debug print - self.strategy.print( - f"rand_prompts:\n {rand_prompts}\n , rand_images:{rand_images}\n , rand_references:{rand_references}\n, rand_labels:{rand_labels}\n " # noqa - ) - - for i, experience in enumerate( - self.experience_maker.make_experience_list( - rand_prompts, - rand_images, - all_videos=rand_videos, - all_references=rand_references, - all_labels=rand_labels, - **self.generate_kwargs - ) - ): - if i == 0: - output = self.tokenizer.batch_decode( - experience.sequences[0].unsqueeze(0), skip_special_tokens=True - ) - self.strategy.print("collect phase: experience.sequences w skip_special_tokens: ", output) - self.strategy.print( - f"collect phase: rand_prompts:\n {rand_prompts[0:2]}\n , rand_images:{rand_images[0:2]}\n , rand_references:{rand_references[0:2]}\n, rand_labels:{rand_labels[0:2]}\n " # noqa - ) - # print all - # self.strategy.print( - # f"rand_prompts:\n {rand_prompts}\n , rand_images:{rand_images}\n , rand_references:{rand_references}\n, rand_labels:{rand_labels}\n " # noqa - # ) + # The helper keeps the rollout loop agnostic to whether the batch is audio or vision. + prompts, images, videos, audios, references, labels = self._unpack_prompt_batch(batch) + experience_list = self._make_experience_list(prompts, images, videos, audios, references, labels) + if not experience_list: + pbar.update() + steps += 1 + continue + for experience in experience_list: self.replay_buffer.append(experience) - self.strategy.report_memory('after replay_buffer ready') + self.strategy.report_memory("after replay_buffer ready") - # Aggregate rollout statistics from replay buffer - # Collect metrics from the rollout/collection phase + # Aggregate rollout statistics from replay buffer before PPO updates clear it. rollout_status = {} if self.replay_buffer.items: all_rewards = [] @@ -408,36 +489,26 @@ def fit( all_response_lengths = [] for item in self.replay_buffer.items: - # Collect rewards from rollout - if hasattr(item, 'info') and item.info is not None and 'reward' in item.info: - all_rewards.append(item.info['reward']) - - # Robust handling of reward_metrics - # 1. Check if info exists - # 2. Check if 'reward_metrics' key exists - # 3. Check if reward_metrics is not None (critical!) + if hasattr(item, "info") and item.info is not None and "reward" in item.info: + all_rewards.append(item.info["reward"]) + if ( - hasattr(item, 'info') and item.info is not None and 'reward_metrics' in item.info - and item.info['reward_metrics'] is not None + hasattr(item, "info") and item.info is not None and "reward_metrics" in item.info + and item.info["reward_metrics"] is not None ): + reward_metrics = item.info["reward_metrics"] + if "format_reward" in reward_metrics: + all_format_rewards.append(reward_metrics["format_reward"]) + if "accuracy_reward" in reward_metrics: + all_accuracy_rewards.append(reward_metrics["accuracy_reward"]) - reward_metrics = item.info['reward_metrics'] - - # Safely extract sub-metrics - if 'format_reward' in reward_metrics: - all_format_rewards.append(reward_metrics['format_reward']) - if 'accuracy_reward' in reward_metrics: - all_accuracy_rewards.append(reward_metrics['accuracy_reward']) - - # Collect response lengths from rollout - if hasattr(item, 'info') and item.info is not None and 'response_length' in item.info: - all_response_lengths.append(item.info['response_length']) + if hasattr(item, "info") and item.info is not None and "response_length" in item.info: + all_response_lengths.append(item.info["response_length"]) - # Compute rollout statistics device = torch.cuda.current_device() if all_rewards: - # [TENSOR-FIX] Handle both tensor lists and scalar lists + # Some reward functions return tensors directly, others scalar values. if isinstance(all_rewards[0], torch.Tensor): rewards_tensor = torch.cat([t.to(device).float() for t in all_rewards]) else: @@ -446,74 +517,58 @@ def fit( rollout_status["rollout_reward_std"] = rewards_tensor.std().item() if all_format_rewards: - # [TENSOR-FIX] Handle both tensor lists and scalar lists - # Issue: all_format_rewards may contain tensors (from reward_metrics), - # but torch.tensor() cannot convert a list of tensors directly. - # Solution: Use torch.cat() for tensor lists, torch.tensor() for scalar lists if isinstance(all_format_rewards[0], torch.Tensor): - # List of tensors: concatenate them format_tensor = torch.cat([t.to(device).float() for t in all_format_rewards]) else: - # List of scalars: convert to tensor format_tensor = torch.tensor(all_format_rewards, dtype=torch.float32, device=device) - mean_format_reward = format_tensor.mean().item() - - # Only display if mean is significantly non-zero if abs(mean_format_reward) > 1e-6: rollout_status["rollout_format_reward"] = mean_format_reward if all_accuracy_rewards: - # [TENSOR-FIX] Handle both tensor lists and scalar lists if isinstance(all_accuracy_rewards[0], torch.Tensor): accuracy_tensor = torch.cat([t.to(device).float() for t in all_accuracy_rewards]) else: accuracy_tensor = torch.tensor(all_accuracy_rewards, dtype=torch.float32, device=device) - mean_accuracy_reward = accuracy_tensor.mean().item() - - # Only display if mean is significantly non-zero if abs(mean_accuracy_reward) > 1e-6: rollout_status["rollout_accuracy_reward"] = mean_accuracy_reward if all_response_lengths: - # [TENSOR-FIX] Handle both tensor lists and scalar lists if isinstance(all_response_lengths[0], torch.Tensor): lengths_tensor = torch.cat([t.to(device).float() for t in all_response_lengths]) else: lengths_tensor = torch.tensor(all_response_lengths, dtype=torch.float32, device=device) - rollout_status["rollout_response_length"] = lengths_tensor.mean().item() - # TODO: Check normalization behavior + # Group-normalized estimators already normalize advantages during experience creation. if self.args.advantage_estimator != "group_norm": self.replay_buffer.normalize("advantages", self.strategy) - self.strategy.report_memory('before train') - + self.strategy.report_memory("before train") status = self.ppo_train(steps) - - self.strategy.report_memory('before clear buffer') + self.strategy.report_memory("before clear buffer") self.replay_buffer.clear() - - self.strategy.report_memory('after train') + self.strategy.report_memory("after train") if "kl" in status: self.kl_ctl.update(status["kl"], args.rollout_batch_size * args.n_samples_per_prompt) - # Update Episode pbar with ROLLOUT statistics (not training statistics!) + # Progress bar reflects rollout quality; wandb/tensorboard will receive both rollout and train metrics. pbar.set_postfix(rollout_status) - - # Logs/checkpoints: save BOTH ROLLOUT and TRAINING statistics to wandb - # [FIX] Merge rollout_status (from inference) and status (from training) - # to ensure wandb logs contain both types of metrics client_states = {"consumed_samples": steps * args.rollout_batch_size} - logs_dict_combined = {**rollout_status, **status} # Merge: rollout first, training second - - self.save_logs_and_checkpoints(args, steps, pbar, logs_dict_combined, client_states, episode=episode) + logs_dict_combined = {**rollout_status, **status} + self.save_logs_and_checkpoints( + args, + steps, + pbar, + logs_dict_combined, + client_states, + episode=episode, + ) pbar.update() - steps = steps + 1 + steps += 1 if self._wandb is not None and self.strategy.is_rank_0(): self._wandb.finish() @@ -522,18 +577,19 @@ def fit( def ppo_train(self, global_steps=0): """ - PPO training loop over the replay buffer. + PPO optimization over the current replay buffer snapshot. - NOTE: This method is not used directly in the main trainer, - as it's overridden by external classes (e.g., lightrft/trainer/spmd_ppo_trainer.py). + NOTE: this method is overridden by the SPMD trainer in the main audio run, + but keeping the base implementation explicit is still useful for non-SPMD execution + and for understanding the reference PPO flow. - :param global_steps: Current global step count, defaults to 0. + :param global_steps: Current global training step. :type global_steps: int - :return: Dictionary of averaged training statistics. + :return: Mean metrics aggregated over all PPO minibatches in the snapshot. :rtype: dict """ torch.cuda.empty_cache() - # Replay buffer may be empty at first, we should rebuild at each training + # Rebuild the dataloader each time because the replay buffer is refreshed after every rollout. dataloader = DataLoader( self.replay_buffer, batch_size=self.replay_buffer.sample_batch_size, @@ -556,15 +612,14 @@ def ppo_train(self, global_steps=0): experience.to_device(device) status = self.training_step(experience, global_steps) - # For DP: weighted mean for KL + # For DP runs, KL is aggregated with response-length weighting. if "kl" in status: status["kl"] *= status["response_length"] status = self.strategy.all_reduce(status) status["kl"] /= status["response_length"] short_status = {} - - # Add core metrics with abbreviations to keep progress bar concise + # Keep progress-bar keys compact while preserving detailed metrics in logs. if "policy_loss" in status: short_status.update({ "pg": status.get("policy_loss"), @@ -575,49 +630,50 @@ def ppo_train(self, global_steps=0): "kl": status.get("kl"), "act_lr": status.get("actor_lr"), }) - if "critic_loss" in status: short_status.update({ "cri": status.get("critic_loss"), "vals": status.get("values"), "cri_lr": status.get("critic_lr"), }) - if "ptx_loss" in status: short_status["ptx"] = status.get("ptx_loss") - - for k, v in status.items(): - if "/" in k: - short_key = k.split('/')[-1] - short_status[short_key] = v + for key, value in status.items(): + if "/" in key: + short_status[key.split("/")[-1]] = value status_list.append(status) pbar.set_postfix(short_status) if status_list: status_mean = status_list[0] - for m in status_list[1:]: - for k, v in m.items(): - status_mean[k] += v - for k in status_mean.keys(): - status_mean[k] /= len(status_list) + for metrics in status_list[1:]: + for key, value in metrics.items(): + status_mean[key] += value + for key in status_mean.keys(): + status_mean[key] /= len(status_list) + torch.cuda.empty_cache() return status_mean - def training_step(self, - experience: ExperienceVL, - global_steps, - entropy_mask: Optional[torch.Tensor] = None) -> Dict[str, float]: + def training_step( + self, + experience: ExperienceVL, + global_steps, + entropy_mask: Optional[torch.Tensor] = None, + ) -> Dict[str, float]: """ - Single training step combining actor and critic updates. + Run one PPO optimization step on a replay-buffer batch. + + Actor updates are applied first and critic updates are added afterwards when a critic exists. - :param experience: Experience batch from replay buffer. + :param experience: Replay-buffer batch containing sequences, masks, rewards, and modality tensors. :type experience: ExperienceVL - :param global_steps: Current global step count. + :param global_steps: Current global step used for actor-freeze logic. :type global_steps: int - :param entropy_mask: Optional mask for high-entropy tokens. + :param entropy_mask: Optional mask for entropy-aware policy loss variants. :type entropy_mask: Optional[torch.Tensor] - :return: Dictionary of training statistics. + :return: Training statistics from actor and critic updates. :rtype: Dict[str, float] """ status = {} @@ -628,34 +684,36 @@ def training_step(self, return status def _validate_qwen_vl_tensors( - self, sequences: torch.Tensor, pixel_values: Optional[torch.Tensor], context: str = "training" + self, + sequences: torch.Tensor, + pixel_values: Optional[torch.Tensor], + context: str = "training", ) -> bool: """ - Validates the consistency between image tokens in sequences and pixel_values features. - - :param sequences: Token sequence tensor. - :type sequences: torch.Tensor - :param pixel_values: Processed pixel values tensor. - :type pixel_values: Optional[torch.Tensor] - :param context: A string indicating where the validation is called from (e.g., "actor_rl", "actor_ptx"). - :type context: str - :return: True if data is consistent, False otherwise. - :rtype: bool - """ + Defensive validation for Qwen-VL style image-token / image-feature consistency. + + This preserves the old skip-on-mismatch safeguard for vision batches. + Audio-language actors bypass it naturally because they do not forward ``pixel_values``. + + :param sequences: Token sequences about to be forwarded through the actor. + :type sequences: torch.Tensor + :param pixel_values: Image features paired with ``sequences``. + :type pixel_values: Optional[torch.Tensor] + :param context: Human-readable call site used in warning logs. + :type context: str + :return: ``True`` when the batch is safe to run, otherwise ``False``. + :rtype: bool + """ if pixel_values is None or pixel_values.numel() == 0: - # This is a text-only batch, no validation needed. return True config = self.strategy.unwrap_model(self.actor.model).config image_token_id = getattr(config, "image_token_id", None) - if image_token_id is None: - # Model does not use special image tokens. return True num_tokens = (sequences == image_token_id).sum().item() num_patches = pixel_values.shape[0] // 4 - if num_tokens != num_patches: self.strategy.print( f"[CRITICAL WARNING] Skipping batch in '{context}'. " @@ -663,47 +721,44 @@ def _validate_qwen_vl_tensors( "This batch will be discarded to prevent a crash." ) return False - return True - def training_step_actor(self, - experience: ExperienceVL, - entropy_mask: Optional[torch.Tensor] = None) -> Dict[str, float]: + def training_step_actor( + self, + experience: ExperienceVL, + entropy_mask: Optional[torch.Tensor] = None, + ) -> Dict[str, float]: """ Actor training step. - :param experience: Experience batch from replay buffer. + Packed and unpacked replay items are normalized into one forward path, and the + actor modality determines whether the model receives vision kwargs or audio kwargs. + + :param experience: Replay-buffer batch for PPO policy optimization. :type experience: ExperienceVL - :return: Dictionary of actor training statistics. + :param entropy_mask: Optional entropy mask forwarded to the policy-loss module. + :type entropy_mask: Optional[torch.Tensor] + :return: Actor-side optimization statistics plus rollout metadata copied from ``experience.info``. :rtype: Dict[str, float] """ self.actor.train() - # TODO: This is a bad indicator to say that data is packed... not supported + # Packed samples concatenate multiple sequences into one row. Unpacked samples stay batched. + # This mirrors the old PPOTrainerVL handling while replacing hard-coded VL kwargs with modality-aware ones. if isinstance(experience.sequences, list): sequences = torch.cat(experience.sequences, dim=0).unsqueeze(0) - - pixel_values = experience.pixel_values - image_grid_thws = experience.image_grid_thws - pixel_values_videos = getattr(experience, "pixel_values_videos", None) - video_grid_thws = getattr(experience, "video_grid_thws", None) - old_action_log_probs = torch.cat(experience.action_log_probs, dim=0).unsqueeze(0) advantages = torch.cat(experience.advantages, dim=0).unsqueeze(0) - num_actions = [v.numel() for v in experience.advantages] - packed_seq_lens = [s.numel() for s in experience.sequences] - attention_mask = torch.cat([torch.full_like(s, i + 1) for i, s in enumerate(experience.sequences)], - dim=0).unsqueeze(0) + num_actions = [value.numel() for value in experience.advantages] + packed_seq_lens = [seq.numel() for seq in experience.sequences] + attention_mask = torch.cat( + [torch.full_like(seq, idx + 1) for idx, seq in enumerate(experience.sequences)], + dim=0, + ).unsqueeze(0) if self.args.use_kl_loss and experience.base_action_log_probs is not None: base_action_log_probs = torch.cat(experience.base_action_log_probs, dim=0).unsqueeze(0) else: sequences = experience.sequences - - pixel_values = experience.pixel_values - image_grid_thws = experience.image_grid_thws - pixel_values_videos = getattr(experience, "pixel_values_videos", None) - video_grid_thws = getattr(experience, "video_grid_thws", None) - old_action_log_probs = experience.action_log_probs advantages = experience.advantages num_actions = experience.action_mask.size(1) @@ -713,36 +768,25 @@ def training_step_actor(self, base_action_log_probs = experience.base_action_log_probs if advantages is not None: - # Log max advantage before clipping for debugging (optional) + # Clipping prevents a few extreme group-normalized values from dominating the PPO step. max_adv = advantages.max().item() if max_adv > 10.0: self.strategy.print(f"[Warning] Huge advantage detected: {max_adv}") advantages = torch.clamp(advantages, min=-10.0, max=10.0) - # [DEFENSIVE CHECK] Validate RL data before actor forward pass - # NOTE: This validation is now primarily done in spmd_ppo_trainer.py BEFORE calling training_step - # to ensure all ranks make the same skip decision. This check remains as a safety fallback. - # If this triggers, it indicates a bug in the pre-validation logic. - if not self._validate_qwen_vl_tensors(sequences, pixel_values, context="actor_rl_update"): + # Actor loss. + # Build modality-aware kwargs from the replay item instead of assuming vision-specific fields. + actor_kwargs = self._build_model_kwargs(experience) + if not self._validate_qwen_vl_tensors( + sequences, + actor_kwargs.get("pixel_values"), + context="actor_rl_update", + ): self.strategy.print( "[CRITICAL ERROR] Validation failed inside training_step_actor. " "This should have been caught by pre-validation in spmd_ppo_trainer.py!" ) - return {} # Emergency fallback - should not normally execute - - # Actor loss - # Build kwargs based on actor's modality - only include supported parameters - candidate_params = { - "pixel_values": pixel_values, - "image_grid_thw": image_grid_thws, - "pixel_values_videos": pixel_values_videos, - "video_grid_thw": video_grid_thws, - } - # Audio-language actors expect audio_values; pipeline stores them in pixel_values slot - if "audio_values" in self._actor_supported_params: - candidate_params["audio_values"] = candidate_params.get("pixel_values") - - actor_kwargs = {key: value for key, value in candidate_params.items() if key in self._actor_supported_params} + return {} action_log_probs, output = self.actor( sequences, @@ -750,15 +794,9 @@ def training_step_actor(self, attention_mask=attention_mask, return_output=True, packed_seq_lens=packed_seq_lens, - **actor_kwargs + **actor_kwargs, ) - # NOTE: Explicit masking in log-space is incorrect - removed - # if experience.action_mask is not None: - # # Setting masked positions to 0 to match old_action_log_probs is WRONG in log-space - # action_log_probs = action_log_probs * experience.action_mask - - # Loss function actor_loss = self.actor_loss_fn( action_log_probs, old_action_log_probs, @@ -769,27 +807,18 @@ def training_step_actor(self, if self.args.use_kl_loss: if self.initial_model is not None: - # TODO(pu): Text-only action mask for KL calculation - kl = compute_approx_kl( action_log_probs, base_action_log_probs, experience.action_mask, kl_estimator=self.args.kl_estimator, ) - - # [Protection measure 2] Per-token KL Clamping - # NOTE: Adding this causes svkng training to not converge - # kl = torch.clamp(kl, min=0.0, max=20.0) - else: kl = torch.zeros_like(action_log_probs, dtype=action_log_probs.dtype, device=action_log_probs.device) if not self.args.packing_samples: kl_mean = masked_mean(kl, experience.action_mask, dim=-1) - # Not supported for packed samples else: - # Convert tensor into list of tensors for easier manipulation within dataset kl = unpacking_samples(kl, num_actions) kl_mean = torch.tensor([each_kl.mean() for each_kl in kl], device=action_log_probs.device) @@ -798,23 +827,26 @@ def training_step_actor(self, else: kl_loss = 0 - # Mixtral auxiliary loss - if self.aux_loss: - aux_loss = output.aux_loss - else: - aux_loss = 0 - + aux_loss = output.aux_loss if self.aux_loss else 0 loss = actor_loss + aux_loss * self.args.aux_loss_coef + kl_loss * self.kl_ctl.value if torch.isnan(loss) or torch.isinf(loss): self.strategy.print("[CRITICAL ERROR] Actor loss is NaN or Inf at step. Skipping update.") self.strategy.print(f" Actor Loss: {actor_loss.item()}") - self.strategy.print(f" KL Loss: {kl_loss.item() if isinstance(kl_loss, torch.Tensor) else kl_loss}") + if isinstance(kl_loss, torch.Tensor): + self.strategy.print(f" KL Loss: {kl_loss.item()}") + else: + self.strategy.print(f" KL Loss: {kl_loss}") self.strategy.backward(loss, self.actor, self.actor_optim) - # PTX loss for supervised fine-tuning + # PTX loss for supervised fine-tuning. + # Audio PTX is intentionally left unsupported here because the old PTX path + # was tightly coupled to vision-style tensors. if self.pretrain_dataloader is not None: + if self._is_audio_actor: + raise NotImplementedError("PTX data path for audio-language actors is not implemented in PPOTrainerVL.") + data = next(self.pretrain_dataloader) inputs = data[1].squeeze(1).to(torch.cuda.current_device()) attention_mask = data[2].squeeze(1).to(torch.cuda.current_device()) @@ -831,17 +863,11 @@ def training_step_actor(self, attention_mask=attention_mask, pixel_values=pixel_values, image_grid_thw=image_grid_thws, - return_output=True + return_output=True, ) ptx_log_probs = output["logits"] - - # Loss function ptx_loss = self.ptx_loss_fn(ptx_log_probs, label) - # Mixtral auxiliary loss - if self.aux_loss: - aux_loss = output.aux_loss - else: - aux_loss = 0 + aux_loss = output.aux_loss if self.aux_loss else 0 loss = ptx_loss + aux_loss * self.args.aux_loss_coef self.strategy.backward(self.ptx_coef * loss, self.actor, self.actor_optim) @@ -850,58 +876,44 @@ def training_step_actor(self, if self.ema_model: self.strategy.moving_average(self.actor, self.ema_model, self.ema_beta, "cuda") - # Status status = {"policy_loss": actor_loss.item(), "actor_lr": self.actor_scheduler.get_last_lr()[0]} - if self.pretrain_dataloader is not None: + if self.pretrain_dataloader is not None and not self._is_audio_actor: status["ptx_loss"] = ptx_loss.item() - # Add ratio and loss component statistics from PolicyLoss for diagnosis - if hasattr(self.actor_loss_fn, 'get_last_stats'): - policy_stats = self.actor_loss_fn.get_last_stats() - status.update(policy_stats) - - # self.strategy.print(f"experience.info:{experience.info}") - - # Robustly handle various data types in experience.info for logging - # Note: We keep all metrics in status dict for internal use (e.g., KL weighting, progress bar) - # but will filter out rollout-only metrics when logging to wandb to avoid duplication - for k, v in experience.info.items(): - # Special handling for KL divergence, which is already a scalar item - if k == "kl": - # KL is often weighted by response length, handle it carefully if it's tensor - if isinstance(v, torch.Tensor): - # This logic assumes 'v' is a tensor of KL values per item in the batch - weighted_kl = (v * + # Add ratio and loss-component diagnostics from PolicyLoss when available. + if hasattr(self.actor_loss_fn, "get_last_stats"): + status.update(self.actor_loss_fn.get_last_stats()) + + # Keep rollout-side info in the status dict so upper layers can log both rollout and train metrics together. + # This keeps the old logging behavior where experience.info remained the single source of rollout metadata. + for key, value in experience.info.items(): + if key == "kl": + if isinstance(value, torch.Tensor): + weighted_kl = (value * experience.info["response_length"]).sum() / experience.info["response_length"].sum() - status[k] = weighted_kl.item() - else: # If it's already a scalar float - status[k] = v + status[key] = weighted_kl.item() + else: + status[key] = value continue - # Handle nested dictionaries like 'reward_metrics' - if isinstance(v, dict): - for sub_k, sub_v in v.items(): - log_key = f"{k}/{sub_k}" - if isinstance(sub_v, torch.Tensor): - status[log_key] = sub_v.mean().item() - elif isinstance(sub_v, list) and sub_v and isinstance(sub_v[0], (int, float)): - status[log_key] = sum(sub_v) / len(sub_v) - elif isinstance(sub_v, (int, float)): - status[log_key] = sub_v + if isinstance(value, dict): + for sub_key, sub_value in value.items(): + log_key = f"{key}/{sub_key}" + if isinstance(sub_value, torch.Tensor): + status[log_key] = sub_value.mean().item() + elif isinstance(sub_value, list) and sub_value and isinstance(sub_value[0], (int, float)): + status[log_key] = sum(sub_value) / len(sub_value) + elif isinstance(sub_value, (int, float)): + status[log_key] = sub_value continue - # General handling for other keys - if isinstance(v, torch.Tensor): - # If it's a tensor, it's safe to call .mean() - status[k] = v.float().mean().item() - elif isinstance(v, list): - # If it's a list, only compute mean if it contains numbers - if v and isinstance(v[0], (int, float)): - status[k] = sum(v) / len(v) - # Otherwise, it's a list of strings or dicts, which cannot be averaged. Skip it. - elif isinstance(v, (int, float)): - # If it's already a scalar number, just use it - status[k] = v + if isinstance(value, torch.Tensor): + status[key] = value.float().mean().item() + elif isinstance(value, list): + if value and isinstance(value[0], (int, float)): + status[key] = sum(value) / len(value) + elif isinstance(value, (int, float)): + status[key] = value return status @@ -909,61 +921,28 @@ def training_step_critic(self, experience: ExperienceVL) -> Dict[str, float]: """ Critic training step. - :param experience: Experience batch from replay buffer. + It uses the same modality-aware kwargs assembly as actor training, so audio and vision + stay consistent during PPO value updates. + + :param experience: Replay-buffer batch for PPO value-function optimization. :type experience: ExperienceVL - :return: Dictionary of critic training statistics. + :return: Critic-side optimization statistics. :rtype: Dict[str, float] """ self.critic.train() - - # Layer 1: Get current GPU device device = torch.cuda.current_device() - # Layer 2: Helper function for robust device placement - def ensure_device_and_contiguous(tensor, name="tensor"): - """ - Ensure tensor is: - 1. On the correct GPU device - 2. Contiguous in memory (required by Triton) - 3. Return None safely if input is None - - :param tensor: Input tensor to process. - :type tensor: torch.Tensor or None - :param name: Name for logging purposes, defaults to "tensor". - :type name: str - :return: Processed tensor or None. - :rtype: torch.Tensor or None - """ - if tensor is None: - return None - - # Move to GPU if not already there - if tensor.device.type != 'cuda' or tensor.device.index != device: - tensor = tensor.to(device) - - # Ensure contiguous memory layout for Triton kernels - if not tensor.is_contiguous(): - tensor = tensor.contiguous() - - return tensor - - # Layer 3: Apply defensive device placement to all multimodal tensors - pixel_values = ensure_device_and_contiguous(experience.pixel_values, "pixel_values") - image_grid_thws = ensure_device_and_contiguous(experience.image_grid_thws, "image_grid_thws") - pixel_values_videos = ensure_device_and_contiguous( - getattr(experience, "pixel_values_videos", None), "pixel_values_videos" - ) - video_grid_thws = ensure_device_and_contiguous(getattr(experience, "video_grid_thws", None), "video_grid_thws") - - # TODO: This is a bad indicator to say that data is packed... + # Match the packed/unpacked normalization used in actor training. if isinstance(experience.sequences, list): sequences = torch.cat(experience.sequences, dim=0).unsqueeze(0) old_values = torch.cat(experience.values, dim=0).unsqueeze(0) returns = torch.cat(experience.returns, dim=0).unsqueeze(0) - num_actions = [v.numel() for v in experience.advantages] - packed_seq_lens = [s.numel() for s in experience.sequences] - attention_mask = torch.cat([torch.full_like(s, i + 1) for i, s in enumerate(experience.sequences)], - dim=0).unsqueeze(0) + num_actions = [value.numel() for value in experience.advantages] + packed_seq_lens = [seq.numel() for seq in experience.sequences] + attention_mask = torch.cat( + [torch.full_like(seq, idx + 1) for idx, seq in enumerate(experience.sequences)], + dim=0, + ).unsqueeze(0) else: sequences = experience.sequences old_values = experience.values @@ -972,162 +951,127 @@ def ensure_device_and_contiguous(tensor, name="tensor"): packed_seq_lens = None attention_mask = experience.attention_mask - # Ensure sequences and attention_mask are also on device and contiguous - sequences = ensure_device_and_contiguous(sequences, "sequences") - attention_mask = ensure_device_and_contiguous(attention_mask, "attention_mask") + sequences = self._ensure_device_and_contiguous(sequences, device) + attention_mask = self._ensure_device_and_contiguous(attention_mask, device) + old_values = self._ensure_device_and_contiguous(old_values, device) + returns = self._ensure_device_and_contiguous(returns, device) + critic_kwargs = self._build_model_kwargs(experience, device=device) - # Critic loss values, output = self.critic( sequences, num_actions=num_actions, attention_mask=attention_mask, - pixel_values=pixel_values, - image_grid_thw=image_grid_thws, - pixel_values_videos=pixel_values_videos, - video_grid_thw=video_grid_thws, return_output=True, packed_seq_lens=packed_seq_lens, + **critic_kwargs, ) - # Loss function + critic_loss = self.critic_loss_fn( values, old_values, returns, action_mask=experience.action_mask, ) - # Mixtral auxiliary loss - if self.aux_loss: - aux_loss = output.aux_loss - else: - aux_loss = 0 + aux_loss = output.aux_loss if self.aux_loss else 0 loss = critic_loss + aux_loss * self.args.aux_loss_coef self.strategy.backward(loss, self.critic, self.critic_optim) self.strategy.optimizer_step(self.critic_optim, self.critic, self.critic_scheduler, name="critic") - # Status - status = { + return { "critic_loss": critic_loss.item(), "values": masked_mean(values, experience.action_mask).item(), "critic_lr": self.critic_scheduler.get_last_lr()[0], } - return status - def save_logs_and_checkpoints(self, args, global_step, step_bar, logs_dict={}, client_states={}, episode=0): + def save_logs_and_checkpoints( + self, + args, + global_step, + step_bar, + logs_dict={}, + client_states={}, + episode=0, + ): """ - Save logs to wandb/tensorboard and save model checkpoints. + Log rollout/train/eval metrics and save checkpoints on the configured schedule. - :param args: Training arguments. + :param args: Runtime training arguments controlling log/eval/save cadence. :type args: Namespace - :param global_step: Current global step. + :param global_step: Current training step. :type global_step: int - :param step_bar: Progress bar object. + :param step_bar: Progress-bar instance for the outer rollout loop. :type step_bar: tqdm - :param logs_dict: Dictionary of metrics to log. Should contain both: - - Rollout statistics (rollout_reward, rollout_response_length, etc.) - from inference/generation phase - - Training statistics (policy_loss, critic_loss, kl, etc.) - from optimization phase - Defaults to {}. + :param logs_dict: Combined metrics from rollout collection and PPO optimization. :type logs_dict: dict - :param client_states: Client state for checkpoint recovery, defaults to {}. + :param client_states: Extra state saved into checkpoints for resume support. :type client_states: dict - :param episode: Current episode number, defaults to 0. + :param episode: Current episode index. :type episode: int """ - - # 1. LOGGING TRAIN & ROLLOUT METRICS if global_step % args.logging_steps == 0: - # Metrics that are already logged in rollout/ namespace should not be duplicated in train/ - ROLLOUT_ONLY_METRICS = {'reward', 'response_length', 'total_length', 'num_actions', 'return'} - ROLLOUT_ONLY_PREFIXES = {'reward_metrics/'} - + # Rollout metrics are logged under their own namespace and should not be duplicated under train/*. + rollout_only_metrics = {"reward", "response_length", "total_length", "num_actions", "return"} + rollout_only_prefixes = {"reward_metrics/"} rollout_metrics = {} train_metrics = {} - for k, v in logs_dict.items(): - if k.startswith('rollout_'): - # Clean key: rollout_reward -> reward - clean_key = k.replace('rollout_', '', 1) - rollout_metrics[clean_key] = v - elif k in ROLLOUT_ONLY_METRICS: + for key, value in logs_dict.items(): + if key.startswith("rollout_"): + rollout_metrics[key.replace("rollout_", "", 1)] = value + elif key in rollout_only_metrics: continue - elif any(k.startswith(prefix) for prefix in ROLLOUT_ONLY_PREFIXES): + elif any(key.startswith(prefix) for prefix in rollout_only_prefixes): continue else: - # Everything else is considered a training metric - train_metrics[k] = v + train_metrics[key] = value - # Wandb Logging if self._wandb is not None and self.strategy.is_rank_0(): all_wandb_logs = {} - - # Add Rollout Metrics - for k, v in rollout_metrics.items(): - all_wandb_logs[f"rollout/{k}"] = v + for key, value in rollout_metrics.items(): + all_wandb_logs[f"rollout/{key}"] = value all_wandb_logs["rollout/global_step"] = global_step all_wandb_logs["rollout/episode"] = episode - # Add Train Metrics - for k, v in train_metrics.items(): - all_wandb_logs[f"train/{k}"] = v + for key, value in train_metrics.items(): + all_wandb_logs[f"train/{key}"] = value all_wandb_logs["train/global_step"] = global_step all_wandb_logs["train/episode"] = episode - # Performance Stats - if self.experience_maker.perf_stats is not None: - for k, v in self.experience_maker.perf_stats.items(): - all_wandb_logs[f"perf/experience_maker/{k}"] = v + # FastExperienceMaker can publish collection-side performance stats opportunistically. + perf_stats = getattr(self.experience_maker, "perf_stats", None) + if perf_stats is not None: + for key, value in perf_stats.items(): + all_wandb_logs[f"perf/experience_maker/{key}"] = value - # Commit Train/Rollout logs with unique system step if all_wandb_logs: self.wandb_log_counter += 1 self._wandb.log(all_wandb_logs, step=self.wandb_log_counter, commit=True) - - # TensorBoard Logging elif self._tensorboard is not None and self.strategy.is_rank_0(): - for k, v in rollout_metrics.items(): - self._tensorboard.add_scalar(f"rollout/{k}", v, global_step) - for k, v in train_metrics.items(): - self._tensorboard.add_scalar(f"train/{k}", v, global_step) + for key, value in rollout_metrics.items(): + self._tensorboard.add_scalar(f"rollout/{key}", value, global_step) + for key, value in train_metrics.items(): + self._tensorboard.add_scalar(f"train/{key}", value, global_step) - # 2. EVALUATION if global_step % args.eval_steps == 0 and self.eval_dataloader is not None: - # Run evaluation + # Eval runs through the same experience maker, but only collects metrics instead of updating PPO state. raw_eval_metrics = self.evaluate(self.eval_dataloader, global_step) - - # Only log if we have results if raw_eval_metrics and self.strategy.is_rank_0(): self.eval_step_counter += 1 - - # Wandb Logging for Eval if self._wandb is not None: eval_logs = {} - for k, v in raw_eval_metrics.items(): - # Remove "eval_" prefix if present to avoid "eval/eval_reward" - clean_key = k.replace("eval_", "") if k.startswith("eval_") else k - eval_logs[f"eval/{clean_key}"] = v - - # Custom X-axis for Eval + for key, value in raw_eval_metrics.items(): + clean_key = key.replace("eval_", "") if key.startswith("eval_") else key + eval_logs[f"eval/{clean_key}"] = value eval_logs["eval/global_step"] = self.eval_step_counter - # Reference to main training step eval_logs["eval/train_step"] = global_step eval_logs["eval/episode"] = episode - - # IMPORTANT: - # Use wandb_log_counter to ensure eval has a unique system step - # This prevents eval metrics from being overwritten by train metrics - # The plots will still use eval/global_step as X-axis due to define_metric self.wandb_log_counter += 1 self._wandb.log(eval_logs, step=self.wandb_log_counter, commit=True) - - # TensorBoard Logging for Eval elif self._tensorboard is not None: - for k, v in raw_eval_metrics.items(): - # Clean key - clean_key = k.replace("eval_", "") if k.startswith("eval_") else k - self._tensorboard.add_scalar(f"eval/{clean_key}", v, global_step) + for key, value in raw_eval_metrics.items(): + clean_key = key.replace("eval_", "") if key.startswith("eval_") else key + self._tensorboard.add_scalar(f"eval/{clean_key}", value, global_step) - # 3. CHECKPOINTING if global_step % args.save_steps == 0: tag = f"global_step{global_step}" self._save_checkpoint(args, tag, client_states) @@ -1136,11 +1080,13 @@ def _save_checkpoint(self, args, tag, client_states): """ Save model checkpoint to disk. - :param args: Training arguments. + This keeps the old DS checkpoint path and the optional rotated HF/LoRA export path. + + :param args: Runtime training arguments containing checkpoint settings. :type args: Namespace - :param tag: Checkpoint tag (e.g., "global_step1000"). + :param tag: Checkpoint tag such as ``global_step1000``. :type tag: str - :param client_states: Client state for checkpoint recovery. + :param client_states: Extra client state persisted for checkpoint resume. :type client_states: dict """ ckpt_path = args.ckpt_path @@ -1155,12 +1101,14 @@ def _save_checkpoint(self, args, tag, client_states): ) if self.critic is not None: self.strategy.save_ckpt( - self.critic, os.path.join(ckpt_path, "_critic"), tag, args.max_ckpt_num, args.max_ckpt_mem + self.critic, + os.path.join(ckpt_path, "_critic"), + tag, + args.max_ckpt_num, + args.max_ckpt_mem, ) - # For LoRA, we ALWAYS save the HF adapter as it is much smaller and more convenient for deployment. if self.save_hf_ckpt or self.is_lora: - # Rotate HF checkpoints if self.strategy.is_rank_0(): os.makedirs(ckpt_path, exist_ok=True) max_num = getattr(args, "max_ckpt_num", 3) @@ -1177,13 +1125,17 @@ def _save_checkpoint(self, args, tag, client_states): def evaluate(self, eval_dataloader, global_step): """ - Evaluate the model on evaluation dataset. + Evaluate the model on evaluation data. + + Evaluation reuses the same experience-maker path as rollout collection, but only aggregates + reward and response-length statistics instead of updating PPO state. - :param eval_dataloader: DataLoader for evaluation data. + :param eval_dataloader: Evaluation dataloader normalized through the same batch-unpack helper + used during training. :type eval_dataloader: DataLoader - :param global_step: Current global step for logging. + :param global_step: Training step associated with this evaluation run. :type global_step: int - :return: Dictionary of evaluation metrics. + :return: Aggregated evaluation metrics. :rtype: dict """ if eval_dataloader is None: @@ -1203,58 +1155,39 @@ def evaluate(self, eval_dataloader, global_step): all_response_lengths = [] num_eval_batches = 0 - # Helper to extract values - def extract_values(val): - if isinstance(val, torch.Tensor): - return val.view(-1).cpu().tolist() - elif isinstance(val, (list, tuple)): - return list(val) - else: - return [float(val)] + def extract_values(value): + # Reward helpers may emit tensors, lists, or scalars depending on the recipe. + if isinstance(value, torch.Tensor): + return value.view(-1).cpu().tolist() + if isinstance(value, (list, tuple)): + return list(value) + return [float(value)] with torch.no_grad(): for batch in eval_dataloader: - if len(batch) == 5: - eval_prompts, eval_images, eval_videos, eval_references, eval_labels = batch - else: - eval_prompts, eval_images, eval_references, eval_labels = batch - eval_videos = None - - # Generate responses using experience maker (but don't train on them) - # We reuse the experience maker but only for generation - # TODO: simplify this logic - for i, experience in enumerate( - self.experience_maker.make_experience_list( - eval_prompts, eval_images, eval_videos, eval_references, eval_labels, **self.generate_kwargs - ) - ): - if i == 0: - output = self.tokenizer.batch_decode( - experience.sequences[0].unsqueeze(0), skip_special_tokens=True - ) - self.strategy.print("eval phase: experience.sequences w skip_special_tokens: ", output) - self.strategy.print( - f"eval phase: eval_prompts:\n {eval_prompts[0:2]}\n , rand_images:{eval_images[0:2]}\n , eval_references:{eval_references[0:2]}\n, eval_labels:{eval_labels[0:2]}\n " # noqa - ) - if hasattr(experience, 'info') and experience.info: + prompts, images, videos, audios, references, labels = self._unpack_prompt_batch(batch) + experience_list = self._make_experience_list(prompts, images, videos, audios, references, labels) + if not experience_list: + continue + + for experience in experience_list: + if hasattr(experience, "info") and experience.info: info = experience.info - if 'reward' in info: - all_rewards.extend(extract_values(info['reward'])) - if 'response_length' in info: - all_response_lengths.extend(extract_values(info['response_length'])) - - if 'reward_metrics' in info: - rm = info['reward_metrics'] - if 'format_reward' in rm: - all_format_rewards.extend(extract_values(rm['format_reward'])) - if 'accuracy_reward' in rm: - all_accuracy_rewards.extend(extract_values(rm['accuracy_reward'])) + if "reward" in info: + all_rewards.extend(extract_values(info["reward"])) + if "response_length" in info: + all_response_lengths.extend(extract_values(info["response_length"])) + if "reward_metrics" in info: + reward_metrics = info["reward_metrics"] + if "format_reward" in reward_metrics: + all_format_rewards.extend(extract_values(reward_metrics["format_reward"])) + if "accuracy_reward" in reward_metrics: + all_accuracy_rewards.extend(extract_values(reward_metrics["accuracy_reward"])) num_eval_batches += 1 if num_eval_batches >= len(eval_dataloader): break - # Compute statistics metrics = {} device = torch.cuda.current_device() @@ -1262,23 +1195,20 @@ def compute_stats(name, values_list): if not values_list: return if isinstance(values_list[0], torch.Tensor): - t = torch.cat([x.to(device).float() for x in values_list]) + tensor = torch.cat([value.to(device).float() for value in values_list]) else: - t = torch.tensor(values_list, dtype=torch.float32, device=device) - metrics[f"{name}_mean"] = t.mean().item() - # metrics[f"{name}_std"] = t.std().item() # Optional + tensor = torch.tensor(values_list, dtype=torch.float32, device=device) + metrics[f"{name}_mean"] = tensor.mean().item() compute_stats("reward", all_rewards) compute_stats("format_reward", all_format_rewards) compute_stats("accuracy_reward", all_accuracy_rewards) compute_stats("response_length", all_response_lengths) - metrics["num_samples"] = len(all_rewards) - # Print results self.strategy.print(f"Evaluation Results (Step {global_step}):") - for k, v in metrics.items(): - self.strategy.print(f" {k}: {v:.4f}") + for key, value in metrics.items(): + self.strategy.print(f" {key}: {value:.4f}") self.strategy.print(f"{'=' * 60}\n") self.actor.train() diff --git a/lightrft/trainer/replay_buffer_utils.py b/lightrft/trainer/replay_buffer_utils.py index 8380d13a..5870b20b 100644 --- a/lightrft/trainer/replay_buffer_utils.py +++ b/lightrft/trainer/replay_buffer_utils.py @@ -77,9 +77,11 @@ class BufferItemVL: sequences: torch.Tensor pixel_values: Optional[torch.Tensor] = None # image pixel processed by HF processor + audio_values: Optional[torch.Tensor] = None # audio features processed by HF processor image_grid_thws: Optional[torch.Tensor] = None # image grid thw pixel_values_videos: Optional[torch.Tensor] = None # video pixel processed by HF processor video_grid_thws: Optional[torch.Tensor] = None # video grid thw + feature_attention_mask: Optional[torch.Tensor] = None # audio feature mask raw_images: Optional[List[Image.Image]] = None # raw images before processing action_log_probs: torch.Tensor = None @@ -282,6 +284,8 @@ def _split_experience_batch_vl(experience: ExperienceVL) -> List: "attention_mask", "action_mask", "action_entropy", + "audio_values", + "feature_attention_mask", ) for key in keys: # Use getattr with default None to handle optional attributes like action_entropy @@ -355,6 +359,13 @@ def _split_experience_batch_vl(experience: ExperienceVL) -> List: batch_kwargs[i]["pixel_values"] = pixel_values[index:index + num_image_tokens] index += num_image_tokens + if getattr(experience, "audio_values", None) is not None: + audio_values = experience.audio_values + vals = torch.unbind(audio_values) if isinstance(audio_values, torch.Tensor) else audio_values + assert batch_size == len(vals), f"audio_values size mismatch: {len(vals)} vs {batch_size}" + for i, v in enumerate(vals): + batch_kwargs[i]["audio_values"] = v + # Split video data if experience.pixel_values_videos is not None: pixel_values_videos = experience.pixel_values_videos @@ -619,6 +630,7 @@ def _make_experience_batch_vl(items: List, packing_samples: bool = False) -> Exp "advantages", "attention_mask", "action_mask", + "feature_attention_mask", ) for key in keys: vals = [getattr(item, key) for item in items] @@ -655,6 +667,9 @@ def _make_experience_batch_vl(items: List, packing_samples: bool = False) -> Exp ] kwargs["pixel_values"] = torch.cat(pixel_values_list, dim=0) if pixel_values_list else None + audio_values_list = [item.audio_values for item in items if getattr(item, "audio_values", None) is not None] + kwargs["audio_values"] = torch.stack(audio_values_list, dim=0) if audio_values_list else None + image_grid_thws_list = [ item.image_grid_thws.unsqueeze(0) if (item.image_grid_thws is not None and item.image_grid_thws.dim() == 1) else item.image_grid_thws diff --git a/lightrft/trainer/utils.py b/lightrft/trainer/utils.py index abbf2931..8ad197f6 100644 --- a/lightrft/trainer/utils.py +++ b/lightrft/trainer/utils.py @@ -42,6 +42,8 @@ def fire_sampling( all_images_num: Optional[List[int]] = None, all_videos: Optional[List] = None, all_videos_num: Optional[List[int]] = None, + all_audios: Optional[List] = None, + all_audio_num: Optional[List[int]] = None, sampling_params: Optional[Union[dict, object]] = None, tokenizer: Optional[Any] = None, ) -> List: @@ -116,8 +118,10 @@ def fire_sampling( all_prompts=all_prompts if is_multimodal else None, all_images=all_images, all_videos=all_videos, + all_audios=all_audios, images_num=all_images_num if is_multimodal else None, videos_num=all_videos_num if is_multimodal else None, + audios_num=all_audio_num if is_multimodal else None, ) # Log first-token top-k frequency distribution @@ -182,8 +186,10 @@ def fire_sampling( all_prompts=all_prompts_rest if is_multimodal else None, all_images=all_images, all_videos=all_videos, + all_audios=all_audios, images_num=all_images_num if is_multimodal else None, videos_num=all_videos_num if is_multimodal else None, + audios_num=all_audio_num if is_multimodal else None, ) # Merge the first token with the remaining tokens From 9e6d2d7f9291988c0efbe4a48db4db3ca3bf3ca5 Mon Sep 17 00:00:00 2001 From: niuyazhe Date: Wed, 15 Apr 2026 18:12:45 +0800 Subject: [PATCH 04/11] fix(nyz): fix sglang output reward bug --- examples/r1_aqa/audio_dataset.py | 2 + examples/r1_aqa/reward_models_utils.py | 23 ++ lightrft/strategy/strategy_base.py | 62 ++++- lightrft/strategy/test_fake_strategy.py | 44 ++++ lightrft/trainer/audio_utils.py | 299 ++++++++++++++++++++++++ lightrft/trainer/modality_utils.py | 20 ++ 6 files changed, 445 insertions(+), 5 deletions(-) create mode 100644 lightrft/trainer/audio_utils.py create mode 100644 lightrft/trainer/modality_utils.py diff --git a/examples/r1_aqa/audio_dataset.py b/examples/r1_aqa/audio_dataset.py index 4e008300..777cd316 100644 --- a/examples/r1_aqa/audio_dataset.py +++ b/examples/r1_aqa/audio_dataset.py @@ -120,6 +120,8 @@ def __getitem__(self, idx: int) -> Tuple[str, Any, str, str]: except Exception as exc: self.strategy.print(f"[WARNING] Chat template failed for idx {idx}: {exc}") prompt_text = self._extract_text_from_messages(prompt_messages) + + prompt_text = prompt_text.replace('', ' ') # ---- 3. Load audio ---- audio_path = data.get(self.audio_path_key, "") diff --git a/examples/r1_aqa/reward_models_utils.py b/examples/r1_aqa/reward_models_utils.py index 7a32a683..96a59815 100644 --- a/examples/r1_aqa/reward_models_utils.py +++ b/examples/r1_aqa/reward_models_utils.py @@ -73,6 +73,15 @@ def accuracy_reward_fn(content: str, solution: str) -> float: reward = 1.0 except Exception: pass + if reward == 0.0: + sol_match = re.search(r"(.*?)", solution) + ground_truth = sol_match.group(1).strip() if sol_match else solution.strip() + student_answer = content.strip() + import torch.distributed as dist + if dist.is_initialized() and dist.get_rank() == 0: + print(f"student_answer: {student_answer}, ground_truth: {ground_truth}") + if student_answer == ground_truth: + reward = 1.0 return reward @@ -134,6 +143,19 @@ def avqa_combined_reward_fn( return total_r, acc_r, fmt_r +def clean_solution(sol: str) -> str: + # <|im_start|>assistantat sea<|im_end|> + """ + Extract the string between <|im_start|>assistant and <|im_end|> tags. + + Example: + input: "<|im_start|>assistantat sea<|im_end|>" + output: "at sea" + """ + import re + # Pattern matches text between <|im_start|>assistant and <|im_end|> + match = re.search(r"<\|im_start\|>assistant(.*?)<\|im_end\|>", sol, re.DOTALL) + return match.group(1).strip() if match else sol.strip() # ============================================================================ # Reward Function (LightRFT interface — called by the trainer) # ============================================================================ @@ -170,6 +192,7 @@ def reward_fn( for i in range(B): sol = queries[i] + sol = clean_solution(sol) gt = refs[i] if i < len(refs) else "" total_r, acc_r, fmt_r = avqa_combined_reward_fn(sol, gt) final_reward[i] = total_r diff --git a/lightrft/strategy/strategy_base.py b/lightrft/strategy/strategy_base.py index 8b6c2c44..90364f76 100644 --- a/lightrft/strategy/strategy_base.py +++ b/lightrft/strategy/strategy_base.py @@ -11,6 +11,7 @@ import random import time import io +import numbers from loguru import logger from abc import ABC, abstractmethod from collections import defaultdict @@ -76,6 +77,48 @@ def _serialize_audio_for_sglang(audio_item: Any, default_sr: int = 16000): return buffer.getvalue() +def _prepare_audio_for_vllm(audio_item: Any, default_sr: int = 16000): + """ + Convert local audio payloads into the form accepted by vLLM. + + vLLM expects waveform-like objects such as ``(audio, sampling_rate)`` tuples, + arrays, lists of floats, or tensors. Unlike SGLang, passing serialized WAV bytes + through ``multi_modal_data["audio"]`` causes the parser to fail. + """ + if audio_item is None: + return None + if isinstance(audio_item, tuple) and len(audio_item) == 2: + audio_array, sr = audio_item + return np.asarray(audio_array, dtype=np.float32), int(sr) + if isinstance(audio_item, list): + if len(audio_item) == 2 and isinstance(audio_item[1], numbers.Number): + audio_array, sr = audio_item + return np.asarray(audio_array, dtype=np.float32), int(sr) + return np.asarray(audio_item, dtype=np.float32) + if isinstance(audio_item, np.ndarray): + return np.asarray(audio_item, dtype=np.float32) + if isinstance(audio_item, torch.Tensor): + return audio_item.detach().cpu() + if isinstance(audio_item, bytes): + audio_array, sr = sf.read(io.BytesIO(audio_item), dtype="float32") + return np.asarray(audio_array, dtype=np.float32), int(sr) + if isinstance(audio_item, str): + if os.path.exists(audio_item): + audio_array, sr = sf.read(audio_item, dtype="float32") + return np.asarray(audio_array, dtype=np.float32), int(sr) + raise TypeError( + "Unsupported vLLM audio payload: string paths or URLs must be loaded into waveform data before rollout." + ) + if isinstance(audio_item, dict): + if "array" in audio_item: + sr = audio_item.get("sampling_rate", audio_item.get("sample_rate", default_sr)) + return np.asarray(audio_item["array"], dtype=np.float32), int(sr) + raise TypeError( + "Unsupported vLLM audio payload dict: expected an 'array' field and optional sampling rate metadata." + ) + raise TypeError(f"Unsupported vLLM audio payload type: {type(audio_item).__name__}") + + class EngineStatus(Enum): """ Enum class for inference engine status. @@ -944,6 +987,7 @@ def _build_multimodal_inputs( videos_num, all_audios=None, audios_num=None, + engine_type: str = "sglang", ): """ Build multimodal inputs for inference engine (vLLM/SGLang). @@ -1002,11 +1046,18 @@ def _build_multimodal_inputs( raw_audio_list = all_audios[i] else: raw_audio_list = all_audios[audio_start_idx:audio_start_idx + audio_num] - # Serialize in one place so the rest of the rollout stack can keep audio payloads - # in their native Python forms. - audio_list = [ - _serialize_audio_for_sglang(audio_item) for audio_item in raw_audio_list if audio_item is not None - ] + if engine_type == "vllm": + audio_list = [ + _prepare_audio_for_vllm(audio_item) for audio_item in raw_audio_list if audio_item is not None + ] + else: + # Serialize in one place so the rest of the rollout stack can keep audio + # payloads in their native Python forms. + audio_list = [ + _serialize_audio_for_sglang(audio_item) + for audio_item in raw_audio_list + if audio_item is not None + ] else: audio_list = [] @@ -1111,6 +1162,7 @@ def gather_and_generate( videos_num=videos_num, all_audios=all_audios, audios_num=audios_num, + engine_type=self.inference_engine_type, ) else: inputs = all_prompt_token_ids diff --git a/lightrft/strategy/test_fake_strategy.py b/lightrft/strategy/test_fake_strategy.py index 25568089..d6736f01 100644 --- a/lightrft/strategy/test_fake_strategy.py +++ b/lightrft/strategy/test_fake_strategy.py @@ -11,6 +11,7 @@ import unittest from unittest.mock import MagicMock +import numpy as np import torch import torch.nn as nn from torch.utils.data import DataLoader, TensorDataset @@ -322,5 +323,48 @@ def __init__(self): self.assertEqual(strategy.config.zpg, 2) +class TestBuildMultimodalInputs(unittest.TestCase): + """Focused tests for engine-specific multimodal payload construction.""" + + def test_keeps_waveform_audio_for_vllm(self): + """vLLM audio inputs should stay waveform-native instead of being WAV bytes.""" + audio = (np.asarray([0.1, -0.2, 0.3], dtype=np.float32), 16000) + + inputs = FakeStrategy._build_multimodal_inputs( + all_prompts=["prompt"], + all_images=None, + images_num=None, + all_videos=None, + videos_num=None, + all_audios=[audio], + audios_num=[1], + engine_type="vllm", + ) + + payload = inputs[0]["multi_modal_data"]["audio"][0] + self.assertIsInstance(payload, tuple) + self.assertEqual(payload[1], 16000) + self.assertTrue(np.array_equal(payload[0], audio[0])) + + def test_serializes_audio_for_sglang(self): + """SGLang audio inputs should still be serialized to WAV bytes.""" + audio = (np.asarray([0.1, -0.2, 0.3], dtype=np.float32), 16000) + + inputs = FakeStrategy._build_multimodal_inputs( + all_prompts=["prompt"], + all_images=None, + images_num=None, + all_videos=None, + videos_num=None, + all_audios=[audio], + audios_num=[1], + engine_type="sglang", + ) + + payload = inputs[0]["multi_modal_data"]["audio"][0] + self.assertIsInstance(payload, bytes) + self.assertTrue(payload.startswith(b"RIFF")) + + if __name__ == "__main__": unittest.main() diff --git a/lightrft/trainer/audio_utils.py b/lightrft/trainer/audio_utils.py new file mode 100644 index 00000000..1a202e26 --- /dev/null +++ b/lightrft/trainer/audio_utils.py @@ -0,0 +1,299 @@ +""" +Utilities for audio-language rollout processing. + +This module keeps audio-specific preprocessing out of the vision-language path. +""" + +from __future__ import annotations + +import inspect +import io +import numbers +from typing import Any, List, Optional, Tuple, Union + +import numpy as np +import soundfile as sf +import torch +from easydict import EasyDict + + +def normalize_audios(raw_audios: List[Any]) -> List[Any]: + """Audio payloads are already normalized by the dataset layer.""" + return raw_audios + + +def is_single_audio_payload(audio_item: Any) -> bool: + """ + Return True when ``audio_item`` represents exactly one audio payload. + + Accepted forms include: + - ``(waveform, sampling_rate)`` + - ``[waveform, sampling_rate]`` + - raw waveform arrays + - engine-ready ``str`` / ``bytes`` / ``dict`` payloads + - one-element wrappers such as ``[(waveform, sr)]`` + """ + if audio_item is None: + return False + if isinstance(audio_item, (str, bytes, dict, np.ndarray)): + return True + if isinstance(audio_item, tuple) and len(audio_item) == 2: + return isinstance(audio_item[1], numbers.Number) + if isinstance(audio_item, list): + if len(audio_item) == 1: + return is_single_audio_payload(audio_item[0]) + if len(audio_item) == 2 and isinstance(audio_item[1], numbers.Number): + return True + return False + + +def canonicalize_audio_payload(audio_item: Any) -> Any: + """ + Normalize single-audio wrappers to a stable payload shape. + + This keeps the audio rollout path permissive about harmless container differences + while still rejecting true multi-audio inputs at a higher level. + """ + if isinstance(audio_item, list) and len(audio_item) == 1 and is_single_audio_payload(audio_item[0]): + return canonicalize_audio_payload(audio_item[0]) + return audio_item + + +def get_audios_num(all_audios: Optional[List[Any]]) -> Optional[List[int]]: + """ + Count audio items per sample. + + Audio RL currently supports zero or one audio payload per prompt in rollout. + """ + if all_audios is None: + return None + counts = [] + for audio in all_audios: + if audio is None: + counts.append(0) + elif is_single_audio_payload(audio): + counts.append(1) + elif isinstance(audio, list): + counts.append(len(audio)) + else: + counts.append(1) + return counts + + +def extract_audio_array(audio_item: Any, default_sr: int = 16000) -> Tuple[np.ndarray, int]: + """Normalize supported audio payloads to ``(waveform, sampling_rate)``.""" + audio_item = canonicalize_audio_payload(audio_item) + if isinstance(audio_item, tuple) and len(audio_item) == 2: + audio_array, sr = audio_item + return np.asarray(audio_array, dtype=np.float32), int(sr) + if isinstance(audio_item, list) and len(audio_item) == 2 and isinstance(audio_item[1], numbers.Number): + audio_array, sr = audio_item + return np.asarray(audio_array, dtype=np.float32), int(sr) + if isinstance(audio_item, np.ndarray): + return np.asarray(audio_item, dtype=np.float32), default_sr + raise TypeError(f"Unsupported audio payload type: {type(audio_item).__name__}") + + +def serialize_audio_for_sglang(audio_item: Any, default_sr: int = 16000) -> Union[str, bytes, dict, None]: + """ + Convert a local audio payload into a SGLang-compatible object. + + SGLang accepts file paths / URLs / bytes, but not ``(waveform, sr)`` tuples directly. + """ + if audio_item is None: + return None + if isinstance(audio_item, (str, bytes, dict)): + return audio_item + + audio_array, sr = extract_audio_array(audio_item, default_sr=default_sr) + buffer = io.BytesIO() + sf.write(buffer, audio_array, sr, format="WAV") + return buffer.getvalue() + + +def normalize_audio_features( + input_features: torch.Tensor, + feature_attention_mask: Optional[torch.Tensor], + expected_mel_len: int = 3000, +) -> Tuple[torch.Tensor, torch.Tensor]: + """ + Normalize Qwen2-Audio features to ``(B, mel_bins, expected_mel_len)``. + """ + if not isinstance(input_features, torch.Tensor): + input_features = torch.as_tensor(input_features) + if feature_attention_mask is not None and not isinstance(feature_attention_mask, torch.Tensor): + feature_attention_mask = torch.as_tensor(feature_attention_mask) + + if input_features.dim() != 3: + raise RuntimeError( + f"Expected 3D audio features, but got shape {tuple(input_features.shape)}. " + "Qwen2-Audio should return (batch, mel_bins, time)." + ) + + if input_features.shape[1] in (80, 128): + normalized = input_features + elif input_features.shape[-1] in (80, 128): + normalized = input_features.transpose(1, 2).contiguous() + else: + raise RuntimeError( + f"Unexpected Qwen2-Audio feature shape {tuple(input_features.shape)}: " + "unable to identify mel-bin dimension." + ) + + current_len = normalized.shape[-1] + target_len = min(current_len, expected_mel_len) + if current_len < expected_mel_len: + normalized = torch.nn.functional.pad(normalized, (0, expected_mel_len - current_len), value=0.0) + elif current_len > expected_mel_len: + normalized = normalized[..., :expected_mel_len] + + if feature_attention_mask is None: + feature_attention_mask = torch.zeros( + normalized.shape[0], expected_mel_len, dtype=torch.long, device=normalized.device + ) + feature_attention_mask[:, :target_len] = 1 + else: + feature_attention_mask = feature_attention_mask.to(device=normalized.device, dtype=torch.long) + if feature_attention_mask.shape[-1] < expected_mel_len: + feature_attention_mask = torch.nn.functional.pad( + feature_attention_mask, (0, expected_mel_len - feature_attention_mask.shape[-1]), value=0 + ) + elif feature_attention_mask.shape[-1] > expected_mel_len: + feature_attention_mask = feature_attention_mask[..., :expected_mel_len] + + return normalized, feature_attention_mask + + +class AudioDataProcessor: + """ + Audio-language rollout preprocessor. + + Unlike the VL processor, audio inputs stay on an explicit audio path: + raw audio payloads are kept for the inference engine, while ``audio_values`` + and ``feature_attention_mask`` are prepared for actor/reference forward. + """ + def __init__(self, tokenizer, processor, prompt_max_len: int): + self.tokenizer = tokenizer + self.processor = processor + self.prompt_max_len = prompt_max_len + + def process_audio_batch( + self, + all_prompts: List[str], + all_audios: List[Any], + all_references: Optional[List[str]], + n_samples_per_prompt: int, + ) -> EasyDict: + N = n_samples_per_prompt + L = len(all_prompts) + if all_audios is None: + all_audios = [None] * L + + all_prompts_text, all_prompts_audio = [], [] + all_audios_valid = [] + text_idx = [] + + for idx, (prompt, audio) in enumerate(zip(all_prompts, all_audios)): + if audio is None: + all_prompts_text.append(prompt) + text_idx.append(idx) + else: + audio = canonicalize_audio_payload(audio) + if isinstance(audio, list) and not is_single_audio_payload(audio): + raise RuntimeError( + "Audio RL rollout currently expects at most one audio payload per prompt. " + f"Received list input for sample {idx}." + ) + all_prompts_audio.append(prompt) + all_audios_valid.append(audio) + + all_prompts_text = sum([[prompt] * N for prompt in all_prompts_text], []) + all_prompts_audio = sum([[prompt] * N for prompt in all_prompts_audio], []) + all_audios_valid = [audio for audio in all_audios_valid for _ in range(N)] + + if all_prompts_text: + inputs_text = self.tokenizer( + all_prompts_text, + max_length=self.prompt_max_len, + truncation=True, + add_special_tokens=False, + ) + all_prompt_token_ids_text = inputs_text["input_ids"] + else: + all_prompt_token_ids_text = [] + + all_prompt_token_ids_audio = [] + all_audio_values = None + all_feature_attention_mask = None + if all_prompts_audio: + proc_sig = inspect.signature(self.processor.__call__) + audio_kwarg = "audio" if "audio" in proc_sig.parameters else "audios" + flat_audios = [extract_audio_array(audio, default_sr=16000)[0] for audio in all_audios_valid] + inputs_audio = self.processor( + text=all_prompts_audio, + **{audio_kwarg: flat_audios}, + add_special_tokens=False, + max_length=self.prompt_max_len, + truncation=True, + padding=True, + return_tensors="pt", + sampling_rate=getattr(self.processor.feature_extractor, "sampling_rate", 16000), + ) + all_prompt_token_ids_audio = inputs_audio["input_ids"].tolist() + all_audio_values = inputs_audio.get("input_features", None) + if all_audio_values is None: + raise RuntimeError( + f"Processor {type(self.processor).__name__} returned no 'input_features'. " + f"Available keys: {list(inputs_audio.keys())}" + ) + all_feature_attention_mask = inputs_audio.get("feature_attention_mask", None) + all_audio_values, all_feature_attention_mask = normalize_audio_features( + all_audio_values, all_feature_attention_mask, expected_mel_len=3000 + ) + + total_samples = L * N + all_prompts_out = [None] * total_samples + all_audios_out = [None] * total_samples + all_prompt_token_ids_out = [None] * total_samples + all_audio_values_out = None + all_feature_attention_mask_out = None + if all_audio_values is not None: + all_audio_values_out = all_audio_values.new_zeros((total_samples, ) + tuple(all_audio_values.shape[1:])) + all_feature_attention_mask_out = all_feature_attention_mask.new_zeros( + (total_samples, ) + tuple(all_feature_attention_mask.shape[1:]) + ) + + text_ptr = 0 + for orig_idx in text_idx: + for n in range(N): + gid = orig_idx * N + n + all_prompts_out[gid] = all_prompts_text[text_ptr] + all_prompt_token_ids_out[gid] = all_prompt_token_ids_text[text_ptr] + text_ptr += 1 + + audio_ptr = 0 + for orig_idx in range(L): + if orig_idx in text_idx: + continue + for n in range(N): + gid = orig_idx * N + n + all_prompts_out[gid] = all_prompts_audio[audio_ptr] + all_audios_out[gid] = all_audios_valid[audio_ptr] + all_prompt_token_ids_out[gid] = all_prompt_token_ids_audio[audio_ptr] + if all_audio_values_out is not None: + all_audio_values_out[gid] = all_audio_values[audio_ptr] + all_feature_attention_mask_out[gid] = all_feature_attention_mask[audio_ptr] + audio_ptr += 1 + + if all_references is not None: + all_references = sum([[ref] * N for ref in all_references], []) + + return EasyDict( + all_prompt_token_ids=all_prompt_token_ids_out, + all_prompts=all_prompts_out, + all_audios=all_audios_out, + all_audio_num=get_audios_num(all_audios_out), + all_audio_values=all_audio_values_out, + all_feature_attention_mask=all_feature_attention_mask_out, + all_references=all_references, + ) diff --git a/lightrft/trainer/modality_utils.py b/lightrft/trainer/modality_utils.py new file mode 100644 index 00000000..5c7ee69f --- /dev/null +++ b/lightrft/trainer/modality_utils.py @@ -0,0 +1,20 @@ +""" +Shared helpers for modality-specific trainer/model wiring. +""" + +from typing import Any, Dict, Set + + +def build_supported_model_kwargs(source: Any, supported_params: Set[str]) -> Dict[str, Any]: + """ + Extract only the multimodal kwargs that the current actor supports. + """ + candidate_params = { + "pixel_values": getattr(source, "pixel_values", None), + "image_grid_thw": getattr(source, "image_grid_thw", getattr(source, "image_grid_thws", None)), + "pixel_values_videos": getattr(source, "pixel_values_videos", None), + "video_grid_thw": getattr(source, "video_grid_thw", getattr(source, "video_grid_thws", None)), + "audio_values": getattr(source, "audio_values", None), + "feature_attention_mask": getattr(source, "feature_attention_mask", None), + } + return {key: value for key, value in candidate_params.items() if key in supported_params} From 54d41b8daf3910fbaab8c535ec874306366c5731 Mon Sep 17 00:00:00 2001 From: niuyazhe Date: Thu, 16 Apr 2026 16:04:46 +0800 Subject: [PATCH 05/11] feature(nyz): add Qwen2.5 Omni audio support --- examples/r1_aqa/eval.py | 41 +++-- examples/r1_aqa/train_colocate.py | 24 +-- lightrft/models/actor_al.py | 246 +++++++++++++++++++++++-- lightrft/models/tests/test_actor_al.py | 141 +++++++++++++- lightrft/models/utils.py | 4 + lightrft/strategy/fsdp/fsdpv2.py | 33 +++- 6 files changed, 434 insertions(+), 55 deletions(-) diff --git a/examples/r1_aqa/eval.py b/examples/r1_aqa/eval.py index 55081335..0eae4270 100644 --- a/examples/r1_aqa/eval.py +++ b/examples/r1_aqa/eval.py @@ -35,6 +35,7 @@ """ import argparse +import inspect import json import os import re @@ -42,6 +43,13 @@ import torch +from lightrft.models.actor_al import ( + AUDIO_MODEL_TYPE_QWEN2_5_OMNI, + create_audio_processor, + get_audio_model_class, + infer_audio_model_type, +) + # --------------------------------------------------------------------------- # Message building (MMAU and MMAR share question/choices + format) @@ -137,16 +145,16 @@ def run_inference_hf( """ Run inference using HuggingFace Transformers (sequential). """ - from transformers import AutoProcessor + model_type = infer_audio_model_type(model_path) try: - from transformers import Qwen2AudioForConditionalGeneration - model = Qwen2AudioForConditionalGeneration.from_pretrained( + model_cls = get_audio_model_class(model_type) + model = model_cls.from_pretrained( model_path, torch_dtype=torch.bfloat16, trust_remote_code=True, ).to("cuda").eval() - except (ImportError, OSError): + except (ImportError, OSError, NotImplementedError): from transformers import AutoModelForCausalLM model = AutoModelForCausalLM.from_pretrained( model_path, @@ -154,7 +162,7 @@ def run_inference_hf( trust_remote_code=True, ).to("cuda").eval() - processor = AutoProcessor.from_pretrained(model_path, trust_remote_code=True) + processor = create_audio_processor(model_path, trust_remote_code=True) try: import librosa @@ -162,6 +170,8 @@ def run_inference_hf( raise ImportError("librosa is required for audio loading: pip install librosa") sr = getattr(processor.feature_extractor, "sampling_rate", 16000) + proc_sig = inspect.signature(processor.__call__) + audio_kwarg = "audio" if "audio" in proc_sig.parameters else "audios" def get_audio_path(sample: Dict) -> str: raw = sample.get("audio_path") or sample.get("audio_id", "") @@ -184,7 +194,7 @@ def get_audio_path(sample: Dict) -> str: audio, _ = librosa.load(audio_path, sr=sr) inputs = processor( text=text, - audios=[audio], + **{audio_kwarg: [audio]}, sampling_rate=sr, return_tensors="pt", padding=True, @@ -195,12 +205,21 @@ def get_audio_path(sample: Dict) -> str: inputs = {k: v.to("cuda") for k, v in inputs.items()} - with torch.no_grad(): - outputs = model.generate( - **inputs, - max_new_tokens=max_new_tokens, - do_sample=False, + generation_kwargs = { + "do_sample": False, + } + if model_type == AUDIO_MODEL_TYPE_QWEN2_5_OMNI: + generation_kwargs.update( + { + "generation_mode": "text", + "thinker_max_new_tokens": max_new_tokens, + } ) + else: + generation_kwargs["max_new_tokens"] = max_new_tokens + + with torch.no_grad(): + outputs = model.generate(**inputs, **generation_kwargs) input_len = inputs["input_ids"].shape[-1] generated = outputs[0][input_len:] diff --git a/examples/r1_aqa/train_colocate.py b/examples/r1_aqa/train_colocate.py index c7e053d5..188aa94b 100644 --- a/examples/r1_aqa/train_colocate.py +++ b/examples/r1_aqa/train_colocate.py @@ -25,7 +25,7 @@ import torch.nn.functional as F from lightrft.utils import add_arguments from lightrft.datasets import SFTDatasetVL -from lightrft.models.actor_al import ActorAL +from lightrft.models.actor_al import ActorAL, create_audio_processor from lightrft.models.actor_language import ActorLanguage from lightrft.strategy import get_strategy from lightrft.trainer.spmd_ppo_trainer import SPMDPPOTrainerVL @@ -153,21 +153,13 @@ def train(args): use_fast=not strategy.args.disable_fast_tokenizer, ) - # Ensure we have the correct Qwen2AudioProcessor (AutoProcessor may - # fall back to a generic text processor that ignores the `audios` kwarg). - try: - from transformers import Qwen2AudioProcessor - if not isinstance(processor, Qwen2AudioProcessor): - strategy.print( - f"[WARN] AutoProcessor loaded {type(processor).__name__}, " - "re-loading as Qwen2AudioProcessor" - ) - processor = Qwen2AudioProcessor.from_pretrained( - args.pretrain, trust_remote_code=True - ) - except ImportError: - strategy.print("[WARN] Qwen2AudioProcessor not available in this transformers version") - assert processor is not None, "Qwen2-Audio processor is required" + processor = create_audio_processor( + args.pretrain, + processor=processor, + print_fn=strategy.print, + ) + + assert processor is not None, "Audio-language processor is required" # ==================== Data Loading ==================== strategy.print(f"Loading prompts dataset from: {args.prompt_data}") diff --git a/lightrft/models/actor_al.py b/lightrft/models/actor_al.py index 10f7faae..f4660bd7 100644 --- a/lightrft/models/actor_al.py +++ b/lightrft/models/actor_al.py @@ -8,12 +8,12 @@ """ import os -from typing import Optional, Tuple, Union +from typing import Any, Callable, Optional, Tuple, Union import torch import torch.distributed as dist import torch.nn as nn -from transformers import Qwen2AudioForConditionalGeneration +from transformers import AutoConfig, Qwen2AudioForConditionalGeneration, Qwen2_5OmniForConditionalGeneration from transformers.integrations.deepspeed import HfDeepSpeedConfig from .actor_modality import ActorModality @@ -25,6 +25,174 @@ ) +AUDIO_MODEL_TYPE_QWEN2_AUDIO = "qwen2_audio" +AUDIO_MODEL_TYPE_QWEN2_5_OMNI = "qwen2_5_omni" + + +def normalize_audio_model_type(model_type: Optional[str]) -> Optional[str]: + """ + Collapse backbone-specific variants into a stable audio model family name. + """ + if model_type in { + AUDIO_MODEL_TYPE_QWEN2_AUDIO, + AUDIO_MODEL_TYPE_QWEN2_5_OMNI, + }: + return model_type + if model_type == "qwen2_5_omni_thinker": + return AUDIO_MODEL_TYPE_QWEN2_5_OMNI + return model_type + + +def infer_audio_model_type(pretrain_or_model: Any) -> Optional[str]: + """ + Infer the audio backbone family from a checkpoint path or a loaded model. + """ + if not isinstance(pretrain_or_model, str): + config = getattr(pretrain_or_model, "config", None) + return normalize_audio_model_type(getattr(config, "model_type", None)) + + try: + config = AutoConfig.from_pretrained(pretrain_or_model, trust_remote_code=True) + model_type = normalize_audio_model_type(getattr(config, "model_type", None)) + if model_type is not None: + return model_type + except Exception: + pass + + lowered = pretrain_or_model.lower() + if "qwen2.5-omni" in lowered or "qwen2_5_omni" in lowered: + return AUDIO_MODEL_TYPE_QWEN2_5_OMNI + if "qwen2-audio" in lowered or "qwen2_audio" in lowered: + return AUDIO_MODEL_TYPE_QWEN2_AUDIO + return None + + +def _resolve_audio_model_name_or_path(pretrain_or_model: Any) -> Optional[str]: + """ + Best-effort resolution of a checkpoint path for processor/model loading. + """ + if isinstance(pretrain_or_model, str): + return pretrain_or_model + + direct_name = getattr(pretrain_or_model, "name_or_path", None) + if direct_name: + return direct_name + + config = getattr(pretrain_or_model, "config", None) + for attr_name in ("_name_or_path", "name_or_path"): + name_or_path = getattr(config, attr_name, None) + if name_or_path: + return name_or_path + + return None + + +def get_audio_model_class(model_type: Optional[str]): + """ + Return the Hugging Face model class for a supported audio-language backbone. + """ + normalized = normalize_audio_model_type(model_type) + if normalized == AUDIO_MODEL_TYPE_QWEN2_AUDIO: + return Qwen2AudioForConditionalGeneration + if normalized == AUDIO_MODEL_TYPE_QWEN2_5_OMNI: + return Qwen2_5OmniForConditionalGeneration + raise NotImplementedError(f"Unsupported audio-language model type: {model_type}") + + +def get_audio_processor_class(model_type: Optional[str]): + """ + Return the Hugging Face processor class for a supported audio-language backbone. + """ + normalized = normalize_audio_model_type(model_type) + if normalized == AUDIO_MODEL_TYPE_QWEN2_AUDIO: + from transformers import Qwen2AudioProcessor + return Qwen2AudioProcessor + if normalized == AUDIO_MODEL_TYPE_QWEN2_5_OMNI: + from transformers import Qwen2_5OmniProcessor + return Qwen2_5OmniProcessor + return None + + +def create_audio_processor( + pretrain_or_model: Any, + processor=None, + trust_remote_code: bool = True, + print_fn: Optional[Callable[[str], None]] = None, + **from_pretrained_kwargs, +): + """ + Create or normalize the audio processor for the given backbone. + + If an existing processor is supplied and already matches the resolved audio + backbone, it is reused as-is. Otherwise the correct backbone-specific + processor is reloaded from the checkpoint path. + """ + model_type = infer_audio_model_type(pretrain_or_model) + try: + processor_cls = get_audio_processor_class(model_type) + except ImportError as exc: + if print_fn is not None: + print_fn(f"[WARN] Failed to import audio processor for {model_type}: {exc}") + processor_cls = None + + if processor_cls is not None and processor is not None and isinstance(processor, processor_cls): + return processor + + source = _resolve_audio_model_name_or_path(pretrain_or_model) + if source is None: + if processor is not None: + return processor + raise ValueError("Unable to resolve a checkpoint path for creating the audio processor.") + + if processor_cls is None: + from transformers import AutoProcessor + if print_fn is not None: + print_fn("[WARN] Falling back to AutoProcessor for audio model inputs.") + return AutoProcessor.from_pretrained( + source, + trust_remote_code=trust_remote_code, + **from_pretrained_kwargs, + ) + + if processor is not None and print_fn is not None: + print_fn( + f"[WARN] AutoProcessor loaded {type(processor).__name__}, " + f"re-loading as {processor_cls.__name__}" + ) + + return processor_cls.from_pretrained( + source, + trust_remote_code=trust_remote_code, + **from_pretrained_kwargs, + ) + + +def get_audio_forward_model(model: Any): + """ + Return the submodule used for token-level logprob forward passes. + + Qwen2.5-Omni generation is wrapped by the full model, while token scoring should + run through its ``thinker`` branch. + """ + model_type = infer_audio_model_type(model) + if model_type == AUDIO_MODEL_TYPE_QWEN2_5_OMNI: + thinker = getattr(model, "thinker", None) + if thinker is None: + raise AttributeError("Qwen2.5-Omni model does not expose a `thinker` module.") + return thinker + return model + + +def get_audio_model_and_type(pretrain_or_model: str, **from_pretrained_kwargs) -> Tuple[Any, str]: + """ + Load a supported audio-language backbone and return ``(model, model_type)``. + """ + model_type = infer_audio_model_type(pretrain_or_model) + model_cls = get_audio_model_class(model_type) + model = model_cls.from_pretrained(pretrain_or_model, **from_pretrained_kwargs) + return model, model_type + + class _AudioEmbedPositions(nn.Module): """FSDP2-safe replacement for ``nn.Embedding`` used in Whisper's audio tower. @@ -109,9 +277,11 @@ def __init__( **kwargs, ) -> None: super().__init__() + self.packing_samples = packing_samples if isinstance(pretrain_or_model, str): self.pretrain_or_model = pretrain_or_model + self.model_type = infer_audio_model_type(pretrain_or_model) attn_implementation = "flash_attention_2" if use_flash_attention_2 else "eager" # Note: dschf is defined in function scope to avoid global effects @@ -121,8 +291,7 @@ def __init__( else: dschf = None # noqa: F841 - # Load Qwen2Audio model - self.model = Qwen2AudioForConditionalGeneration.from_pretrained( + self.model, self.model_type = get_audio_model_and_type( pretrain_or_model, trust_remote_code=True, attn_implementation=attn_implementation, @@ -144,12 +313,14 @@ def __init__( # https://github.com/huggingface/transformers/issues/26877 # Use `model.generate(use_cache=True)` instead.` - self.model.config.use_cache = False - - # packing samples using Flash Attention 2 - self.packing_samples = packing_samples + if hasattr(self.model.config, "use_cache"): + self.model.config.use_cache = False + forward_model = get_audio_forward_model(self.model) + if hasattr(forward_model.config, "use_cache"): + forward_model.config.use_cache = False else: self.model = pretrain_or_model + self.model_type = infer_audio_model_type(pretrain_or_model) self.pretrain_or_model = pretrain_or_model.config.model_type # ------------------------------------------------------------------ @@ -169,7 +340,8 @@ def __init__( # The Whisper encoder is small (~12 layers), so using eager # attention has negligible impact on overall training throughput. # ------------------------------------------------------------------ - audio_tower = getattr(self.model, "audio_tower", None) or getattr(self.model, "audio_encoder", None) + forward_model = get_audio_forward_model(self.model) + audio_tower = getattr(forward_model, "audio_tower", None) or getattr(forward_model, "audio_encoder", None) if audio_tower is not None: # Fix 1: embed_positions if hasattr(audio_tower, "embed_positions") and isinstance(audio_tower.embed_positions, nn.Embedding): @@ -184,13 +356,43 @@ def __init__( if hasattr(module, "_attn_implementation"): module._attn_implementation = "eager" # Also patch the config so any lazily-constructed layers use eager - audio_cfg = getattr(self.model.config, "audio_config", None) + audio_cfg = getattr(forward_model.config, "audio_config", None) if audio_cfg is not None: audio_cfg._attn_implementation = "eager" print("[ActorAL] Set audio_tower attention to 'eager' for FSDP2 compat") print("pretrain_or_model: ", self.pretrain_or_model) + def get_fsdp_target_model(self) -> nn.Module: + """ + Return the concrete module FSDP should shard, optimize, and checkpoint for actor training. + + The actor wrapper intentionally keeps ``self.model`` as the full Hugging Face + object so inference-time APIs such as ``generate()`` continue to behave like + the original checkpoint. However, RL training does not always optimize that + whole object. + + Examples: + - ``Qwen2-Audio``: the trainable language/audio path is the root model + itself, so FSDP should shard ``self.model`` directly. + - ``Qwen2.5-Omni`` during PPO/GRPO actor training: token-level log-prob + computation runs through ``self.model.thinker(...)``. The sibling + branches such as ``talker`` and ``token2wav`` are generation-only for + speech output and are not used in the actor loss. + - ``Qwen2.5-Omni`` during text generation: we still call + ``self.model.generate(...)`` on the full root object, but that does not + mean FSDP should wrap the full root for training. + + Returning the wrong target here is not just inefficient; it can change FSDP + behavior materially. In practice, wrapping the full Omni root caused FSDP2 + to traverse branches that actor training never uses, and that led to invalid + nested mesh composition when ``fully_shard`` tried to apply its mesh layout. + Returning ``thinker`` keeps sharding aligned with the actual forward path + used by ``ActorAL.forward()`` and with the parameter set seen by the actor + optimizer. + """ + return get_audio_forward_model(self.model) + @torch.no_grad() def generate( self, @@ -254,7 +456,11 @@ def generate( generate_args["input_features"] = input_features generate_args["feature_attention_mask"] = feature_attention_mask - if kwargs.get("max_new_tokens", None): + if self.model_type == AUDIO_MODEL_TYPE_QWEN2_5_OMNI: + generate_args["generation_mode"] = "text" + if kwargs.get("max_new_tokens", None): + generate_args["thinker_max_new_tokens"] = kwargs.get("max_new_tokens") + elif kwargs.get("max_new_tokens", None): generate_args["max_new_tokens"] = kwargs.get("max_new_tokens") if kwargs.get("max_length", None): generate_args["max_length"] = kwargs.get("max_length") @@ -359,9 +565,17 @@ def forward( attention_mask=attention_mask, pad_token_id=pad_token_id, ) - position_ids = attention_mask.long().cumsum(-1) - 1 - position_ids.masked_fill_(attention_mask == 0, 1) + if self.model_type == AUDIO_MODEL_TYPE_QWEN2_5_OMNI: + # Let Omni thinker build its own multimodal 3D RoPE positions via + # get_rope_index(...). Passing the usual 1D cumsum position_ids here + # would bypass that path and treat audio tokens like plain text. + position_ids = None + else: + position_ids = attention_mask.long().cumsum(-1) - 1 + position_ids.masked_fill_(attention_mask == 0, 1) else: + if self.model_type == AUDIO_MODEL_TYPE_QWEN2_5_OMNI: + raise NotImplementedError("packing_samples is not supported for Qwen2.5-Omni audio actors.") # convert attention_mask to position_ids position_ids = reset_position_ids(attention_mask) # explicitly ignore attention_mask for packing_samples @@ -369,6 +583,8 @@ def forward( # Pipeline passes audio as audio_values; Qwen2Audio expects input_features. input_features = audio_values + forward_model = get_audio_forward_model(self.model) + forward_config = forward_model.config model_kwargs = { "input_ids": sequences, @@ -394,7 +610,7 @@ def forward( # same expanded sequences, so the log-prob *ratio* used for # the policy gradient is still consistent. # ---------------------------------------------------------- - audio_token_id = getattr(self.model.config, "audio_token_id", None) + audio_token_id = getattr(forward_config, "audio_token_id", None) has_audio_placeholder = (audio_token_id is not None and (sequences == audio_token_id).any().item()) if has_audio_placeholder: @@ -418,7 +634,7 @@ def forward( model_kwargs["feature_attention_mask"] = feature_attention_mask # else: audio_token_id absent → text-only forward (see comment above) - output = self.model(**model_kwargs) + output = forward_model(**model_kwargs) if num_actions is None: # default assert return_output diff --git a/lightrft/models/tests/test_actor_al.py b/lightrft/models/tests/test_actor_al.py index a0a1c619..fb6372ef 100644 --- a/lightrft/models/tests/test_actor_al.py +++ b/lightrft/models/tests/test_actor_al.py @@ -8,13 +8,15 @@ """ from unittest.mock import Mock, patch +import importlib.util import os +import pathlib import pytest +import sys import torch +import types -# Add the lightrft package to the path - -from lightrft.models import ActorAL +from lightrft.models.actor_al import ActorAL, AUDIO_MODEL_TYPE_QWEN2_5_OMNI, create_audio_processor class TestActorAL: @@ -39,20 +41,52 @@ def mock_output(self): "logits": torch.randn(2, 10, 32000) # batch_size=2, seq_len=10, vocab_size=32000 } + @pytest.fixture + def mock_omni_config(self): + config = Mock() + config.model_type = "qwen2_5_omni" + config.use_cache = True + config.pad_token_id = 0 + return config + + @pytest.fixture + def mock_omni_thinker_config(self): + config = Mock() + config.model_type = "qwen2_5_omni_thinker" + config.audio_token_id = 151646 + config.use_cache = True + return config + @pytest.fixture def mock_model(self, mock_config, mock_output): """Set up mock model fixture.""" model = Mock() model.config = mock_config + model.audio_tower = None + model.audio_encoder = None model.generate.return_value = torch.randint(0, 32000, (2, 15)) # batch_size=2, seq_len=15 model.return_value = mock_output return model - @patch('lightrft.models.actor_al.Qwen2AudioForConditionalGeneration') - def test_actor_al_initialization(self, mock_qwen2_audio, mock_model): + @pytest.fixture + def mock_omni_model(self, mock_omni_config, mock_omni_thinker_config, mock_output): + model = Mock() + model.config = mock_omni_config + model.generate.return_value = torch.randint(0, 32000, (2, 15)) + + thinker = Mock() + thinker.config = mock_omni_thinker_config + thinker.audio_tower = None + thinker.audio_encoder = None + thinker.return_value = mock_output + model.thinker = thinker + return model + + @patch('lightrft.models.actor_al.get_audio_model_and_type') + def test_actor_al_initialization(self, mock_get_audio_model, mock_model): """Test ActorAL initialization with mock model.""" # Set up mock - mock_qwen2_audio.from_pretrained.return_value = mock_model + mock_get_audio_model.return_value = (mock_model, "qwen2_audio") # Initialize ActorAL actor = ActorAL( @@ -197,6 +231,101 @@ def test_print_trainable_parameters(self, mock_model): actor.print_trainable_parameters() mock_model.print_trainable_parameters.assert_called_once() + def test_forward_with_qwen2_5_omni_routes_to_thinker(self, mock_omni_model): + actor = ActorAL(pretrain_or_model=mock_omni_model, packing_samples=False) + + sequences = torch.randint(0, 32000, (2, 10)) + attention_mask = torch.ones(2, 10) + + with patch('lightrft.models.actor_al.log_probs_from_logits') as mock_log_probs: + mock_log_probs.return_value = torch.randn(2, 9) + result = actor.forward( + sequences=sequences, + num_actions=4, + attention_mask=attention_mask, + audio_values=None, + ) + + assert isinstance(result, torch.Tensor) + assert actor.model_type == AUDIO_MODEL_TYPE_QWEN2_5_OMNI + mock_omni_model.thinker.assert_called_once() + _, kwargs = mock_omni_model.thinker.call_args + assert kwargs["position_ids"] is None + assert kwargs["attention_mask"] is attention_mask + + def test_generate_with_qwen2_5_omni_uses_text_mode(self, mock_omni_model): + actor = ActorAL(pretrain_or_model=mock_omni_model, packing_samples=False) + + input_ids = torch.randint(0, 32000, (2, 5)) + sequences, attention_mask, action_mask = actor.generate( + input_ids=input_ids, + max_new_tokens=12, + temperature=0.8, + do_sample=True, + eos_token_id=2, + pad_token_id=0, + ) + + assert sequences.shape[0] == 2 + assert attention_mask.shape[0] == 2 + assert action_mask.shape[0] == 2 + + _, kwargs = mock_omni_model.generate.call_args + assert kwargs["generation_mode"] == "text" + assert kwargs["thinker_max_new_tokens"] == 12 + + def test_get_fsdp_target_model_uses_root_for_qwen2_audio(self, mock_model): + actor = ActorAL(pretrain_or_model=mock_model, packing_samples=False) + assert actor.get_fsdp_target_model() is mock_model + + def test_get_fsdp_target_model_uses_thinker_for_qwen2_5_omni(self, mock_omni_model): + actor = ActorAL(pretrain_or_model=mock_omni_model, packing_samples=False) + assert actor.get_fsdp_target_model() is mock_omni_model.thinker + + @patch("lightrft.models.actor_al.get_audio_processor_class") + @patch("lightrft.models.actor_al.infer_audio_model_type") + def test_create_audio_processor_reuses_matching_instance( + self, + mock_infer_audio_model_type, + mock_get_audio_processor_class, + ): + dummy_processor_cls = type("DummyProcessor", (), {}) + existing_processor = dummy_processor_cls() + + mock_infer_audio_model_type.return_value = "qwen2_audio" + mock_get_audio_processor_class.return_value = dummy_processor_cls + + assert create_audio_processor("test_model_path", processor=existing_processor) is existing_processor + + @patch("lightrft.models.actor_al.get_audio_processor_class") + @patch("lightrft.models.actor_al.infer_audio_model_type") + def test_create_audio_processor_reloads_mismatched_instance( + self, + mock_infer_audio_model_type, + mock_get_audio_processor_class, + ): + reloaded_processor = Mock() + dummy_processor_cls = type("DummyProcessor", (), {}) + dummy_processor_cls.from_pretrained = Mock(return_value=reloaded_processor) + print_fn = Mock() + + mock_infer_audio_model_type.return_value = "qwen2_5_omni" + mock_get_audio_processor_class.return_value = dummy_processor_cls + + result = create_audio_processor( + "test_model_path", + processor=Mock(), + trust_remote_code=True, + print_fn=print_fn, + ) + + assert result is reloaded_processor + dummy_processor_cls.from_pretrained.assert_called_once_with( + "test_model_path", + trust_remote_code=True, + ) + print_fn.assert_called_once() + class TestActorALWithRealData: """Test cases for ActorAL with real model and data (if available).""" diff --git a/lightrft/models/utils.py b/lightrft/models/utils.py index d4a10f6f..5780e618 100644 --- a/lightrft/models/utils.py +++ b/lightrft/models/utils.py @@ -61,12 +61,16 @@ def find_all_linear_modules(model: "nn.Module", freeze_vision_tower: bool) -> Li forbidden.add("multi_modal_projector") elif model_type in ["qwen2_vl", "qwen2_5_vl"]: forbidden.add("merger") + elif model_type == "qwen2_5_omni": + forbidden.update({"talker", "token2wav"}) if freeze_vision_tower: if model_type in ["mllama"]: forbidden.add("vision_model") elif model_type in ["qwen2_vl", "qwen2_5_vl"]: forbidden.add("visual") + elif model_type == "qwen2_5_omni": + forbidden.add("visual") else: forbidden.add("vision_tower") diff --git a/lightrft/strategy/fsdp/fsdpv2.py b/lightrft/strategy/fsdp/fsdpv2.py index b8b2c99c..9620ffbd 100755 --- a/lightrft/strategy/fsdp/fsdpv2.py +++ b/lightrft/strategy/fsdp/fsdpv2.py @@ -69,6 +69,7 @@ "Qwen2VLVisionBlock", "Qwen2_5_VLVisionBlock", "Qwen2_5_VLDecoderLayer", + "Qwen2_5OmniDecoderLayer", "Qwen2DecoderLayer", "LlamaDecoderLayer", # for DeepSeek-R1-Distill-Llama-70B "DeepseekDecoderLayer", @@ -77,9 +78,28 @@ vit_transformer_cls_names = [ "Qwen2VLVisionBlock", "Qwen2_5_VLVisionBlock", + "Qwen2_5OmniVisionEncoder", + "Qwen2_5OmniAudioEncoder", ] +def _get_fsdp_training_target(model: nn.Module) -> nn.Module: + """ + Resolve the module FSDP should shard/optimize. + + Actor wrappers may keep extra inference-only branches on ``model.model``; when + available, prefer the actor-provided FSDP target instead of assuming the full + wrapped backbone should be sharded. + """ + if not is_actor(model): + return model + + get_target = getattr(model, "get_fsdp_target_model", None) + if callable(get_target): + return get_target() + return model.model + + class FSDPV2Strategy(StrategyBase): """ The strategy for training with PyTorch's Fully Sharded Data Parallel V2. @@ -162,8 +182,7 @@ def create_optimizer(self, model, **kwargs) -> Optimizer: >>> optimizer = strategy.create_optimizer(model, lr=1e-4, weight_decay=0.01) """ - if is_actor(model): - model = model.model + model = _get_fsdp_training_target(model) # group params by (dtype, dtensor shard size, weight_dacay) to avoid error in clip_grad and opt.step self.grouped_params = group_parameters_for_optimizer_dtensor(model, kwargs["weight_decay"]) # Convert the grouped parameters into the final format for the optimizer @@ -219,8 +238,7 @@ def optimizer_step( """ self.cur_step[name] += 1 if self.cur_step[name] == self.accumulated_gradient: - if is_actor(model): - model = model.model + model = _get_fsdp_training_target(model) grad_norms = [] for param_group in self.grouped_params.values(): @@ -314,7 +332,7 @@ def _fsdp_init_model(self, model, is_training, shard_size=-1, reshard_after_forw naive_mp_training = self.use_naive_opt and is_training - model_to_wrap = model.model if is_actor(model) else model + model_to_wrap = _get_fsdp_training_target(model) if isinstance(model_to_wrap, FSDPModule): return model @@ -322,7 +340,7 @@ def _fsdp_init_model(self, model, is_training, shard_size=-1, reshard_after_forw self.report_memory("before FSDP2 wrap model pos2") # this is not sufficient enough, for example, it will only return Qwen2DecoderLayer for qwen2 - default_transformer_cls_names_to_wrap = getattr(model_to_wrap, "_no_split_modules", []) + default_transformer_cls_names_to_wrap = list(getattr(model_to_wrap, "_no_split_modules", [])) # so we add some manual rules transformer_cls_names_to_wrap = default_transformer_cls_names_to_wrap @@ -336,7 +354,8 @@ def _fsdp_init_model(self, model, is_training, shard_size=-1, reshard_after_forw # we either keep vision model in full state, or keep it in FSDP's root module. # below we keep vit in root module to avoid stuck for cls_name in vit_transformer_cls_names: - transformer_cls_names_to_wrap.remove(cls_name) + if cls_name in transformer_cls_names_to_wrap: + transformer_cls_names_to_wrap.remove(cls_name) transformer_cls_to_wrap = list() # noqa vit_transformer_cls = list() # noqa From d6d527bcad833484fa6a032ce1145e85658e39ea Mon Sep 17 00:00:00 2001 From: niuyazhe Date: Fri, 17 Apr 2026 16:22:12 +0800 Subject: [PATCH 06/11] fix(nyz): add training convergence version --- .gitignore | 1 + .../r1_aqa/run_grpo_r1_aqa_qwen2_audio_7b.sh | 51 +++-- examples/r1_aqa/train_colocate.py | 20 +- lightrft/models/actor_al.py | 203 ++++++++++++++++-- lightrft/models/loss.py | 93 +++++++- lightrft/models/tests/test_actor_al.py | 182 ++++++++++++++-- lightrft/strategy/fsdp/fsdpv2.py | 42 ++++ lightrft/trainer/fast_exp_maker.py | 21 +- lightrft/trainer/spmd_ppo_trainer.py | 4 + 9 files changed, 557 insertions(+), 60 deletions(-) diff --git a/.gitignore b/.gitignore index 33d356e2..6964940f 100644 --- a/.gitignore +++ b/.gitignore @@ -1221,3 +1221,4 @@ build/* examples/math_benchmarks/eval_results/ .llmconfig.yaml tb/* +debug*.py diff --git a/examples/r1_aqa/run_grpo_r1_aqa_qwen2_audio_7b.sh b/examples/r1_aqa/run_grpo_r1_aqa_qwen2_audio_7b.sh index 7f7a42b6..a1e64ff1 100644 --- a/examples/r1_aqa/run_grpo_r1_aqa_qwen2_audio_7b.sh +++ b/examples/r1_aqa/run_grpo_r1_aqa_qwen2_audio_7b.sh @@ -6,24 +6,21 @@ # with rule-based rewards, faithfully migrating the R1-AQA training pipeline. # # -# Migration from R1-AQA: -# R1-AQA num_generations=8 → N_SAMPLES=8 -# R1-AQA temperature=1.0 → TEMPERATURE=1.0 -# R1-AQA max_prompt_length=512 → PROMPT_MAX_LEN=512 -# R1-AQA per_device_batch=1 → MICRO_TRAIN=1 -# R1-AQA grad_accum=2 → TBS adjusted -# R1-AQA DeepSpeed ZeRO3 → --zero_stage 3 -# ################################################################################ # Part 1: User Configuration # ################################################################################ # --- Model and Dataset Paths --- -# Qwen2-Audio-7B-Instruct base model +# Qwen2-Audio-7B-Instruct/Qwen2.5-Omni-7B base model PATH_TO_YOUR_BASE_MODEL="" -# Path to the preprocessed AVQA dataset (output of data_preprocess/avqa.py) +# Path to the cleaned AVQA dataset directory. +# Recommended workflow: +# 1. Build parquet with examples/r1_aqa/data_preprocess/avqa.py +# 2. Clean missing/broken audio rows with +# examples/r1_aqa/data_preprocess/clean_audio_dataset.py +# 3. Point this variable to the cleaned output directory PATH_TO_YOUR_AVQA_DATASET="" # --- Experiment and Logging --- @@ -42,7 +39,7 @@ export WANDB_PROJECT="LightRFT-R1-AQA" # --- GRPO Settings (from R1-AQA) --- GROUP_METHOD="normal" -N_SAMPLES=4 # num_generations reduced from 8→4 to save memory +N_SAMPLES=8 # num_generations per prompt EPISODE=10 # Number of training episodes WARMUP=0.03 # Learning rate warmup ratio TEMPERATURE=1.0 # Sampling temperature (R1-AQA default: 1.0) @@ -50,10 +47,10 @@ TEMPERATURE=1.0 # Sampling temperature (R1-AQA default: 1.0) # --- Batch Size Configuration --- # Constraint: train_batch_size >= rollout_batch_size * n_samples_per_prompt # Reduced for single-GPU memory constraints (140 GiB GPU). -RBS=4 # Rollout Batch Size (reduced from 128 to fit in memory) -TBS=16 # Train Batch Size (RBS * N_SAMPLES = 4 * 4 = 16) -MICRO_ROLLOUT=1 # Micro rollout batch size per GPU (reduced from 2) -MICRO_TRAIN=1 # Micro train batch size per GPU (R1-AQA: per_device=1) +RBS=64 # Rollout Batch Size +TBS=512 # Train Batch Size +MICRO_ROLLOUT=4 # Micro rollout batch size per GPU +MICRO_TRAIN=4 # Micro train batch size per GPU # --- Learning and Model Settings --- KL=0.01 # KL divergence coefficient @@ -123,8 +120,6 @@ torchrun \ --fsdp \ --use_kl_loss \ --flash_attn \ - --rm_use_engine \ - --reward_pretrain "{}" \ --save_path "results/${EXPERIMENT_NAME}/${SAVE_MODEL_NAME}" \ --ckpt_path "results/${EXPERIMENT_NAME}/${SAVE_MODEL_NAME}" \ $( [ -n "${CKPT_PATH_LOCAL}" ] && echo "--ckpt_path_local ${CKPT_PATH_LOCAL}" ) \ @@ -152,13 +147,14 @@ torchrun \ --gradient_checkpointing \ --save_steps ${SAVE_STEPS} \ --max_ckpt_num 3 \ - --engine_type sglang \ + --engine_type vllm \ --engine_mem_util ${ENGINE_MEM_UTIL} \ --engine_tp_size $ENGINE_TP \ --enable_engine_sleep \ --l2 1.0e-2 \ --adam_offload \ - --use_tensorboard "tb/r1-aqa-baseline" \ + --mixed_mm_data \ + --use_wandb "${WANDB_API_KEY}" \ --wandb_project "${WANDB_PROJECT}" \ --wandb_run_name "${WANDB_RUN_NAME}" \ 2>&1 | tee "rft_logs/${EXPERIMENT_NAME}/node${NODE_RANK}_${current_time}.log" @@ -175,16 +171,25 @@ torchrun \ # --audio_dir data/AVQA/audios \ # # --local_save_dir /path/to/preprocessed/avqa_lightrft # # # -# Step 2: Configure the Script # +# Step 2: Clean Missing / Broken Audio Rows # +# Strongly recommended before training. This avoids distributed hangs caused # +# by prompts that still contain audio placeholders while their audio files # +# are missing on disk. # +# # +# python examples/r1_aqa/data_preprocess/clean_audio_dataset.py \ # +# --input_dataset /path/to/preprocessed/avqa_lightrft \ # +# --output_dir /path/to/preprocessed/avqa_lightrft_clean # +# # +# Step 3: Configure the Script # # Edit "Part 1: User Configuration" above: # # - Set PATH_TO_YOUR_BASE_MODEL (Qwen2-Audio-7B-Instruct) # -# - Set PATH_TO_YOUR_AVQA_DATASET # +# - Set PATH_TO_YOUR_AVQA_DATASET to the cleaned dataset directory # # - Set GPU count in GPU_PER_NODE # # # -# Step 3: Run Training # +# Step 4: Run Training # # bash examples/r1_aqa/run_grpo_r1_aqa_qwen2_audio_7b.sh # # # -# Step 4: Evaluate on MMAU Test-Mini # +# Step 5: Evaluate on MMAU Test-Mini # # python examples/r1_aqa/eval_mmau.py \ # # --model_path results/lightrft-r1-aqa-grpo-training/... \ # # --data_file /path/to/mmau-test-mini.json \ # diff --git a/examples/r1_aqa/train_colocate.py b/examples/r1_aqa/train_colocate.py index 188aa94b..7cd96493 100644 --- a/examples/r1_aqa/train_colocate.py +++ b/examples/r1_aqa/train_colocate.py @@ -25,7 +25,12 @@ import torch.nn.functional as F from lightrft.utils import add_arguments from lightrft.datasets import SFTDatasetVL -from lightrft.models.actor_al import ActorAL, create_audio_processor +from lightrft.models.actor_al import ( + AUDIO_MODEL_TYPE_QWEN2_5_OMNI, + ActorAL, + create_audio_processor, + infer_audio_model_type, +) from lightrft.models.actor_language import ActorLanguage from lightrft.strategy import get_strategy from lightrft.trainer.spmd_ppo_trainer import SPMDPPOTrainerVL @@ -37,6 +42,17 @@ from audio_dataset import AudioPromptDataset +def _validate_inference_engine(args) -> None: + if args.text_only: + return + + model_type = infer_audio_model_type(args.pretrain) + if model_type == AUDIO_MODEL_TYPE_QWEN2_5_OMNI and args.engine_type == "sglang": + raise ValueError( + "Qwen2.5-Omni audio demo currently requires --engine_type vllm; sglang is not supported." + ) + + def train(args): """ Main training function for R1-AQA GRPO with co-located rule-based rewards. @@ -52,6 +68,8 @@ def train(args): 8. Run training loop via SPMDPPOTrainerVL 9. Save final model """ + _validate_inference_engine(args) + # ==================== Strategy ==================== strategy = get_strategy(args) diff --git a/lightrft/models/actor_al.py b/lightrft/models/actor_al.py index f4660bd7..b059013a 100644 --- a/lightrft/models/actor_al.py +++ b/lightrft/models/actor_al.py @@ -501,6 +501,7 @@ def forward( pixel_values_videos: Optional[torch.Tensor] = None, video_grid_thw: Optional[torch.Tensor] = None, return_output=False, + return_aligned_inputs: bool = False, packed_seq_lens: Optional[list[int]] = None, audio_values: Optional[torch.Tensor] = None, feature_attention_mask: Optional[torch.Tensor] = None, @@ -531,6 +532,11 @@ def forward( :type video_grid_thw: Optional[torch.Tensor] :param return_output: Whether to return the full model output along with log probs :type return_output: bool + :param return_aligned_inputs: Whether to additionally return the exact ``input_ids`` and + ``attention_mask`` that were fed into the backbone after audio placeholder alignment. + This is primarily used during rollout so replay batches can reuse the identical + token layout instead of reconstructing it from the raw engine output. + :type return_aligned_inputs: bool :param packed_seq_lens: Sequence lengths for packed samples :type packed_seq_lens: Optional[list[int]] :param audio_values: Preprocessed audio features (mel-spectrogram from pipeline) @@ -565,14 +571,7 @@ def forward( attention_mask=attention_mask, pad_token_id=pad_token_id, ) - if self.model_type == AUDIO_MODEL_TYPE_QWEN2_5_OMNI: - # Let Omni thinker build its own multimodal 3D RoPE positions via - # get_rope_index(...). Passing the usual 1D cumsum position_ids here - # would bypass that path and treat audio tokens like plain text. - position_ids = None - else: - position_ids = attention_mask.long().cumsum(-1) - 1 - position_ids.masked_fill_(attention_mask == 0, 1) + position_ids = None else: if self.model_type == AUDIO_MODEL_TYPE_QWEN2_5_OMNI: raise NotImplementedError("packing_samples is not supported for Qwen2.5-Omni audio actors.") @@ -605,35 +604,86 @@ def forward( # appears in the token sequence, so the model's merge step # would fail with a shape-mismatch error. # - # When the placeholder is absent we fall back to a text-only - # forward. Both the actor AND the reference model see the - # same expanded sequences, so the log-prob *ratio* used for - # the policy gradient is still consistent. + # When the rollout engine expands the prompt differently from + # the local processor, we rewrite the prompt-side placeholder + # run to the audio tower's expected token count instead of + # dropping audio conditioning. # ---------------------------------------------------------- audio_token_id = getattr(forward_config, "audio_token_id", None) has_audio_placeholder = (audio_token_id is not None and (sequences == audio_token_id).any().item()) - if has_audio_placeholder: + if has_audio_placeholder or (not self.packing_samples and feature_attention_mask is not None): input_features, feature_attention_mask = self._prepare_audio_features( input_features, feature_attention_mask=feature_attention_mask, sequences=sequences, audio_token_id=audio_token_id, ) + original_audio_token_counts = ( + (sequences == audio_token_id).sum(dim=1) + if audio_token_id is not None + else None + ) + expected_audio_token_counts = self._infer_audio_output_token_counts( + forward_model, + feature_attention_mask, + ) + if ( + not self.packing_samples + and audio_token_id is not None + and expected_audio_token_counts is not None + ): + sequences, attention_mask = self._align_audio_placeholder_counts( + sequences=sequences, + attention_mask=attention_mask, + audio_token_id=audio_token_id, + expected_audio_token_counts=expected_audio_token_counts, + pad_token_id=pad_token_id, + num_actions=(num_actions if isinstance(num_actions, int) else None), + ) + actual_audio_token_counts = ( + (sequences == audio_token_id).sum(dim=1) + if audio_token_id is not None + else None + ) + if ( + actual_audio_token_counts is not None + and expected_audio_token_counts is not None + and torch.any(actual_audio_token_counts > expected_audio_token_counts) + ): + raise RuntimeError( + "Audio placeholder alignment failed before Qwen2.5-Omni merge: " + f"actual={actual_audio_token_counts.tolist()} " + f"expected={expected_audio_token_counts.tolist()}" + ) if os.environ.get("LIGHTRFT_AUDIO_DEBUG", "0") == "1": rank = dist.get_rank() if dist.is_initialized() else 0 print( f"[ActorAL][rank={rank}] sequences={tuple(sequences.shape)} " f"audio_values={tuple(input_features.shape)} " f"feature_attention_mask={tuple(feature_attention_mask.shape)} " - f"audio_token_count={(sequences == audio_token_id).sum(dim=1).tolist()} " - f"feature_len={feature_attention_mask.sum(dim=1).tolist()}", + f"audio_token_count={actual_audio_token_counts.tolist() if actual_audio_token_counts is not None else None} " + f"original_audio_token_count=" + f"{original_audio_token_counts.tolist() if original_audio_token_counts is not None else None} " + f"feature_len={feature_attention_mask.sum(dim=1).tolist()} " + f"expected_audio_token_count=" + f"{expected_audio_token_counts.tolist() if expected_audio_token_counts is not None else None} " + f"audio_merge=True", flush=True, ) + model_kwargs["input_ids"] = sequences + model_kwargs["attention_mask"] = attention_mask model_kwargs["input_features"] = input_features model_kwargs["feature_attention_mask"] = feature_attention_mask - # else: audio_token_id absent → text-only forward (see comment above) + # else: no placeholder token and no feature mask to infer expected count from + if not self.packing_samples and self.model_type != AUDIO_MODEL_TYPE_QWEN2_5_OMNI: + position_ids = attention_mask.long().cumsum(-1) - 1 + position_ids.masked_fill_(attention_mask == 0, 1) + model_kwargs["position_ids"] = position_ids + + sequences = model_kwargs["input_ids"] + attention_mask = model_kwargs["attention_mask"] output = forward_model(**model_kwargs) if num_actions is None: # default @@ -654,6 +704,8 @@ def forward( offset += seq_len action_log_probs = torch.cat(action_log_probs, dim=1) + if return_output and return_aligned_inputs: + return (action_log_probs, output, sequences, attention_mask) if return_output: return (action_log_probs, output) else: @@ -714,6 +766,125 @@ def _prepare_audio_features( return input_features, feature_attention_mask + @staticmethod + def _infer_audio_output_token_counts( + forward_model: nn.Module, + feature_attention_mask: Optional[torch.Tensor], + ) -> Optional[torch.Tensor]: + """ + Infer how many audio placeholder tokens the backbone expects per sample. + + External rollout engines can expand audio placeholders differently from the + local processor. If the local audio tower would emit a different number of + encoder states than the number of ``audio_token_id`` slots present in + ``input_ids``, the subsequent masked scatter would fail on CUDA. + """ + if feature_attention_mask is None: + return None + + audio_tower = getattr(forward_model, "audio_tower", None) + get_output_lengths = getattr(audio_tower, "_get_feat_extract_output_lengths", None) + if get_output_lengths is None: + return None + + feature_lengths = feature_attention_mask.to(dtype=torch.long).sum(dim=1) + _, output_lengths = get_output_lengths(feature_lengths) + return output_lengths.to(device=feature_attention_mask.device, dtype=torch.long) + + @staticmethod + def _align_audio_placeholder_counts( + sequences: torch.Tensor, + attention_mask: Optional[torch.Tensor], + audio_token_id: int, + expected_audio_token_counts: Optional[torch.Tensor], + pad_token_id: int, + num_actions: Optional[int] = None, + ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: + """ + Rewrite prompt-side audio placeholders to the expected count per sample. + + The sequence returned by the rollout engine may under/over-expand the + audio placeholder block compared with the local HF audio tower. + + This alignment must stay stable when the same sample is replayed in a + different PPO micro-batch. Relying on batch-level ``num_actions`` to + infer the prompt/response split is unsafe here because replay batches can + have a different max response length than the rollout batch that first + produced the sample. Instead, rewrite the first contiguous audio-token + block directly from the active tokens so the aligned ``input_ids`` are + batch-invariant. + """ + if attention_mask is None or expected_audio_token_counts is None: + return sequences, attention_mask + + adjusted_rows = [] + batch_size = sequences.size(0) + + for row_idx in range(batch_size): + active_tokens = sequences[row_idx, attention_mask[row_idx].bool()] + active_len = int(active_tokens.numel()) + if active_len == 0: + adjusted_rows.append(active_tokens) + continue + + expected_count = max(0, int(expected_audio_token_counts[row_idx].item())) + audio_positions = torch.nonzero(active_tokens == audio_token_id, as_tuple=False).flatten() + + if audio_positions.numel() == 0: + if expected_count == 0: + row_tokens = active_tokens.clone() + else: + new_audio_block = active_tokens.new_full((expected_count,), audio_token_id) + row_tokens = torch.cat((new_audio_block, active_tokens), dim=0) + else: + block_start = int(audio_positions[0].item()) + block_end = block_start + while block_end + 1 < active_len and int(active_tokens[block_end + 1].item()) == audio_token_id: + block_end += 1 + + actual_count = block_end - block_start + 1 + if actual_count == expected_count: + row_tokens = active_tokens.clone() + else: + prompt_prefix = active_tokens[:block_start] + prompt_suffix = active_tokens[block_end + 1:] + new_audio_block = active_tokens.new_full((expected_count,), audio_token_id) + row_tokens = torch.cat((prompt_prefix, new_audio_block, prompt_suffix), dim=0) + + # If the response contains stray audio placeholders, keep the sequence + # length stable and replace surplus placeholders from the end. + all_audio_positions = torch.nonzero(row_tokens == audio_token_id, as_tuple=False).flatten() + total_count = int(all_audio_positions.numel()) + if total_count > expected_count: + for pos in all_audio_positions.flip(0)[:total_count - expected_count]: + row_tokens[pos] = pad_token_id + + adjusted_rows.append(row_tokens) + + max_active_len = max((int(row.numel()) for row in adjusted_rows), default=0) + target_len = max(sequences.size(1), max_active_len) + + aligned_sequences = torch.full( + (batch_size, target_len), + pad_token_id, + dtype=sequences.dtype, + device=sequences.device, + ) + aligned_attention_mask = torch.zeros( + (batch_size, target_len), + dtype=attention_mask.dtype, + device=attention_mask.device, + ) + + for row_idx, row_tokens in enumerate(adjusted_rows): + row_len = int(row_tokens.numel()) + if row_len == 0: + continue + aligned_sequences[row_idx, -row_len:] = row_tokens + aligned_attention_mask[row_idx, -row_len:] = 1 + + return aligned_sequences, aligned_attention_mask + def gradient_checkpointing_enable(self, gradient_checkpointing_kwargs={"use_reentrant": False}): """ Enable gradient checkpointing to reduce memory usage during training. diff --git a/lightrft/models/loss.py b/lightrft/models/loss.py index d593cb93..2faccc10 100644 --- a/lightrft/models/loss.py +++ b/lightrft/models/loss.py @@ -176,6 +176,66 @@ def __init__( self.use_dapo = use_dapo self.use_cpg_loss = use_cpg_loss self.high_entropy_token_ratio = high_entropy_token_ratio + self._last_stats: dict[str, float] = {} + + @staticmethod + def _stats_over_mask(values: torch.Tensor, mask: Optional[torch.Tensor], prefix: str) -> dict[str, float]: + if mask is None: + selected = values.reshape(-1) + else: + selected = values.masked_select(mask.to(dtype=torch.bool)) + + if selected.numel() == 0: + return { + f"{prefix}_mean": 0.0, + f"{prefix}_min": 0.0, + f"{prefix}_max": 0.0, + } + + selected = selected.detach().float() + return { + f"{prefix}_mean": selected.mean().item(), + f"{prefix}_min": selected.min().item(), + f"{prefix}_max": selected.max().item(), + } + + def get_last_stats(self) -> dict[str, float]: + return dict(self._last_stats) + + def _update_last_stats( + self, + *, + stats_mask: torch.Tensor, + advantages: torch.Tensor, + logprob_delta: torch.Tensor, + token_loss: torch.Tensor, + log_probs: Optional[torch.Tensor] = None, + old_log_probs: Optional[torch.Tensor] = None, + ratio: Optional[torch.Tensor] = None, + ) -> None: + valid_token_count = float(stats_mask.sum().item()) + stats = { + "policy/valid_tokens": valid_token_count, + **self._stats_over_mask(advantages, stats_mask, "policy/adv"), + **self._stats_over_mask(logprob_delta, stats_mask, "policy/logprob_delta"), + **self._stats_over_mask(token_loss, stats_mask, "policy/token_loss"), + } + + if ratio is not None: + denom = max(valid_token_count, 1.0) + clipped_high = (ratio > 1 + self.clip_eps).masked_select(stats_mask) + clipped_low = (ratio < 1 - self.clip_eps).masked_select(stats_mask) + stats.update( + { + "policy/clipfrac_high": float(clipped_high.numel() / denom), + "policy/clipfrac_low": float(clipped_low.numel() / denom), + **self._stats_over_mask(log_probs, stats_mask, "policy/logprob"), + **self._stats_over_mask(old_log_probs, stats_mask, "policy/old_logprob"), + **self._stats_over_mask(ratio, stats_mask, "policy/ratio"), + } + ) + + self._last_stats = stats def forward( self, @@ -236,21 +296,44 @@ def forward( else: # No entropy masking, use action_mask only final_mask = action_mask + + stats_mask = final_mask + if stats_mask is None: + stats_mask = torch.ones_like(log_probs, dtype=torch.bool) + else: + stats_mask = stats_mask.to(dtype=torch.bool) + + logprob_delta = log_probs - old_log_probs if self.use_cpg_loss: clipped_log_probs = torch.where( advantages > 0, torch.clamp(log_probs, max=torch.log(torch.tensor(1 + self.clip_eps)) + old_log_probs), torch.clamp(log_probs, min=torch.log(torch.tensor(1 - self.clip_eps)) + old_log_probs) ) - loss = -clipped_log_probs * advantages - loss = masked_mean(loss, final_mask, dim=-1).mean() + token_loss = -clipped_log_probs * advantages + loss = masked_mean(token_loss, final_mask, dim=-1).mean() + self._update_last_stats( + stats_mask=stats_mask, + advantages=advantages, + logprob_delta=logprob_delta, + token_loss=token_loss, + ) return loss # PPO loss - ratio = (log_probs - old_log_probs).exp() + ratio = logprob_delta.exp() surr1 = ratio * advantages surr2 = ratio.clamp(1 - self.clip_eps, 1 + self.clip_eps) * advantages - loss = -torch.min(surr1, surr2) - loss = masked_mean(loss, final_mask, dim=-1).mean() + token_loss = -torch.min(surr1, surr2) + loss = masked_mean(token_loss, final_mask, dim=-1).mean() + self._update_last_stats( + stats_mask=stats_mask, + advantages=advantages, + logprob_delta=logprob_delta, + token_loss=token_loss, + log_probs=log_probs, + old_log_probs=old_log_probs, + ratio=ratio, + ) return loss diff --git a/lightrft/models/tests/test_actor_al.py b/lightrft/models/tests/test_actor_al.py index fb6372ef..7a965f03 100644 --- a/lightrft/models/tests/test_actor_al.py +++ b/lightrft/models/tests/test_actor_al.py @@ -8,24 +8,15 @@ """ from unittest.mock import Mock, patch -import importlib.util import os -import pathlib import pytest -import sys import torch -import types from lightrft.models.actor_al import ActorAL, AUDIO_MODEL_TYPE_QWEN2_5_OMNI, create_audio_processor class TestActorAL: """Test cases for ActorAL (Audio Language) class.""" - @pytest.fixture - def device(self): - """Set up device fixture.""" - return torch.device("cuda" if torch.cuda.is_available() else "cpu") - @pytest.fixture def mock_config(self): """Set up mock config fixture.""" @@ -76,12 +67,41 @@ def mock_omni_model(self, mock_omni_config, mock_omni_thinker_config, mock_outpu thinker = Mock() thinker.config = mock_omni_thinker_config - thinker.audio_tower = None + thinker.audio_tower = Mock() + thinker.audio_tower._get_feat_extract_output_lengths.return_value = ( + torch.tensor([4, 8]), + torch.tensor([1, 2]), + ) thinker.audio_encoder = None thinker.return_value = mock_output model.thinker = thinker return model + def _make_actor(self, model, *, packing_samples=False): + actor = ActorAL(pretrain_or_model=model, packing_samples=packing_samples) + actor.packing_samples = packing_samples + return actor + + def _make_omni_forward_inputs(self, *, sequences, attention_mask, feature_attention_mask): + sequences = torch.tensor(sequences) + return { + "sequences": sequences, + "attention_mask": torch.tensor(attention_mask), + "audio_values": torch.randn(sequences.size(0), 80, 10), + "feature_attention_mask": torch.tensor(feature_attention_mask), + } + + def _run_mocked_omni_forward(self, actor, *, num_actions=2, **forward_kwargs): + sequences = forward_kwargs["sequences"] + with patch("lightrft.models.actor_al.log_probs_from_logits") as mock_log_probs: + mock_log_probs.return_value = torch.randn(sequences.size(0), sequences.size(1) - 1) + result = actor.forward( + num_actions=num_actions, + **forward_kwargs, + ) + + return result, actor.model.thinker.call_args.kwargs + @patch('lightrft.models.actor_al.get_audio_model_and_type') def test_actor_al_initialization(self, mock_get_audio_model, mock_model): """Test ActorAL initialization with mock model.""" @@ -232,7 +252,7 @@ def test_print_trainable_parameters(self, mock_model): mock_model.print_trainable_parameters.assert_called_once() def test_forward_with_qwen2_5_omni_routes_to_thinker(self, mock_omni_model): - actor = ActorAL(pretrain_or_model=mock_omni_model, packing_samples=False) + actor = self._make_actor(mock_omni_model) sequences = torch.randint(0, 32000, (2, 10)) attention_mask = torch.ones(2, 10) @@ -253,6 +273,144 @@ def test_forward_with_qwen2_5_omni_routes_to_thinker(self, mock_omni_model): assert kwargs["position_ids"] is None assert kwargs["attention_mask"] is attention_mask + @pytest.mark.parametrize( + ( + "output_lengths", + "feature_lengths", + "expected_audio_counts", + "expected_attention_mask", + ), + [ + ([8, 4], [2, 1], [2, 1], [[0, 1, 1, 1, 1, 1], [0, 0, 1, 1, 1, 1]]), + ([4, 8], [1, 2], [1, 2], [[0, 0, 1, 1, 1, 1], [0, 1, 1, 1, 1, 1]]), + ], + ) + def test_forward_with_qwen2_5_omni_aligns_audio_placeholders( + self, + mock_omni_model, + output_lengths, + feature_lengths, + expected_audio_counts, + expected_attention_mask, + ): + actor = self._make_actor(mock_omni_model) + mock_omni_model.thinker.audio_tower._get_feat_extract_output_lengths.return_value = ( + torch.tensor(output_lengths), + torch.tensor(feature_lengths), + ) + audio_token_id = mock_omni_model.thinker.config.audio_token_id + + inputs = self._make_omni_forward_inputs( + sequences=[ + [0, 0, audio_token_id, 11, 12, 13], + [0, audio_token_id, audio_token_id, 21, 22, 23], + ], + attention_mask=[ + [0, 0, 1, 1, 1, 1], + [0, 1, 1, 1, 1, 1], + ], + feature_attention_mask=[ + [1, 1, 1, 1], + [1, 1, 1, 1, 1, 1, 1, 1], + ], + ) + + result, kwargs = self._run_mocked_omni_forward(actor, **inputs) + assert isinstance(result, torch.Tensor) + assert "input_features" in kwargs + assert "feature_attention_mask" in kwargs + assert (kwargs["input_ids"] == audio_token_id).sum(dim=1).tolist() == expected_audio_counts + assert kwargs["attention_mask"].tolist() == expected_attention_mask + + def test_forward_with_qwen2_5_omni_trims_response_audio_placeholders(self, mock_omni_model): + actor = self._make_actor(mock_omni_model) + mock_omni_model.thinker.audio_tower._get_feat_extract_output_lengths.return_value = ( + torch.tensor([8, 8]), + torch.tensor([2, 2]), + ) + + audio_token_id = mock_omni_model.thinker.config.audio_token_id + inputs = self._make_omni_forward_inputs( + sequences=[ + [0, 0, audio_token_id, audio_token_id, 11, audio_token_id], + [0, audio_token_id, audio_token_id, 21, 22, audio_token_id], + ], + attention_mask=[ + [0, 1, 1, 1, 1, 1], + [0, 1, 1, 1, 1, 1], + ], + feature_attention_mask=[ + [1, 1, 1, 1, 1, 1, 1, 1], + [1, 1, 1, 1, 1, 1, 1, 1], + ], + ) + + result, kwargs = self._run_mocked_omni_forward(actor, **inputs) + assert isinstance(result, torch.Tensor) + assert (kwargs["input_ids"] == audio_token_id).sum(dim=1).tolist() == [2, 2] + + def test_forward_can_return_aligned_inputs_for_audio_replay(self, mock_omni_model): + actor = self._make_actor(mock_omni_model) + mock_omni_model.thinker.audio_tower._get_feat_extract_output_lengths.return_value = ( + torch.tensor([8]), + torch.tensor([2]), + ) + + audio_token_id = mock_omni_model.thinker.config.audio_token_id + inputs = self._make_omni_forward_inputs( + sequences=[[0, 0, audio_token_id, 11, 12, 13]], + attention_mask=[[0, 0, 1, 1, 1, 1]], + feature_attention_mask=[[1, 1, 1, 1, 1, 1, 1, 1]], + ) + + ( + action_log_probs, + model_output, + aligned_sequences, + aligned_attention_mask, + ), _ = self._run_mocked_omni_forward( + actor, + return_output=True, + return_aligned_inputs=True, + **inputs, + ) + + assert isinstance(action_log_probs, torch.Tensor) + assert model_output is not None + assert (aligned_sequences == audio_token_id).sum(dim=1).tolist() == [2] + assert aligned_attention_mask.tolist() == [[0, 1, 1, 1, 1, 1]] + + def test_align_audio_placeholders_is_batch_invariant_to_num_actions(self, mock_omni_model): + actor = self._make_actor(mock_omni_model) + audio_token_id = mock_omni_model.thinker.config.audio_token_id + + sequences = torch.tensor( + [[11, 12, 13, 14, audio_token_id, audio_token_id, audio_token_id, audio_token_id, 31, 41, 42, 43]] + ) + attention_mask = torch.ones_like(sequences) + expected_counts = torch.tensor([5]) + + aligned_short, mask_short = actor._align_audio_placeholder_counts( + sequences=sequences, + attention_mask=attention_mask, + audio_token_id=audio_token_id, + expected_audio_token_counts=expected_counts, + pad_token_id=0, + num_actions=3, + ) + aligned_long, mask_long = actor._align_audio_placeholder_counts( + sequences=sequences, + attention_mask=attention_mask, + audio_token_id=audio_token_id, + expected_audio_token_counts=expected_counts, + pad_token_id=0, + num_actions=6, + ) + + assert torch.equal(aligned_short, aligned_long) + assert torch.equal(mask_short, mask_long) + assert (aligned_short == audio_token_id).sum(dim=1).tolist() == [5] + def test_generate_with_qwen2_5_omni_uses_text_mode(self, mock_omni_model): actor = ActorAL(pretrain_or_model=mock_omni_model, packing_samples=False) @@ -279,7 +437,7 @@ def test_get_fsdp_target_model_uses_root_for_qwen2_audio(self, mock_model): assert actor.get_fsdp_target_model() is mock_model def test_get_fsdp_target_model_uses_thinker_for_qwen2_5_omni(self, mock_omni_model): - actor = ActorAL(pretrain_or_model=mock_omni_model, packing_samples=False) + actor = self._make_actor(mock_omni_model) assert actor.get_fsdp_target_model() is mock_omni_model.thinker @patch("lightrft.models.actor_al.get_audio_processor_class") diff --git a/lightrft/strategy/fsdp/fsdpv2.py b/lightrft/strategy/fsdp/fsdpv2.py index 9620ffbd..9cba06f9 100755 --- a/lightrft/strategy/fsdp/fsdpv2.py +++ b/lightrft/strategy/fsdp/fsdpv2.py @@ -100,6 +100,27 @@ def _get_fsdp_training_target(model: nn.Module) -> nn.Module: return model.model +def _collect_floating_param_dtypes(module: nn.Module) -> dict[torch.dtype, list[str]]: + """ + Collect floating-point parameter dtypes for diagnostics before FSDP wrapping. + """ + dtype_to_names: dict[torch.dtype, list[str]] = defaultdict(list) + for name, param in module.named_parameters(): + if param is None or not torch.is_floating_point(param): + continue + dtype_to_names[param.dtype].append(name) + return dict(dtype_to_names) + + +def _format_dtype_summary(dtype_to_names: dict[torch.dtype, list[str]], limit: int = 8) -> str: + parts = [] + for dtype, names in dtype_to_names.items(): + shown = names[:limit] + suffix = "" if len(names) <= limit else f" ... (+{len(names) - limit} more)" + parts.append(f"{dtype}: {shown}{suffix}") + return "; ".join(parts) + + class FSDPV2Strategy(StrategyBase): """ The strategy for training with PyTorch's Fully Sharded Data Parallel V2. @@ -337,6 +358,27 @@ def _fsdp_init_model(self, model, is_training, shard_size=-1, reshard_after_forw if isinstance(model_to_wrap, FSDPModule): return model + floating_param_dtypes = _collect_floating_param_dtypes(model_to_wrap) + if len(floating_param_dtypes) > 1: + summary = _format_dtype_summary(floating_param_dtypes) + if self.bf16: + self.print( + "[FSDP] Detected mixed floating parameter dtypes before sharding; " + f"casting to bfloat16. {summary}" + ) + model_to_wrap = model_to_wrap.to(torch.bfloat16) + floating_param_dtypes = _collect_floating_param_dtypes(model_to_wrap) + if len(floating_param_dtypes) > 1: + raise RuntimeError( + "Failed to normalize model parameter dtypes before FSDP wrap. " + f"Remaining dtypes: {_format_dtype_summary(floating_param_dtypes)}" + ) + else: + raise RuntimeError( + "Mixed floating parameter dtypes before FSDP wrap with bf16 disabled. " + f"{summary}" + ) + self.report_memory("before FSDP2 wrap model pos2") # this is not sufficient enough, for example, it will only return Qwen2DecoderLayer for qwen2 diff --git a/lightrft/trainer/fast_exp_maker.py b/lightrft/trainer/fast_exp_maker.py index 179dfa07..1a5fa350 100644 --- a/lightrft/trainer/fast_exp_maker.py +++ b/lightrft/trainer/fast_exp_maker.py @@ -1674,17 +1674,32 @@ def _make_experience_list_by_model( Timer.start(' actor_logprob') # Check if we need to compute entropy for high-entropy token filtering need_entropy = hasattr(self.actor, 'high_entropy_token_ratio') and self.actor.high_entropy_token_ratio > 0.0 + # Qwen2.5-Omni may rewrite audio placeholder positions to match the local HF audio tower's + # expected token count. Replay must reuse that aligned layout verbatim, otherwise current + # and old/reference logprobs are scored on different tokenizations of the same sample. + should_capture_aligned_audio_inputs = "audio_values" in self._actor_supported_params for output in outputs: - if need_entropy: - # Request full output to get action_entropy - action_log_probs, model_output = self.actor( + if need_entropy or should_capture_aligned_audio_inputs: + actor_forward_result = self.actor( output.sequences, output.num_actions, output.attention_mask, packed_seq_lens=output.packed_seq_lens, return_output=True, + return_aligned_inputs=should_capture_aligned_audio_inputs, **output.inputs_extra_kwargs ) + if should_capture_aligned_audio_inputs: + # Persist the aligned ids/mask from rollout so policy/ref/critic all replay + # exactly the same post-alignment inputs during PPO updates. + action_log_probs, model_output, aligned_sequences, aligned_attention_mask = actor_forward_result + output.sequences = aligned_sequences + output.attention_mask = aligned_attention_mask + if aligned_attention_mask is not None: + output.total_length = aligned_attention_mask.float().sum(dim=-1) + else: + action_log_probs, model_output = actor_forward_result + output.action_log_probs = action_log_probs # Extract action_entropy if available if "action_entropy" in model_output: diff --git a/lightrft/trainer/spmd_ppo_trainer.py b/lightrft/trainer/spmd_ppo_trainer.py index fd4a91ac..e91fb6cc 100644 --- a/lightrft/trainer/spmd_ppo_trainer.py +++ b/lightrft/trainer/spmd_ppo_trainer.py @@ -275,6 +275,10 @@ def ppo_train(self, global_steps=0): # Currently using this rewritten ppo_train "kl": status["kl"], # KL divergence "act_lr": status["actor_lr"], # actor learning rate } + if "policy/ratio_max" in status: + short_status["rmax"] = status["policy/ratio_max"] + if "policy/logprob_delta_max" in status: + short_status["dmax"] = status["policy/logprob_delta_max"] if "critic_loss" in status: short_status["cri"] = status["critic_loss"] From 9cf3360660d48deb320215c84cf705663a2b40eb Mon Sep 17 00:00:00 2001 From: niuyazhe Date: Fri, 17 Apr 2026 16:25:14 +0800 Subject: [PATCH 07/11] style(nyz): correct format --- examples/r1_aqa/audio_dataset.py | 5 +- examples/r1_aqa/data_preprocess/avqa.py | 34 +++--- .../data_preprocess/clean_audio_dataset.py | 4 +- examples/r1_aqa/eval.py | 101 +++++++++--------- examples/r1_aqa/reward_models_utils.py | 7 +- examples/r1_aqa/train_colocate.py | 89 ++++++++------- lightrft/models/actor_al.py | 34 ++---- lightrft/models/loss.py | 16 ++- lightrft/models/tests/test_actor_al.py | 6 +- lightrft/strategy/test_fake_strategy.py | 1 - 10 files changed, 148 insertions(+), 149 deletions(-) diff --git a/examples/r1_aqa/audio_dataset.py b/examples/r1_aqa/audio_dataset.py index 777cd316..69eca3e1 100644 --- a/examples/r1_aqa/audio_dataset.py +++ b/examples/r1_aqa/audio_dataset.py @@ -24,11 +24,11 @@ import librosa from torch.utils.data import Dataset - # ============================================================================ # Audio Loading # ============================================================================ + def load_audio(audio_path: str, sr: int = 16000) -> Tuple[Any, int]: """Load an audio file as ``(waveform, sampling_rate)``.""" return librosa.load(audio_path, sr=sr) @@ -66,7 +66,6 @@ class AudioPromptDataset(Dataset): - ``audio_payload`` is kept as raw waveform + sampling rate for rollout-side processing - ``reference`` and ``label`` are passed through to reward computation """ - def __init__( self, dataset, @@ -120,7 +119,7 @@ def __getitem__(self, idx: int) -> Tuple[str, Any, str, str]: except Exception as exc: self.strategy.print(f"[WARNING] Chat template failed for idx {idx}: {exc}") prompt_text = self._extract_text_from_messages(prompt_messages) - + prompt_text = prompt_text.replace('', ' ') # ---- 3. Load audio ---- diff --git a/examples/r1_aqa/data_preprocess/avqa.py b/examples/r1_aqa/data_preprocess/avqa.py index 4913341a..d4aea003 100644 --- a/examples/r1_aqa/data_preprocess/avqa.py +++ b/examples/r1_aqa/data_preprocess/avqa.py @@ -47,11 +47,11 @@ import datasets - # --------------------------------------------------------------------------- # Prompt template — faithfully ported from R1-AQA src/dataset/dataset.py # --------------------------------------------------------------------------- + def build_prompt_and_solution( obj: Dict[str, Any], audio_dir: Optional[str] = None, @@ -89,21 +89,23 @@ def build_prompt_and_solution( "and final answer in ." ) else: - question_template = ( - f"{question_text} {choice_str} " - "Output the final answer in ." - ) + question_template = (f"{question_text} {choice_str} " + "Output the final answer in .") # Chat-format prompt with audio content type (Qwen2-Audio format) - prompt = [ - { - "role": "user", - "content": [ - {"type": "audio", "audio_url": audio_path}, - {"type": "text", "text": question_template}, - ], - } - ] + prompt = [{ + "role": "user", + "content": [ + { + "type": "audio", + "audio_url": audio_path + }, + { + "type": "text", + "text": question_template + }, + ], + }] # Correct answer string answer_str = multi_choice[answer_idx] @@ -238,9 +240,7 @@ def preprocess_avqa( if __name__ == "__main__": - parser = argparse.ArgumentParser( - description="Preprocess AVQA dataset (R1-AQA format) for LightRFT training" - ) + parser = argparse.ArgumentParser(description="Preprocess AVQA dataset (R1-AQA format) for LightRFT training") parser.add_argument( "--input_jsonl", required=True, diff --git a/examples/r1_aqa/data_preprocess/clean_audio_dataset.py b/examples/r1_aqa/data_preprocess/clean_audio_dataset.py index 324b79e0..47fdfcdd 100644 --- a/examples/r1_aqa/data_preprocess/clean_audio_dataset.py +++ b/examples/r1_aqa/data_preprocess/clean_audio_dataset.py @@ -60,9 +60,7 @@ def clean_split(dataset, verify_decode: bool) -> tuple[list[int], list[dict[str, if verify_decode: ok, error = can_decode_audio(audio_path) if not ok: - dropped_rows.append( - {"index": idx, "audio_path": audio_path, "reason": "decode_error", "error": error} - ) + dropped_rows.append({"index": idx, "audio_path": audio_path, "reason": "decode_error", "error": error}) continue keep_indices.append(idx) diff --git a/examples/r1_aqa/eval.py b/examples/r1_aqa/eval.py index 0eae4270..319884f1 100644 --- a/examples/r1_aqa/eval.py +++ b/examples/r1_aqa/eval.py @@ -50,11 +50,11 @@ infer_audio_model_type, ) - # --------------------------------------------------------------------------- # Message building (MMAU and MMAR share question/choices + format) # --------------------------------------------------------------------------- + def build_message(obj_dict: Dict, audio_dir: Optional[str] = None) -> list: """ Build the chat message for MMAU or MMAR evaluation. @@ -66,10 +66,8 @@ def build_message(obj_dict: Dict, audio_dir: Optional[str] = None) -> list: :return: Chat messages list. """ choice_str = f"Please choose the answer from the following options: {obj_dict['choices']}." - question_template = ( - f"{obj_dict['question']} {choice_str} " - "Output the final answer in ." - ) + question_template = (f"{obj_dict['question']} {choice_str} " + "Output the final answer in .") # MMAU uses audio_id; MMAR uses audio_path (e.g. ./audio/xxx.wav) raw_path = obj_dict.get("audio_path") or obj_dict.get("audio_id", "") @@ -79,15 +77,19 @@ def build_message(obj_dict: Dict, audio_dir: Optional[str] = None) -> list: else: audio_path = raw_path - message = [ - { - "role": "user", - "content": [ - {"type": "audio", "audio_url": audio_path}, - {"type": "text", "text": question_template}, - ], - } - ] + message = [{ + "role": "user", + "content": [ + { + "type": "audio", + "audio_url": audio_path + }, + { + "type": "text", + "text": question_template + }, + ], + }] return message @@ -109,6 +111,7 @@ def extract_answer(output_str: str) -> str: # MMAR official evaluation uses token-based string_match; optional local use # --------------------------------------------------------------------------- + def _string_match_mmar(answer: str, prediction: str, choices: List[str]) -> bool: """ MMAR evaluation.py string_match: tokenize and check answer ⊆ prediction @@ -135,6 +138,7 @@ def tokenize(text: str): # Inference # --------------------------------------------------------------------------- + def run_inference_hf( model_path: str, data: List[Dict], @@ -185,9 +189,7 @@ def get_audio_path(sample: Dict) -> str: print(f"Processing {i + 1}/{len(data)}...") message = build_message(sample, audio_dir) - text = processor.apply_chat_template( - message, tokenize=False, add_generation_prompt=True - ) + text = processor.apply_chat_template(message, tokenize=False, add_generation_prompt=True) audio_path = get_audio_path(sample) try: @@ -209,12 +211,10 @@ def get_audio_path(sample: Dict) -> str: "do_sample": False, } if model_type == AUDIO_MODEL_TYPE_QWEN2_5_OMNI: - generation_kwargs.update( - { - "generation_mode": "text", - "thinker_max_new_tokens": max_new_tokens, - } - ) + generation_kwargs.update({ + "generation_mode": "text", + "thinker_max_new_tokens": max_new_tokens, + }) else: generation_kwargs["max_new_tokens"] = max_new_tokens @@ -271,16 +271,16 @@ def get_audio_path(sample: Dict) -> str: all_inputs = [] for sample in data: message = build_message(sample, audio_dir) - text = processor.apply_chat_template( - message, tokenize=False, add_generation_prompt=True - ) + text = processor.apply_chat_template(message, tokenize=False, add_generation_prompt=True) audio_path = get_audio_path(sample) try: audio, _ = librosa.load(audio_path, sr=sr) all_inputs.append({ "prompt": text, - "multi_modal_data": {"audio": [(audio, sr)]}, + "multi_modal_data": { + "audio": [(audio, sr)] + }, }) except Exception as e: print(f"[WARNING] Audio load failed for {audio_path}: {e}") @@ -296,6 +296,7 @@ def get_audio_path(sample: Dict) -> str: # Main # --------------------------------------------------------------------------- + def load_data(data_file: str, benchmark: str) -> List[Dict]: """Load MMAU (JSON array) or MMAR (JSONL) data.""" with open(data_file, "r", encoding="utf-8") as f: @@ -307,27 +308,19 @@ def load_data(data_file: str, benchmark: str) -> List[Dict]: def main(): - parser = argparse.ArgumentParser( - description="MMAU and MMAR evaluation for R1-AQA" + parser = argparse.ArgumentParser(description="MMAU and MMAR evaluation for R1-AQA") + parser.add_argument( + "--benchmark", type=str, required=True, choices=["mmau", "mmar"], help="Benchmark: mmau or mmar" + ) + parser.add_argument("--model_path", type=str, required=True, help="Path to the trained model (HF format)") + parser.add_argument( + "--data_file", type=str, required=True, help="Path to benchmark data (MMAU: .json, MMAR: .jsonl)" ) - parser.add_argument("--benchmark", type=str, required=True, - choices=["mmau", "mmar"], - help="Benchmark: mmau or mmar") - parser.add_argument("--model_path", type=str, required=True, - help="Path to the trained model (HF format)") - parser.add_argument("--data_file", type=str, required=True, - help="Path to benchmark data (MMAU: .json, MMAR: .jsonl)") - parser.add_argument("--audio_dir", type=str, default=None, - help="Base directory for audio files") - parser.add_argument("--out_file", type=str, required=True, - help="Output file for evaluation results") - parser.add_argument("--batch_size", type=int, default=32, - help="Batch size for inference") - parser.add_argument("--max_new_tokens", type=int, default=1024, - help="Maximum new tokens to generate") - parser.add_argument("--engine", type=str, default="hf", - choices=["hf", "vllm"], - help="Inference engine: hf or vllm") + parser.add_argument("--audio_dir", type=str, default=None, help="Base directory for audio files") + parser.add_argument("--out_file", type=str, required=True, help="Output file for evaluation results") + parser.add_argument("--batch_size", type=int, default=32, help="Batch size for inference") + parser.add_argument("--max_new_tokens", type=int, default=1024, help="Maximum new tokens to generate") + parser.add_argument("--engine", type=str, default="hf", choices=["hf", "vllm"], help="Inference engine: hf or vllm") args = parser.parse_args() # Load data @@ -339,13 +332,19 @@ def main(): print(f"Running inference with {args.engine} engine...") if args.engine == "vllm": all_outputs = run_inference_vllm( - args.model_path, data, args.audio_dir, - args.batch_size, args.max_new_tokens, + args.model_path, + data, + args.audio_dir, + args.batch_size, + args.max_new_tokens, ) else: all_outputs = run_inference_hf( - args.model_path, data, args.audio_dir, - args.batch_size, args.max_new_tokens, + args.model_path, + data, + args.audio_dir, + args.batch_size, + args.max_new_tokens, ) # Prediction key per benchmark diff --git a/examples/r1_aqa/reward_models_utils.py b/examples/r1_aqa/reward_models_utils.py index 96a59815..472f0bb4 100644 --- a/examples/r1_aqa/reward_models_utils.py +++ b/examples/r1_aqa/reward_models_utils.py @@ -27,11 +27,11 @@ import torch - # ============================================================================ # R1-AQA Accuracy Reward (ported from src/utils/rewards.py) # ============================================================================ + def accuracy_reward_fn(content: str, solution: str) -> float: """ R1-AQA accuracy reward function. @@ -90,6 +90,7 @@ def accuracy_reward_fn(content: str, solution: str) -> float: # R1-AQA Format Reward (ported from src/utils/rewards.py) # ============================================================================ + def format_reward_fn(content: str, enable_think: bool = False) -> float: """ R1-AQA format reward function. @@ -118,6 +119,7 @@ def format_reward_fn(content: str, enable_think: bool = False) -> float: # Combined Reward (per-sample) # ============================================================================ + def avqa_combined_reward_fn( sol: str, gt: str, @@ -156,10 +158,13 @@ def clean_solution(sol: str) -> str: # Pattern matches text between <|im_start|>assistant and <|im_end|> match = re.search(r"<\|im_start\|>assistant(.*?)<\|im_end\|>", sol, re.DOTALL) return match.group(1).strip() if match else sol.strip() + + # ============================================================================ # Reward Function (LightRFT interface — called by the trainer) # ============================================================================ + def reward_fn( model_reward_list: List[torch.Tensor], labels: Sequence[str], diff --git a/examples/r1_aqa/train_colocate.py b/examples/r1_aqa/train_colocate.py index 7cd96493..e88201af 100644 --- a/examples/r1_aqa/train_colocate.py +++ b/examples/r1_aqa/train_colocate.py @@ -48,9 +48,7 @@ def _validate_inference_engine(args) -> None: model_type = infer_audio_model_type(args.pretrain) if model_type == AUDIO_MODEL_TYPE_QWEN2_5_OMNI and args.engine_type == "sglang": - raise ValueError( - "Qwen2.5-Omni audio demo currently requires --engine_type vllm; sglang is not supported." - ) + raise ValueError("Qwen2.5-Omni audio demo currently requires --engine_type vllm; sglang is not supported.") def train(args): @@ -145,9 +143,7 @@ def train(args): ) if args.fsdp: shard_size = ( - args.initial_model_shard_size - if args.initial_model_shard_size is not None - else strategy.world_size + args.initial_model_shard_size if args.initial_model_shard_size is not None else strategy.world_size ) initial_model = strategy.prepare_model(initial_model, is_training=False, shard_size=shard_size) strategy.offload_model(initial_model) @@ -209,19 +205,28 @@ def train(args): if eval_data_path: strategy.print(f"Loading evaluation dataset from {eval_data_path}") eval_data = blending_datasets( - eval_data_path, "1.0", strategy, args.seed, - return_eval=False, train_split=args.eval_split, + eval_data_path, + "1.0", + strategy, + args.seed, + return_eval=False, + train_split=args.eval_split, ) if len(eval_data) > 0: eval_data = eval_data.select(range(min(args.max_eval_samples, len(eval_data)))) eval_dataset = AudioPromptDataset( - eval_data, tokenizer, processor, args.prompt_max_len, strategy, + eval_data, + tokenizer, + processor, + args.prompt_max_len, + strategy, input_template=args.input_template, ) eval_dataloader = strategy.setup_dataloader( eval_dataset, args.rollout_batch_size // strategy.world_size, - False, False, + False, + False, collate_fn=eval_dataset.collate_fn, ) strategy.print(f"Evaluation dataset: {len(eval_dataset)} samples") @@ -233,7 +238,8 @@ def train(args): prompts_dataloader = strategy.setup_dataloader( prompts_dataset, args.rollout_batch_size // strategy.world_size, - True, True, + True, + True, collate_fn=prompts_dataset.collate_fn, ) @@ -255,9 +261,7 @@ def train(args): (critic, critic_optim, critic_scheduler), reward_models, initial_model, - ) = strategy.prepare_models_and_optimizers( - actor, critic, reward_models, initial_model, args, max_steps - ) + ) = strategy.prepare_models_and_optimizers(actor, critic, reward_models, initial_model, args, max_steps) strategy.print(reward_models) @@ -269,8 +273,10 @@ def train(args): consumed_samples = 0 if args.load_checkpoint and os.path.exists(os.path.join(args.ckpt_path, "_actor")): _, states = strategy.load_ckpt( - actor.model, os.path.join(args.ckpt_path, "_actor"), - optimizer=actor_optim, scheduler=actor_scheduler, + actor.model, + os.path.join(args.ckpt_path, "_actor"), + optimizer=actor_optim, + scheduler=actor_scheduler, ) consumed_samples = states["consumed_samples"] strategy.print(f"Loaded checkpoint: {args.ckpt_path}, consumed_samples: {consumed_samples}") @@ -361,10 +367,13 @@ def train(args): parser = argparse.ArgumentParser() # Engine - parser.add_argument("--engine_type", type=str, default="vllm", - help="Inference engine: vllm or sglang") - parser.add_argument("--text_only", action="store_true", default=False, - help="Text-only mode (no multimodal). Default False for audio tasks.") + parser.add_argument("--engine_type", type=str, default="vllm", help="Inference engine: vllm or sglang") + parser.add_argument( + "--text_only", + action="store_true", + default=False, + help="Text-only mode (no multimodal). Default False for audio tasks." + ) # Checkpoint parser.add_argument("--save_path", type=str, default="./ckpt") @@ -396,10 +405,10 @@ def train(args): parser.add_argument("--micro_rollout_batch_size", type=int, default=8) parser.add_argument("--max_epochs", type=int, default=1) # R1-AQA default: max_prompt_length=512 - parser.add_argument("--prompt_max_len", type=int, default=512, - help="Max tokens for each prompt (R1-AQA default: 512)") - parser.add_argument("--generate_max_len", type=int, default=1024, - help="Max tokens to generate") + parser.add_argument( + "--prompt_max_len", type=int, default=512, help="Max tokens for each prompt (R1-AQA default: 512)" + ) + parser.add_argument("--generate_max_len", type=int, default=1024, help="Max tokens to generate") parser.add_argument("--max_len", type=int, default=None) parser.add_argument("--max_samples", type=int, default=1000000) parser.add_argument("--max_norm", type=float, default=1.0) @@ -422,8 +431,12 @@ def train(args): parser.add_argument("--freeze_prefix", action="store_true", default=False) parser.add_argument("--freezing_actor_steps", type=int, default=-1) # R1-AQA default: num_generations=8 - parser.add_argument("--n_samples_per_prompt", type=int, default=8, - help="Number of responses per prompt in GRPO (R1-AQA default: 8)") + parser.add_argument( + "--n_samples_per_prompt", + type=int, + default=8, + help="Number of responses per prompt in GRPO (R1-AQA default: 8)" + ) parser.add_argument("--save_value_network", action="store_true", default=False) # R1-AQA default: lr not explicitly set, using 1e-6 as reasonable default parser.add_argument("--actor_learning_rate", type=float, default=1e-6) @@ -431,9 +444,9 @@ def train(args): parser.add_argument("--lr_warmup_ratio", type=float, default=0.03) parser.add_argument("--kl_target", type=float, default=None) parser.add_argument("--init_kl_coef", type=float, default=0.01) - parser.add_argument("--kl_estimator", type=str, default="k3", - choices=["k1", "k2", "k3"], - help="GRPO uses k3 as KL estimator") + parser.add_argument( + "--kl_estimator", type=str, default="k3", choices=["k1", "k2", "k3"], help="GRPO uses k3 as KL estimator" + ) parser.add_argument("--adam_betas", type=float, nargs=2, default=(0.9, 0.95)) # Reward/Advantage Norm/Clip @@ -467,10 +480,13 @@ def train(args): parser.add_argument("--initial_model_shard_size", type=int, default=None) # Advantage estimator - parser.add_argument("--advantage_estimator", type=str, - choices=["gae", "reinforce", "rloo", "reinforce_baseline", "group_norm", "cpgd", "reinforce++"], - default="group_norm", - help="Advantage estimation method. R1-AQA uses GRPO = group_norm") + parser.add_argument( + "--advantage_estimator", + type=str, + choices=["gae", "reinforce", "rloo", "reinforce_baseline", "group_norm", "cpgd", "reinforce++"], + default="group_norm", + help="Advantage estimation method. R1-AQA uses GRPO = group_norm" + ) parser.add_argument("--use_kl_loss", action="store_true", default=False) # LoRA @@ -510,8 +526,7 @@ def train(args): parser.add_argument("--wandb_org", type=str, default=None) parser.add_argument("--wandb_group", type=str, default=None) parser.add_argument("--wandb_project", type=str, default="lightrft_r1_aqa") - parser.add_argument("--wandb_run_name", type=str, - default="r1_aqa_%s" % datetime.now().strftime("%m%dT%H:%M")) + parser.add_argument("--wandb_run_name", type=str, default="r1_aqa_%s" % datetime.now().strftime("%m%dT%H:%M")) # TensorBoard parser.add_argument("--use_tensorboard", type=str, default=None) @@ -539,9 +554,7 @@ def train(args): args.critic_pretrain = args.pretrain if args.advantage_estimator in ["rloo", "reinforce_baseline", "group_norm"]: - assert args.n_samples_per_prompt > 1, ( - f"{args.advantage_estimator} requires n_samples_per_prompt > 1" - ) + assert args.n_samples_per_prompt > 1, (f"{args.advantage_estimator} requires n_samples_per_prompt > 1") if args.use_kl_loss: if args.kl_estimator not in ["k2", "k3"]: diff --git a/lightrft/models/actor_al.py b/lightrft/models/actor_al.py index b059013a..09566cfd 100644 --- a/lightrft/models/actor_al.py +++ b/lightrft/models/actor_al.py @@ -24,7 +24,6 @@ reset_position_ids, ) - AUDIO_MODEL_TYPE_QWEN2_AUDIO = "qwen2_audio" AUDIO_MODEL_TYPE_QWEN2_5_OMNI = "qwen2_5_omni" @@ -155,10 +154,8 @@ def create_audio_processor( ) if processor is not None and print_fn is not None: - print_fn( - f"[WARN] AutoProcessor loaded {type(processor).__name__}, " - f"re-loading as {processor_cls.__name__}" - ) + print_fn(f"[WARN] AutoProcessor loaded {type(processor).__name__}, " + f"re-loading as {processor_cls.__name__}") return processor_cls.from_pretrained( source, @@ -619,19 +616,14 @@ def forward( sequences=sequences, audio_token_id=audio_token_id, ) - original_audio_token_counts = ( - (sequences == audio_token_id).sum(dim=1) - if audio_token_id is not None - else None - ) + original_audio_token_counts = ((sequences == audio_token_id).sum(dim=1) + if audio_token_id is not None else None) expected_audio_token_counts = self._infer_audio_output_token_counts( forward_model, feature_attention_mask, ) if ( - not self.packing_samples - and audio_token_id is not None - and expected_audio_token_counts is not None + not self.packing_samples and audio_token_id is not None and expected_audio_token_counts is not None ): sequences, attention_mask = self._align_audio_placeholder_counts( sequences=sequences, @@ -641,14 +633,10 @@ def forward( pad_token_id=pad_token_id, num_actions=(num_actions if isinstance(num_actions, int) else None), ) - actual_audio_token_counts = ( - (sequences == audio_token_id).sum(dim=1) - if audio_token_id is not None - else None - ) + actual_audio_token_counts = ((sequences == audio_token_id).sum(dim=1) + if audio_token_id is not None else None) if ( - actual_audio_token_counts is not None - and expected_audio_token_counts is not None + actual_audio_token_counts is not None and expected_audio_token_counts is not None and torch.any(actual_audio_token_counts > expected_audio_token_counts) ): raise RuntimeError( @@ -662,7 +650,7 @@ def forward( f"[ActorAL][rank={rank}] sequences={tuple(sequences.shape)} " f"audio_values={tuple(input_features.shape)} " f"feature_attention_mask={tuple(feature_attention_mask.shape)} " - f"audio_token_count={actual_audio_token_counts.tolist() if actual_audio_token_counts is not None else None} " + f"audio_token_count={actual_audio_token_counts.tolist() if actual_audio_token_counts is not None else None} " # noqa f"original_audio_token_count=" f"{original_audio_token_counts.tolist() if original_audio_token_counts is not None else None} " f"feature_len={feature_attention_mask.sum(dim=1).tolist()} " @@ -834,7 +822,7 @@ def _align_audio_placeholder_counts( if expected_count == 0: row_tokens = active_tokens.clone() else: - new_audio_block = active_tokens.new_full((expected_count,), audio_token_id) + new_audio_block = active_tokens.new_full((expected_count, ), audio_token_id) row_tokens = torch.cat((new_audio_block, active_tokens), dim=0) else: block_start = int(audio_positions[0].item()) @@ -848,7 +836,7 @@ def _align_audio_placeholder_counts( else: prompt_prefix = active_tokens[:block_start] prompt_suffix = active_tokens[block_end + 1:] - new_audio_block = active_tokens.new_full((expected_count,), audio_token_id) + new_audio_block = active_tokens.new_full((expected_count, ), audio_token_id) row_tokens = torch.cat((prompt_prefix, new_audio_block, prompt_suffix), dim=0) # If the response contains stray audio placeholders, keep the sequence diff --git a/lightrft/models/loss.py b/lightrft/models/loss.py index 2faccc10..62c4671b 100644 --- a/lightrft/models/loss.py +++ b/lightrft/models/loss.py @@ -225,15 +225,13 @@ def _update_last_stats( denom = max(valid_token_count, 1.0) clipped_high = (ratio > 1 + self.clip_eps).masked_select(stats_mask) clipped_low = (ratio < 1 - self.clip_eps).masked_select(stats_mask) - stats.update( - { - "policy/clipfrac_high": float(clipped_high.numel() / denom), - "policy/clipfrac_low": float(clipped_low.numel() / denom), - **self._stats_over_mask(log_probs, stats_mask, "policy/logprob"), - **self._stats_over_mask(old_log_probs, stats_mask, "policy/old_logprob"), - **self._stats_over_mask(ratio, stats_mask, "policy/ratio"), - } - ) + stats.update({ + "policy/clipfrac_high": float(clipped_high.numel() / denom), + "policy/clipfrac_low": float(clipped_low.numel() / denom), + **self._stats_over_mask(log_probs, stats_mask, "policy/logprob"), + **self._stats_over_mask(old_log_probs, stats_mask, "policy/old_logprob"), + **self._stats_over_mask(ratio, stats_mask, "policy/ratio"), + }) self._last_stats = stats diff --git a/lightrft/models/tests/test_actor_al.py b/lightrft/models/tests/test_actor_al.py index 7a965f03..6624d74e 100644 --- a/lightrft/models/tests/test_actor_al.py +++ b/lightrft/models/tests/test_actor_al.py @@ -384,9 +384,9 @@ def test_align_audio_placeholders_is_batch_invariant_to_num_actions(self, mock_o actor = self._make_actor(mock_omni_model) audio_token_id = mock_omni_model.thinker.config.audio_token_id - sequences = torch.tensor( - [[11, 12, 13, 14, audio_token_id, audio_token_id, audio_token_id, audio_token_id, 31, 41, 42, 43]] - ) + sequences = torch.tensor([[ + 11, 12, 13, 14, audio_token_id, audio_token_id, audio_token_id, audio_token_id, 31, 41, 42, 43 + ]]) attention_mask = torch.ones_like(sequences) expected_counts = torch.tensor([5]) diff --git a/lightrft/strategy/test_fake_strategy.py b/lightrft/strategy/test_fake_strategy.py index d6736f01..519dcc9a 100644 --- a/lightrft/strategy/test_fake_strategy.py +++ b/lightrft/strategy/test_fake_strategy.py @@ -325,7 +325,6 @@ def __init__(self): class TestBuildMultimodalInputs(unittest.TestCase): """Focused tests for engine-specific multimodal payload construction.""" - def test_keeps_waveform_audio_for_vllm(self): """vLLM audio inputs should stay waveform-native instead of being WAV bytes.""" audio = (np.asarray([0.1, -0.2, 0.3], dtype=np.float32), 16000) From 3b93d57f9ac37f07adfafb728367b383f247a58c Mon Sep 17 00:00:00 2001 From: niuyazhe Date: Fri, 17 Apr 2026 20:27:55 +0800 Subject: [PATCH 08/11] fix(nyz): fix qwen2 audio compatibility bugs --- examples/r1_aqa/audio_dataset.py | 27 ------------------------ examples/r1_aqa/data_preprocess/avqa.py | 5 +++-- examples/r1_aqa/eval.py | 4 +++- examples/r1_aqa/reward_models_utils.py | 3 --- lightrft/strategy/fsdp/fsdpv2.py | 26 +++++++++++------------ lightrft/strategy/vllm_utils/__init__.py | 6 ++++++ 6 files changed, 25 insertions(+), 46 deletions(-) diff --git a/examples/r1_aqa/audio_dataset.py b/examples/r1_aqa/audio_dataset.py index 69eca3e1..f678e2ba 100644 --- a/examples/r1_aqa/audio_dataset.py +++ b/examples/r1_aqa/audio_dataset.py @@ -16,7 +16,6 @@ from __future__ import annotations -import copy import json import os from typing import Any, Dict, List, Optional, Tuple @@ -34,29 +33,6 @@ def load_audio(audio_path: str, sr: int = 16000) -> Tuple[Any, int]: return librosa.load(audio_path, sr=sr) -def sanitize_qwen2_audio_messages(messages) -> List[Dict[str, Any]]: - """ - Remove misleading keys before applying the Qwen2-Audio chat template. - - Some parquet rows keep ``audio_url=None`` on text segments, and the upstream - template interprets the presence of that key as an audio placeholder. - """ - sanitized = copy.deepcopy(messages) - for message in sanitized: - content = message.get("content") - if not isinstance(content, list): - continue - for segment in content: - if not isinstance(segment, dict): - continue - segment_type = segment.get("type") - if segment_type == "text": - segment.pop("audio_url", None) - elif segment_type == "audio": - segment.pop("text", None) - return sanitized - - class AudioPromptDataset(Dataset): """ PyTorch dataset for the R1-AQA audio prompt format. @@ -109,7 +85,6 @@ def __getitem__(self, idx: int) -> Tuple[str, Any, str, str]: prompt_messages = [{"role": "user", "content": prompt_messages}] # ---- 2. Render via processor's chat template ---- - prompt_messages = sanitize_qwen2_audio_messages(prompt_messages) try: prompt_text = self.processor.apply_chat_template( prompt_messages, @@ -120,8 +95,6 @@ def __getitem__(self, idx: int) -> Tuple[str, Any, str, str]: self.strategy.print(f"[WARNING] Chat template failed for idx {idx}: {exc}") prompt_text = self._extract_text_from_messages(prompt_messages) - prompt_text = prompt_text.replace('', ' ') - # ---- 3. Load audio ---- audio_path = data.get(self.audio_path_key, "") audio_data = None diff --git a/examples/r1_aqa/data_preprocess/avqa.py b/examples/r1_aqa/data_preprocess/avqa.py index d4aea003..c5a38bc1 100644 --- a/examples/r1_aqa/data_preprocess/avqa.py +++ b/examples/r1_aqa/data_preprocess/avqa.py @@ -82,15 +82,16 @@ def build_prompt_and_solution( # Build question template (matches R1-AQA) choice_str = f"Please choose the answer from the following options: {multi_choice}." + # There should be a space between and if enable_think: question_template = ( f"{question_text} {choice_str} " "Output the thinking process in " - "and final answer in ." + "and final answer in ." ) else: question_template = (f"{question_text} {choice_str} " - "Output the final answer in .") + "Output the final answer in .") # Chat-format prompt with audio content type (Qwen2-Audio format) prompt = [{ diff --git a/examples/r1_aqa/eval.py b/examples/r1_aqa/eval.py index 319884f1..6911e178 100644 --- a/examples/r1_aqa/eval.py +++ b/examples/r1_aqa/eval.py @@ -279,7 +279,9 @@ def get_audio_path(sample: Dict) -> str: all_inputs.append({ "prompt": text, "multi_modal_data": { - "audio": [(audio, sr)] + # vLLM Qwen2-Audio expects a single audio sample as a scalar + # ``(waveform, sampling_rate)`` payload, not ``[(waveform, sr)]``. + "audio": (audio, sr) }, }) except Exception as e: diff --git a/examples/r1_aqa/reward_models_utils.py b/examples/r1_aqa/reward_models_utils.py index 472f0bb4..f705e0de 100644 --- a/examples/r1_aqa/reward_models_utils.py +++ b/examples/r1_aqa/reward_models_utils.py @@ -77,9 +77,6 @@ def accuracy_reward_fn(content: str, solution: str) -> float: sol_match = re.search(r"(.*?)", solution) ground_truth = sol_match.group(1).strip() if sol_match else solution.strip() student_answer = content.strip() - import torch.distributed as dist - if dist.is_initialized() and dist.get_rank() == 0: - print(f"student_answer: {student_answer}, ground_truth: {ground_truth}") if student_answer == ground_truth: reward = 1.0 diff --git a/lightrft/strategy/fsdp/fsdpv2.py b/lightrft/strategy/fsdp/fsdpv2.py index 9cba06f9..9f9143d2 100755 --- a/lightrft/strategy/fsdp/fsdpv2.py +++ b/lightrft/strategy/fsdp/fsdpv2.py @@ -74,8 +74,8 @@ "LlamaDecoderLayer", # for DeepSeek-R1-Distill-Llama-70B "DeepseekDecoderLayer", ] - -vit_transformer_cls_names = [ +# multi-modal modules +mm_module_cls_names = [ "Qwen2VLVisionBlock", "Qwen2_5_VLVisionBlock", "Qwen2_5OmniVisionEncoder", @@ -394,25 +394,25 @@ def _fsdp_init_model(self, model, is_training, shard_size=-1, reshard_after_forw # Note:if we have mixed multi-modal data across DP ranks # (e.g. some ranks pure text, other ranks contains images) # we either keep vision model in full state, or keep it in FSDP's root module. - # below we keep vit in root module to avoid stuck - for cls_name in vit_transformer_cls_names: + # below we keep multi-modal modules in root module to avoid stuck + for cls_name in mm_module_cls_names: if cls_name in transformer_cls_names_to_wrap: transformer_cls_names_to_wrap.remove(cls_name) transformer_cls_to_wrap = list() # noqa - vit_transformer_cls = list() # noqa + mm_module_cls = list() # noqa for layer_class in transformer_cls_names_to_wrap: transformer_cls = get_module_class_from_name(model_to_wrap, layer_class) if transformer_cls is not None: transformer_cls_to_wrap.append(transformer_cls) - # Note: in this way, we keep vit in full state by passing no_shard_mesh - # this is less memory efficient compared to keep vit in root module - # vit_transformer_cls = list() - # for layer_class in vit_transformer_cls_names: + # Note: in this way, we keep multi-modal modules in full state by passing no_shard_mesh + # this is less memory efficient compared to keep multi-modal modules in root module + # mm_module_cls = list() + # for layer_class in mm_module_cls_names: # transformer_cls = get_module_class_from_name(model_to_wrap, layer_class) # if transformer_cls is not None: - # vit_transformer_cls.append(transformer_cls) + # mm_module_cls.append(transformer_cls) if len(transformer_cls_to_wrap) == 0: self.print("len(transformer_cls_to_wrap)=0", model_to_wrap) @@ -444,13 +444,13 @@ def _fsdp_init_model(self, model, is_training, shard_size=-1, reshard_after_forw # fsdp_kwargs_no_shard = fsdp_kwargs.copy() # fsdp_kwargs_no_shard['mesh'] = no_shard_mesh # fsdp_kwargs_no_shard['reshard_after_forward'] = True - # fsdp_kwargs_vit = fsdp_kwargs_no_shard if self.no_shard_vit else fsdp_kwargs + # fsdp_kwargs_mm = fsdp_kwargs_no_shard if self.no_shard_mm else fsdp_kwargs for cls_to_wrap in transformer_cls_to_wrap: for module in model_to_wrap.modules(): if isinstance(module, cls_to_wrap): - # if cls_to_wrap in vit_transformer_cls: - # fully_shard(module, **fsdp_kwargs_vit) + # if cls_to_wrap in mm_module_cls: + # fully_shard(module, **fsdp_kwargs_mm) fully_shard(module, **fsdp_kwargs) if not self.args.fused_linear_logprob: diff --git a/lightrft/strategy/vllm_utils/__init__.py b/lightrft/strategy/vllm_utils/__init__.py index fc9e67e6..f30aba6a 100644 --- a/lightrft/strategy/vllm_utils/__init__.py +++ b/lightrft/strategy/vllm_utils/__init__.py @@ -83,6 +83,7 @@ def get_vllm_engine_for_rollout(args: Any): mem_util=args.engine_mem_util, max_model_len=args.prompt_max_len + args.generate_max_len, enable_sleep=args.enable_engine_sleep, + seed=getattr(args, "seed", 42), **kwargs, ) return vllm_engine @@ -95,6 +96,7 @@ def get_vllm_engine( mem_util: float = 0.5, max_model_len: int = 4096, enable_sleep: bool = True, + seed: int = 42, **kwargs: Any ): """ @@ -116,6 +118,9 @@ def get_vllm_engine( :type max_model_len: int :param enable_sleep: Whether to enable sleep mode for memory efficiency. Defaults to True. :type enable_sleep: bool + :param seed: Random seed forwarded to vLLM. Required by newer vLLM releases when + using the external launcher backend so workers share the same sampling config. + :type seed: int :param kwargs: Additional keyword arguments passed to the LLM constructor. :type kwargs: Any @@ -160,6 +165,7 @@ def get_vllm_engine( worker_cls="lightrft.strategy.vllm_utils.vllm_worker_wrap_no_ray.WorkerWrap", enable_sleep_mode=enable_sleep, max_model_len=max_model_len, + seed=seed, # enforce_eager=True, **kwargs, ) From f17fac757916a508f3391945ed9d518e00b4fb67 Mon Sep 17 00:00:00 2001 From: niuyazhe Date: Sat, 18 Apr 2026 10:36:43 +0800 Subject: [PATCH 09/11] polish(nyz): add qwen2-audio sglang pipeline --- examples/r1_aqa/audio_dataset.py | 17 +++++++ examples/r1_aqa/reward_models_utils.py | 6 --- .../r1_aqa/run_grpo_r1_aqa_qwen2_audio_7b.sh | 4 +- examples/r1_aqa/train_colocate.py | 18 ++++--- lightrft/trainer/ppo_trainer_vl.py | 49 +++++++++++++++++-- lightrft/trainer/spmd_ppo_trainer.py | 24 +++++---- 6 files changed, 90 insertions(+), 28 deletions(-) diff --git a/examples/r1_aqa/audio_dataset.py b/examples/r1_aqa/audio_dataset.py index f678e2ba..17dcd038 100644 --- a/examples/r1_aqa/audio_dataset.py +++ b/examples/r1_aqa/audio_dataset.py @@ -83,6 +83,7 @@ def __getitem__(self, idx: int) -> Tuple[str, Any, str, str]: prompt_messages = json.loads(prompt_messages) except (json.JSONDecodeError, TypeError): prompt_messages = [{"role": "user", "content": prompt_messages}] + prompt_messages = self._drop_none_fields(prompt_messages) # ---- 2. Render via processor's chat template ---- try: @@ -132,3 +133,19 @@ def _extract_text_from_messages(messages) -> str: if isinstance(segment, dict) and segment.get("type") == "text": texts.append(segment.get("text", "")) return " ".join(texts) + + @staticmethod + def _drop_none_fields(obj): + """ + Remove ``None`` values from nested prompt content before chat templating. + + The parquet loader materializes a union of nested content keys, so a text + block may arrive as ``{"type": "text", "text": "...", "audio_url": None}``. + Qwen2-Audio's default chat template checks key existence instead of value, + which would misclassify that text block as a second audio placeholder. + """ + if isinstance(obj, list): + return [AudioPromptDataset._drop_none_fields(item) for item in obj] + if isinstance(obj, dict): + return {key: AudioPromptDataset._drop_none_fields(value) for key, value in obj.items() if value is not None} + return obj diff --git a/examples/r1_aqa/reward_models_utils.py b/examples/r1_aqa/reward_models_utils.py index f705e0de..ac18dab9 100644 --- a/examples/r1_aqa/reward_models_utils.py +++ b/examples/r1_aqa/reward_models_utils.py @@ -73,12 +73,6 @@ def accuracy_reward_fn(content: str, solution: str) -> float: reward = 1.0 except Exception: pass - if reward == 0.0: - sol_match = re.search(r"(.*?)", solution) - ground_truth = sol_match.group(1).strip() if sol_match else solution.strip() - student_answer = content.strip() - if student_answer == ground_truth: - reward = 1.0 return reward diff --git a/examples/r1_aqa/run_grpo_r1_aqa_qwen2_audio_7b.sh b/examples/r1_aqa/run_grpo_r1_aqa_qwen2_audio_7b.sh index a1e64ff1..a3d97e5c 100644 --- a/examples/r1_aqa/run_grpo_r1_aqa_qwen2_audio_7b.sh +++ b/examples/r1_aqa/run_grpo_r1_aqa_qwen2_audio_7b.sh @@ -13,6 +13,8 @@ # --- Model and Dataset Paths --- # Qwen2-Audio-7B-Instruct/Qwen2.5-Omni-7B base model +# Qwen2-Audio-7B-Instruct can only work with sglang engine, not vllm +# Qwen2.5-Omni-7B can only work with vllm engine, not sglang PATH_TO_YOUR_BASE_MODEL="" # Path to the cleaned AVQA dataset directory. @@ -147,7 +149,7 @@ torchrun \ --gradient_checkpointing \ --save_steps ${SAVE_STEPS} \ --max_ckpt_num 3 \ - --engine_type vllm \ + --engine_type sglang \ --engine_mem_util ${ENGINE_MEM_UTIL} \ --engine_tp_size $ENGINE_TP \ --enable_engine_sleep \ diff --git a/examples/r1_aqa/train_colocate.py b/examples/r1_aqa/train_colocate.py index e88201af..2f06b897 100644 --- a/examples/r1_aqa/train_colocate.py +++ b/examples/r1_aqa/train_colocate.py @@ -26,7 +26,7 @@ from lightrft.utils import add_arguments from lightrft.datasets import SFTDatasetVL from lightrft.models.actor_al import ( - AUDIO_MODEL_TYPE_QWEN2_5_OMNI, + AUDIO_MODEL_TYPE_QWEN2_AUDIO, ActorAL, create_audio_processor, infer_audio_model_type, @@ -43,12 +43,18 @@ def _validate_inference_engine(args) -> None: - if args.text_only: - return - model_type = infer_audio_model_type(args.pretrain) - if model_type == AUDIO_MODEL_TYPE_QWEN2_5_OMNI and args.engine_type == "sglang": - raise ValueError("Qwen2.5-Omni audio demo currently requires --engine_type vllm; sglang is not supported.") + if model_type == AUDIO_MODEL_TYPE_QWEN2_AUDIO: + expected_engine = "sglang" + else: + expected_engine = "vllm" + + if args.engine_type != expected_engine: + model_name = model_type or "non-qwen2-audio" + raise ValueError( + f"Model type `{model_name}` requires --engine_type {expected_engine}, " + f"but got --engine_type {args.engine_type}." + ) def train(args): diff --git a/lightrft/trainer/ppo_trainer_vl.py b/lightrft/trainer/ppo_trainer_vl.py index d34b412e..0c158df3 100644 --- a/lightrft/trainer/ppo_trainer_vl.py +++ b/lightrft/trainer/ppo_trainer_vl.py @@ -723,6 +723,49 @@ def _validate_qwen_vl_tensors( return False return True + def _validate_multimodal_training_batch( + self, + experience: ExperienceVL, + context: str = "training", + ) -> bool: + """ + Validate replay batches before forwarding multimodal actors. + + Vision batches keep the existing image-token consistency check. Audio batches + additionally reject replay rows whose attention mask is entirely zero, because + Qwen2-Audio cannot infer a valid padding side from a batch that mixes empty + rows with normal left-padded rows. + """ + if not self._validate_qwen_vl_tensors( + experience.sequences, + getattr(experience, "pixel_values", None), + context=context, + ): + return False + + if not self._is_audio_actor: + return True + + attention_mask = getattr(experience, "attention_mask", None) + if attention_mask is None or attention_mask.ndim != 2: + self.strategy.print( + f"[CRITICAL WARNING] Skipping batch in '{context}'. " + "Audio replay batch is missing a valid 2D attention_mask." + ) + return False + + active_lengths = attention_mask.long().sum(dim=-1) + invalid_rows = torch.nonzero(active_lengths <= 0, as_tuple=False).flatten().tolist() + if invalid_rows: + self.strategy.print( + f"[CRITICAL WARNING] Skipping batch in '{context}'. " + f"Audio replay batch contains empty attention_mask rows at indices {invalid_rows}. " + "This points to a degenerate rollout/replay sample rather than a missing audio file." + ) + return False + + return True + def training_step_actor( self, experience: ExperienceVL, @@ -777,11 +820,7 @@ def training_step_actor( # Actor loss. # Build modality-aware kwargs from the replay item instead of assuming vision-specific fields. actor_kwargs = self._build_model_kwargs(experience) - if not self._validate_qwen_vl_tensors( - sequences, - actor_kwargs.get("pixel_values"), - context="actor_rl_update", - ): + if not self._validate_multimodal_training_batch(experience, context="actor_rl_update"): self.strategy.print( "[CRITICAL ERROR] Validation failed inside training_step_actor. " "This should have been caught by pre-validation in spmd_ppo_trainer.py!" diff --git a/lightrft/trainer/spmd_ppo_trainer.py b/lightrft/trainer/spmd_ppo_trainer.py index e91fb6cc..d49cfb15 100644 --- a/lightrft/trainer/spmd_ppo_trainer.py +++ b/lightrft/trainer/spmd_ppo_trainer.py @@ -221,16 +221,20 @@ def ppo_train(self, global_steps=0): # Currently using this rewritten ppo_train # Step 1: Each rank validates its local data should_skip_local = False - if self.VLM and hasattr(self, '_validate_qwen_vl_tensors'): - # Call the same validation logic used in training_step_actor - sequences = experience.sequences - pixel_values = experience.pixel_values - - # Validate before any forward pass - is_valid = self._validate_qwen_vl_tensors( - sequences, pixel_values, context="pre_training_validation" - ) - should_skip_local = not is_valid + if self.VLM: + if hasattr(self, "_validate_multimodal_training_batch"): + is_valid = self._validate_multimodal_training_batch( + experience, context="pre_training_validation" + ) + should_skip_local = not is_valid + elif hasattr(self, "_validate_qwen_vl_tensors"): + # Backward-compatible fallback for older trainer implementations. + sequences = experience.sequences + pixel_values = experience.pixel_values + is_valid = self._validate_qwen_vl_tensors( + sequences, pixel_values, context="pre_training_validation" + ) + should_skip_local = not is_valid # Step 2: Synchronize skip decision across all ranks via all_reduce # This ensures all ranks agree on whether to skip, preventing execution divergence From 7ef597a8225a3319d6e967462b64b934387e6b47 Mon Sep 17 00:00:00 2001 From: PaParaZz1 Date: Sat, 25 Apr 2026 12:29:45 +0800 Subject: [PATCH 10/11] polish(nyz): docs and details --- examples/r1_aqa/README.md | 5 +- examples/r1_aqa/README_zh.md | 233 +++++++++++++++++++++++++ examples/r1_aqa/audio_dataset.py | 46 ++++- examples/r1_aqa/eval.py | 2 +- examples/r1_aqa/reward_models_utils.py | 1 - lightrft/models/loss.py | 46 +++++ lightrft/trainer/ppo_trainer_vl.py | 32 +++- 7 files changed, 357 insertions(+), 8 deletions(-) create mode 100644 examples/r1_aqa/README_zh.md diff --git a/examples/r1_aqa/README.md b/examples/r1_aqa/README.md index b08e1252..86ba8a5b 100644 --- a/examples/r1_aqa/README.md +++ b/examples/r1_aqa/README.md @@ -1,5 +1,7 @@ # R1-AQA on LightRFT: Audio Question Answering with GRPO +**English** | [中文](README_zh.md) + This example migrates [R1-AQA](https://github.com/xiaomi-research/r1-aqa) (Audio Question Answering via GRPO on Qwen2-Audio) into the [LightRFT](https://github.com/opendilab/LightRFT) training framework. ## Overview @@ -18,7 +20,8 @@ examples/r1_aqa/ ├── train_colocate.py # GRPO training entry point ├── eval.py # Evaluation script (e.g., MMAU-style tests) ├── run_grpo_r1_aqa_qwen2_audio_7b.sh # Training launch script -└── README.md # This file +├── README.md # This file +└── README_zh.md # Chinese version ``` ## Quick Start diff --git a/examples/r1_aqa/README_zh.md b/examples/r1_aqa/README_zh.md new file mode 100644 index 00000000..a2e174cd --- /dev/null +++ b/examples/r1_aqa/README_zh.md @@ -0,0 +1,233 @@ +# LightRFT 上的 R1-AQA：使用 GRPO 的音频问答 + +[English](README.md) | **中文** + +本示例将 [R1-AQA](https://github.com/xiaomi-research/r1-aqa)（基于 Qwen2-Audio 的音频问答 GRPO 训练）迁移到 [LightRFT](https://github.com/opendilab/LightRFT) 训练框架中。 + +## 概述 + +R1-AQA 将 Group Relative Policy Optimization（GRPO）应用到 Qwen2-Audio-7B-Instruct 上，用于音频问答任务。训练使用 AVQA 数据集上的规则奖励（准确率 + 格式奖励）。这个 LightRFT 示例在保留核心训练流程的同时，复用了 LightRFT 的分布式训练基础设施、GRPO 实现和奖励处理系统。 + +## 文件结构 + +``` +examples/r1_aqa/ +├── data_preprocess/ +│ ├── avqa.py # 将 R1-AQA JSONL 转成 LightRFT parquet +│ └── clean_audio_dataset.py # 删除音频缺失或无法读取的样本 +├── audio_dataset.py # 音频数据集与多模态输入封装 +├── reward_models_utils.py # 规则奖励（准确率 + 格式） +├── train_colocate.py # GRPO 训练入口 +├── eval.py # 评测脚本（例如 MMAU 风格测试） +├── run_grpo_r1_aqa_qwen2_audio_7b.sh # 训练启动脚本 +├── README.md # 英文版 +└── README_zh.md # 本文件 +``` + +## 快速开始 + +### 前置依赖 + +```bash +# 核心依赖（通常已随 LightRFT 安装） +pip install transformers torch deepspeed + +# 音频依赖 +pip install librosa soundfile + +# 可选：用于符号化答案校验的 math_verify +pip install math_verify +``` + +### 第 1 步：准备 AVQA 数据集 + +首先获取 R1-AQA 使用的 AVQA 训练数据（JSONL 格式）。原始 AVQA 数据如何转换，可参考 [R1-AQA README](https://github.com/xiaomi-research/r1-aqa)。 + +JSONL 文件中每一行应是一个 JSON 对象，字段类似： + +```json +{ + "id": 183, + "question_text": "What happened in the video?", + "multi_choice": ["motorboat", "Yacht consignment", "Sailboat set sail", "Consignment car"], + "answer": 1, + "dataset_name": "AVQA", + "audio_path": "path/to/-HG3Omg_89c_30.wav" +} +``` + +你也可以直接从 https://huggingface.co/datasets/Joysw909/AVQA 下载数据： + +```bash +huggingface-cli download --repo-type dataset --resume-download Joysw909/AVQA --local-dir path/to/AVQA +cd path/to/AVQA +mkdir -p all_audios +# 将各个 VGG 目录下的音频复制到 all_audios +cp VGG10000/* VGG20000/* VGG30000/* VGG40000/* all_audios/ 2>/dev/null || true +``` + +转换为 LightRFT 使用的格式： + +```bash +python examples/r1_aqa/data_preprocess/avqa.py \\ + --input_jsonl path/to/AVQA/train_r1aqa_line.json \\ + --audio_dir path/to/AVQA/all_audios \\ + --local_save_dir ./avqa_lightrft +``` + +### 第 2 步：清理缺失或损坏的音频样本 + +在训练前，强烈建议先对 parquet 数据做一次清理。在分布式 GRPO 训练中，如果某些样本的 prompt 仍然包含音频占位符，但其 `audio_path` 指向的文件已缺失，那么某些 rank 可能会走文本分支，而其他 rank 仍然走音频分支，后续常见表现就是 actor forward 阶段卡住。 + +运行： + +```bash +python examples/r1_aqa/data_preprocess/clean_audio_dataset.py \\ + --input_dataset ./avqa_lightrft \\ + --output_dir ./avqa_lightrft_clean +``` + +如果想做更严格的校验： + +```bash +python examples/r1_aqa/data_preprocess/clean_audio_dataset.py \\ + --input_dataset ./avqa_lightrft \\ + --output_dir ./avqa_lightrft_clean \\ + --verify_decode +``` + +该脚本会输出： + +- `train.parquet`：仅保留有效音频样本的清洗后数据 +- `train.dropped.jsonl`：被丢弃样本的记录，包含原始数据索引、`audio_path` 和原因 + +推荐工作流： + +1. 先运行一次 `avqa.py` 生成 parquet 数据集。 +2. 再对该 parquet 目录运行一次 `clean_audio_dataset.py`。 +3. 训练时使用清理后的输出目录，而不是原始 parquet 目录。 + +### 第 3 步：配置并启动训练 + +先编辑脚本，填入你的路径： + +```bash +# 在 run_grpo_r1_aqa_qwen2_audio_7b.sh 中： +PATH_TO_YOUR_BASE_MODEL="Qwen/Qwen2-Audio-7B-Instruct" +PATH_TO_YOUR_AVQA_DATASET="/path/to/your/avqa_lightrft_clean" +``` + +启动训练： + +```bash +bash examples/r1_aqa/run_grpo_r1_aqa_qwen2_audio_7b.sh +``` + +### 第 4 步：在 MMAU / MMAR 上评测 + +```bash +# MMAU (test-mini) +python examples/r1_aqa/eval.py \ + --benchmark mmau \ + --model_path results/lightrft-r1-aqa-grpo-training// \ + --data_file /path/to/mmau-test-mini.json \ + --audio_dir /path/to/mmau/audio \ + --out_file results/res_mmau_mini.json + +# 运行 MMAU 官方评测脚本 +python /path/to/mmau/evaluation.py --input results/res_mmau_mini.json + + +# MMAR +python examples/r1_aqa/eval.py \ + --benchmark mmar \ + --model_path results/lightrft-r1-aqa-grpo-training// \ + --data_file /path/to/MMAR-meta.jsonl \ + --audio_dir /path/to/mmar/audio \ + --out_file results/res_mmar.jsonl + +# 运行 MMAR 官方评测脚本 +python /path/to/mmar/code/evaluation.py --input results/res_mmar.jsonl +``` + +## Batch Size 约束 + +LightRFT 对 GRPO 有如下 batch size 关系约束： + +``` +train_batch_size >= rollout_batch_size × n_samples_per_prompt +``` + +R1-AQA 默认配置（`n_samples=8`）下示例： + +| 配置 | rollout_batch_size | n_samples | train_batch_size | 合法？ | +|---|---|---|---|---| +| 默认 | 16 | 8 | 128 | 128 >= 16×8=128 ✓ | +| 最小 | 4 | 4 | 32 | 32 >= 4×4=16 ✓ | +| 单卡 | 4 | 4 | 16 | 16 >= 4×4=16 ✓ | + +## 常见问题 + +### 1. 找不到音频路径 + +确保预处理脚本中的 `audio_dir` 指向实际存放 `.wav` 文件的目录。JSONL 中的音频路径既可以是相对路径，也可以是绝对路径。 + +如果训练日志显示各个 rank 的音频样本数不一致，例如某个 rank 打印出的 `<|AUDIO|>` prompt 数量更少，或者成功加载的音频数比其他 rank 少，先清理 parquet 数据，再使用清理后的目录进行训练： + +```bash +python examples/r1_aqa/data_preprocess/clean_audio_dataset.py \\ + --input_dataset /path/to/avqa_lightrft \\ + --output_dir /path/to/avqa_lightrft_clean +``` + +然后把 `PATH_TO_YOUR_AVQA_DATASET` 更新为清理后的输出目录。 + +### 2. 显存 / OOM + +- 减小 `MICRO_TRAIN` 和 `MICRO_ROLLOUT`（例如设为 1） +- 减小 `N_SAMPLES`（例如从 8 改成 4） +- 开启 `--gradient_checkpointing` 和 `--adam_offload` +- 调低 `ENGINE_MEM_UTIL`（例如设为 0.4） + +### 3. 推理引擎问题 + +- Qwen2-Audio 需要支持音频模型的 vLLM 或 SGLang +- 检查你的 vLLM 版本是否支持 `Qwen2AudioForConditionalGeneration` +- 如果使用 SGLang，确认已经具备音频多模态支持 + +### 4. MMAU 输出字段不匹配 + +评测脚本输出的是 `model_prediction`，这与 MMAU 期望的字段名一致。如果你使用自定义评测脚本，请确认输出字段名是否匹配。 + +### 5. Think Mode + +R1-AQA 支持可选的 `` 模式。启用方式如下： + +```bash +# 在数据预处理阶段： +python examples/r1_aqa/data_preprocess/avqa.py --enable_think ... +``` + +奖励函数会自动兼容两种模式。当 `enable_think=True` 时，格式奖励还会额外检查 `...` 标签。 + +## 设计说明 + +### 1. 奖励求和，而不是加权 + +R1-AQA 直接将准确率奖励和格式奖励相加（最大值为 2.0）；而 LightRFT 中 GSM8K/Geo3K 的实现使用加权组合（`0.9×accuracy + 0.1×format`，最大值为 1.0）。这里保留 R1-AQA 的求和方式，以确保奖励信号与原实现一致。GRPO 的归一化过程会处理这部分量纲差异。 + +### 2. 原生音频 rollout 路径 + +音频 RL 现在在 LightRFT 核心代码里走专门的 rollout 路径： + +- 原始音频负载保留在生成侧，并以 `audio_data` 的形式传给 SGLang +- 处理后的 mel 特征会显式保存在 `audio_values` 中 +- Qwen2-Audio 的特征掩码会显式保存在 `feature_attention_mask` 中 + +### 3. ActorAL（音频语言 Actor） + +Qwen2-Audio 使用的是 `Qwen2AudioForConditionalGeneration`（而不是 `AutoModelForVision2Seq`），其 forward 也需要 `audio_values`，而不是 `pixel_values` + `image_grid_thw`。因此这里使用 `lightrft.models.actor_al` 中的 `ActorAL`，它原生支持 Qwen2-Audio 所需的参数接口。 + +### 4. Chat Template + +R1-AQA 会把音频 URL 以 `{"type": "audio", "audio_url": path}` 的形式嵌入到 chat message 的 content 中。这里保留这一格式，并使用 Qwen2-Audio processor 的 `apply_chat_template` 将其转换成带有音频占位符的正确 token 格式。 diff --git a/examples/r1_aqa/audio_dataset.py b/examples/r1_aqa/audio_dataset.py index 17dcd038..25d462f8 100644 --- a/examples/r1_aqa/audio_dataset.py +++ b/examples/r1_aqa/audio_dataset.py @@ -29,7 +29,13 @@ def load_audio(audio_path: str, sr: int = 16000) -> Tuple[Any, int]: - """Load an audio file as ``(waveform, sampling_rate)``.""" + """ + Load an audio file as ``(waveform, sampling_rate)``. + + :param audio_path: Path to the audio file on disk. + :param sr: Target sampling rate used when decoding the audio file. + :return: Tuple of ``(waveform, sampling_rate)``. + """ return librosa.load(audio_path, sr=sr) @@ -51,6 +57,16 @@ def __init__( strategy, input_template: Optional[str] = None, ): + """ + Initialize the R1-AQA audio prompt dataset wrapper. + + :param dataset: Underlying dataset object that stores prompt/audio metadata. + :param tokenizer: Tokenizer kept for compatibility with the LightRFT dataset interface. + :param processor: Multimodal processor used to render the chat template and expose audio config. + :param max_length: Maximum sequence length tracked by the dataset interface. + :param strategy: Training strategy object used for config lookup and logging. + :param input_template: Optional example-side input template; unused in this dataset. + """ super().__init__() self.dataset = dataset self.tokenizer = tokenizer @@ -71,9 +87,20 @@ def __init__( self.target_sr = getattr(processor.feature_extractor, "sampling_rate", 16000) def __len__(self) -> int: + """ + Return the number of samples in the wrapped dataset. + + :return: Total sample count. + """ return len(self.dataset) def __getitem__(self, idx: int) -> Tuple[str, Any, str, str]: + """ + Load and normalize one training sample. + + :param idx: Sample index in the wrapped dataset. + :return: Tuple of ``(prompt_text, audio_data, reference, label)``. + """ data = self.dataset[idx] # ---- 1. Extract prompt (chat messages with audio content) ---- @@ -112,13 +139,23 @@ def __getitem__(self, idx: int) -> Tuple[str, Any, str, str]: return prompt_text, audio_data, reference, label def collate_fn(self, batch: List[Tuple]) -> Tuple[List, List, List, List]: - """Keep prompts/audios/references/labels as plain Python lists for the rollout stack.""" + """ + Keep prompts, audios, references, and labels as plain Python lists for rollout. + + :param batch: Batch of dataset items produced by ``__getitem__``. + :return: Tuple of lists ``(prompts, audios, refs, labels)``. + """ prompts, audios, refs, labels = zip(*batch) return list(prompts), list(audios), list(refs), list(labels) @staticmethod def _extract_text_from_messages(messages) -> str: - """Fallback text extraction used when the upstream chat template fails.""" + """ + Extract fallback text content when the upstream chat template fails. + + :param messages: Chat message payload in the processor input format. + :return: Concatenated text segments extracted from the message content. + """ texts = [] for msg in messages: if not isinstance(msg, dict): @@ -143,6 +180,9 @@ def _drop_none_fields(obj): block may arrive as ``{"type": "text", "text": "...", "audio_url": None}``. Qwen2-Audio's default chat template checks key existence instead of value, which would misclassify that text block as a second audio placeholder. + + :param obj: Nested list/dict structure that may contain ``None`` values. + :return: Structure with ``None`` entries removed recursively. """ if isinstance(obj, list): return [AudioPromptDataset._drop_none_fields(item) for item in obj] diff --git a/examples/r1_aqa/eval.py b/examples/r1_aqa/eval.py index 6911e178..fdf451df 100644 --- a/examples/r1_aqa/eval.py +++ b/examples/r1_aqa/eval.py @@ -67,7 +67,7 @@ def build_message(obj_dict: Dict, audio_dir: Optional[str] = None) -> list: """ choice_str = f"Please choose the answer from the following options: {obj_dict['choices']}." question_template = (f"{obj_dict['question']} {choice_str} " - "Output the final answer in .") + "Output the final answer in .") # MMAU uses audio_id; MMAR uses audio_path (e.g. ./audio/xxx.wav) raw_path = obj_dict.get("audio_path") or obj_dict.get("audio_id", "") diff --git a/examples/r1_aqa/reward_models_utils.py b/examples/r1_aqa/reward_models_utils.py index ac18dab9..dde3d593 100644 --- a/examples/r1_aqa/reward_models_utils.py +++ b/examples/r1_aqa/reward_models_utils.py @@ -137,7 +137,6 @@ def avqa_combined_reward_fn( def clean_solution(sol: str) -> str: - # <|im_start|>assistantat sea<|im_end|> """ Extract the string between <|im_start|>assistant and <|im_end|> tags. diff --git a/lightrft/models/loss.py b/lightrft/models/loss.py index 62c4671b..3c7ab96d 100644 --- a/lightrft/models/loss.py +++ b/lightrft/models/loss.py @@ -163,6 +163,9 @@ class PolicyLoss(nn.Module): - PPO: https://arxiv.org/abs/1707.06347 - CPGD: https://arxiv.org/abs/2505.12504 - High-Entropy Token Filtering: https://arxiv.org/abs/2506.01939 + + :ivar dict[str, float] _last_stats: Cached statistics from the most recent + ``forward`` call for trainer-side logging and debugging. """ def __init__( self, @@ -180,6 +183,20 @@ def __init__( @staticmethod def _stats_over_mask(values: torch.Tensor, mask: Optional[torch.Tensor], prefix: str) -> dict[str, float]: + """ + Summarize masked tensor values as mean/min/max scalars. + + :param values: Tensor whose selected elements should be summarized. + :type values: torch.Tensor + :param mask: Optional boolean-like mask that selects valid elements in ``values``. + If ``None``, all elements are used. + :type mask: Optional[torch.Tensor] + :param prefix: Prefix used to build the output metric keys. + :type prefix: str + :return: Dictionary with ``{prefix}_mean``, ``{prefix}_min``, and + ``{prefix}_max`` entries. + :rtype: dict[str, float] + """ if mask is None: selected = values.reshape(-1) else: @@ -200,6 +217,12 @@ def _stats_over_mask(values: torch.Tensor, mask: Optional[torch.Tensor], prefix: } def get_last_stats(self) -> dict[str, float]: + """ + Return statistics captured during the most recent policy loss computation. + + :return: Shallow copy of the cached scalar metrics dictionary. + :rtype: dict[str, float] + """ return dict(self._last_stats) def _update_last_stats( @@ -213,6 +236,26 @@ def _update_last_stats( old_log_probs: Optional[torch.Tensor] = None, ratio: Optional[torch.Tensor] = None, ) -> None: + """ + Refresh cached policy diagnostics for the latest forward pass. + + :param stats_mask: Boolean mask indicating which tokens contribute to the statistics. + :type stats_mask: torch.Tensor + :param advantages: Advantage tensor associated with the current minibatch. + :type advantages: torch.Tensor + :param logprob_delta: Difference ``log_probs - old_log_probs`` for each token. + :type logprob_delta: torch.Tensor + :param token_loss: Per-token policy loss values before masked reduction. + :type token_loss: torch.Tensor + :param log_probs: Optional current-policy log probabilities used for PPO diagnostics. + :type log_probs: Optional[torch.Tensor] + :param old_log_probs: Optional old-policy log probabilities used for PPO diagnostics. + :type old_log_probs: Optional[torch.Tensor] + :param ratio: Optional PPO importance-sampling ratio used to compute clip fractions. + :type ratio: Optional[torch.Tensor] + :return: ``None``. Metrics are stored in ``self._last_stats``. + :rtype: None + """ valid_token_count = float(stats_mask.sum().item()) stats = { "policy/valid_tokens": valid_token_count, @@ -268,6 +311,9 @@ def forward( :returns: Scalar policy loss averaged over valid (and optionally high-entropy) tokens. :rtype: torch.Tensor + The method also stores token-level summary statistics from the current call + in ``self._last_stats`` so the trainer can log them via ``get_last_stats()``. + **Masking Strategy:** The final mask is computed as: diff --git a/lightrft/trainer/ppo_trainer_vl.py b/lightrft/trainer/ppo_trainer_vl.py index 52c2aed4..0a2b34a1 100644 --- a/lightrft/trainer/ppo_trainer_vl.py +++ b/lightrft/trainer/ppo_trainer_vl.py @@ -265,6 +265,8 @@ def __init__( wandb.define_metric("rollout/*", step_metric="rollout/global_step") wandb.define_metric("train/global_step") wandb.define_metric("train/*", step_metric="train/global_step") + # eval/* uses its own counter, allowing it to be plotted sequentially + # even if evaluations happen rarely wandb.define_metric("eval/global_step") wandb.define_metric("eval/*", step_metric="eval/global_step") @@ -423,12 +425,18 @@ def fit( ) # Calculate number of rollouts per episode. - # Regardless of the TBS/RBS relationship, rollout count should depend on total sample volume, - # not on how those samples are internally split across optimizer steps. + # Regardless of TBS and RBS relationship, rollout count should be determined by "total data / rollout size". + # Numerator (num_update_steps * train_batch_size) equals "total samples planned for this episode". + # Denominator (rollout_batch_size * n_samples) equals "samples produced per rollout". + # This calculation ensures data collection volume is constant. + # When TBS=64, num_update_steps is naturally twice as large as when TBS=128. + # Substituting into formula: (2N * 0.5T) / R = (N * T) / R. + # Conclusion: Rollout count unchanged, but internal update loop count doubles due to smaller TBS. num_rollouts_per_episodes = ( num_update_steps_per_episodes * args.train_batch_size // args.max_epochs // args.rollout_batch_size // args.n_samples_per_prompt ) + # Safeguard to prevent num_rollouts_per_episodes from being 0 if num_rollouts_per_episodes == 0: # Use ceil as a safeguard when integer division would otherwise drop a fractional rollout. num_rollouts_per_episodes = math.ceil( @@ -489,6 +497,10 @@ def fit( all_response_lengths = [] for item in self.replay_buffer.items: + # Robust handling of reward_metrics + # 1. Check if info exists + # 2. Check if 'reward_metrics' key exists + # 3. Check if reward_metrics is not None (critical!) if hasattr(item, "info") and item.info is not None and "reward" in item.info: all_rewards.append(item.info["reward"]) @@ -517,6 +529,10 @@ def fit( rollout_status["rollout_reward_std"] = rewards_tensor.std().item() if all_format_rewards: + # [TENSOR-FIX] Handle both tensor lists and scalar lists + # Issue: all_format_rewards may contain tensors (from reward_metrics), + # but torch.tensor() cannot convert a list of tensors directly. + # Solution: Use torch.cat() for tensor lists, torch.tensor() for scalar lists if isinstance(all_format_rewards[0], torch.Tensor): format_tensor = torch.cat([t.to(device).float() for t in all_format_rewards]) else: @@ -554,6 +570,10 @@ def fit( # Progress bar reflects rollout quality; wandb/tensorboard will receive both rollout and train metrics. pbar.set_postfix(rollout_status) + + # Logs/checkpoints: save BOTH ROLLOUT and TRAINING statistics to wandb + # [FIX] Merge rollout_status (from inference) and status (from training) + # to ensure wandb logs contain both types of metrics client_states = {"consumed_samples": steps * args.rollout_batch_size} logs_dict_combined = {**rollout_status, **status} self.save_logs_and_checkpoints( @@ -703,11 +723,13 @@ def _validate_qwen_vl_tensors( :rtype: bool """ if pixel_values is None or pixel_values.numel() == 0: + # This is a text-only batch, no validation needed. return True config = self.strategy.unwrap_model(self.actor.model).config image_token_id = getattr(config, "image_token_id", None) if image_token_id is None: + # Model does not use special image tokens. return True num_tokens = (sequences == image_token_id).sum().item() @@ -850,6 +872,9 @@ def training_step_actor( experience.action_mask, kl_estimator=self.args.kl_estimator, ) + # [Protection measure 2] Per-token KL Clamping + # NOTE: Adding this causes svkng training to not converge + # kl = torch.clamp(kl, min=0.0, max=20.0) else: kl = torch.zeros_like(action_log_probs, dtype=action_log_probs.dtype, device=action_log_probs.device) @@ -1081,6 +1106,9 @@ def save_logs_and_checkpoints( all_wandb_logs[f"perf/experience_maker/{key}"] = value if all_wandb_logs: + # Use wandb_log_counter to ensure eval has a unique system step + # This prevents eval metrics from being overwritten by train metrics + # The plots will still use eval/global_step as X-axis due to define_metric self.wandb_log_counter += 1 self._wandb.log(all_wandb_logs, step=self.wandb_log_counter, commit=True) elif self._tensorboard is not None and self.strategy.is_rank_0(): From f90b893daed9f90e0992a59949358bce093b8ad2 Mon Sep 17 00:00:00 2001 From: PaParaZz1 Date: Sat, 25 Apr 2026 13:09:12 +0800 Subject: [PATCH 11/11] refactor(nyz): polish data layout materialize --- lightrft/trainer/modality_utils.py | 18 ++ lightrft/trainer/ppo_trainer_vl.py | 281 ++++++++++++++++++++++++----- 2 files changed, 255 insertions(+), 44 deletions(-) diff --git a/lightrft/trainer/modality_utils.py b/lightrft/trainer/modality_utils.py index 5c7ee69f..9779f7e2 100644 --- a/lightrft/trainer/modality_utils.py +++ b/lightrft/trainer/modality_utils.py @@ -8,6 +8,24 @@ def build_supported_model_kwargs(source: Any, supported_params: Set[str]) -> Dict[str, Any]: """ Extract only the multimodal kwargs that the current actor supports. + + This helper is the bridge between the trainer's generic replay objects and the + modality-specific model signatures. + + Example:: + + # Vision-language model + supported_params = {"pixel_values", "image_grid_thw"} + kwargs = build_supported_model_kwargs(experience, supported_params) + # -> {"pixel_values": experience.pixel_values, "image_grid_thw": experience.image_grid_thws} + + # Audio-language model + supported_params = {"audio_values", "feature_attention_mask"} + kwargs = build_supported_model_kwargs(experience, supported_params) + # -> {"audio_values": experience.audio_values, "feature_attention_mask": experience.feature_attention_mask} + + Keeping this mapping in one place avoids trainer call sites accidentally mixing + image fields and audio fields during future refactors. """ candidate_params = { "pixel_values": getattr(source, "pixel_values", None), diff --git a/lightrft/trainer/ppo_trainer_vl.py b/lightrft/trainer/ppo_trainer_vl.py index 0a2b34a1..5b9e16d3 100644 --- a/lightrft/trainer/ppo_trainer_vl.py +++ b/lightrft/trainer/ppo_trainer_vl.py @@ -301,10 +301,35 @@ def _ensure_device_and_contiguous(value, device): value = value.contiguous() return value + @staticmethod + def _cache_identity(value): + """ + Build a lightweight identity signature for replay-cache invalidation. + + Lists are represented by the identity of their elements so packed-sample caches + are refreshed if a new list of tensors replaces the old one. + """ + if isinstance(value, list): + return tuple(id(item) for item in value) + return id(value) + def _build_model_kwargs(self, source, device: Optional[int] = None) -> Dict[str, Any]: """ Select and optionally relocate only the multimodal kwargs supported by the current actor modality. + Example:: + + # Vision-language actor + kwargs = self._build_model_kwargs(experience) + # -> {"pixel_values": ..., "image_grid_thw": ..., "pixel_values_videos": ..., "video_grid_thw": ...} + + # Audio-language actor + kwargs = self._build_model_kwargs(experience) + # -> {"audio_values": ..., "feature_attention_mask": ...} + + The trainer therefore keeps one forward call-site while still preserving the original + modality-specific tensor names expected by each model family. + :param source: Replay item or mapping containing candidate multimodal tensors. :type source: Any :param device: Optional CUDA device index used to normalize tensor placement. @@ -324,6 +349,16 @@ def _unpack_prompt_batch(self, batch): Audio example datasets still produce a 4-field batch, but the second field now maps to ``audios`` instead of overloading the image slot. + Example:: + + # Vision dataset collate + batch = (prompts, images, references, labels) + # -> (prompts, images, None, None, references, labels) + + # Audio dataset collate + batch = (prompts, audios, references, labels) + # -> (prompts, None, None, audios, references, labels) + :param batch: Raw batch emitted by the prompt dataloader. :type batch: tuple or list :return: Tuple of ``(prompts, images, videos, audios, references, labels)`` used by rollout code. @@ -339,6 +374,180 @@ def _unpack_prompt_batch(self, batch): return prompts, modality_inputs, None, None, references, labels raise ValueError(f"Unsupported prompt batch format with {len(batch)} fields.") + def _materialize_replay_batch_layout(self, experience: ExperienceVL): + """ + Normalize replay-buffer sequence layout for validation and forward passes. + + The replay buffer uses two layouts: + + - padded batches: ``experience.sequences`` is already a ``(B, S)`` tensor and the + original ``attention_mask`` can be forwarded directly; + - packed batches: ``experience.sequences`` is a Python list of per-sample tensors, + so we concatenate them into the single packed row expected by actor/critic forward. + + The result is cached on the ``experience`` object because the same replay batch is + typically consumed three times in one PPO step: + + 1. pre-validation in ``SPMDPPOTrainer``, + 2. actor update, + 3. critic update. + + Without caching, each of those stages would repeat the same ``torch.cat`` work for + packed samples. + + Example:: + + packed input: + sequences = [tensor([11, 12]), tensor([21, 22, 23])] + + returned layout: + sequences = tensor([[11, 12, 21, 22, 23]]) + attention_mask = tensor([[1, 1, 2, 2, 2]]) + packed_seq_lens = [2, 3] + num_actions = [...] + + The synthetic packed attention mask mirrors the one used by the actual forward path, + so validation and training reason about exactly the same token layout. + + :param experience: Replay-buffer item or minibatch. + :type experience: ExperienceVL + :return: Dictionary containing cached normalized replay layout fields. + :rtype: Dict[str, Any] + """ + cache = getattr(experience, "_ppo_trainer_vl_cache", None) + cache_signature = ( + self._cache_identity(experience.sequences), + self._cache_identity(experience.attention_mask), + self._cache_identity(experience.advantages), + self._cache_identity(experience.action_mask), + ) + if cache is None: + cache = {"signature": cache_signature} + setattr(experience, "_ppo_trainer_vl_cache", cache) + elif cache.get("signature") != cache_signature: + cache = {"signature": cache_signature} + setattr(experience, "_ppo_trainer_vl_cache", cache) + + layout = cache.get("layout") + if layout is not None: + return layout + + is_packed = isinstance(experience.sequences, list) + if is_packed: + packed_seq_lens = [seq.numel() for seq in experience.sequences] + sequences = torch.cat(experience.sequences, dim=0).unsqueeze(0) + attention_mask = torch.cat( + [torch.full_like(seq, idx + 1) for idx, seq in enumerate(experience.sequences)], + dim=0, + ).unsqueeze(0) + num_actions = [value.numel() for value in experience.advantages] + else: + packed_seq_lens = None + sequences = experience.sequences + attention_mask = experience.attention_mask + num_actions = experience.action_mask.size(1) + + layout = { + "is_packed": is_packed, + "sequences": sequences, + "attention_mask": attention_mask, + "packed_seq_lens": packed_seq_lens, + "num_actions": num_actions, + } + cache["layout"] = layout + return layout + + def _materialize_policy_training_inputs(self, experience: ExperienceVL): + """ + Normalize replay-buffer fields needed by the actor PPO update. + + This helper keeps the packed-vs-padded branching in one place. Without it, + ``training_step_actor`` would need to partially unpack sequence layout in one helper + and then still manually unpack log-probs / advantages / KL references inline, which + makes the control flow harder to scan. + + Example:: + + packed replay item: + experience.sequences = [s0, s1] + experience.action_log_probs = [lp0, lp1] + experience.advantages = [adv0, adv1] + + normalized actor inputs: + sequences.shape == (1, total_len) + old_action_log_probs.shape == (1, total_actions) + advantages.shape == (1, total_actions) + num_actions == [len(adv0), len(adv1)] + + :param experience: Replay-buffer batch for actor optimization. + :type experience: ExperienceVL + :return: Dictionary of normalized actor inputs. + :rtype: Dict[str, Any] + """ + layout = self._materialize_replay_batch_layout(experience) + sequences = layout["sequences"] + attention_mask = layout["attention_mask"] + packed_seq_lens = layout["packed_seq_lens"] + num_actions = layout["num_actions"] + base_action_log_probs = None + + if layout["is_packed"]: + old_action_log_probs = torch.cat(experience.action_log_probs, dim=0).unsqueeze(0) + advantages = torch.cat(experience.advantages, dim=0).unsqueeze(0) + if self.args.use_kl_loss and experience.base_action_log_probs is not None: + base_action_log_probs = torch.cat(experience.base_action_log_probs, dim=0).unsqueeze(0) + else: + old_action_log_probs = experience.action_log_probs + advantages = experience.advantages + if self.args.use_kl_loss and experience.base_action_log_probs is not None: + base_action_log_probs = experience.base_action_log_probs + + return { + "sequences": sequences, + "attention_mask": attention_mask, + "packed_seq_lens": packed_seq_lens, + "old_action_log_probs": old_action_log_probs, + "advantages": advantages, + "num_actions": num_actions, + "base_action_log_probs": base_action_log_probs, + } + + def _materialize_value_training_inputs(self, experience: ExperienceVL): + """ + Normalize replay-buffer fields needed by the critic PPO update. + + The critic uses the same packed sequence layout as the actor, but different target + tensors (`values` / `returns`). Keeping this in a dedicated helper avoids having + `training_step_critic` repeat the same packed-sample branching that already exists + in the actor path. + + :param experience: Replay-buffer batch for critic optimization. + :type experience: ExperienceVL + :return: Dictionary of normalized critic inputs. + :rtype: Dict[str, Any] + """ + layout = self._materialize_replay_batch_layout(experience) + sequences = layout["sequences"] + attention_mask = layout["attention_mask"] + packed_seq_lens = layout["packed_seq_lens"] + num_actions = layout["num_actions"] + + if layout["is_packed"]: + old_values = torch.cat(experience.values, dim=0).unsqueeze(0) + returns = torch.cat(experience.returns, dim=0).unsqueeze(0) + else: + old_values = experience.values + returns = experience.returns + + return { + "sequences": sequences, + "attention_mask": attention_mask, + "packed_seq_lens": packed_seq_lens, + "old_values": old_values, + "returns": returns, + "num_actions": num_actions, + } + def _make_experience_list(self, prompts, images, videos, audios, references, labels): """ Shared rollout helper used by both training and evaluation. @@ -755,9 +964,15 @@ def _validate_multimodal_training_batch( additionally reject replay rows whose attention mask is entirely zero, because Qwen2-Audio cannot infer a valid padding side from a batch that mixes empty rows with normal left-padded rows. + + Packed replay samples are normalized first so validation sees the same concatenated + sequence layout that actor/critic forward will later consume. """ + layout = self._materialize_replay_batch_layout(experience) + sequences = layout["sequences"] + attention_mask = layout["attention_mask"] if not self._validate_qwen_vl_tensors( - experience.sequences, + sequences, getattr(experience, "pixel_values", None), context=context, ): @@ -766,8 +981,7 @@ def _validate_multimodal_training_batch( if not self._is_audio_actor: return True - attention_mask = getattr(experience, "attention_mask", None) - if attention_mask is None or attention_mask.ndim != 2: + if not isinstance(attention_mask, torch.Tensor) or attention_mask.ndim != 2: self.strategy.print( f"[CRITICAL WARNING] Skipping batch in '{context}'. " "Audio replay batch is missing a valid 2D attention_mask." @@ -806,29 +1020,14 @@ def training_step_actor( """ self.actor.train() - # Packed samples concatenate multiple sequences into one row. Unpacked samples stay batched. - # This mirrors the old PPOTrainerVL handling while replacing hard-coded VL kwargs with modality-aware ones. - if isinstance(experience.sequences, list): - sequences = torch.cat(experience.sequences, dim=0).unsqueeze(0) - old_action_log_probs = torch.cat(experience.action_log_probs, dim=0).unsqueeze(0) - advantages = torch.cat(experience.advantages, dim=0).unsqueeze(0) - num_actions = [value.numel() for value in experience.advantages] - packed_seq_lens = [seq.numel() for seq in experience.sequences] - attention_mask = torch.cat( - [torch.full_like(seq, idx + 1) for idx, seq in enumerate(experience.sequences)], - dim=0, - ).unsqueeze(0) - if self.args.use_kl_loss and experience.base_action_log_probs is not None: - base_action_log_probs = torch.cat(experience.base_action_log_probs, dim=0).unsqueeze(0) - else: - sequences = experience.sequences - old_action_log_probs = experience.action_log_probs - advantages = experience.advantages - num_actions = experience.action_mask.size(1) - packed_seq_lens = None - attention_mask = experience.attention_mask - if self.args.use_kl_loss and experience.base_action_log_probs is not None: - base_action_log_probs = experience.base_action_log_probs + actor_inputs = self._materialize_policy_training_inputs(experience) + sequences = actor_inputs["sequences"] + attention_mask = actor_inputs["attention_mask"] + packed_seq_lens = actor_inputs["packed_seq_lens"] + old_action_log_probs = actor_inputs["old_action_log_probs"] + advantages = actor_inputs["advantages"] + num_actions = actor_inputs["num_actions"] + base_action_log_probs = actor_inputs["base_action_log_probs"] if advantages is not None: # Clipping prevents a few extreme group-normalized values from dominating the PPO step. @@ -839,6 +1038,10 @@ def training_step_actor( # Actor loss. # Build modality-aware kwargs from the replay item instead of assuming vision-specific fields. + # Example outputs: + # - vision actor -> {"pixel_values": ..., "image_grid_thw": ...} + # - audio actor -> {"audio_values": ..., "feature_attention_mask": ...} + # The call site therefore stays identical even though the underlying model signatures differ. actor_kwargs = self._build_model_kwargs(experience) if not self._validate_multimodal_training_batch(experience, context="actor_rl_update"): self.strategy.print( @@ -994,29 +1197,19 @@ def training_step_critic(self, experience: ExperienceVL) -> Dict[str, float]: self.critic.train() device = torch.cuda.current_device() - # Match the packed/unpacked normalization used in actor training. - if isinstance(experience.sequences, list): - sequences = torch.cat(experience.sequences, dim=0).unsqueeze(0) - old_values = torch.cat(experience.values, dim=0).unsqueeze(0) - returns = torch.cat(experience.returns, dim=0).unsqueeze(0) - num_actions = [value.numel() for value in experience.advantages] - packed_seq_lens = [seq.numel() for seq in experience.sequences] - attention_mask = torch.cat( - [torch.full_like(seq, idx + 1) for idx, seq in enumerate(experience.sequences)], - dim=0, - ).unsqueeze(0) - else: - sequences = experience.sequences - old_values = experience.values - returns = experience.returns - num_actions = experience.action_mask.size(1) - packed_seq_lens = None - attention_mask = experience.attention_mask + critic_inputs = self._materialize_value_training_inputs(experience) + sequences = critic_inputs["sequences"] + attention_mask = critic_inputs["attention_mask"] + packed_seq_lens = critic_inputs["packed_seq_lens"] + old_values = critic_inputs["old_values"] + returns = critic_inputs["returns"] + num_actions = critic_inputs["num_actions"] sequences = self._ensure_device_and_contiguous(sequences, device) attention_mask = self._ensure_device_and_contiguous(attention_mask, device) old_values = self._ensure_device_and_contiguous(old_values, device) returns = self._ensure_device_and_contiguous(returns, device) + # Reuse the same modality selection logic as actor training so the two updates stay aligned. critic_kwargs = self._build_model_kwargs(experience, device=device) values, output = self.critic(