Add files using upload-large-folder tool

config/config.py

@@ -0,0 +1,193 @@
+import os
+import torch
+
+from data_utils.paths import OUTPUTS_DIR, project_path
+
+# ====== Model Configuration ======
+MODEL_CONFIG = {
+    "pretrained_model_path": "llava-hf/llava-onevision-qwen2-0.5b-ov-hf",  # two-stage grpo
+    "use_flash_attention_2": True,
+    "torch_dtype": "bfloat16",
+}
+
+# ====== Training Configuration ======
+TRAINING_CONFIG = {
+    "task": 'chart',
+    "num_gpus": 8,  # 使用的 GPU 数量
+    "num_client": 8,  # 并发客户端数量，通常与 GPU 数量相同
+    # RL阶段的参数 (根据原脚本的rl_args)
+    "dyme_args": {
+        "output_dir": os.path.join(OUTPUTS_DIR, "dyme-k-8"),
+        "logging_steps": 1,
+        "num_generations": 8,  # RL 阶段可以生成多个响应进行比较
+        "max_completion_length": 200,
+        "temperature": 0.8,
+        "repetition_penalty": 1.15,
+        "per_device_train_batch_size": 2,
+        "gradient_accumulation_steps": 16,
+        "num_train_epochs": 10,
+        "learning_rate": 8e-5,
+        "bf16": True,  # 使用 bf16 而不是 fp16
+        "gradient_checkpointing": False,
+        "ddp_find_unused_parameters": False,
+        "max_grad_norm": 1.0,
+        "save_strategy": "epoch",
+        "weight_decay": 0.01,
+        "warmup_steps": 0,
+        "beta": 0.0,  # DyMETrainer requires beta=0 (no reference-model KL path yet)
+        "loss_type": 'grpo',  # GRPO specific
+        "seed": 42,
+    },
+    "sft_args": {
+        "output_dir": os.path.join(OUTPUTS_DIR, "sft-chart-llava_cot"),
+        "logging_steps": 1,
+        "per_device_train_batch_size": 2,
+        "gradient_accumulation_steps": 4,
+        "num_train_epochs": 10,
+        "learning_rate": 1e-5,
+        "bf16": True,  # 使用 bf16 而不是 fp16
+        "gradient_checkpointing": False,
+        "ddp_find_unused_parameters": False,
+        "max_grad_norm": 1.0,
+        # "save_steps": 100,
+        "save_strategy": "epoch",
+        "weight_decay": 0.01,
+        "warmup_steps": 0,
+        "seed": 42,
+        "remove_unused_columns": False
+    },
+    "grpo_args":{
+        "output_dir": os.path.join(OUTPUTS_DIR, "grpo-chart-llava-beta"),
+        "logging_steps": 1,
+        "num_generations": 4,  # RL 阶段可以生成多个响应进行比较
+        "max_completion_length": 576,
+        "max_prompt_length": None,
+        "per_device_train_batch_size": 2,
+        "gradient_accumulation_steps": 4,
+        "num_train_epochs": 10,
+        "learning_rate": 1e-5,
+        "bf16": True,  # 使用 bf16 而不是 fp16
+        "gradient_checkpointing": False,
+        "ddp_find_unused_parameters": False,
+        "max_grad_norm": 1.0,
+        "save_strategy": "epoch",
+        "weight_decay": 0.01,
+        "warmup_steps": 0,
+        "beta": 0.04,  # GRPO specific
+        "loss_type": 'grpo',  # GRPO specific
+        "seed": 42,
+    }
+
+}
+
+RL_CONFIG = {
+    "answer_flag": "Answer:",
+    "end_flag": "<|im_end|>"
+}
+
+# OPSD / TriMode training (disabled by default -> original DyME behavior)
+DYME_OPSD_CONFIG = {
+    "enabled": False,
+    "mode": "dyme",
+    "privileged_profile": "hybrid",
+    "privileged_providers": ["text"],
+    "privileged_image": {
+        "mode": "single",
+        "crop_strategy": "bbox_then_center",
+        "bbox_coord": "normalized",
+        "margin_ratio": 0.25,
+    },
+    "privileged_debug": {
+        "save_images": True,
+        "image_subdir": "logs/images",
+        "max_samples_per_detail": 2,
+    },
+    "gate": {
+        "correct_threshold": 0.5,
+        "teacher_recoverable": "privileged_available",
+        "recoverable_tau": 0.5,
+        "use_edge_mask": False,
+        # TriMode: OPSD only on completions that are correct + well-formatted (not whole group).
+        "per_completion_opsd": True,
+        "require_format_for_opsd": True,
+        "skip_degenerate_for_opsd": True,
+    },
+    "loss": {
+        "beta": 0.5,
+        "opsd_weight": 2.0,
+        "grpo_weight": 1.0,
+        "sft_weight": 1.0,
+    },
+    # Optional: [format, context, accuracy] — default uniform; antidegen uses [0.5, 1.5, 1.0]
+    "reward_weights": [1.0, 1.0, 1.0],
+    "debug": {
+        # Full weak-signal diagnostic bundle every N global steps (rank 0). 0 = off.
+        "detail_every": 10,
+        # Lightweight [OPSD-PROBE] on every (re)generate (rank 0). Independent of OPSD-DEBUG.
+        "probe_on_generate": False,
+        "probe_sample_count": 4,
+        # Deep generate diagnostics ([OPSD-GENDBG]) — prompt tail, first-token logits, model context.
+        "probe_first_token_logits": True,
+        "probe_prompt_tail_tokens": 16,
+        "probe_log_model_context": True,
+        "health_monitor": {
+            "enabled": True,
+            "window": 20,
+            "log_on_generate": True,
+            "log_every_step": True,
+            "log_detail_bundle": True,
+            "log_alerts_immediately": True,
+            "metrics_every_step": True,
+        },
+    },
+}
+
+# ====== Client Configuration for Reward Calculation ======
+CLIENT_CONFIG = {
+    "client_type": "openai",  # 客户端主机地址
+    "api_key": "none",  # 客户端主机
+    "api_base": "http://127.0.0.1:%s/v1",  # 客户端，如果是本地服务需要预留端口
+    "timeout": 60,  # 请求超时时间
+    "model_id": "Qwen/Qwen2.5-14B-Instruct-AWQ",  # 使用的模型ID
+    "init_port": 23333, # 或者none代表在线服务
+    "num_server": 8
+}
+
+# ====== Dataset Configuration ======
+DATASET_CONFIG = {
+    "train_dataset": project_path("data/chartqa/train_medium_vf_full.json"),
+    # 训练数据路径
+    "eval_dataset": "HuggingFaceM4/ChartQA",  # 验证数据路径
+}
+
+# ====== DePlot offline preprocessing (ChartQA F2) ======
+DEPLOT_CONFIG = {
+    "enabled": True,
+    "model_id": "google/deplot",
+    "batch_size": 8,
+    "max_new_tokens": 384,
+    "cache_path": project_path("data/chartqa/deplot_cache.json"),
+    "prompt": "Generate underlying data table of the figure below:",
+}
+
+# ====== Full Configuration ======
+CONFIG = {
+    "model": MODEL_CONFIG,
+    "training": TRAINING_CONFIG,
+    "rl": RL_CONFIG,
+    "opsd": DYME_OPSD_CONFIG,
+    "client": CLIENT_CONFIG,
+    "dataset": DATASET_CONFIG,
+    "deplot": DEPLOT_CONFIG,
+}
+
+# Save configuration to a file for reference
+def save_config(config, config_path="./config.json"):
+    import json
+    with open(config_path, "w") as f:
+        json.dump(config, f, indent=4)
+
+# Example usage to save config
+if __name__ == "__main__":
+    save_config(CONFIG)
+

config/config_change.py

@@ -0,0 +1,79 @@
+import os
+import torch
+
+# ====== Model Configuration ======
+MODEL_CONFIG = {
+    "pretrained_model_path": "llava-hf/llava-onevision-qwen2-0.5b-ov-hf",  # two-stage grpo
+    "use_flash_attention_2": True,
+    "torch_dtype": "bfloat16",
+}
+
+# ====== Training Configuration ======
+TRAINING_CONFIG = {
+    "task": 'chart',
+    "num_gpus": 8, 
+    "num_client": 8,  
+
+    "dyme_args": {
+        "output_dir": '/path/to/dyme-llavaov-chart-change3',
+        "logging_steps": 1,
+        "num_generations": 4,  
+        "max_completion_length": 200,
+        "per_device_train_batch_size": 1,
+        "gradient_accumulation_steps": 16, # norm 16; 20 for sft hybrid
+        "num_train_epochs": 10,
+        "learning_rate": 8e-5,
+        "bf16": True,  
+        "gradient_checkpointing": False,
+        "ddp_find_unused_parameters": False,
+        "max_grad_norm": 1.0,
+        "save_strategy": "epoch",
+        "weight_decay": 0.01,
+        "warmup_steps": 0,
+        "beta": 0.0,  # GRPO specific
+        "loss_type": 'grpo',  # GRPO specific
+        "seed": 42,
+    },
+}
+
+RL_CONFIG = {
+    "answer_flag": "Answer:",
+    "end_flag": "<|im_end|>"
+}
+
+# ====== Client Configuration for Reward Calculation ======
+CLIENT_CONFIG = {
+    "client_type": "openai", 
+    "api_key": "none",  # 客户端主机
+    "api_base": "http://127.0.0.1:%s/v1",  
+    "timeout": 60,  
+    "model_id": "Qwen/Qwen2.5-14B-Instruct-AWQ",
+    "init_port": 23333, 
+    "num_server": 8
+}
+
+# ====== Dataset Configuration ======
+DATASET_CONFIG = {
+    "train_dataset": "/path/to/data/chartqa_output/json/train_new_prerefine.json",
+    "eval_dataset": "HuggingFaceM4/ChartQA",  
+}
+
+# ====== Full Configuration ======
+CONFIG = {
+    "model": MODEL_CONFIG,
+    "training": TRAINING_CONFIG,
+    "rl": RL_CONFIG,
+    "client": CLIENT_CONFIG,
+    "dataset": DATASET_CONFIG,
+}
+
+# Save configuration to a file for reference
+def save_config(config, config_path="./config.json"):
+    import json
+    with open(config_path, "w") as f:
+        json.dump(config, f, indent=4)
+
+# Example usage to save config
+if __name__ == "__main__":
+    save_config(CONFIG)
+

config/config_trimode_antidegen.py

@@ -0,0 +1,61 @@
+"""
+TriMode ChartQA config with anti-degeneration overrides (log-validated).
+
+Based on train_trimode_4gpu_20260610_173637.log analysis:
+- decoding: shorter max length, lower temperature, higher repetition penalty
+- training: lower LR + warmup to mitigate step-1 collapse
+- OPSD gate: require_format_for_opsd=False to raise opsd_mask coverage
+- reward_weights: emphasize continuous context F1 for group reward spread
+"""
+import os
+
+import config.config_trimode as trimode
+from data_utils.paths import OUTPUTS_DIR
+
+MODEL_CONFIG = dict(trimode.MODEL_CONFIG)
+
+TRAINING_CONFIG = {
+    **trimode.TRAINING_CONFIG,
+    "dyme_args": {
+        **trimode.TRAINING_CONFIG["dyme_args"],
+        "output_dir": os.environ.get(
+            "DYME_OUTPUT_DIR",
+            os.path.join(OUTPUTS_DIR, "dyme-trimode-antidegen"),
+        ),
+        "max_completion_length": 150,
+        "temperature": 0.7,
+        "repetition_penalty": 1.25,
+        "learning_rate": 5e-5,
+        "warmup_steps": 50,
+    },
+}
+
+_reward_weights_raw = os.environ.get("DYME_REWARD_WEIGHTS", "0.5,1.5,1.0")
+try:
+    _reward_weights = [float(x.strip()) for x in _reward_weights_raw.split(",") if x.strip()]
+    if len(_reward_weights) != 3:
+        raise ValueError("expected 3 weights")
+except ValueError:
+    _reward_weights = [0.5, 1.5, 1.0]
+
+_require_format_raw = os.environ.get("DYME_OPSD_REQUIRE_FORMAT", "0").strip().lower()
+_require_format_for_opsd = _require_format_raw not in ("0", "false", "no", "off")
+
+DYME_OPSD_CONFIG = {
+    **trimode.DYME_OPSD_CONFIG,
+    "gate": {
+        **trimode.DYME_OPSD_CONFIG.get("gate", {}),
+        "require_format_for_opsd": _require_format_for_opsd,
+        "skip_degenerate_for_opsd": True,
+    },
+    "reward_weights": _reward_weights,
+}
+
+CONFIG = {
+    "model": MODEL_CONFIG,
+    "training": TRAINING_CONFIG,
+    "rl": trimode.RL_CONFIG,
+    "opsd": DYME_OPSD_CONFIG,
+    "client": trimode.CLIENT_CONFIG,
+    "dataset": trimode.DATASET_CONFIG,
+}

config/loader.py

@@ -0,0 +1,66 @@
+import importlib
+import importlib.util
+import os
+from typing import Any
+
+_CONFIG_ALIASES: dict[str, str] = {
+    "norm": "config.config",
+    "trimode": "config.config_trimode",
+    "trimode_antidegen": "config.config_trimode_antidegen",
+    "rlsd": "config.config_rlsd_chartqa",
+    "rlsd_chartqa": "config.config_rlsd_chartqa",
+    "opd_7b": "config.config_opd_7b_chartqa",
+    "opd_7b_chartqa": "config.config_opd_7b_chartqa",
+    "llavacot": "config.config_llavacot",
+    "low": "config.config_low",
+    "aok": "config.config_aok",
+    "change": "config.config_change",
+    "7b": "config.config_7B",
+    "llm": "config.config_llm",
+}
+
+
+def _resolve_config_path(config_arg: str) -> str:
+    if os.path.isfile(config_arg):
+        return os.path.abspath(config_arg)
+
+    project_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+    candidate = os.path.join(project_root, config_arg)
+    if os.path.isfile(candidate):
+        return os.path.abspath(candidate)
+
+    raise FileNotFoundError(f"Config file not found: {config_arg}")
+
+
+def load_config(config_arg: str) -> dict[str, Any]:
+    """
+    Load CONFIG from a Python config file path or shorthand alias.
+
+    Examples:
+        load_config("config/config.py")
+        load_config("config/config_trimode.py")
+        load_config("norm")
+        load_config("trimode")
+    """
+    if config_arg.endswith(".py"):
+        path = _resolve_config_path(config_arg)
+        module_name = f"_dyme_runtime_config_{abs(hash(path)) % 10 ** 8}"
+        spec = importlib.util.spec_from_file_location(module_name, path)
+        if spec is None or spec.loader is None:
+            raise ImportError(f"Cannot load config module from {path}")
+        module = importlib.util.module_from_spec(spec)
+        spec.loader.exec_module(module)
+        if not hasattr(module, "CONFIG"):
+            raise ValueError(f"{path} must define a CONFIG dict")
+        return module.CONFIG
+
+    module_path = _CONFIG_ALIASES.get(config_arg)
+    if module_path is None:
+        raise ValueError(
+            f"Unknown config alias '{config_arg}'. "
+            f"Use a .py path (e.g. config/config.py) or one of: {', '.join(sorted(_CONFIG_ALIASES))}"
+        )
+    module = importlib.import_module(module_path)
+    if not hasattr(module, "CONFIG"):
+        raise ValueError(f"{module_path} must define a CONFIG dict")
+    return module.CONFIG

configs/deepspeed/zero2_bf16.json

@@ -0,0 +1,40 @@
+{
+  "bf16": {
+    "enabled": "auto"
+  },
+  "optimizer": {
+    "type": "AdamW",
+    "params": {
+      "lr": "auto",
+      "betas": "auto",
+      "eps": "auto",
+      "weight_decay": "auto",
+      "torch_adam": true,
+      "adam_w_mode": true
+    }
+  },
+  "scheduler": {
+    "type": "WarmupDecayLR",
+    "params": {
+      "warmup_min_lr": "auto",
+      "warmup_max_lr": "auto",
+      "warmup_num_steps": "auto",
+      "total_num_steps": "auto"
+    }
+  },
+  "zero_optimization": {
+    "stage": 2,
+    "allgather_partitions": true,
+    "allgather_bucket_size": 200000000.0,
+    "overlap_comm": true,
+    "reduce_scatter": true,
+    "reduce_bucket_size": "auto",
+    "contiguous_gradients": true
+  },
+  "gradient_accumulation_steps": "auto",
+  "gradient_clipping": "auto",
+  "train_batch_size": "auto",
+  "train_micro_batch_size_per_gpu": "auto",
+  "steps_per_print": 2000,
+  "wall_clock_breakdown": false
+}

data_utils/aokvqa/download.py

@@ -0,0 +1,279 @@
+import os
+import json
+import base64
+import io
+import multiprocessing
+from functools import partial
+from PIL import Image
+from datasets import load_dataset
+from tqdm import tqdm
+from openai import OpenAI, BadRequestError
+import time
+
+from data_utils.paths import AOKVQA_DIR
+
+# Ensure PIL can handle large images
+Image.MAX_IMAGE_PIXELS = None
+
+# --- API and MLLM prompt configuration ---
+
+API_PORTS = list(range(23333, 23333 + 8))
+API_URL_TEMPLATE = "http://127.0.0.1:{port}/v1/"
+
+# 2. Define prompts used to obtain visual_fact
+VISUAL_FACT_SYSTEM_PROMPT = """
+You are a helpful assistant that analyzes images and provides visual facts.
+Your response MUST be a single, valid JSON object.
+The JSON object should contain:
+1. "description": A detailed and accurate description of the image.
+2. "objects": A list of key objects, including their name, attributes, and approximate position in the image.
+
+Example format:
+{
+  "description": "A person riding a bicycle on a city street.... (detailed description here)",
+  "objects": [
+    {"name": "person", "attributes": ["wearing helmet", "blue shirt"], "position": "center"},
+    {"name": "bicycle", "attributes": ["red", "mountain bike"], "position": "center"},
+    {"name": "street", "attributes": ["asphalt", "wet"], "position": "bottom"}
+  ]
+}
+"""
+
+VISUAL_FACT_USER_PROMPT = """
+Analyze the attached image and provide the visual facts in the required JSON format.
+For context, the user will be asked this question about the image (do not answer the question, just use it for context):
+"{question}"
+"""
+
+
+def encode_image_to_base64(pil_image):
+    """Convert a PIL Image object to a Base64-encoded string"""
+    buffered = io.BytesIO()
+    if pil_image.mode == "RGBA" or "transparency" in pil_image.info:
+        pil_image.save(buffered, format="PNG")
+        mime_type = "image/png"
+    else:
+        pil_image.save(buffered, format="JPEG")
+        mime_type = "image/jpeg"
+
+    img_str = base64.b64encode(buffered.getvalue()).decode('utf-8')
+    return f"data:{mime_type};base64,{img_str}"
+
+
+def get_visual_fact(api_url, pil_image, question):
+    """
+    Call an external MLLM API to obtain visual_fact.
+    """
+    try:
+        client = OpenAI(base_url=api_url, api_key="DUMMY_KEY")
+        image_url = encode_image_to_base64(pil_image)
+
+        messages = [
+            {"role": "system", "content": VISUAL_FACT_SYSTEM_PROMPT},
+            {
+                "role": "user",
+                "content": [
+                    {"type": "image_url", "image_url": {"url": image_url,}},
+                    {"type": "text", "text": VISUAL_FACT_USER_PROMPT.format(question=question)}
+                ]
+            }
+        ]
+
+        try:
+            response = client.chat.completions.create(
+                model="Qwen/Qwen2.5-VL-32B-Instruct-AWQ",
+                messages=messages,
+                max_tokens=1024,
+                temperature=0.0,
+            )
+            response_content = response.choices[0].message.content
+            return response_content
+
+        except (BadRequestError, Exception):
+            response = client.chat.completions.create(
+                model="Qwen/Qwen2.5-VL-32B-Instruct-AWQ",
+                messages=messages,
+                max_tokens=1024,
+                temperature=0.0,
+            )
+            response_content = response.choices[0].message.content
+
+            if response_content.startswith("```json"):
+                response_content = response_content[7:].strip("` \n")
+
+            return response_content
+
+    except json.JSONDecodeError as e:
+        print(f"!! JSON parse failed: {e}. API: {api_url}. Raw response: {response_content[:100]}...")
+        return {"error": "Failed to parse JSON response", "raw_response": response_content}
+    except Exception as e:
+        print(f"!! API call failed: {e}. API: {api_url}")
+        time.sleep(1)
+        return {"error": f"API call failed: {str(e)}"}
+
+
+# --- (This is the modified function) ---
+def process_example_worker(example_with_index, split, image_output_dir, api_ports_list, fetch_visual_facts):
+    """
+    Worker function for multiprocessing.
+    Process a single example: save image and (optionally) call the API.
+    """
+    i, example = example_with_index
+
+    try:
+        # 1. Extract metadata
+        question = example["question"]
+        pil_image = example["image"]
+
+        # --- (A) Extract new fields ---
+        choices_list = example.get("choices")
+        correct_idx = example.get("correct_choice_idx")
+        direct_answers_list = example.get("direct_answers")
+        rationales_list = example.get("rationales")  # A-OKVQA has a 'rationales' field
+
+        # --- (B) Determine "answer" (prefer choice as requested) ---
+        answer = None
+        if choices_list and correct_idx is not None and 0 <= correct_idx < len(choices_list):
+            answer = choices_list[correct_idx]
+        elif direct_answers_list:
+            answer = direct_answers_list[0]  # Fallback to direct_answers
+
+        # --- (C) Determine "hint" (the longest rationale) ---
+        hint = None
+        if rationales_list:
+            try:
+                # Ensure rationales_list is a non-empty list of strings
+                if isinstance(rationales_list, list) and len(rationales_list) > 0:
+                    string_rationales = [r for r in rationales_list if isinstance(r, str)]
+                    if string_rationales:
+                        hint = max(string_rationales, key=len)  # Choose the longest string
+            except Exception as e:
+                print(f"!! Warning: Failed to compute 'hint' (index {i}). Error: {e}")
+                pass  # hint remains None
+
+        # --- (D) End of extraction ---
+        # 2. Generate and save image filename
+        generated_filename = f"{split}_{i:07d}.png"
+        image_save_path = os.path.join(image_output_dir, generated_filename)
+
+        # 3. Save image (I/O operation)
+        if not os.path.exists(image_save_path):
+            pil_image.save(image_save_path)
+
+        # 4. (Key step) Get Visual Fact
+        visual_fact_data = None
+        if fetch_visual_facts:
+            # Use API ports in a round-robin manner according to index 'i'
+            port_to_use = api_ports_list[i % len(api_ports_list)]
+            api_url = API_URL_TEMPLATE.format(port=port_to_use)
+
+            visual_fact_data = get_visual_fact(api_url, pil_image, question)
+
+        # 5. Build metadata dict (using the new format you specified)
+        metadata = {
+            "question": question,
+            "question_wo_prompt": question,
+            "answer": answer,
+            "choices": choices_list,  # Store the full options list
+            "image": image_save_path,
+            "visual_fact": visual_fact_data,
+            "hint": hint
+        }
+
+        return metadata
+
+    except Exception as e:
+        print(f"!! Fatal error in worker (index {i}): {e}")
+        return None  # The main process will filter out None
+
+
+def save_aokvqa_with_facts(base_output_dir=None, fetch_visual_facts=None):
+    """
+    Load A-OKVQA, save images with multiprocessing, and optionally obtain visual_facts for the training split.
+
+    Set env FETCH_VISUAL_FACTS=1 to call local MLLM APIs; default is images-only download.
+    """
+    if base_output_dir is None:
+        base_output_dir = AOKVQA_DIR
+    if fetch_visual_facts is None:
+        fetch_visual_facts = os.environ.get("FETCH_VISUAL_FACTS", "0") == "1"
+
+    # 1. Define output directories
+    base_dir_abs = os.path.abspath(base_output_dir)
+    image_output_dir = os.path.join(base_dir_abs, "images")
+    json_output_dir = os.path.join(base_dir_abs, "json")
+
+    # 2. Create directories
+    os.makedirs(image_output_dir, exist_ok=True)
+    os.makedirs(json_output_dir, exist_ok=True)
+
+    print(f"Images will be saved to (absolute path): {image_output_dir}")
+    print(f"JSON will be saved to (absolute path): {json_output_dir}")
+    print(f"Using {len(API_PORTS)} API ports: {API_PORTS}")
+
+    # 3. Load dataset
+    print("Loading HuggingFaceM4/A-OKVQA dataset...")
+    try:
+        dataset = load_dataset("HuggingFaceM4/A-OKVQA")
+        # Use a small subset for testing
+        # dataset['train'] = dataset['train'].select(range(100))
+    except Exception as e:
+        print(f"Failed to load dataset: {e}")
+        return
+
+    print(f"Dataset splits: {list(dataset.keys())}")
+
+    # 4. Define number of worker processes
+    num_workers = 64
+    print(f"Starting {num_workers} worker processes...")
+
+    # 5. Iterate over each split
+    for split in dataset.keys():
+
+        fetch_facts = fetch_visual_facts and (split == 'train')
+
+        if fetch_facts:
+            print(f"\n--- Processing split {split} (will call MLLM API) ---")
+        else:
+            print(f"\n--- Processing split {split} (images only, visual_fact=None) ---")
+
+        metadata_list = []
+
+        worker_func = partial(
+            process_example_worker,
+            split=split,
+            image_output_dir=image_output_dir,
+            api_ports_list=API_PORTS,
+            fetch_visual_facts=fetch_facts
+        )
+
+        tasks = list(enumerate(dataset[split]))
+        total_count = len(tasks)
+
+        # 6. Use multiprocessing pool
+        with multiprocessing.Pool(processes=num_workers) as pool:
+            for result in tqdm(pool.imap_unordered(worker_func, tasks), total=total_count, desc=f"Processing {split}"):
+                if result:
+                    metadata_list.append(result)
+
+        print(f"  Split {split} finished. Success: {len(metadata_list)} / {total_count}.")
+
+        # 7. (Optional) Sort
+        metadata_list.sort(key=lambda x: x['image'])
+
+        # 8. Write JSON file
+        json_filename = os.path.join(json_output_dir, f"{split}.json")
+        print(f"Saving {len(metadata_list)} metadata entries to {json_filename}...")
+
+        with open(json_filename, 'w', encoding='utf-8') as f:
+            json.dump(metadata_list, f, indent=4, ensure_ascii=False)
+
+    print(f"\n--- All processing completed! ---")
+    print(f"All image files saved in: '{image_output_dir}'")
+    print(f"All JSON files saved in: '{json_output_dir}'")
+
+
+if __name__ == "__main__":
+    multiprocessing.set_start_method("spawn", force=True)
+
+    save_aokvqa_with_facts()

data_utils/aokvqa/evaluator.py

@@ -0,0 +1,21 @@
+# --- New: A-OKVQA evaluation function ---
+def eval_aokvqa_direct(prediction, ground_truth_list):
+    """
+    Simple VQA accuracy: returns 1 if the prediction is in the direct_answers list, otherwise 0.
+    Matching is case-insensitive.
+    """
+    if not ground_truth_list:  # Handle empty ground truth list
+        return 0.0
+
+    # Clean predicted answer
+    pred_cleaned = prediction.lower().strip().strip('.').strip()
+
+    # Create a cleaned set of ground truth answers
+    gt_cleaned_set = set(ans.lower().strip().strip('.').strip() for ans in ground_truth_list)
+
+    # Check if prediction is in the set
+    if pred_cleaned in gt_cleaned_set:
+        return 1.0
+
+    # (Optional) You can add more sophisticated VQA matching logic, but simple "in" check is the baseline
+    return 0.0

data_utils/aokvqa/patch_aokvqa.py

@@ -0,0 +1,101 @@
+import os
+import json
+from datasets import load_dataset
+from tqdm import tqdm
+import shutil
+
+
+def patch_json_with_direct_answers(base_output_dir):
+    """
+    Load existing JSON files and the original A-OKVQA dataset,
+    and add the 'direct_answers' field back into the JSON files.
+
+    This script assumes that the entries in your generated JSON files
+    (e.g., train.json) are ordered according to the original dataset indices.
+    Your original script's
+    `metadata_list.sort(key=lambda x: x['image'])` ensures this alignment.
+    """
+
+    json_output_dir = os.path.join(base_output_dir, "json")
+    print(f"JSON directory to patch: {json_output_dir}")
+
+    # 1. Load the original dataset
+    print("Loading HuggingFaceM4/A-OKVQA dataset (metadata only)...")
+    try:
+        # Specify a cache directory to avoid repeated downloads
+        cache_dir = os.path.join(base_output_dir, ".cache")
+        os.makedirs(cache_dir, exist_ok=True)
+        dataset = load_dataset("HuggingFaceM4/A-OKVQA", cache_dir=cache_dir)
+    except Exception as e:
+        print(f"Failed to load original dataset: {e}")
+        return
+
+    print(f"Available splits in dataset: {list(dataset.keys())}")
+
+    # 2. Iterate over each split
+    for split in dataset.keys():
+        json_filename = os.path.join(json_output_dir, f"{split}.json")
+
+        if not os.path.exists(json_filename):
+            print(f"!! Warning: {json_filename} not found, skipping split '{split}'.")
+            continue
+
+        print(f"\n--- Patching split {split} ({json_filename}) ---")
+
+        # 3. Load existing (incomplete) JSON data
+        try:
+            with open(json_filename, 'r', encoding='utf-8') as f:
+                generated_data_list = json.load(f)
+            print(f"  Loaded {len(generated_data_list)} processed entries.")
+        except Exception as e:
+            print(f"  !! Error: Failed to load {json_filename}: {e}")
+            continue
+
+        # 4. Load original split data
+        original_split_data = dataset[split]
+        print(f"  Loaded {len(original_split_data)} original entries.")
+
+        # 5. Sanity check (ensure counts match)
+        if len(generated_data_list) != len(original_split_data):
+            print(f"  !! Critical error: Mismatch in number of entries!")
+            print(f"  JSON ({split}.json) has {len(generated_data_list)} records.")
+            print(f"  Original dataset ('{split}') has {len(original_split_data)} records.")
+            print(f"  Skipping this split.")
+            continue
+
+        # 6. Core logic: merge data using zip (relying on aligned ordering)
+        # Your 'image' field (e.g., "train_0000001.png") ensures that
+        # the order of 'generated_data_list' matches 'original_split_data'.
+
+        print(f"  Merging 'direct_answers'...")
+        for generated_metadata, original_example in \
+                tqdm(zip(generated_data_list, original_split_data),
+                     total=len(generated_data_list),
+                     desc=f"Merging {split}"):
+            # Add (or overwrite) the missing field
+            generated_metadata['direct_answers'] = original_example.get('direct_answers')
+
+        # 7. Backup and overwrite
+        backup_filename = os.path.join(json_output_dir, f"{split}.backup.json")
+        try:
+            if not os.path.exists(backup_filename):  # Backup only once
+                shutil.copyfile(json_filename, backup_filename)
+                print(f"  Backed up original file to {backup_filename}")
+            else:
+                print(f"  Backup file {backup_filename} already exists, will overwrite {json_filename} directly")
+        except Exception as e:
+            print(f"  !! Warning: Failed to create backup: {e}. Will overwrite directly.")
+
+        print(f"  Writing {len(generated_data_list)} updated metadata entries back to {json_filename}...")
+        with open(json_filename, 'w', encoding='utf-8') as f:
+            # generated_data_list has already been modified in memory
+            json.dump(generated_data_list, f, indent=4, ensure_ascii=False)
+
+    print("\n--- Patching completed! ---")
+
+
+if __name__ == "__main__":
+    from data_utils.paths import AOKVQA_DIR
+
+    print(f"Target root directory: {AOKVQA_DIR}")
+    patch_json_with_direct_answers(base_output_dir=AOKVQA_DIR)

data_utils/chart/pre_refinement.py

@@ -0,0 +1,134 @@
+# process_json_distributed.py
+
+import json
+import multiprocessing
+import os
+from copy import deepcopy
+from tqdm import tqdm
+import sys
+
+from data_utils.paths import CHARTQA_JSON_DIR, PROJECT_ROOT
+sys.path.append(PROJECT_ROOT)
+from client_utils.openai_api import OpenAIClient
+from data_utils.chart.prompts import prompt_refine
+from data_utils.commom_util import prompt_ic
+
+class ContextRefiner:
+
+    def __init__(self, CLIENT_CONFIG, gpu_id=0):
+        self.refine_templetes = ["""Goal: [State the user's objective, e.g., Find the year with the highest sales]
+Observation: [List key data points from the chart, e.g., 2020: 150, 2021: 200, 2022: 180]
+Reasoning: [State the logical step, e.g., Compare the values. 200 is the maximum.]
+Conclusion: [Draw the conclusion, e.g., The year with the highest sales was 2021.]
+"""]
+        if CLIENT_CONFIG['client_type'] == 'openai':
+
+            self.client = OpenAIClient(config=CLIENT_CONFIG)
+        else:
+            raise ValueError(f"Client type '{CLIENT_CONFIG['client_type']}' not supported.")
+
+    def refine_hint(self, question: str, hint: str, reference_answer: str, task: str):
+        if not hint:
+            return hint
+        system_prompt = None
+        if 'chart' in task:
+            system_prompt = 'You are a seasoned professional in the field of chart analysis...'
+        else:
+            raise Exception('Unknown expert task')
+        try:
+            in_context_example = self.client.get_completion(prompt_ic % hint, system_prompt=system_prompt,
+                                                            max_tokens=5000)
+            if 'chart' in task:
+                evaluation_prompt = prompt_refine % (in_context_example, question, reference_answer,
+                                                     self.refine_templetes[0])
+                output = self.client.get_completion(evaluation_prompt, system_prompt=system_prompt, max_tokens=1000)
+                return output
+            else:
+                raise ValueError(f"Task '{task}' not supported for thinking reward.")
+        except Exception as e:
+            print(f"Error occurred while processing '{question}': {e}")
+            return hint
+
+
+refiner_instance = None
+
+
+def worker_initializer(base_client_config):
+    global refiner_instance
+
+    # Key change: get the unique ID of the current worker process (starting from 1)
+    # This is the variable we use to simulate gpu_id
+    worker_id = multiprocessing.current_process()._identity[0] - 1
+
+    # Create a deep copy of the config to avoid interference between processes
+    worker_config = deepcopy(base_client_config)
+
+    # Key change: implement your port calculation logic
+    if worker_config.get('init_port') is not None and worker_config.get('num_server') is not None:
+        num_server = int(worker_config['num_server'])
+        # server_id decides which port to use
+        server_id = worker_id % num_server
+        port = worker_config['init_port'] + server_id
+
+        # Format api_base to assign a fixed port for this process
+        worker_config['api_base'] = worker_config['api_base'] % str(port)
+
+        print(f"Process {os.getpid()} (Worker-{worker_id}) initializing... connecting to {worker_config['api_base']}")
+    else:
+        print(f"Process {os.getpid()} (Worker-{worker_id}) initializing... using default api_base")
+
+    # Use the customized config for this specific process to create the instance
+    refiner_instance = ContextRefiner(worker_config, gpu_id=worker_id)
+
+
+def process_item_worker(item):
+    """Function executed by a single worker process (unchanged)"""
+    global refiner_instance
+    if refiner_instance is None:
+        raise Exception("Refiner has not been initialized in the worker process!")
+
+    new_hint = refiner_instance.refine_hint(
+        question=item['question'],
+        hint=item['hint'],
+        reference_answer=item['answer'],
+        task='chart'
+    )
+    item['hint'] = new_hint
+    return item
+
+
+# ---------------- Main logic ----------------
+def main():
+    # Configuration that contains port and server count information
+    from config import CLIENT_CONFIG
+    input_filename = os.path.join(CHARTQA_JSON_DIR, 'train.json')
+    output_filename = os.path.join(CHARTQA_JSON_DIR, 'train_new_prerefine.json')
+
+    NUM_PROCESSES = 64
+    print(f"Using {NUM_PROCESSES} processes and distributing requests to {CLIENT_CONFIG['num_server']} servers...")
+
+    try:
+        with open(input_filename, 'r', encoding='utf-8') as f:
+            data = json.load(f)
+    except FileNotFoundError:
+        print(f"Error: input file '{input_filename}' not found.")
+        return
+
+    processed_data = []
+
+    # Key point: pass the base configuration to the initializer of each process
+    with multiprocessing.Pool(processes=NUM_PROCESSES, initializer=worker_initializer,
+                              initargs=(CLIENT_CONFIG,)) as pool:
+        with tqdm(total=len(data), desc="Processing JSON in parallel") as pbar:
+            for result in pool.imap_unordered(process_item_worker, data):
+                processed_data.append(result)
+                pbar.update(1)
+
+    with open(output_filename, 'w', encoding='utf-8') as f:
+        json.dump(processed_data, f, ensure_ascii=False, indent=4)
+
+    print(f"\nProcessing completed! Results saved to '{output_filename}'.")
+
+
+if __name__ == "__main__":
+    main()

data_utils/chart/prompts.py

@@ -0,0 +1,114 @@
+prompt_thinking_reward = """
+Given
+<IC>: the data of an image
+<Q>: a question
+<A>: a reference answer
+<R>: a reasoning text
+
+Goal:
+Assess whether <R> correctly and reasonably uses visible data in <IC> to justify that the correct answer to <Q> is <A>. Rate the quality as low / medium / high according to:
+(a) low: Does not use data from <IC> at all, or the language is not fluent/natural, or it fails to indicate the answer to <Q> is <A>, or the structure is not good.
+(b) medium: Uses data from <IC> and is written fluently, but the reasoning is overly brief or insufficiently clear.
+(c) high: Uses data from <IC> and is written fluently; Has a very strong CoT structure (e.g. clear and logical steps); the reasoning progresses step by step with depth, each step is correct and reasonable; the data from <IC> appears exactly where it should; overall, the reasoning text provides very strong support that the answer to <Q> is <A>.
+
+Example:
+<IC>: [
+  {"object": "bar", "attributes": ["~120k", "Q4"], "label": "Product A"},
+  {"object": "bar", "attributes": ["~150k", "Q4"], "label": "Product B"},
+  {"object": "bar", "attributes": ["~90k", "Q4"], "label": "Product C"},
+  {"title": "Quarterly Revenue"}
+]
+<Q>: Which product has the highest revenue in Q4?
+<A>: product b
+<R>:
+    [Extraction] Reads Q4 bar heights: A ~120k, B ~150k, C ~90k.
+    [Calculation] Compares values: B > A and B > C.
+    [Conclusion] Therefore, Product B is highest, matching the answer "product b".
+
+<Output>: medium
+
+According to the requirements and examples above, score the input into three categories. Please give me the result directly without any explanation or description.
+
+<IC>: %s
+<Q>: %s
+<A>: %s
+<R>: %s
+<Output>: 
+"""
+
+
+prompt_refine = """Given:
+<IC>: the data of an image
+<Q>: a question
+<A>: a reference answer
+<T>: a writing template
+
+Goal:
+Transform the visual information in <IC> into a textualized data description and incorporate it into a smooth, natural explanation that reasons why the correct answer to <Q> is <A>, using the format and tone defined by <T>.
+
+Example:
+<IC>: [
+    {"Year": 2011, "Favorable": 0, "Unfavorable": 3.1},
+    {"Year": 2012, "Favorable": 56, "Unfavorable": 38.0},
+    {"Year": 2013, "Favorable": 0, "Unfavorable": 0.0},
+    {"Year": 2014, "Favorable": 51, "Unfavorable": 48.0},
+    {"Year": 2015, "Favorable": 0, "Unfavorable": 53.0}
+]
+<Q>: In which year the value was 51?
+<A>: 2014
+<T>:
+Goal: [State the user's objective, e.g., Find the year with the highest sales]
+Observation: [List key data points from the chart, e.g., 2020: 150, 2021: 200, 2022: 180]
+Reasoning: [State the logical step, e.g., Compare the values. 200 is the maximum.]
+Conclusion: [Draw the conclusion, e.g., The year with the highest sales was 2021.]
+
+<Output>:
+Goal: Find the year in which the 'Favorable' value was 51.
+Observation: The data shows the 'Favorable' values for each year are: 2011: 0, 2012: 56, 2013: 0, 2014: 51, and 2015: 0.
+Reasoning: Scanning the 'Favorable' column for the number 51 leads to the corresponding year in that row.
+Conclusion: The value 51 occurred in the year 2014.
+
+Now, according to the requirements and the examples above, convert my input into the target reasoning text. Please give me the result directly without any explanation or description.
+
+<IC>: %s
+<Q>: %s
+<A>: %s
+<T>: %s
+<Output>:
+"""
+
+prompt_template = """Analyze the preceding text (which is a Chain of Thought, or "CoT").
+
+**Your Task:**
+1.  **Evaluate Structure and Logic:** First, determine if the CoT is **well-structured** and **logical**. A "well-structured" CoT contains clear, distinct, and labeled reasoning steps (e.g., "Goal:", "Observation:", "Reasoning:", "Conclusion:").
+2.  **Extract Template:** If (and only if) the CoT is **well-structured**, extract a **generic reasoning template** from it. In this template, use brackets `[ ]` to describe the general purpose of each step (as shown in Example 1).
+3.  **Output None:** If the CoT is **unstructured** (e.g., it reads like a single conversational paragraph), lacks clear steps, or is logically unsound, you must output **None** (as shown in Example 2).
+
+---
+
+**Example 1: [Well-Structured CoT]**
+
+**Input:**
+"Goal: Find the lowest value of the red graph.\nObservation: The data for the 'Rep/Lean Rep' category across the years are: 2018: 72, 2019: 70, and 2020: 77.\nReasoning: Comparing the values, the minimum value is 70.\nConclusion: The lowest value of the red graph is 70."
+
+**Output:**
+Goal: [State the user's objective, e.g., Find the year with the highest sales]
+Observation: [List key data points from the chart, e.g., 2020: 150, 2021: 200, 2022: 180]
+Reasoning: [State the logical step, e.g., Compare the values. 200 is the maximum.]
+Conclusion: [Draw the conclusion, e.g., The year with the highest sales was 2021.]
+
+---
+
+**Example 2: [Poorly-Structured / Conversational CoT]**
+
+**Input:**
+I'm trying to figure out the year when the unfavorable view reaches its highest point. Looking at the data, I see that the values for each year are pretty low until 2016, where it jumps to 55. But the peak doesn't happen until 2017, when the value spikes to 64. So, it seems like the unfavorable view really hits its maximum in 2017.
+
+**Output:**
+None
+
+Now, according to the requirements and the examples above, give me the result directly without any explanation or description.
+**Input:**
+%s
+**Output:**
+"""

data_utils/clevr/clevr_processor.py

@@ -0,0 +1,461 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+from accelerate.utils import gather_object
+
+try:
+    from transformers import activations
+
+    activations.PytorchGELUTanh = activations.GELUTanh
+except ImportError:
+    print("Note: Unable to apply PytorchGELUTanh patch. If you encounter an ImportError, please check the transformers version.")
+# --- End of patch ---
+
+import os
+import shutil
+import json
+import re
+import numpy as np
+from PIL import Image
+from tqdm import tqdm
+from datasets import load_dataset
+from transformers import Qwen2_5_VLForConditionalGeneration, AutoProcessor
+from accelerate import Accelerator
+
+# (qwen_vl_utils import and fallback remain unchanged)
+try:
+    from qwen_vl_utils import process_vision_info
+except ImportError:
+    print("Warning: Failed to import 'qwen_vl_utils.process_vision_info'.")
+
+
+    def process_vision_info(messages):
+        images = []
+        for msg in messages:
+            if msg['role'] == 'user':
+                for content in msg['content']:
+                    if content['type'] == 'image':
+                        images.append(content['image'])
+        return images, None  # Return (images, videos)
+
+Image.MAX_IMAGE_PIXELS = None
+
+# (RexOmni dependency and DummyRex remain unchanged)
+try:
+    from rex_omni import RexOmniWrapper
+except ImportError:
+    print("Warning: 'from rex_omni import RexOmniWrapper' failed.")
+    print("Using a dummy RexOmniWrapper (DummyRex) for testing only.")
+
+
+    class DummyRex:
+        def __init__(self, *args, **kwargs):
+            print("INFO: DUMMY: Using DummyRex detector.")
+
+        def inference(self, images, task, categories, **kwargs):
+            print("INFO: DUMMY: DummyRex returning fake center boxes.")
+
+            # Batch-supporting Dummy
+            results = []
+
+            # Ensure images is a list
+            if not isinstance(images, list):
+                images = [images]
+
+            for img in images:
+                if isinstance(img, Image.Image):
+                    w, h = img.size
+                else:
+                    w, h = 800, 600
+                x0, y0 = w * 0.25, h * 0.25
+                x1, y1 = w * 0.75, h * 0.75
+                results.append({"extracted_predictions": {"anything": [{"type": "box", "coords": [x0, y0, x1, y1]}]}})
+            return results
+
+
+    RexOmniWrapper = DummyRex
+
+
+def _strip_tags(text, tag_name):
+    # (unchanged)
+    if not isinstance(text, str):
+        text = str(text)
+    text = re.sub(rf'<{tag_name}>', '', text, flags=re.IGNORECASE)
+    text = re.sub(rf'</{tag_name}>', '', text, flags=re.IGNORECASE)
+    return text.strip()
+
+
+# --- Core VQA helper functions (moved to global scope) ---
+
+def _crop_and_expand_box(image, box, padding_pixels=20):
+    # (unchanged)
+    x0, y0, x1, y1 = [int(c) for c in box]
+    img_w, img_h = image.size
+    x0_new = max(0, x0 - padding_pixels)
+    y0_new = max(0, y0 - padding_pixels)
+    x1_new = min(img_w, x1 + padding_pixels)
+    y1_new = min(img_h, y1 + padding_pixels)
+    return image.crop((x0_new, y0_new, x1_new, y1_new))
+
+
+# --- ★★★ Optimization 1: Change VQA queries to batched processing ★★★ ---
+def _query_qwen_vl_BATCH(crop_images_list, model, processor, accelerator):
+    """
+    Use Qwen-VL to query cropped image patches in batch and return a list of JSON strings.
+    """
+    if not crop_images_list:
+        return []
+
+    prompt = """This is an object from a CLEVR scene. Analyze the primary object in the image.
+Respond *strictly* with a JSON list (containing one dictionary) in the following format:
+[
+    {"object": "object_name", "attributes": ["attr1", "attr2"]}
+]
+- "object": The shape of the object (e.g., "sphere", "cube", "cylinder", "cone").
+- "attributes": A list of visual attributes (e.g., ["blue", "large", "metal", "shiny", "rubber"]).
+Provide only the JSON list:"""
+
+    # 1. Create messages for each image in the batch
+    template_messages = [
+        {
+            "role": "user",
+            "content": [
+                {"type": "image", "image": "placeholder.jpg"},  # Placeholder
+                {"type": "text", "text": prompt},
+            ],
+        }
+    ]
+
+    try:
+        # 2. Generate chat prompt text once
+        chat_prompt_text = processor.apply_chat_template(
+            template_messages, tokenize=False, add_generation_prompt=True
+        )
+
+        num_crops = len(crop_images_list)
+        batch_text = [chat_prompt_text] * num_crops
+        batch_images = crop_images_list
+
+        unwrapped_model = accelerator.unwrap_model(model)
+
+        # 3. Use text list and image list for batched processing
+        inputs = processor(
+            text=batch_text,
+            images=batch_images,
+            padding=True,  # Important: enable padding to handle batching
+            return_tensors="pt",
+        ).to(unwrapped_model.device)
+
+        # 4. Batch generation
+        generated_ids = unwrapped_model.generate(**inputs, max_new_tokens=256, do_sample=False)
+
+        generated_ids_trimmed = [
+            out_ids[len(in_ids):] for in_ids, out_ids in zip(inputs.input_ids, generated_ids)
+        ]
+
+        # 5. Batch decode
+        output_texts_list = processor.batch_decode(
+            generated_ids_trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False
+        )
+
+        return output_texts_list
+
+    except Exception as e:
+        print(f"Qwen-VL batched inference failed: {e}")
+        # Return a list of empty JSON placeholders to avoid zip errors
+        return ["[]"] * len(crop_images_list)
+
+
+def _parse_qwen_json(response_text):
+    # (unchanged)
+    try:
+        match = re.search(r'```json\s*(\[.*\])\s*```', response_text, re.DOTALL)
+        if match:
+            json_str = match.group(1)
+            return json.loads(json_str)
+        match = re.search(r'(\[.*\])', response_text, re.DOTALL)
+        if match:
+            json_str = match.group(1)
+            return json.loads(json_str)
+        return []
+    except json.JSONDecodeError:
+        print(f"Failed to parse JSON: {response_text}")
+        return []
+    except Exception as e:
+        print(f"Unknown error occurred while parsing JSON: {e}")
+        return []
+
+
+def _load_and_preprocess_data(base_output_dir, image_output_dir):
+    # (unchanged)
+    print("Loading 'MMInstruction/Clevr_CoGenT_TrainA_R1'...")
+    try:
+        dataset = load_dataset("MMInstruction/Clevr_CoGenT_TrainA_R1", split='train')
+    except Exception as e:
+        print(f"Failed to load dataset 'MMInstruction/Clevr_CoGenT_TrainA_R1': {e}")
+        return []
+
+    # Only use the first few samples for testing (still 100 here)
+    # dataset = dataset.select(range(100))
+    print(f"Loaded {len(dataset)} samples.")
+
+    job_list = []
+    print("Preprocessing data (saving images and parsing text)...")
+    for i, example in enumerate(tqdm(dataset, desc="Preprocessing progress")):
+        prompt = example['problem']
+        hint = _strip_tags(example['thinking'], 'think')
+        answer = _strip_tags(example['solution'], 'answer')
+
+        image = example['image']
+        if not isinstance(image, Image.Image):
+            print(f"Warning: sample {i} is not a PIL image, skipped.")
+            continue
+
+        image_filename = f"clevr_cogent_trainA_r1_{i:07d}.jpg"
+        destination_image_path = os.path.join(image_output_dir, image_filename)
+
+        try:
+            os.makedirs(os.path.dirname(destination_image_path), exist_ok=True)
+            if not os.path.exists(destination_image_path):
+                image.convert("RGB").save(destination_image_path, "JPEG")
+        except Exception as e:
+            print(f"Warning: failed to save image for sample {i}, skipped. Error: {e}")
+            continue
+
+        job_list.append({
+            "prompt": prompt,
+            "answer": answer,
+            "hint": hint,
+            "destination_image_path": destination_image_path
+        })
+
+    print(f"Successfully preprocessed {len(job_list)} items.")
+
+    job_list_path = os.path.join(base_output_dir, "job_list.json")
+    with open(job_list_path, 'w', encoding='utf-8') as f:
+        json.dump(job_list, f)
+
+    print(f"Job list saved to: {job_list_path}")
+    return job_list
+
+
+def main():
+    # (1. Initialize Accelerator - unchanged)
+    accelerator = Accelerator()
+
+    # (0. Define configuration - unchanged)
+    MODEL_CONFIGS = {
+        "rex_path": "IDEA-Research/Rex-Omni",
+        "qwen_path": "Qwen/Qwen2.5-VL-32B-Instruct-AWQ"
+    }
+    OUTPUT_DIR = "/path/to/data/clevr_cogent_output"
+    IMAGE_OUTPUT_DIR = os.path.join(OUTPUT_DIR, "images")
+    JSON_OUTPUT_DIR = os.path.join(OUTPUT_DIR, "json")
+
+    os.makedirs(IMAGE_OUTPUT_DIR, exist_ok=True)
+    os.makedirs(JSON_OUTPUT_DIR, exist_ok=True)
+
+    # (2. Preprocessing (run only on main process) - unchanged)
+    job_list_path = os.path.join(OUTPUT_DIR, "job_list.json")
+    if accelerator.is_main_process:
+        print("Main process [Pre-processing]: loading and preprocessing data...")
+        _load_and_preprocess_data(OUTPUT_DIR, IMAGE_OUTPUT_DIR)
+
+    # (3. Synchronization - unchanged)
+    accelerator.wait_for_everyone()
+
+    # (4. Load and distribute jobs - unchanged)
+    if not accelerator.is_main_process:
+        print(f"Process {accelerator.process_index}: loading job_list.json...")
+    try:
+        with open(job_list_path, 'r', encoding='utf-8') as f:
+            all_jobs = json.load(f)
+    except Exception as e:
+        print(f"Process {accelerator.process_index} failed to load job_list.json: {e}")
+        return
+    total_jobs = len(all_jobs)
+    num_processes = accelerator.num_processes
+    jobs_per_process = total_jobs // num_processes
+    start_index = accelerator.process_index * jobs_per_process
+    end_index = (accelerator.process_index + 1) * jobs_per_process
+    if accelerator.is_last_process:
+        end_index = total_jobs
+    my_jobs = all_jobs[start_index:end_index]
+    print(f"[Process {accelerator.process_index}]:"
+          f" assigned {len(my_jobs)} jobs (indices from {start_index} to {end_index}).")
+
+    # (5. Load models (each process loads its own copy) - unchanged)
+    try:
+        try:
+            from transformers import activations
+            activations.PytorchGELUTanh = activations.GELUTanh
+        except ImportError:
+            pass
+
+        print(f"[Process {accelerator.process_index}]: loading RexOmni...")
+        rex_model = RexOmniWrapper(
+            model_path=MODEL_CONFIGS['rex_path'],
+            backend="transformers",
+            max_tokens=2048,
+            temperature=0.0,
+        )
+
+        print(f"[Process {accelerator.process_index}]: loading Qwen-VL...")
+        qwen_model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
+            MODEL_CONFIGS['qwen_path'],
+            torch_dtype="float16",
+            device_map="cuda",
+            attn_implementation="flash_attention_2"
+        )
+        qwen_processor = AutoProcessor.from_pretrained(MODEL_CONFIGS['qwen_path'])
+
+        # Note: RexOmniWrapper (if not Dummy) may need .to(accelerator.device),
+        # but Qwen-VL already has its device specified via device_map="cuda".
+        # accelerator.prepare is still a good way to manage models.
+        qwen_model, rex_model = accelerator.prepare(qwen_model, rex_model)
+
+        print(f"[Process {accelerator.process_index}]: models loaded.")
+    except Exception as e:
+        print(f"[Process {accelerator.process_index}]: failed to load models: {e}")
+        import traceback
+        traceback.print_exc()
+        return
+
+    # 1. Define a batch size (Batch Size)
+    REX_BATCH_SIZE = 16  # <-- Adjust this value according to your VRAM
+
+    print(f"[Process {accelerator.process_index}]:"
+          f" starting to process {len(my_jobs)} jobs with Rex batch size {REX_BATCH_SIZE}.")
+
+    processed_metadata_list = []
+
+    # 2. Modify main loop: iterate with step REX_BATCH_SIZE
+    for i in tqdm(range(0, len(my_jobs), REX_BATCH_SIZE),
+                  desc=f"Worker {accelerator.process_index} batch progress",
+                  disable=not accelerator.is_main_process):
+
+        # 3. Prepare jobs and images for this batch
+        batch_jobs = my_jobs[i: i + REX_BATCH_SIZE]
+        batch_images = []
+        batch_image_paths = []  # for debugging
+
+        valid_jobs_in_batch = []
+
+        for job in batch_jobs:
+            try:
+                img_path = job['destination_image_path']
+                batch_image_paths.append(img_path)
+                batch_images.append(Image.open(img_path).convert("RGB"))
+                valid_jobs_in_batch.append(job)  # only jobs with successfully loaded images are valid
+            except Exception as e:
+                print(f"[Process {accelerator.process_index}]:"
+                      f" failed to load image {img_path}: {e}, skipping this image in this batch.")
+                # We do not add the image or job, keeping batch_images and valid_jobs_in_batch in sync
+
+        if not batch_images:  # if all images in this batch failed to load
+            continue
+
+        try:
+            # 4. ★ Key: run RexOmni in batch
+            # (we only pass successfully loaded images)
+            all_rex_results = rex_model.inference(
+                images=batch_images,  # pass the image list
+                task="detection",
+                categories=["anything"]
+            )
+
+            # 5. Iterate over results in this batch
+            # all_rex_results length should equal batch_images (and valid_jobs_in_batch)
+            if len(all_rex_results) != len(valid_jobs_in_batch):
+                print(f"[Process {accelerator.process_index}]: Warning: RexOmni "
+                      f"returned {len(all_rex_results)} results, but "
+                      f"{len(valid_jobs_in_batch)} inputs were provided. Skipping this batch.")
+                continue
+
+            for job, image, rex_result in zip(valid_jobs_in_batch, batch_images, all_rex_results):
+
+                predictions = rex_result["extracted_predictions"]
+                detected_boxes = predictions.get("anything", [])
+
+                visual_facts = []
+                crops_to_process = []
+                box_coords_list = []
+
+                # 6. Collect all crops to be processed (from this image)
+                for annotation in detected_boxes:
+                    if annotation.get("type") == "box" and len(annotation.get("coords", [])) == 4:
+                        coords = annotation["coords"]
+                        crop_image = _crop_and_expand_box(image, coords)
+                        crops_to_process.append(crop_image)
+                        box_coords_list.append(coords)
+
+                # 7. Batch VQA (logic remains the same, still batch *per image* crops)
+                if crops_to_process:
+                    json_str_list = _query_qwen_vl_BATCH(
+                        crops_to_process, qwen_model, qwen_processor, accelerator
+                    )
+
+                    # 8. Iterate over batched results and parse (logic unchanged)
+                    for json_str, coords in zip(json_str_list, box_coords_list):
+                        json_obj_list = _parse_qwen_json(json_str)
+                        if json_obj_list:
+                            try:
+                                obj_dict = json_obj_list[0]
+                                obj_dict["bounding_box"] = [round(c, 2) for c in coords]
+                                visual_facts.append(obj_dict)
+                            except (IndexError, TypeError, KeyError) as e:
+                                print(f"[Process {accelerator.process_index}]: "
+                                      f"Error while parsing batched result: {e} | JSON: {json_str}")
+
+                # 9. Aggregate results for this job (logic unchanged)
+                processed_metadata_list.append({
+                    "question": job['prompt'],
+                    "answer": job['answer'],
+                    "question_wo_prompt": job['prompt'],
+                    "hint": job['hint'],
+                    "image": job['destination_image_path'],
+                    "visual_fact": visual_facts
+                })
+                # --- End of inner loop logic ---
+
+        except Exception as e:
+            print(f"[Process {accelerator.process_index}]: "
+                  f"Error while processing batch {i // REX_BATCH_SIZE} (images {batch_image_paths}): {e}")
+            import traceback
+            traceback.print_exc()
+
+    # --- End of loop ---
+
+    print(f"[Process {accelerator.process_index}]:"
+          f" process finished, handled {len(processed_metadata_list)} items.")
+
+    # (7. Gather all results - unchanged)
+    print(f"[Process {accelerator.process_index}]: gathering results...")
+    all_results_list_of_lists = gather_object(processed_metadata_list)
+
+    # (8. Save (only on main process) - ★★★ using fixed GATHER logic ★★★)
+    if accelerator.is_main_process:
+        print("Main process [Saving]: aggregating and saving all results...")
+
+        # --- Key fix ---
+        # gather_object already returns a flattened list of dictionaries (List[dict]).
+        final_metadata_list = all_results_list_of_lists
+        # --- End of fix ---
+
+        json_filename = os.path.join(JSON_OUTPUT_DIR, "clevr_cogent_trainA_r1_processed.json")
+
+        # Verify the count
+        print(f"Total number of aggregated items: {len(final_metadata_list)}")
+        if len(final_metadata_list) > 0:
+            print(f"Type of first item: {type(final_metadata_list[0])}")
+
+        print(f"\nSaving {len(final_metadata_list)} metadata entries to {json_filename}...")
+        with open(json_filename, 'w', encoding='utf-8') as f:
+            json.dump(final_metadata_list, f, indent=4, ensure_ascii=False)
+
+        print(f"\n--- Processing completed! ---")
+        print(f"All image files have been saved in: '{IMAGE_OUTPUT_DIR}'")
+        print(f"Final JSON file has been saved in: '{JSON_OUTPUT_DIR}'")
+
+
+if __name__ == "__main__":
+    main()

data_utils/clevr/run_single_test.py

@@ -0,0 +1,194 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+from transformers import activations
+activations.PytorchGELUTanh = activations.GELUTanh
+import os
+import json
+from PIL import Image
+from datasets import load_dataset
+from transformers import Qwen2_5_VLForConditionalGeneration, AutoProcessor
+
+# Import model wrapper
+try:
+    from rex_omni import RexOmniWrapper
+except ImportError:
+    # Import DummyRex matching the one in clevr_processor.py
+    print("Warning: 'from rex_omni import RexOmniWrapper' failed.")
+    print("Using a dummy RexOmniWrapper (DummyRex) for testing only.")
+
+
+    class DummyRex:
+        def __init__(self, *args, **kwargs):
+            print("INFO: DUMMY: Using DummyRex detector.")
+
+        def inference(self, images, task, categories, **kwargs):
+            print("INFO: DUMMY: DummyRex returning a fake center box.")
+            if isinstance(images, Image.Image):
+                w, h = images.size
+            else:
+                w, h = 800, 600
+            x0, y0 = w * 0.25, h * 0.25
+            x1, y1 = w * 0.75, h * 0.75
+            return [{"extracted_predictions": {"anything": [{"type": "box", "coords": [x0, y0, x1, y1]}]}}]
+
+
+    RexOmniWrapper = DummyRex
+
+try:
+    from qwen_vl_utils import process_vision_info
+except ImportError:
+    print("Warning: Failed to import 'qwen_vl_utils.process_vision_info'.")
+
+
+    def process_vision_info(messages):
+        images = []
+        for msg in messages:
+            if msg['role'] == 'user':
+                for content in msg['content']:
+                    if content['type'] == 'image':
+                        images.append(content['image'])
+        return images, None
+
+from clevr_processor import ClevrFactExtractor, _strip_tags
+
+
+def run_test(configs, paths, gpu_id=0, sample_index=0):
+    """
+    Run the test pipeline on a single sample.
+    """
+    print("--- Starting single-run test (CoGenT) ---")
+
+    # --- 1. Set environment and load models ---
+    os.environ["CUDA_VISIBLE_DEVICES"] = str(gpu_id)
+    print(f"Set CUDA_VISIBLE_DEVICES={gpu_id}")
+
+    try:
+        print(f"Loading RexOmni... ({configs['rex_path']})")
+        rex_model = RexOmniWrapper(
+            model_path=configs['rex_path'],
+            backend="transformers",
+            max_tokens=2048,
+            temperature=0.0,
+        )
+
+        print(f"Loading Qwen-VL... ({configs['qwen_path']})")
+        qwen_model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
+            configs['qwen_path'],
+            torch_dtype="float16",
+            device_map="cuda:0",
+            attn_implementation="flash_attention_2"
+        )
+        qwen_processor = AutoProcessor.from_pretrained(configs['qwen_path'])
+
+        print("Models loaded.")
+    except Exception as e:
+        print(f"Failed to load models: {e}")
+        return
+
+    print("Loading dataset metadata...")
+    try:
+        dataset = load_dataset("MMInstruction/Clevr_CoGenT_TrainA_R1", split='train', streaming=True)
+        example_iter = iter(dataset)
+        for _ in range(sample_index + 1):
+            example = next(example_iter)
+
+    except Exception as e:
+        print(f"Failed to load or filter dataset: {e}")
+        return
+
+    print(f"Processing sample {sample_index}...")
+
+    try:
+        # 1. Preprocessing
+        prompt = example['problem']
+        hint = _strip_tags(example['thinking'], 'think')
+        answer = _strip_tags(example['solution'], 'answer')
+        image = example['image'].convert("RGB")  # Get PIL image and convert to RGB
+
+        # Save image for testing
+        destination_image_path = os.path.join(paths['output_dir'], "images", f"test_sample_{sample_index}.jpg")
+        os.makedirs(os.path.dirname(destination_image_path), exist_ok=True)
+        image.save(destination_image_path, "JPEG")
+        print(f"Loaded and saved test image: {destination_image_path}")
+
+        # --- Stage 1: RexOmni detection ---
+        print("Running RexOmni detection...")
+        rex_results = rex_model.inference(images=image, task="detection", categories=["anything"])
+        predictions = rex_results[0]["extracted_predictions"]
+        detected_boxes = predictions.get("anything", [])
+        print(f"RexOmni detected {len(detected_boxes)} 'anything' boxes.")
+
+        visual_facts = []
+
+        # --- Stage 2: Qwen-VL VQA ---
+        for i, annotation in enumerate(detected_boxes):
+            if annotation.get("type") == "box" and len(annotation.get("coords", [])) == 4:
+
+                coords = annotation["coords"]
+                print(f"  Processing box {i}: {coords}")
+
+                crop_image = ClevrFactExtractor._crop_and_expand_box(image, coords)
+
+                # Save cropped image for debugging
+                crop_filename = f"./test_crop_{sample_index}_{i}.jpg"
+                crop_image.save(crop_filename)
+                print(f"    -> Saved cropped image for inspection: {crop_filename}")
+
+                json_str = ClevrFactExtractor._query_qwen_vl(
+                    crop_image, qwen_model, qwen_processor
+                )
+
+                json_obj_list = ClevrFactExtractor._parse_qwen_json(json_str)
+
+                if json_obj_list:
+                    obj_dict = json_obj_list[0]
+                    obj_dict["bounding_box"] = [round(c, 2) for c in coords]
+                    visual_facts.append(obj_dict)
+                    print(f"    -> Qwen-VL result: {obj_dict}")
+                else:
+                    print(f"    -> Qwen-VL did not return valid JSON.")
+
+        # --- 4. Print final result ---
+        final_result = {
+            "prompt": prompt,
+            "answer": answer,
+            "hint": hint,
+            "image": destination_image_path,
+            "visual_fact": visual_facts
+        }
+
+        print("\n" + "=" * 30)
+        print("--- Single test result ---")
+        print(json.dumps(final_result, indent=4, ensure_ascii=False))
+        print("=" * 30 + "\n")
+
+    except Exception as e:
+        print(f"Error while processing sample {sample_index}: {e}")
+        import traceback
+        traceback.print_exc()
+
+
+if __name__ == "__main__":
+    # --- 1. Model configs ---
+    MODEL_CONFIGS = {
+        "rex_path": "IDEA-Research/Rex-Omni",
+        "qwen_path": "Qwen/Qwen2.5-VL-32B-Instruct-AWQ"
+    }
+
+    # --- 2. Paths config ---
+    PATHS = {
+        # !! Change this to the directory where you want to save images and JSON !!
+        "output_dir": "./clevr_cogent_output"
+    }
+
+    # --- 3. Test parameters ---
+    GPU_ID_TO_USE = 0
+    SAMPLE_INDEX_TO_TEST = 0  # Test the first CLEVR sample
+
+    # --- 4. Run test ---
+    run_test(
+        configs=MODEL_CONFIGS,
+        paths=PATHS,
+        gpu_id=GPU_ID_TO_USE,
+        sample_index=SAMPLE_INDEX_TO_TEST
+    )

data_utils/paths.py

@@ -0,0 +1,58 @@
+"""Project-relative data paths for DyME."""
+import os
+
+PROJECT_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))
+
+DATA_DIR = os.path.join(PROJECT_ROOT, "data")
+DATA_IMAGES_DIR = os.path.join(DATA_DIR, "images")
+
+CHARTQA_DIR = os.path.join(DATA_IMAGES_DIR, "chartqa")
+CHARTQA_IMAGES_DIR = os.path.join(CHARTQA_DIR, "images")
+CHARTQA_JSON_DIR = os.path.join(CHARTQA_DIR, "json")
+
+AOKVQA_DIR = os.path.join(DATA_IMAGES_DIR, "aokvqa")
+AOKVQA_IMAGES_DIR = os.path.join(AOKVQA_DIR, "images")
+AOKVQA_JSON_DIR = os.path.join(AOKVQA_DIR, "json")
+
+OUTPUTS_DIR = os.path.join(PROJECT_ROOT, "outputs")
+
+# Legacy absolute prefixes in JSON files -> canonical project directories
+_LEGACY_PREFIX_MAP = [
+    ("/chartqa_output/", CHARTQA_DIR + os.sep),
+    ("/path/to/chartqa_output/", CHARTQA_DIR + os.sep),
+    ("/path/to/data/chartqa_output/", CHARTQA_DIR + os.sep),
+    ("/path/to/data/aokvqa/", AOKVQA_DIR + os.sep),
+]
+
+
+def project_path(*parts: str) -> str:
+    return os.path.join(PROJECT_ROOT, *parts)
+
+
+def resolve_image_path(path: str) -> str:
+    """Resolve a stored image path to an existing file under the project tree."""
+    if not path:
+        return path
+    if os.path.exists(path):
+        return path
+
+    candidates = []
+    for old, new in _LEGACY_PREFIX_MAP:
+        if old in path:
+            candidates.append(path.replace(old, new))
+
+    if not os.path.isabs(path):
+        candidates.append(os.path.join(PROJECT_ROOT, path))
+
+    basename = os.path.basename(path)
+    candidates.extend([
+        os.path.join(CHARTQA_IMAGES_DIR, basename),
+        os.path.join(AOKVQA_IMAGES_DIR, basename),
+        os.path.join(PROJECT_ROOT, "chartqa_output", "images", basename),
+    ])
+
+    for candidate in candidates:
+        if os.path.exists(candidate):
+            return os.path.abspath(candidate)
+
+    return path

default_config_zero2.yaml

@@ -0,0 +1,20 @@
+compute_environment: LOCAL_MACHINE
+debug: false
+deepspeed_config:
+  deepspeed_config_file: configs/deepspeed/zero2_bf16.json
+  zero3_init_flag: false
+distributed_type: DEEPSPEED
+downcast_bf16: 'no'
+enable_cpu_affinity: false
+machine_rank: 0
+main_training_function: main
+num_machines: 1
+# ZeRO-2 student sharding (HF Accelerate + DeepSpeed JSON). Pair with DYME_TEACHER_DEVICE_MAP=auto.
+# Do NOT set mixed_precision / zero_stage here when using deepspeed_config_file (Accelerate requirement).
+num_processes: 2
+rdzv_backend: static
+same_network: true
+tpu_env: []
+tpu_use_cluster: false
+tpu_use_sudo: false
+use_cpu: false

main_7B.py

@@ -0,0 +1,192 @@
+# train_grpo.py
+"""
+Main script for training a Llava-based model using the custom MyGRPOTrainer.
+
+This script handles:
+1. Configuration loading.
+2. Initialization of Weights & Biases (wandb) and Hugging Face Accelerate.
+3. Loading the model (with PEFT/LoRA) and processor.
+4. Preparing the training and evaluation datasets.
+5. Setting up and running the GPRO trainer.
+"""
+
+import os
+from functools import partial
+from typing import Dict, Any
+
+import torch
+import wandb
+from accelerate import Accelerator
+from datasets import Dataset, load_dataset
+from transformers import AutoProcessor, Qwen2_5_VLForConditionalGeneration
+from trl import GRPOConfig
+from peft import LoraConfig, get_peft_model, TaskType  # NEW: Import PEFT modules
+
+from config.config_7B import CONFIG
+from data_utils.commom_util import collate_fn, define_task_data_func
+from trainer.DyMETrainer_7B import DyMETrainer
+from reward_utils.checker import RewardCalculator, RewardCalculatorLocal
+from reward_utils.refiner import ContextRefiner, ContextRefinerLocal
+
+
+def setup_accelerator_and_wandb(bf16) -> Accelerator:
+    """
+    Initializes Weights & Biases and the Hugging Face Accelerator.
+
+    Returns:
+        Accelerator: The configured accelerator instance.
+    """
+    wandb_key = os.environ.get("WANDB_API_KEY")
+    if wandb_key:
+        wandb.login(key=wandb_key)
+    if bf16:
+        accelerator = Accelerator(mixed_precision="bf16", log_with="wandb")
+    else:
+        accelerator = Accelerator(log_with="wandb")
+    return accelerator
+
+
+# NEW: Updated function to accept peft_config
+def load_model_and_processor(model_config: Dict[str, Any], peft_config: Dict[str, Any]):
+    """
+    Loads the pre-trained vision-language model, applies PEFT/LoRA, and loads its processor.
+
+    Args:
+        model_config (Dict[str, Any]): Configuration dictionary for the model.
+        peft_config (Dict[str, Any]): Configuration dictionary for PEFT/LoRA.
+
+    Returns:
+        Tuple[PeftModel, PreTrainedProcessor]: The loaded PEFT-enabled model and processor.
+    """
+    model_id = model_config['pretrained_model_path']
+
+    model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
+        model_id,
+        torch_dtype=getattr(torch, model_config['torch_dtype']),
+        attn_implementation='flash_attention_2' if model_config['use_flash_attention_2'] else 'sdpa',
+        low_cpu_mem_usage=True,
+    )
+
+    # Freeze the vision tower to save memory and computation
+    model.model.visual.requires_grad_(False)
+
+    # NEW: Create and apply LoRA configuration
+    print("Applying LoRA configuration...")
+    lora_config = peft_config
+
+    model = get_peft_model(model, lora_config)
+
+    # NEW: Print trainable parameters to verify LoRA is active
+    # This should show a very small percentage of trainable parameters.
+    model.print_trainable_parameters()
+
+    processor = AutoProcessor.from_pretrained(model_id)
+    processor.tokenizer.padding_side = "left"
+    # image_token_id = processor.tokenizer.additional_special_tokens_ids[
+    #     processor.tokenizer.additional_special_tokens.index("<image>")]
+
+    return model, processor
+
+
+def prepare_datasets(task: str, dataset_config: Dict[str, Any]) -> (Dataset, Dataset):
+    """
+    Prepares the training and evaluation datasets based on the specified task.
+
+    Args:
+        task (str): The name of the task (e.g., 'chartqa').
+        dataset_config (Dict[str, Any]): Configuration for datasets.
+
+    Returns:
+        Tuple[Dataset, Dataset]: The training and evaluation datasets.
+    """
+    data_func = define_task_data_func(task)
+
+    # Create training dataset
+    train_data_list = data_func(json_path=dataset_config['train_dataset'])
+    train_dataset = Dataset.from_list(train_data_list)
+
+    # Create evaluation dataset
+    if 'chart' in task:
+        eval_dataset = load_dataset(dataset_config['eval_dataset'])['test']
+        # Note: You can uncomment the line below for quick testing/debugging.
+        # eval_dataset = eval_dataset.select(range(1000, 1100))
+    else:
+        # Extend this section for other tasks if needed in the future.
+        raise NotImplementedError(f"Task '{task}' is not supported for evaluation in this script.")
+
+    return train_dataset, eval_dataset
+
+
+def main():
+    """
+    Main function to orchestrate the model training pipeline.
+    """
+
+    # 1. Load Configurations
+    model_config = CONFIG['model']
+    training_config = CONFIG['training']
+    rl_config = CONFIG['rl']
+    client_config = CONFIG['client']
+    dataset_config = CONFIG['dataset']
+    peft_config = LoraConfig(
+        r=64,
+        lora_alpha=128,
+        lora_dropout=0.05,
+        target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"],
+        use_rslora=True,
+        bias="none",
+        task_type="CAUSAL_LM",
+    )
+
+    task = training_config['task']
+
+    # 2. Setup Environment
+    accelerator = setup_accelerator_and_wandb(bf16=training_config['dyme_args']['bf16'])
+    device_id = accelerator.process_index
+
+    # 3. Initialize Model and Processor
+    # NEW: Pass peft_config to the loading function
+    model, processor = load_model_and_processor(model_config, peft_config)
+
+    # 4. Prepare Datasets
+    train_dataset, eval_dataset = prepare_datasets(task, dataset_config)
+
+    # 5. Initialize Reward Calculator
+
+    checker = RewardCalculatorLocal(rl_config, client_config.copy(), gpu_id=device_id)
+    refiner = ContextRefinerLocal(rl_config, client_config.copy(), gpu_id=device_id)
+    # 6. Define Training Arguments
+    training_args = GRPOConfig(**training_config['dyme_args'])
+
+    collate_fn_with_processor = partial(collate_fn, processor=processor)
+    # 7. Initialize the Trainer
+    dyme_trainer = DyMETrainer(
+        model=model,  # NEW: This is now a PeftModel
+        checker=checker,
+        refiner=refiner,
+        args=training_args,
+        train_dataset=train_dataset,
+        eval_dataset=eval_dataset,
+        processing_class=processor,
+        processing_func=collate_fn_with_processor,
+        task_name=task,
+        end_flag=rl_config['end_flag'],
+    )
+
+    # 8. Start Training
+    dyme_trainer.train()
+
+    output_dir = training_args.output_dir
+    output_dir = os.path.join(output_dir, "final_checkpoint")
+
+    # NEW: save_model will now save only the LoRA adapter weights
+    dyme_trainer.save_model(output_dir)
+
+    if accelerator.is_main_process:
+        processor.save_pretrained(output_dir)
+        # NEW: The saved model is just the adapter, not the full model.
+        print(f"LoRA adapter and processor saved to {output_dir}")
+
+
+if __name__ == "__main__":
+    main()

opsd_utils/__init__.py

@@ -0,0 +1,18 @@
+from opsd_utils.constants import MODE_GRPO, MODE_OPSD, MODE_SFT
+from opsd_utils.mode_router import route_prompt_modes, expand_modes_to_completions
+from opsd_utils.opsd_loss import compute_vlm_opsd_loss
+from opsd_utils.privileged import build_privileged_context
+from opsd_utils.recoverability import estimate_recoverable_flags
+from opsd_utils.prompt_builder import build_teacher_prompt_batch
+
+__all__ = [
+    "MODE_GRPO",
+    "MODE_OPSD",
+    "MODE_SFT",
+    "route_prompt_modes",
+    "expand_modes_to_completions",
+    "compute_vlm_opsd_loss",
+    "build_privileged_context",
+    "estimate_recoverable_flags",
+    "build_teacher_prompt_batch",
+]

opsd_utils/constants.py

@@ -0,0 +1,40 @@
+MODE_GRPO = 0
+MODE_OPSD = 1
+MODE_SFT = 2
+
+DEFAULT_OPSD_CONFIG = {
+    "enabled": False,
+    "mode": "dyme",
+    "privileged_profile": "hybrid",
+    "privileged_providers": ["text"],
+    "privileged_image": {
+        "mode": "single",
+        "crop_strategy": "bbox_then_center",
+        "bbox_coord": "normalized",
+        "margin_ratio": 0.25,
+    },
+    "privileged_debug": {
+        "save_images": True,
+        "image_subdir": "logs/images",
+        "max_samples_per_detail": 2,
+    },
+    "gate": {
+        "correct_threshold": 0.5,
+        "teacher_recoverable": "privileged_available",
+        "recoverable_tau": 0.5,
+        "recoverable_without_privilege": False,
+        "use_edge_mask": False,
+        "per_completion_opsd": True,
+        "require_format_for_opsd": True,
+        "skip_degenerate_for_opsd": True,
+    },
+    "loss": {
+        "beta": 0.5,
+        "opsd_weight": 2.0,
+        "grpo_weight": 1.0,
+        "sft_weight": 1.0,
+        "acc_gate": True,
+    },
+    "text_include_gold": True,
+    "reward_weights": [1.0, 1.0, 1.0],
+}

opsd_utils/debug_log.py

@@ -0,0 +1,426 @@
+"""Verbose debug logging for the OPSD / TriMode training pipeline."""
+from __future__ import annotations
+
+import json
+import os
+import time
+import traceback
+from contextlib import contextmanager
+from typing import Any, Optional
+
+_DEBUG_ENABLED = False
+_DETAIL_EVERY = 10
+_PROBE_ON_GENERATE = False
+_PROBE_FIRST_TOKEN_LOGITS = True
+_PROBE_PROMPT_TAIL_TOKENS = 16
+_PROBE_LOG_MODEL_CONTEXT = True
+_HEALTH_MONITOR_ENABLED = True
+_HEALTH_LOG_ON_GENERATE = True
+_HEALTH_LOG_EVERY_STEP = True
+_HEALTH_LOG_DETAIL_BUNDLE = True
+_HEALTH_LOG_ALERTS_IMMEDIATELY = True
+_RANK = 0
+_WORLD_SIZE = 1
+_STEP_LABEL = "init"
+_DETAIL_STEP: Optional[int] = None
+_CALL_COUNTER = 0
+
+MODE_NAMES = {0: "GRPO", 1: "OPSD", 2: "SFT"}
+
+
+def _env_debug_enabled() -> bool:
+    return os.environ.get("DYME_OPSD_DEBUG", "").strip().lower() in ("1", "true", "yes", "on")
+
+
+def _env_detail_every() -> int:
+    raw = os.environ.get("DYME_OPSD_DETAIL_EVERY", "").strip()
+    if not raw:
+        return 10
+    try:
+        return max(0, int(raw))
+    except ValueError:
+        return 10
+
+
+def _env_probe_on_generate() -> Optional[bool]:
+    raw = os.environ.get("DYME_OPSD_PROBE_ON_GENERATE", "").strip().lower()
+    if not raw:
+        return None
+    return raw in ("1", "true", "yes", "on")
+
+
+def _env_probe_first_token_logits() -> Optional[bool]:
+    raw = os.environ.get("DYME_OPSD_PROBE_FIRST_TOKEN_LOGITS", "").strip().lower()
+    if not raw:
+        return None
+    return raw in ("1", "true", "yes", "on")
+
+
+def _env_probe_prompt_tail_tokens() -> Optional[int]:
+    raw = os.environ.get("DYME_OPSD_PROBE_PROMPT_TAIL_TOKENS", "").strip()
+    if not raw:
+        return None
+    try:
+        return max(1, int(raw))
+    except ValueError:
+        return 16
+
+
+def _env_probe_log_model_context() -> Optional[bool]:
+    raw = os.environ.get("DYME_OPSD_PROBE_LOG_MODEL_CONTEXT", "").strip().lower()
+    if not raw:
+        return None
+    return raw in ("1", "true", "yes", "on")
+
+
+def _env_health_monitor_enabled() -> Optional[bool]:
+    raw = os.environ.get("DYME_OPSD_HEALTH_MONITOR", "").strip().lower()
+    if not raw:
+        return None
+    return raw in ("1", "true", "yes", "on")
+
+
+def configure(
+    *,
+    enabled: Optional[bool] = None,
+    detail_every: Optional[int] = None,
+    probe_on_generate: Optional[bool] = None,
+    probe_first_token_logits: Optional[bool] = None,
+    probe_prompt_tail_tokens: Optional[int] = None,
+    probe_log_model_context: Optional[bool] = None,
+    health_monitor_enabled: Optional[bool] = None,
+    health_log_on_generate: Optional[bool] = None,
+    health_log_every_step: Optional[bool] = None,
+    health_log_detail_bundle: Optional[bool] = None,
+    health_log_alerts_immediately: Optional[bool] = None,
+    rank: Optional[int] = None,
+    world_size: Optional[int] = None,
+) -> bool:
+    """Configure global OPSD debug logging. Returns whether debug is enabled."""
+    global _DEBUG_ENABLED, _DETAIL_EVERY, _PROBE_ON_GENERATE
+    global _PROBE_FIRST_TOKEN_LOGITS, _PROBE_PROMPT_TAIL_TOKENS, _PROBE_LOG_MODEL_CONTEXT
+    global _HEALTH_MONITOR_ENABLED, _HEALTH_LOG_ON_GENERATE, _HEALTH_LOG_EVERY_STEP
+    global _HEALTH_LOG_DETAIL_BUNDLE, _HEALTH_LOG_ALERTS_IMMEDIATELY
+    global _RANK, _WORLD_SIZE
+    if enabled is None:
+        enabled = _env_debug_enabled()
+    _DEBUG_ENABLED = bool(enabled)
+    if detail_every is not None:
+        _DETAIL_EVERY = max(0, int(detail_every))
+    elif _env_detail_every() != 10 or os.environ.get("DYME_OPSD_DETAIL_EVERY"):
+        _DETAIL_EVERY = _env_detail_every()
+    env_probe = _env_probe_on_generate()
+    if probe_on_generate is not None:
+        _PROBE_ON_GENERATE = bool(probe_on_generate)
+    elif env_probe is not None:
+        _PROBE_ON_GENERATE = env_probe
+    env_first_logits = _env_probe_first_token_logits()
+    if probe_first_token_logits is not None:
+        _PROBE_FIRST_TOKEN_LOGITS = bool(probe_first_token_logits)
+    elif env_first_logits is not None:
+        _PROBE_FIRST_TOKEN_LOGITS = env_first_logits
+    env_tail = _env_probe_prompt_tail_tokens()
+    if probe_prompt_tail_tokens is not None:
+        _PROBE_PROMPT_TAIL_TOKENS = max(1, int(probe_prompt_tail_tokens))
+    elif env_tail is not None:
+        _PROBE_PROMPT_TAIL_TOKENS = env_tail
+    env_model_ctx = _env_probe_log_model_context()
+    if probe_log_model_context is not None:
+        _PROBE_LOG_MODEL_CONTEXT = bool(probe_log_model_context)
+    elif env_model_ctx is not None:
+        _PROBE_LOG_MODEL_CONTEXT = env_model_ctx
+    env_health = _env_health_monitor_enabled()
+    if health_monitor_enabled is not None:
+        _HEALTH_MONITOR_ENABLED = bool(health_monitor_enabled)
+    elif env_health is not None:
+        _HEALTH_MONITOR_ENABLED = env_health
+    if health_log_on_generate is not None:
+        _HEALTH_LOG_ON_GENERATE = bool(health_log_on_generate)
+    if health_log_every_step is not None:
+        _HEALTH_LOG_EVERY_STEP = bool(health_log_every_step)
+    if health_log_detail_bundle is not None:
+        _HEALTH_LOG_DETAIL_BUNDLE = bool(health_log_detail_bundle)
+    if health_log_alerts_immediately is not None:
+        _HEALTH_LOG_ALERTS_IMMEDIATELY = bool(health_log_alerts_immediately)
+    if rank is not None:
+        _RANK = rank
+    if world_size is not None:
+        _WORLD_SIZE = world_size
+    return _DEBUG_ENABLED
+
+
+def detail_every() -> int:
+    return _DETAIL_EVERY
+
+
+def probe_on_generate() -> bool:
+    return _PROBE_ON_GENERATE
+
+
+def probe_first_token_logits() -> bool:
+    return _PROBE_FIRST_TOKEN_LOGITS
+
+
+def probe_prompt_tail_tokens() -> int:
+    return _PROBE_PROMPT_TAIL_TOKENS
+
+
+def probe_log_model_context() -> bool:
+    return _PROBE_LOG_MODEL_CONTEXT
+
+
+def should_log_probe() -> bool:
+    """True when lightweight per-generate probe should run (rank 0 only)."""
+    return _PROBE_ON_GENERATE and _RANK == 0
+
+
+def should_log_detail(global_step: Optional[int]) -> bool:
+    """True when a full diagnostic bundle should be emitted (rank 0 only)."""
+    if _DETAIL_EVERY <= 0 or _RANK != 0:
+        return False
+    if global_step is None:
+        return False
+    return int(global_step) % _DETAIL_EVERY == 0
+
+
+def is_enabled() -> bool:
+    return _DEBUG_ENABLED
+
+
+def set_step_label(label: str) -> None:
+    global _STEP_LABEL
+    _STEP_LABEL = label
+
+
+def set_detail_step(global_step: Optional[int]) -> None:
+    global _DETAIL_STEP
+    _DETAIL_STEP = global_step
+
+
+def get_detail_step() -> Optional[int]:
+    return _DETAIL_STEP
+
+
+def _next_call_id(stage: str) -> str:
+    global _CALL_COUNTER
+    _CALL_COUNTER += 1
+    return f"{_CALL_COUNTER}:{stage}"
+
+
+def _fmt(value: Any, max_len: int = 240) -> str:
+    if value is None:
+        return "None"
+    try:
+        import torch
+
+        if isinstance(value, torch.Tensor):
+            return (
+                f"Tensor(shape={tuple(value.shape)}, dtype={value.dtype}, "
+                f"device={value.device}, numel={value.numel()})"
+            )
+    except ImportError:
+        pass
+    if isinstance(value, (list, tuple)):
+        if len(value) > 12:
+            head = ", ".join(_fmt(v, max_len=40) for v in value[:6])
+            return f"[{head}, ... +{len(value) - 6} more, total={len(value)}]"
+        return "[" + ", ".join(_fmt(v, max_len=40) for v in value) + "]"
+    if isinstance(value, dict):
+        text = json.dumps(value, ensure_ascii=False, default=str)
+    else:
+        text = repr(value)
+    if len(text) > max_len:
+        return text[: max_len - 3] + "..."
+    return text
+
+
+def _prefix(stage: str, call_id: Optional[str] = None) -> str:
+    cid = call_id or _next_call_id(stage)
+    ts = time.strftime("%Y-%m-%d %H:%M:%S")
+    return f"[OPSD-DEBUG][{ts}][rank={_RANK}/{_WORLD_SIZE}][step={_STEP_LABEL}][{cid}]"
+
+
+def log(stage: str, msg: str, **fields: Any) -> None:
+    if not _DEBUG_ENABLED:
+        return
+    call_id = _next_call_id(stage)
+    extra = ""
+    if fields:
+        extra = " | " + " | ".join(f"{k}={_fmt(v)}" for k, v in fields.items())
+    print(f"{_prefix(stage, call_id)} {msg}{extra}", flush=True)
+
+
+def _detail_prefix(global_step: int, section: str) -> str:
+    ts = time.strftime("%Y-%m-%d %H:%M:%S")
+    return (
+        f"[OPSD-DETAIL][{ts}][rank={_RANK}/{_WORLD_SIZE}]"
+        f"[step={global_step}][every={_DETAIL_EVERY}][{section}]"
+    )
+
+
+def log_detail_banner(global_step: int, title: str) -> None:
+    if not should_log_detail(global_step):
+        return
+    bar = "=" * 20
+    print(f"{_detail_prefix(global_step, 'BANNER')} {bar} {title} {bar}", flush=True)
+
+
+def log_detail(section: str, msg: str, global_step: Optional[int] = None, **fields: Any) -> None:
+    """Full-detail diagnostic line (periodic, rank 0). Independent of verbose OPSD-DEBUG."""
+    step = global_step if global_step is not None else _DETAIL_STEP
+    if step is None or isinstance(step, str):
+        return
+    if not should_log_detail(step):
+        return
+    extra = ""
+    if fields:
+        extra = " | " + " | ".join(f"{k}={_fmt(v, max_len=800)}" for k, v in fields.items())
+    print(f"{_detail_prefix(step, section)} {msg}{extra}", flush=True)
+
+
+def _probe_prefix(section: str) -> str:
+    ts = time.strftime("%Y-%m-%d %H:%M:%S")
+    step = _DETAIL_STEP if _DETAIL_STEP is not None else "?"
+    return (
+        f"[OPSD-PROBE][{ts}][rank={_RANK}/{_WORLD_SIZE}]"
+        f"[global_step={step}][{_STEP_LABEL}][{section}]"
+    )
+
+
+def log_probe(section: str, msg: str, **fields: Any) -> None:
+    """Lightweight per-generate diagnostic (rank 0). Independent of OPSD-DEBUG verbosity."""
+    if not should_log_probe():
+        return
+    extra = ""
+    if fields:
+        extra = " | " + " | ".join(f"{k}={_fmt(v, max_len=1200)}" for k, v in fields.items())
+    print(f"{_probe_prefix(section)} {msg}{extra}", flush=True)
+
+
+def _gendbg_prefix(section: str) -> str:
+    ts = time.strftime("%Y-%m-%d %H:%M:%S")
+    step = _DETAIL_STEP if _DETAIL_STEP is not None else "?"
+    return (
+        f"[OPSD-GENDBG][{ts}][rank={_RANK}/{_WORLD_SIZE}]"
+        f"[global_step={step}][{_STEP_LABEL}][{section}]"
+    )
+
+
+def should_log_gendbg() -> bool:
+    """True when deep generate diagnostics should run (rank 0 only)."""
+    return _PROBE_ON_GENERATE and _RANK == 0
+
+
+def health_monitor_enabled() -> bool:
+    return _HEALTH_MONITOR_ENABLED and _RANK == 0
+
+
+def should_log_health_on_generate() -> bool:
+    return health_monitor_enabled() and _HEALTH_LOG_ON_GENERATE and should_log_probe()
+
+
+def should_log_health_every_step() -> bool:
+    return health_monitor_enabled() and _HEALTH_LOG_EVERY_STEP
+
+
+def should_log_health_detail_bundle() -> bool:
+    return health_monitor_enabled() and _HEALTH_LOG_DETAIL_BUNDLE
+
+
+def should_log_health_alerts_immediately() -> bool:
+    return health_monitor_enabled() and _HEALTH_LOG_ALERTS_IMMEDIATELY
+
+
+def _health_prefix(section: str, global_step: Optional[int] = None) -> str:
+    ts = time.strftime("%Y-%m-%d %H:%M:%S")
+    step = global_step if global_step is not None else (_DETAIL_STEP if _DETAIL_STEP is not None else "?")
+    return f"[OPSD-HEALTH][{ts}][rank={_RANK}/{_WORLD_SIZE}][global_step={step}][{section}]"
+
+
+def log_health(section: str, msg: str, global_step: Optional[int] = None, **fields: Any) -> None:
+    """L1/L2/L4 health lines (rank 0)."""
+    if not health_monitor_enabled():
+        return
+    extra = ""
+    if fields:
+        extra = " | " + " | ".join(f"{k}={_fmt(v, max_len=1200)}" for k, v in fields.items())
+    print(f"{_health_prefix(section, global_step)} {msg}{extra}", flush=True)
+
+
+def log_health_detail_banner(global_step: int, title: str) -> None:
+    if not should_log_health_detail_bundle() or not should_log_detail(global_step):
+        return
+    bar = "=" * 20
+    print(f"{_detail_prefix(global_step, 'health')} {bar} {title} {bar}", flush=True)
+
+
+def log_health_detail(section: str, msg: str, global_step: int, **fields: Any) -> None:
+    """L3 periodic health bundle (rank 0, same cadence as OPSD-DETAIL)."""
+    if not should_log_health_detail_bundle() or not should_log_detail(global_step):
+        return
+    extra = ""
+    if fields:
+        extra = " | " + " | ".join(f"{k}={_fmt(v, max_len=1200)}" for k, v in fields.items())
+    print(f"{_detail_prefix(global_step, section)} {msg}{extra}", flush=True)
+
+
+def log_gendbg(section: str, msg: str, **fields: Any) -> None:
+    """Deep per-generate diagnostic (rank 0). Uses [OPSD-GENDBG] prefix."""
+    if not should_log_gendbg():
+        return
+    extra = ""
+    if fields:
+        extra = " | " + " | ".join(f"{k}={_fmt(v, max_len=1200)}" for k, v in fields.items())
+    print(f"{_gendbg_prefix(section)} {msg}{extra}", flush=True)
+
+
+def log_config(stage: str, title: str, config: dict[str, Any]) -> None:
+    if not _DEBUG_ENABLED:
+        return
+    log(stage, title, config=_fmt(config, max_len=2000))
+
+
+def log_sync_point(stage: str, msg: str, **fields: Any) -> None:
+    """Mark a point where all ranks must reach before a collective call."""
+    log(stage, f"[SYNC_POINT] {msg}", **fields)
+
+
+def log_mode_summary(stage: str, prompt_modes: list[int], completion_modes: Optional[list[int]] = None) -> None:
+    if not _DEBUG_ENABLED:
+        return
+    prompt_names = [MODE_NAMES.get(m, str(m)) for m in prompt_modes]
+    fields: dict[str, Any] = {"prompt_modes": prompt_names}
+    if completion_modes is not None:
+        counts = {name: completion_modes.count(code) for code, name in MODE_NAMES.items()}
+        fields["completion_mode_counts"] = counts
+    log(stage, "mode routing summary", **fields)
+
+
+def log_tensor(stage: str, name: str, tensor: Any) -> None:
+    if not _DEBUG_ENABLED:
+        return
+    log(stage, f"tensor `{name}`", value=_fmt(tensor))
+
+
+def log_exception(stage: str, msg: str, exc: BaseException) -> None:
+    if not _DEBUG_ENABLED:
+        return
+    tb = traceback.format_exc()
+    log(stage, msg, error=repr(exc), traceback=_fmt(tb, max_len=4000))
+
+
+@contextmanager
+def timed(stage: str, msg: str = "", **fields: Any):
+    if not _DEBUG_ENABLED:
+        yield
+        return
+    t0 = time.perf_counter()
+    log(stage, f"START {msg}", **fields)
+    try:
+        yield
+    except Exception as exc:
+        log_exception(stage, f"FAILED {msg}", exc)
+        raise
+    finally:
+        elapsed = time.perf_counter() - t0
+        log(stage, f"END {msg}", elapsed_s=f"{elapsed:.4f}")

opsd_utils/deepspeed_utils.py

@@ -0,0 +1,141 @@
+"""Helpers for DeepSpeed + Accelerate launch detection."""
+from __future__ import annotations
+
+import json
+import os
+import re
+from pathlib import Path
+from typing import Optional
+
+
+def _project_root() -> Path:
+    return Path(__file__).resolve().parents[1]
+
+
+def resolve_accelerate_config_path(config_name: Optional[str] = None) -> Optional[Path]:
+    candidates: list[str] = []
+    if config_name:
+        candidates.append(str(config_name).strip())
+    for env_key in ("ACCELERATE_CONFIG", "ACCELERATE_CONFIG_FILE"):
+        val = os.environ.get(env_key, "").strip()
+        if val:
+            candidates.append(val)
+    for raw in candidates:
+        if not raw:
+            continue
+        path = Path(raw)
+        if not path.is_file():
+            path = _project_root() / raw
+        if path.is_file():
+            return path
+    return None
+
+
+def uses_deepspeed_json_file(config_name: Optional[str] = None) -> bool:
+    """True when Accelerate loads DeepSpeed settings from an external JSON file."""
+    path = resolve_accelerate_config_path(config_name)
+    if path is None:
+        return False
+    return "deepspeed_config_file" in path.read_text(encoding="utf-8")
+
+
+def _yaml_get_str(path: Path, key: str) -> Optional[str]:
+    pattern = re.compile(rf"^{re.escape(key)}\s*:\s*(.+?)\s*$", re.IGNORECASE)
+    for line in path.read_text(encoding="utf-8").splitlines():
+        m = pattern.match(line.strip())
+        if m:
+            return m.group(1).strip().strip("'\"")
+    return None
+
+
+def is_deepspeed_accelerate_config(config_name: Optional[str] = None) -> bool:
+    path = resolve_accelerate_config_path(config_name)
+    if path is None:
+        return False
+    dist = (_yaml_get_str(path, "distributed_type") or "").upper()
+    return dist == "DEEPSPEED"
+
+
+def deepspeed_zero_stage(config_name: Optional[str] = None) -> Optional[int]:
+    path = resolve_accelerate_config_path(config_name)
+    if path is None:
+        return None
+    text = path.read_text(encoding="utf-8")
+    m = re.search(r"zero_stage\s*:\s*(\d+)", text, re.IGNORECASE)
+    if m:
+        return int(m.group(1))
+    m = re.search(r"deepspeed_config_file\s*:\s*(\S+)", text, re.IGNORECASE)
+    if not m:
+        return None
+    json_path = Path(m.group(1).strip().strip("'\""))
+    if not json_path.is_file():
+        json_path = _project_root() / json_path
+    if not json_path.is_file():
+        return None
+    ds_json = json.loads(json_path.read_text(encoding="utf-8"))
+    stage = (ds_json.get("zero_optimization") or {}).get("stage")
+    return int(stage) if stage is not None else None
+
+
+def should_colocate_teacher_with_student(device_map: Optional[str] = None) -> bool:
+    """True when frozen teacher should sit on the same GPU as the trainable student."""
+    raw = (device_map or os.environ.get("DYME_TEACHER_DEVICE_MAP", "")).strip().lower()
+    if raw in ("same", "colocate", "local"):
+        return True
+    if os.environ.get("DYME_DEEPSPEED_COLOCATE", "").strip().lower() in ("1", "true", "yes", "on"):
+        return True
+    if is_deepspeed_accelerate_config() and raw in ("", "auto"):
+        return True
+    return False
+
+
+def gradient_checkpointing_enable_kwargs(config_name: Optional[str] = None) -> Optional[dict]:
+    """
+    Kwargs for ``model.gradient_checkpointing_enable``.
+
+    DeepSpeed ZeRO-1/2 + reentrant checkpointing runs backward twice per segment and
+    hits: "parameter ... has already been reduced".
+    """
+    if not is_deepspeed_accelerate_config(config_name):
+        return None
+    override = os.environ.get("DYME_GRADIENT_CHECKPOINTING_USE_REENTRANT", "").strip().lower()
+    if override in ("1", "true", "yes", "on"):
+        return {"use_reentrant": True}
+    if override in ("0", "false", "no", "off"):
+        return {"use_reentrant": False}
+    return {"use_reentrant": False}
+
+
+def deepspeed_requires_single_student_forward(config_name: Optional[str] = None) -> bool:
+    """
+    DeepSpeed ZeRO-1/2 cannot reduce gradients when the student runs multiple
+    forwards in one backward (GRPO micro-chunks + OPSD loop).
+    """
+    stage = deepspeed_zero_stage(config_name)
+    return stage is not None and stage <= 2
+
+
+def should_disable_gradient_checkpointing(config_name: Optional[str] = None) -> bool:
+    """Gradient checkpointing also triggers double reduction under ZeRO-1/2."""
+    return deepspeed_requires_single_student_forward(config_name)
+
+
+def student_forward_chunk_size(
+    batch_size: int,
+    has_vision: bool,
+    config_name: Optional[str] = None,
+) -> int:
+    """
+    Micro-batch size for student forwards in ``_get_per_token_logps``.
+
+    Under ZeRO-1/2 we must use one forward per backward (full local batch by default).
+    Override with ``DYME_STUDENT_FORWARD_CHUNK`` only if you accept ZeRO-3+ or OOM risk.
+    """
+    if not has_vision:
+        return batch_size
+    if not deepspeed_requires_single_student_forward(config_name):
+        return 1
+    override = os.environ.get("DYME_STUDENT_FORWARD_CHUNK", "").strip()
+    if override.isdigit():
+        return max(1, min(batch_size, int(override)))
+    return batch_size

opsd_utils/diagnostics.py

@@ -0,0 +1,1351 @@
+"""Periodic full-detail diagnostics for weak reward / gradient signals."""
+from __future__ import annotations
+
+import os
+import re
+from typing import Any, Optional
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from opsd_utils import debug_log as opsd_debug
+from opsd_utils.vocab_align import align_cross_model_logits
+
+PAREN_TOKEN_ID = 340
+
+# Reuse logits from the OPSD loss path instead of running extra forwards.
+_OPSD_JSD_DETAIL_CAPTURE: dict[str, Any] = {
+    "active": False,
+    "global_step": None,
+    "target_indices": set(),
+    "entries": [],
+    "skipped_memory": False,
+    "skip_reason": "",
+    "max_samples": 2,
+}
+
+
+def _detail_min_free_gib() -> float:
+    raw = os.environ.get("DYME_OPSD_DETAIL_MIN_FREE_GB", "").strip()
+    if not raw:
+        return 4.0
+    try:
+        return max(0.0, float(raw))
+    except ValueError:
+        return 4.0
+
+
+def cuda_free_gib(device: Optional[torch.device | int] = None) -> Optional[float]:
+    if not torch.cuda.is_available():
+        return None
+    try:
+        if device is None:
+            free_bytes, _ = torch.cuda.mem_get_info()
+        else:
+            dev = torch.device(device) if not isinstance(device, torch.device) else device
+            with torch.cuda.device(dev):
+                free_bytes, _ = torch.cuda.mem_get_info()
+        return free_bytes / (1024**3)
+    except Exception:
+        return None
+
+
+def check_detail_cuda_memory(
+    min_free_gib: Optional[float] = None,
+    device: Optional[torch.device | int] = None,
+) -> tuple[bool, str, Optional[float]]:
+    """Return (ok, reason, free_gib). Skips heavy detail work when GPU headroom is low."""
+    threshold = _detail_min_free_gib() if min_free_gib is None else max(0.0, float(min_free_gib))
+    if not torch.cuda.is_available():
+        return True, "", None
+    free_gib = cuda_free_gib(device)
+    if free_gib is None:
+        return True, "", None
+    if free_gib < threshold:
+        return (
+            False,
+            f"cuda_free_gib={free_gib:.2f} < min_free_gib={threshold:.2f}",
+            free_gib,
+        )
+    return True, "", free_gib
+
+
+def begin_opsd_jsd_detail_capture(
+    global_step: int,
+    opsd_indices: list[int],
+    max_samples: int = 2,
+) -> None:
+    """Prepare to record JSD stats during OPSD loss (no extra model forwards)."""
+    _OPSD_JSD_DETAIL_CAPTURE.update(
+        active=False,
+        global_step=global_step,
+        target_indices=set(),
+        entries=[],
+        skipped_memory=False,
+        skip_reason="",
+        max_samples=max(1, int(max_samples)),
+    )
+    if not opsd_debug.should_log_detail(global_step) or not opsd_indices:
+        return
+
+    ok, reason, free_gib = check_detail_cuda_memory()
+    if not ok:
+        _OPSD_JSD_DETAIL_CAPTURE["skipped_memory"] = True
+        _OPSD_JSD_DETAIL_CAPTURE["skip_reason"] = reason
+        opsd_debug.log_detail(
+            "opsd_jsd",
+            "skip JSD detail capture (CUDA memory guard)",
+            global_step=global_step,
+            reason=reason,
+            cuda_free_gib=free_gib,
+            min_free_gib=_detail_min_free_gib(),
+        )
+        return
+
+    _OPSD_JSD_DETAIL_CAPTURE["active"] = True
+    _OPSD_JSD_DETAIL_CAPTURE["target_indices"] = set(opsd_indices[: _OPSD_JSD_DETAIL_CAPTURE["max_samples"]])
+
+
+def maybe_capture_opsd_jsd_detail(
+    *,
+    global_idx: int,
+    student_logits: torch.Tensor,
+    teacher_logits: torch.Tensor,
+    completion_mask: torch.Tensor,
+    completion_ids: torch.Tensor,
+    beta: float,
+    tokenizer: Any = None,
+    student_prompt_len: Optional[int] = None,
+    teacher_prompt_len: Optional[int] = None,
+) -> None:
+    """Record token-level JSD stats from logits already computed in the loss path."""
+    capture = _OPSD_JSD_DETAIL_CAPTURE
+    if not capture["active"] or global_idx not in capture["target_indices"]:
+        return
+
+    try:
+        with torch.no_grad():
+            s_logits, t_logits = align_cross_model_logits(
+                student_logits.detach(),
+                teacher_logits.detach(),
+            )
+            stats = jsd_token_stats(s_logits, t_logits, completion_mask.float(), beta=beta)
+        stats["sample_index"] = global_idx
+        if student_prompt_len is not None:
+            stats["student_prompt_len"] = int(student_prompt_len)
+        if teacher_prompt_len is not None:
+            stats["teacher_prompt_len"] = int(teacher_prompt_len)
+        if tokenizer is not None:
+            decoded = tokenizer.decode(
+                completion_ids[0][completion_mask[0].bool()],
+                skip_special_tokens=True,
+            )
+            stats["completion_text"] = _preview_text(decoded)
+        capture["entries"].append(stats)
+    except RuntimeError as exc:
+        if "out of memory" in str(exc).lower():
+            if torch.cuda.is_available():
+                torch.cuda.empty_cache()
+            opsd_debug.log_detail(
+                "opsd_jsd",
+                f"skip JSD detail for sample[{global_idx}] (OOM during stats)",
+                global_step=capture.get("global_step"),
+                error=repr(exc),
+            )
+            return
+        raise
+
+
+def _tensor_stats(t: torch.Tensor, name: str) -> dict[str, Any]:
+    if t is None or not isinstance(t, torch.Tensor) or t.numel() == 0:
+        return {name: "empty"}
+    with torch.no_grad():
+        flat = t.detach().float().reshape(-1)
+        return {
+            f"{name}/shape": tuple(t.shape),
+            f"{name}/mean": float(flat.mean().item()),
+            f"{name}/std": float(flat.std(unbiased=False).item()) if flat.numel() > 1 else 0.0,
+            f"{name}/min": float(flat.min().item()),
+            f"{name}/max": float(flat.max().item()),
+            f"{name}/abs_mean": float(flat.abs().mean().item()),
+        }
+
+
+def _preview_text(text: str, max_len: int = 320) -> str:
+    text = (text or "").replace("\n", "\\n")
+    if len(text) > max_len:
+        return text[: max_len - 3] + "..."
+    return text or "<EMPTY>"
+
+
+def _generation_config_summary(generation_config: Any) -> dict[str, Any]:
+    if generation_config is None:
+        return {}
+    keys = (
+        "max_new_tokens",
+        "do_sample",
+        "temperature",
+        "top_p",
+        "top_k",
+        "min_p",
+        "repetition_penalty",
+        "eos_token_id",
+        "pad_token_id",
+        "bos_token_id",
+    )
+    out: dict[str, Any] = {}
+    for k in keys:
+        if hasattr(generation_config, k):
+            v = getattr(generation_config, k)
+            out[k] = v.tolist() if isinstance(v, torch.Tensor) else v
+    return out
+
+
+def _slice_generate_inputs(batch: dict[str, Any], index: int, batch_size: int) -> dict[str, Any]:
+    """Take one row from batched generate/forward tensors (VLM-safe)."""
+    sliced: dict[str, Any] = {}
+    for key, value in batch.items():
+        if isinstance(value, torch.Tensor) and value.dim() >= 1 and value.size(0) == batch_size:
+            sliced[key] = value[index : index + 1]
+        else:
+            sliced[key] = value
+    return sliced
+
+
+def _last_valid_prompt_positions(prompt_mask: torch.Tensor) -> torch.Tensor:
+    """Return index of last valid (non-pad) prompt token per row; supports left padding."""
+    with torch.no_grad():
+        seq_len = prompt_mask.size(1)
+        positions = torch.arange(seq_len, device=prompt_mask.device).expand_as(prompt_mask)
+        masked = torch.where(prompt_mask.bool(), positions, torch.full_like(positions, -1))
+        return masked.max(dim=1).values
+
+
+def _max_same_token_run(token_ids: list[int]) -> tuple[int, Optional[int]]:
+    """Longest run of identical consecutive token ids."""
+    if not token_ids:
+        return 0, None
+    best_run = 1
+    best_tok = token_ids[0]
+    run = 1
+    for i in range(1, len(token_ids)):
+        if token_ids[i] == token_ids[i - 1]:
+            run += 1
+            if run > best_run:
+                best_run = run
+                best_tok = token_ids[i]
+        else:
+            run = 1
+    return best_run, best_tok
+
+
+def _detect_single_token_repeat(token_ids: list[int], min_run: int = 8) -> bool:
+    return _max_same_token_run(token_ids)[0] >= min_run
+
+
+def _detect_char_repeat(text: str, min_run: int = 6) -> bool:
+    """Detect consecutive repeated characters (CJK or ASCII), e.g. 其其其."""
+    if not text or len(text) < min_run:
+        return False
+    run = 1
+    for i in range(1, len(text)):
+        if text[i] == text[i - 1] and not text[i].isspace():
+            run += 1
+            if run >= min_run:
+                return True
+        else:
+            run = 1
+    return False
+
+
+def _count_char_repeat_samples(completions: Optional[list[str]], min_run: int = 6) -> int:
+    if not completions:
+        return 0
+    return sum(1 for c in completions if _detect_char_repeat(c or "", min_run=min_run))
+
+
+def _detect_repeat_loop(token_ids: list[int], min_repeats: int = 4, ngram: int = 3) -> bool:
+    if len(token_ids) < ngram * min_repeats:
+        return False
+    last_start = len(token_ids) - ngram * min_repeats + 1
+    for start in range(max(0, last_start)):
+        gram = token_ids[start : start + ngram]
+        repeats = 1
+        pos = start + ngram
+        while pos + ngram <= len(token_ids) and token_ids[pos : pos + ngram] == gram:
+            repeats += 1
+            pos += ngram
+        if repeats >= min_repeats:
+            return True
+    return False
+
+
+def _detect_degeneration(
+    token_ids: list[int],
+    text: str,
+    *,
+    answer_flag: str = "Answer:",
+    min_single_token_run: int = 8,
+    require_answer_flag: bool = True,
+) -> tuple[bool, list[str]]:
+    """Heuristics for repetitive / format-broken completions."""
+    reasons: list[str] = []
+    if _detect_single_token_repeat(token_ids, min_run=min_single_token_run):
+        run_len, tok = _max_same_token_run(token_ids)
+        reasons.append(f"SINGLE_TOKEN_REPEAT(run={run_len},tok={tok})")
+    if _detect_repeat_loop(token_ids):
+        reasons.append("NGRAM_REPEAT")
+    if len(token_ids) >= 16:
+        unique_ratio = len(set(token_ids)) / len(token_ids)
+        if unique_ratio < 0.12:
+            reasons.append(f"LOW_UNIQUE_RATIO({unique_ratio:.3f})")
+    if require_answer_flag:
+        answer_count = len(re.findall(f"(?i){re.escape(answer_flag)}", text or ""))
+        if answer_count != 1:
+            reasons.append(f"ANSWER_FLAG_COUNT({answer_count})")
+    if _detect_char_repeat(text or ""):
+        reasons.append("CHAR_REPEAT")
+    return bool(reasons), reasons
+
+
+def _count_answer_flag(text: str, answer_flag: str = "Answer:") -> int:
+    return len(re.findall(f"(?i){re.escape(answer_flag)}", text or ""))
+
+
+def is_degenerate_completion(
+    token_ids: list[int],
+    text: str,
+    *,
+    answer_flag: str = "Answer:",
+    min_single_token_run: int = 8,
+    require_answer_flag: bool = True,
+) -> bool:
+    """Return True when completion looks like a repetition / format-broken sample."""
+    is_deg, _ = _detect_degeneration(
+        token_ids,
+        text,
+        answer_flag=answer_flag,
+        min_single_token_run=min_single_token_run,
+        require_answer_flag=require_answer_flag,
+    )
+    return is_deg
+
+
+def _count_paren_then_eos(
+    completion_ids: torch.Tensor,
+    completion_mask: torch.Tensor,
+    eos_id: Optional[int],
+) -> int:
+    if eos_id is None or completion_ids.size(0) == 0:
+        return 0
+    count = 0
+    with torch.no_grad():
+        lengths = completion_mask.sum(dim=1)
+        for i in range(completion_ids.size(0)):
+            eff = int(lengths[i].item())
+            if eff <= 0:
+                continue
+            first = int(completion_ids[i, 0].item())
+            if first != PAREN_TOKEN_ID:
+                continue
+            if eff <= 2:
+                count += 1
+            elif eff >= 2 and int(completion_ids[i, 1].item()) == eos_id:
+                count += 1
+    return count
+
+
+def summarize_generate_probe_stats(
+    completion_ids: torch.Tensor,
+    completion_mask: torch.Tensor,
+    is_eos: torch.Tensor,
+    eos_id: Optional[int],
+    completions: Optional[list[str]] = None,
+    answer_flag: str = "Answer:",
+    max_completion_length: Optional[int] = None,
+) -> dict[str, Any]:
+    with torch.no_grad():
+        lengths = completion_mask.sum(dim=1).float()
+        has_eos = is_eos.any(dim=1)
+    paren_then_eos = _count_paren_then_eos(completion_ids, completion_mask, eos_id)
+    repeat_loop = 0
+    degenerate_count = 0
+    max_run_lengths: list[float] = []
+    unique_ratios: list[float] = []
+    answer_flag_ok = 0
+    clipped_count = 0
+    degenerate_format_count = 0
+    degenerate_repeat_count = 0
+    format_without_thinking = 0
+    for i in range(completion_ids.size(0)):
+        eff = int(lengths[i].item())
+        if eff <= 0:
+            continue
+        ids = completion_ids[i, :eff].tolist()
+        text = completions[i] if completions and i < len(completions) else ""
+        if _detect_repeat_loop(ids) or _detect_single_token_repeat(ids):
+            repeat_loop += 1
+        run_len, _ = _max_same_token_run(ids)
+        max_run_lengths.append(float(run_len))
+        unique_ratios.append(len(set(ids)) / max(len(ids), 1))
+        if _count_answer_flag(text, answer_flag) == 1:
+            answer_flag_ok += 1
+            thinking = (text or "").lower().split(answer_flag.lower())[0]
+            if len(thinking.strip()) < 8:
+                format_without_thinking += 1
+        else:
+            degenerate_format_count += 1
+        if max_completion_length is not None and eff >= max_completion_length - 1:
+            clipped_count += 1
+        is_deg, reasons = _detect_degeneration(ids, text, answer_flag=answer_flag)
+        if is_deg:
+            degenerate_count += 1
+        non_flag = [r for r in reasons if not r.startswith("ANSWER_FLAG")]
+        if non_flag:
+            degenerate_repeat_count += 1
+    char_repeat_count = _count_char_repeat_samples(completions)
+    n = max(completion_ids.size(0), 1)
+    return {
+        "effective_tokens_mean": float(lengths.mean().item()),
+        "char_repeat_count": char_repeat_count,
+        "one_token_count": int((lengths == 1).sum().item()),
+        "paren_then_eos_count": paren_then_eos,
+        "repeat_loop_count": repeat_loop,
+        "eos_terminated_rate": float(has_eos.float().mean().item()),
+        "degenerate_count": degenerate_count,
+        "degenerate_rate": degenerate_count / n,
+        "degenerate_rate_format": degenerate_format_count / n,
+        "degenerate_rate_repeat": degenerate_repeat_count / n,
+        "format_without_thinking_rate": format_without_thinking / n,
+        "max_token_run_mean": float(sum(max_run_lengths) / len(max_run_lengths)) if max_run_lengths else 0.0,
+        "unique_token_ratio_mean": float(sum(unique_ratios) / len(unique_ratios)) if unique_ratios else 0.0,
+        "answer_flag_exactly_once_rate": answer_flag_ok / n,
+        "clipped_count": clipped_count,
+        "clipped_rate": clipped_count / n,
+    }
+
+
+def log_generate_context(
+    *,
+    global_step: int,
+    trainer_step: Optional[int],
+    generate_call_index: int,
+    model: Any,
+    model_wrapped: Any,
+    gradient_checkpointing: bool,
+    generation_config: Any,
+    is_fsdp_enabled: bool,
+    generate_runs_under_no_grad: bool,
+) -> None:
+    if not opsd_debug.should_log_gendbg() or not opsd_debug.probe_log_model_context():
+        return
+
+    opsd_debug.set_detail_step(global_step)
+    dropout_in_train = sum(
+        1 for m in model.modules() if isinstance(m, nn.Dropout) and m.training
+    )
+    opsd_debug.log_gendbg(
+        "context",
+        "generate model context",
+        generate_call_index=generate_call_index,
+        trainer_step=trainer_step,
+        global_step=global_step,
+        model_training=bool(getattr(model, "training", None)),
+        model_wrapped_training=bool(getattr(model_wrapped, "training", None)),
+        gradient_checkpointing=gradient_checkpointing,
+        generation_use_cache=getattr(generation_config, "use_cache", None),
+        is_fsdp_enabled=is_fsdp_enabled,
+        generate_runs_under_no_grad=generate_runs_under_no_grad,
+        dropout_modules_in_train=dropout_in_train,
+        generation_config=_generation_config_summary(generation_config),
+    )
+
+
+def log_prompt_tail_probe(
+    *,
+    global_step: int,
+    trainer_step: Optional[int],
+    generate_call_index: int,
+    prompt_ids: torch.Tensor,
+    prompt_mask: torch.Tensor,
+    tokenizer: Any,
+    sample_count: int = 4,
+    tail_tokens: Optional[int] = None,
+) -> None:
+    if not opsd_debug.should_log_gendbg():
+        return
+
+    opsd_debug.set_detail_step(global_step)
+    n_tail = tail_tokens if tail_tokens is not None else opsd_debug.probe_prompt_tail_tokens()
+    last_pos = _last_valid_prompt_positions(prompt_mask)
+    n = min(sample_count, prompt_ids.size(0))
+
+    for i in range(n):
+        end = int(last_pos[i].item()) + 1
+        start = max(0, end - n_tail)
+        tail_ids = prompt_ids[i, start:end].tolist()
+        eff_len = int(prompt_mask[i].sum().item())
+        tail_decode = tokenizer.decode(tail_ids, skip_special_tokens=False)
+        opsd_debug.log_gendbg(
+            "prompt_tail",
+            f"sample[{i}]",
+            generate_call_index=generate_call_index,
+            trainer_step=trainer_step,
+            prompt_effective_len=eff_len,
+            last_valid_idx=end - 1,
+            prompt_tail_token_ids=tail_ids,
+            prompt_tail_decode=_preview_text(tail_decode, 400),
+        )
+
+
+def summarize_first_token_logits_stats(
+    p_greedy_values: list[float],
+    p_eos_values: list[float],
+    entropy_values: list[float],
+    p_answer_values: Optional[list[float]] = None,
+) -> dict[str, float]:
+    """Aggregate first-token logit probe scalars across samples."""
+    def _mean(vals: list[float]) -> float:
+        return float(sum(vals) / len(vals)) if vals else 0.0
+
+    out = {
+        "p_greedy_first": _mean(p_greedy_values),
+        "p_eos_first": _mean(p_eos_values),
+        "entropy_first": _mean(entropy_values),
+    }
+    if p_answer_values:
+        out["p_answer_first"] = _mean(p_answer_values)
+    return out
+
+
+def answer_first_token_id(tokenizer: Any) -> Optional[int]:
+    for piece in ("Answer", "Answer:"):
+        ids = tokenizer.encode(piece, add_special_tokens=False)
+        if ids:
+            return int(ids[0])
+    return None
+
+
+def log_first_token_logits_probe(
+    *,
+    global_step: int,
+    trainer_step: Optional[int],
+    generate_call_index: int,
+    unwrapped_model: Any,
+    prompt_inputs_generate: dict[str, Any],
+    prompt_mask: torch.Tensor,
+    tokenizer: Any,
+    sample_count: int = 4,
+) -> dict[str, Any]:
+    """Forward once before generate; return greedy ids and aggregated logit stats."""
+    greedy_by_sample: dict[int, int] = {}
+    p_greedy_vals: list[float] = []
+    p_eos_vals: list[float] = []
+    entropy_vals: list[float] = []
+    p_answer_vals: list[float] = []
+    if not opsd_debug.should_log_gendbg() or not opsd_debug.probe_first_token_logits():
+        return {
+            "greedy_by_sample": greedy_by_sample,
+            **summarize_first_token_logits_stats([], [], [], []),
+        }
+
+    opsd_debug.set_detail_step(global_step)
+    eos_id = getattr(tokenizer, "eos_token_id", None)
+    answer_tid = answer_first_token_id(tokenizer)
+    forward_inputs = {k: v for k, v in prompt_inputs_generate.items() if k != "labels"}
+    batch_size = prompt_mask.size(0)
+    n = min(sample_count, batch_size)
+    last_pos = _last_valid_prompt_positions(prompt_mask)
+
+    for i in range(n):
+        try:
+            sample_inputs = _slice_generate_inputs(forward_inputs, i, batch_size)
+            with torch.no_grad():
+                outputs = unwrapped_model(**sample_inputs, use_cache=False)
+                logits = outputs.logits
+                pos = int(last_pos[i].item())
+                next_logits = logits[0, pos, :].float()
+                probs = F.softmax(next_logits, dim=-1)
+                greedy_id = int(next_logits.argmax().item())
+                greedy_by_sample[i] = greedy_id
+                entropy = float(-(probs * (probs + 1e-12).log()).sum().item())
+                fields: dict[str, Any] = {
+                    "generate_call_index": generate_call_index,
+                    "trainer_step": trainer_step,
+                    "last_prompt_idx": pos,
+                    "greedy_token_id": greedy_id,
+                    "p_greedy": float(probs[greedy_id].item()),
+                    "entropy": entropy,
+                    "probe_mode": "per_sample_forward",
+                }
+                if eos_id is not None:
+                    fields["p_eos"] = float(probs[eos_id].item())
+                if answer_tid is not None and answer_tid < probs.size(0):
+                    fields["p_answer_first"] = float(probs[answer_tid].item())
+                    p_answer_vals.append(fields["p_answer_first"])
+                if PAREN_TOKEN_ID < probs.size(0):
+                    fields["p_token_340"] = float(probs[PAREN_TOKEN_ID].item())
+                topk = torch.topk(probs, k=min(5, probs.size(0)))
+                fields["top5"] = [
+                    (int(tid), float(p)) for tid, p in zip(topk.indices.tolist(), topk.values.tolist())
+                ]
+                p_greedy_vals.append(fields["p_greedy"])
+                entropy_vals.append(entropy)
+                if eos_id is not None:
+                    p_eos_vals.append(fields["p_eos"])
+                opsd_debug.log_gendbg("first_token_logits", f"sample[{i}]", **fields)
+        except Exception as exc:
+            opsd_debug.log_gendbg(
+                "first_token_logits",
+                f"FAILED forward for sample[{i}]",
+                generate_call_index=generate_call_index,
+                error=repr(exc),
+            )
+            if torch.cuda.is_available():
+                torch.cuda.empty_cache()
+    logits_stats = summarize_first_token_logits_stats(
+        p_greedy_vals, p_eos_vals, entropy_vals, p_answer_vals
+    )
+    return {"greedy_by_sample": greedy_by_sample, **logits_stats}
+
+
+def log_first_token_logits_match(
+    *,
+    generate_call_index: int,
+    completion_ids: torch.Tensor,
+    greedy_by_sample: dict[int, int],
+    sample_count: int = 4,
+) -> None:
+    """Compare pre-generate greedy next token vs actual first generated token."""
+    if not opsd_debug.should_log_gendbg() or not opsd_debug.probe_first_token_logits():
+        return
+
+    n = min(sample_count, completion_ids.size(0))
+    for i in range(n):
+        if i not in greedy_by_sample or completion_ids.size(1) == 0:
+            continue
+        greedy_id = greedy_by_sample[i]
+        actual_id = int(completion_ids[i, 0].item())
+        opsd_debug.log_gendbg(
+            "first_token_match",
+            f"sample[{i}]",
+            generate_call_index=generate_call_index,
+            greedy_token_id=greedy_id,
+            actual_first_token=actual_id,
+            greedy_matches_actual=(greedy_id == actual_id),
+        )
+
+
+def log_generate_delta(
+    *,
+    generate_call_index: int,
+    current_stats: dict[str, Any],
+    previous_stats: Optional[dict[str, Any]],
+) -> None:
+    if not opsd_debug.should_log_gendbg():
+        return
+
+    fields: dict[str, Any] = {
+        "generate_call_index": generate_call_index,
+        "current": current_stats,
+    }
+    if previous_stats is not None:
+        prev_idx = previous_stats.get("generate_call_index", generate_call_index - 1)
+        fields["prev_generate_call_index"] = prev_idx
+        for key in (
+            "effective_tokens_mean",
+            "one_token_count",
+            "paren_then_eos_count",
+            "repeat_loop_count",
+            "eos_terminated_rate",
+            "degenerate_count",
+            "degenerate_rate",
+            "clipped_rate",
+            "answer_flag_exactly_once_rate",
+        ):
+            cur = current_stats.get(key)
+            prev = previous_stats.get(key)
+            if cur is not None and prev is not None:
+                if isinstance(cur, (int, float)) and isinstance(prev, (int, float)):
+                    fields[f"delta_{key}"] = cur - prev
+    opsd_debug.log_gendbg("delta", "generate stats vs previous regenerate", **fields)
+
+
+def log_generate_probe(
+    *,
+    global_step: int,
+    trainer_step: Optional[int],
+    prompt_length: int,
+    prompt_completion_ids: torch.Tensor,
+    completion_ids: torch.Tensor,
+    completion_mask: torch.Tensor,
+    is_eos: torch.Tensor,
+    eos_idx: torch.Tensor,
+    completions: list[str],
+    tokenizer: Any,
+    generation_config: Any,
+    max_completion_length: int,
+    num_generations: int,
+    sample_count: int = 4,
+    generate_call_index: Optional[int] = None,
+    answer_flag: str = "Answer:",
+    source: str = "generate",
+) -> dict[str, Any]:
+    """Emit [OPSD-PROBE] on every (re)generate — catches 1-token / empty completions early."""
+    stats = summarize_generate_probe_stats(
+        completion_ids,
+        completion_mask,
+        is_eos,
+        getattr(tokenizer, "eos_token_id", None),
+        completions=completions,
+        answer_flag=answer_flag,
+        max_completion_length=max_completion_length,
+    )
+    if generate_call_index is not None:
+        stats["generate_call_index"] = generate_call_index
+
+    if not opsd_debug.should_log_probe():
+        return stats
+
+    opsd_debug.set_detail_step(global_step)
+    tok = tokenizer
+    eos_id = getattr(tok, "eos_token_id", None)
+    pad_id = getattr(tok, "pad_token_id", None)
+    bos_id = getattr(tok, "bos_token_id", None)
+
+    with torch.no_grad():
+        lengths = completion_mask.sum(dim=1).float()
+        has_eos = is_eos.any(dim=1)
+        raw_gen_len = completion_ids.size(1)
+
+    opsd_debug.log_probe(
+        "generate",
+        "sft cold-start GT batch" if source == "sft_cold_start" else "raw generate summary",
+        trainer_step=trainer_step,
+        global_step=global_step,
+        source=source,
+        generate_call_index=generate_call_index,
+        prompt_length=prompt_length,
+        prompt_completion_shape=tuple(prompt_completion_ids.shape),
+        completion_ids_shape=tuple(completion_ids.shape),
+        raw_gen_tokens=raw_gen_len,
+        max_completion_length=max_completion_length,
+        num_generations=num_generations,
+        batch_size=completion_ids.size(0),
+        effective_tokens_mean=stats["effective_tokens_mean"],
+        effective_tokens_min=float(lengths.min().item()),
+        effective_tokens_max=float(lengths.max().item()),
+        zero_length_count=int((lengths == 0).sum().item()),
+        one_token_count=stats["one_token_count"],
+        paren_then_eos_count=stats["paren_then_eos_count"],
+        repeat_loop_count=stats["repeat_loop_count"],
+        eos_terminated_rate=stats["eos_terminated_rate"],
+        degenerate_count=stats["degenerate_count"],
+        degenerate_rate=stats["degenerate_rate"],
+        max_token_run_mean=stats["max_token_run_mean"],
+        unique_token_ratio_mean=stats["unique_token_ratio_mean"],
+        answer_flag_exactly_once_rate=stats["answer_flag_exactly_once_rate"],
+        clipped_count=stats["clipped_count"],
+        clipped_rate=stats["clipped_rate"],
+        tokenizer_eos_id=eos_id,
+        tokenizer_pad_id=pad_id,
+        tokenizer_bos_id=bos_id,
+        generation_config=_generation_config_summary(generation_config),
+    )
+
+    suspicious: list[str] = []
+    n = min(sample_count, len(completions))
+    for i in range(n):
+        eff = int(lengths[i].item())
+        eidx = int(eos_idx[i].item())
+        ids_head = completion_ids[i, : min(16, completion_ids.size(1))].tolist()
+        ids_all = completion_ids[i].tolist()
+        decode_skip = completions[i]
+        decode_keep = tok.decode(completion_ids[i], skip_special_tokens=False)
+        first_tok = int(completion_ids[i, 0].item()) if completion_ids.size(1) > 0 else None
+        first_is_eos = first_tok is not None and first_tok == eos_id
+        flags: list[str] = []
+        patterns: list[str] = []
+        if eff <= 0:
+            flags.append("ZERO_LEN")
+        if eff == 1:
+            flags.append("ONE_TOKEN")
+        if not (decode_skip or "").strip():
+            flags.append("EMPTY_DECODE")
+        if first_is_eos:
+            flags.append("FIRST_IS_EOS")
+        if eff <= 2 and first_tok == PAREN_TOKEN_ID:
+            patterns.append("PAREN_THEN_EOS")
+        if eff > 0:
+            ids_eff = completion_ids[i, :eff].tolist()
+            if _detect_repeat_loop(ids_eff):
+                patterns.append("REPEAT_LOOP")
+            if _detect_single_token_repeat(ids_eff):
+                run_len, repeat_tok_id = _max_same_token_run(ids_eff)
+                patterns.append(f"SINGLE_TOKEN_REPEAT({run_len}x{repeat_tok_id})")
+            is_deg, deg_reasons = _detect_degeneration(ids_eff, decode_skip, answer_flag=answer_flag)
+            if is_deg:
+                patterns.extend(deg_reasons)
+        if flags:
+            suspicious.append(f"sample[{i}]:{','.join(flags)}")
+
+        opsd_debug.log_probe(
+            "generate",
+            f"sample[{i}]",
+            group=i // num_generations if num_generations else 0,
+            effective_tokens=eff,
+            eos_idx=eidx,
+            has_eos=bool(has_eos[i].item()),
+            first_token_id=first_tok,
+            first_is_eos=first_is_eos,
+            token_ids_head=ids_head,
+            token_ids_all=ids_all if len(ids_all) <= 32 else ids_head + ["..."],
+            decode_skip_special=_preview_text(decode_skip, 400),
+            decode_keep_special=_preview_text(decode_keep, 400),
+            flags=flags or None,
+            patterns=patterns or None,
+        )
+
+    if suspicious:
+        opsd_debug.log_probe(
+            "generate",
+            "ALERT suspicious completions",
+            count=len(suspicious),
+            samples=suspicious,
+            hint="check eos_token_id, max_new_tokens, model.generate output, or training-time unwrap/FSDP",
+        )
+
+    if stats.get("degenerate_rate", 0) >= 0.25:
+        opsd_debug.log_probe(
+            "generate",
+            "ALERT repetition / format degeneration",
+            degenerate_count=stats["degenerate_count"],
+            degenerate_rate=stats["degenerate_rate"],
+            clipped_rate=stats["clipped_rate"],
+            answer_flag_exactly_once_rate=stats["answer_flag_exactly_once_rate"],
+            max_token_run_mean=stats["max_token_run_mean"],
+            hint="typical RL collapse: raise repetition_penalty, lower temperature, or shorten max_completion_length",
+        )
+
+    return stats
+
+
+def log_cross_rank_generate_summary(
+    *,
+    accelerator: Any,
+    one_token_count: int,
+    effective_tokens_mean: float,
+    generate_call_index: int,
+) -> None:
+    """Gather per-rank generate stats; all ranks must call, log on rank 0 only."""
+    if not opsd_debug.probe_on_generate():
+        return
+
+    from accelerate.utils import gather as accel_gather
+
+    device = accelerator.device
+    world_size = accelerator.num_processes
+    local = torch.tensor(
+        [float(one_token_count), effective_tokens_mean],
+        dtype=torch.float32,
+        device=device,
+    )
+    # gather_for_metrics on 1D [2] concatenates ranks into [world_size*2]; use gather + view instead.
+    gathered = accel_gather(local.unsqueeze(0))
+    if not opsd_debug.should_log_gendbg():
+        return
+
+    if gathered.dim() == 1:
+        gathered = gathered.view(world_size, -1)
+    elif gathered.size(0) != world_size and gathered.numel() == world_size * 2:
+        gathered = gathered.reshape(world_size, 2)
+    one_tokens = gathered[:, 0].tolist()
+    means = gathered[:, 1].tolist()
+    opsd_debug.log_gendbg(
+        "cross_rank",
+        "generate summary across ranks",
+        generate_call_index=generate_call_index,
+        world_size=world_size,
+        one_token_count_per_rank=one_tokens,
+        effective_tokens_mean_per_rank=means,
+        one_token_count_total=int(sum(one_tokens)),
+    )
+
+
+def log_routed_completion_probe(
+    *,
+    global_step: int,
+    trainer_step: Optional[int],
+    raw_completion_shape: tuple[int, ...],
+    final_completion_ids: torch.Tensor,
+    final_completion_mask: torch.Tensor,
+    opsd_mask_list: list[bool],
+    sample_count: int = 4,
+    tokenizer: Any,
+    sft_replaced_list: Optional[list[bool]] = None,
+    raw_completion_ids: Optional[torch.Tensor] = None,
+) -> None:
+    """Compare raw vs post-routing padded completion shapes (detect routing truncation)."""
+    if not opsd_debug.should_log_probe():
+        return
+
+    with torch.no_grad():
+        final_lengths = final_completion_mask.sum(dim=1).float()
+
+    opsd_debug.log_probe(
+        "route",
+        "post-routing completion shapes",
+        trainer_step=trainer_step,
+        global_step=global_step,
+        raw_completion_shape=raw_completion_shape,
+        final_completion_shape=tuple(final_completion_ids.shape),
+        final_mask_shape=tuple(final_completion_mask.shape),
+        final_effective_tokens_mean=float(final_lengths.mean().item()),
+        final_effective_tokens_min=float(final_lengths.min().item()),
+        final_effective_tokens_max=float(final_lengths.max().item()),
+        opsd_mask_true=sum(opsd_mask_list),
+        opsd_mask_false=len(opsd_mask_list) - sum(opsd_mask_list),
+        sft_replaced_count=sum(sft_replaced_list) if sft_replaced_list else None,
+    )
+
+    n = min(sample_count, final_completion_ids.size(0))
+    for i in range(n):
+        eff = int(final_lengths[i].item())
+        head = final_completion_ids[i, : min(12, final_completion_ids.size(1))].tolist()
+        text = tokenizer.decode(
+            final_completion_ids[i][final_completion_mask[i].bool()],
+            skip_special_tokens=True,
+        )
+        raw_head = None
+        raw_matches_final = None
+        if raw_completion_ids is not None and i < raw_completion_ids.size(0):
+            raw_eff = min(int(raw_completion_ids.size(1)), eff)
+            raw_head = raw_completion_ids[i, : min(12, raw_eff)].tolist()
+            raw_matches_final = raw_head == head if raw_head and head else None
+        opsd_debug.log_probe(
+            "route",
+            f"routed_sample[{i}]",
+            opsd_mask=opsd_mask_list[i] if i < len(opsd_mask_list) else None,
+            sft_replaced=sft_replaced_list[i] if sft_replaced_list and i < len(sft_replaced_list) else None,
+            effective_tokens=eff,
+            token_ids_head=head,
+            raw_token_ids_head=raw_head,
+            raw_matches_routed_head=raw_matches_final,
+            decode=_preview_text(text, 300),
+        )
+
+
+def jsd_token_stats(
+    student_logits: torch.Tensor,
+    teacher_logits: torch.Tensor,
+    mask: torch.Tensor,
+    beta: float = 0.5,
+    top_k: int = 5,
+) -> dict[str, Any]:
+    """Token-level JSD breakdown without building the graph."""
+    with torch.no_grad():
+        student_log_probs = F.log_softmax(student_logits, dim=-1)
+        teacher_log_probs = F.log_softmax(teacher_logits, dim=-1)
+        student_probs = student_log_probs.exp()
+        teacher_probs = teacher_log_probs.exp()
+
+        if beta == 0:
+            jsd = F.kl_div(student_log_probs, teacher_log_probs, reduction="none", log_target=True)
+        elif beta == 1:
+            jsd = F.kl_div(teacher_log_probs, student_log_probs, reduction="none", log_target=True)
+        else:
+            beta_t = torch.tensor(beta, dtype=student_log_probs.dtype, device=student_log_probs.device)
+            mixture_log_probs = torch.logsumexp(
+                torch.stack([student_log_probs + torch.log1p(-beta_t), teacher_log_probs + torch.log(beta_t)]),
+                dim=0,
+            )
+            kl_teacher = F.kl_div(mixture_log_probs, teacher_log_probs, reduction="none", log_target=True)
+            kl_student = F.kl_div(mixture_log_probs, student_log_probs, reduction="none", log_target=True)
+            jsd = beta_t * kl_teacher + (1 - beta_t) * kl_student
+
+        per_token_jsd = jsd.sum(dim=-1)
+        m = mask.float()
+        valid = m.sum().clamp(min=1.0)
+        per_token_jsd_masked = per_token_jsd * m
+
+        # Top-1 agreement on completion tokens
+        s_top = student_logits.argmax(dim=-1)
+        t_top = teacher_logits.argmax(dim=-1)
+        agree = ((s_top == t_top).float() * m).sum() / valid
+
+        # Mean L2 distance of log-probs on gold completion tokens (if provided elsewhere)
+        logprob_l2 = ((student_log_probs - teacher_log_probs) ** 2).sum(dim=-1)
+        logprob_l2_masked = (logprob_l2 * m).sum() / valid
+
+        # Entropy gap
+        s_ent = -(student_probs * student_log_probs).sum(dim=-1)
+        t_ent = -(teacher_probs * teacher_log_probs).sum(dim=-1)
+        ent_gap = ((s_ent - t_ent).abs() * m).sum() / valid
+
+        stats: dict[str, Any] = {
+            "jsd_per_token_mean": float((per_token_jsd_masked.sum() / valid).item()),
+            "jsd_per_token_max": float(per_token_jsd[m.bool()].max().item()) if m.any() else 0.0,
+            "jsd_per_token_min": float(per_token_jsd[m.bool()].min().item()) if m.any() else 0.0,
+            "top1_agreement": float(agree.item()),
+            "logprob_l2_mean": float(logprob_l2_masked.item()),
+            "entropy_gap_mean": float(ent_gap.item()),
+            "mask_valid_tokens": int(valid.item()),
+        }
+
+        if m.any():
+            idx = int((per_token_jsd * m).argmax().item())
+            pos = idx % per_token_jsd.size(-1)
+            stats["max_jsd_token_pos"] = pos
+            stats["max_jsd_value"] = float(per_token_jsd.reshape(-1)[idx].item())
+            s_topk = torch.topk(student_probs[0, pos], k=min(top_k, student_probs.size(-1)))
+            t_topk = torch.topk(teacher_probs[0, pos], k=min(top_k, teacher_probs.size(-1)))
+            stats["student_topk_at_max_jsd"] = [
+                (int(i), float(p)) for i, p in zip(s_topk.indices.tolist(), s_topk.values.tolist())
+            ]
+            stats["teacher_topk_at_max_jsd"] = [
+                (int(i), float(p)) for i, p in zip(t_topk.indices.tolist(), t_topk.values.tolist())
+            ]
+        return stats
+
+
+def log_generation_diagnostics(
+    *,
+    global_step: int,
+    completions: list[str],
+    completion_ids: torch.Tensor,
+    completion_mask: torch.Tensor,
+    is_eos: torch.Tensor,
+    max_completion_length: int,
+    num_generations: int,
+    sample_count: int = 4,
+) -> None:
+    if not opsd_debug.should_log_detail(global_step):
+        return
+
+    opsd_debug.log_detail_banner(global_step, "GENERATION & COMPLETION")
+
+    with torch.no_grad():
+        lengths = completion_mask.sum(dim=1).float()
+        has_eos = is_eos.any(dim=1)
+        eos_rate = has_eos.float().mean().item()
+        clipped = (~has_eos).float().mean().item()
+        all_pad_or_zero = (lengths == 0).float().mean().item()
+        at_max_len = (lengths >= max_completion_length - 1).float().mean().item()
+
+    fields: dict[str, Any] = {
+        "batch_size": completion_ids.size(0),
+        "num_generations": num_generations,
+        "completion_max_len": completion_ids.size(1),
+        "effective_tokens_mean": float(lengths.mean().item()),
+        "effective_tokens_min": float(lengths.min().item()),
+        "effective_tokens_max": float(lengths.max().item()),
+        "eos_terminated_rate": eos_rate,
+        "clipped_no_eos_rate": clipped,
+        "zero_length_rate": all_pad_or_zero,
+        "at_max_length_rate": at_max_len,
+    }
+    opsd_debug.log_detail("generation", "completion mask summary", **fields)
+
+    n = min(sample_count, len(completions))
+    for i in range(n):
+        gen_group = i // num_generations if num_generations else 0
+        opsd_debug.log_detail(
+            "generation",
+            f"sample[{i}]",
+            group=gen_group,
+            effective_tokens=int(lengths[i].item()),
+            has_eos=bool(has_eos[i].item()),
+            text=_preview_text(completions[i]),
+            raw_token_head=completion_ids[i, :12].tolist(),
+        )
+
+
+def log_reward_diagnostics(
+    *,
+    global_step: int,
+    format_rewards: torch.Tensor,
+    acc_rewards: torch.Tensor,
+    context_rewards: torch.Tensor,
+    all_rewards: torch.Tensor,
+    advantages: torch.Tensor,
+    reward_weights: torch.Tensor,
+    num_generations: int,
+    answers: Optional[list[str]] = None,
+    completions: Optional[list[str]] = None,
+    sample_count: int = 4,
+) -> None:
+    if not opsd_debug.should_log_detail(global_step):
+        return
+
+    opsd_debug.log_detail_banner(global_step, "REWARD & ADVANTAGE")
+
+    w = reward_weights.detach().float()
+    weighted = (
+        format_rewards * w[0] + context_rewards * w[1] + acc_rewards * w[2]
+    )
+
+    fields: dict[str, Any] = {
+        "reward_weights": w.tolist(),
+        "format_sum": float(format_rewards.sum().item()),
+        "acc_sum": float(acc_rewards.sum().item()),
+        "context_sum": float(context_rewards.sum().item()),
+        "total_sum": float(all_rewards.sum().item()),
+        "format_zero_rate": float((format_rewards == 0).float().mean().item()),
+        "acc_zero_rate": float((acc_rewards == 0).float().mean().item()),
+        "context_zero_rate": float((context_rewards == 0).float().mean().item()),
+        "weighted_mean": float(weighted.mean().item()),
+        "weighted_std": float(weighted.std(unbiased=False).item()) if weighted.numel() > 1 else 0.0,
+    }
+    adv_flat = advantages.reshape(-1)
+    fields.update(
+        {
+            "advantage_mean": float(adv_flat.mean().item()),
+            "advantage_std": float(adv_flat.std(unbiased=False).item()) if adv_flat.numel() > 1 else 0.0,
+            "advantage_zero_rate": float((adv_flat.abs() < 1e-8).float().mean().item()),
+            "advantage_min": float(adv_flat.min().item()),
+            "advantage_max": float(adv_flat.max().item()),
+        }
+    )
+    opsd_debug.log_detail("reward", "aggregate reward stats", **fields)
+
+    n = min(sample_count, format_rewards.numel())
+    for i in range(n):
+        g = i // num_generations if num_generations else 0
+        extra: dict[str, Any] = {
+            "group": g,
+            "format": float(format_rewards.view(-1)[i].item()),
+            "acc": float(acc_rewards.view(-1)[i].item()),
+            "context": float(context_rewards.view(-1)[i].item()),
+            "weighted": float(weighted.view(-1)[i].item()),
+            "advantage": float(adv_flat.view(-1)[i].item()) if i < adv_flat.numel() else None,
+        }
+        if answers and i < len(answers):
+            extra["gold_answer"] = _preview_text(str(answers[i // num_generations]), 80)
+        if completions and i < len(completions):
+            extra["completion"] = _preview_text(completions[i], 160)
+        opsd_debug.log_detail("reward", f"per_sample[{i}]", **extra)
+
+
+def log_routing_diagnostics(
+    *,
+    global_step: int,
+    opsd_active: bool,
+    opsd_mask_list: list[bool],
+    has_correct: torch.Tensor,
+    completion_modes: Optional[list[int]] = None,
+    recoverable_flags: Optional[list[bool]] = None,
+    completion_advantages: Optional[torch.Tensor] = None,
+    completion_mask: Optional[torch.Tensor] = None,
+) -> None:
+    if not opsd_debug.should_log_detail(global_step):
+        return
+
+    opsd_debug.log_detail_banner(global_step, "OPSD ROUTING & MASK")
+
+    opsd_true = sum(opsd_mask_list)
+    fields: dict[str, Any] = {
+        "opsd_active": opsd_active,
+        "opsd_mask_true": opsd_true,
+        "opsd_mask_false": len(opsd_mask_list) - opsd_true,
+        "opsd_mask_ratio": opsd_true / max(len(opsd_mask_list), 1),
+        "has_correct": has_correct.tolist() if hasattr(has_correct, "tolist") else has_correct,
+    }
+    if recoverable_flags is not None:
+        fields["recoverable_flags"] = recoverable_flags
+    if completion_modes is not None:
+        from opsd_utils.debug_log import MODE_NAMES
+
+        counts = {MODE_NAMES.get(c, str(c)): completion_modes.count(c) for c in set(completion_modes)}
+        fields["completion_mode_counts"] = counts
+
+    if completion_advantages is not None and completion_mask is not None:
+        with torch.no_grad():
+            pos = ((completion_advantages > 0) & (completion_mask > 0)).float().sum(dim=1)
+            neg = ((completion_advantages < 0) & (completion_mask > 0)).float().sum(dim=1)
+            zero_adv = ((completion_advantages.abs() < 1e-8) & (completion_mask > 0)).float().sum(dim=1)
+            fields["adv_pos_tokens_mean"] = float(pos.mean().item())
+            fields["adv_neg_tokens_mean"] = float(neg.mean().item())
+            fields["adv_zero_tokens_mean"] = float(zero_adv.mean().item())
+
+    opsd_debug.log_detail("routing", "mode routing summary", **fields)
+
+
+def log_loss_diagnostics(
+    *,
+    global_step: int,
+    grpo_loss: torch.Tensor,
+    per_token_logps: torch.Tensor,
+    old_per_token_logps: torch.Tensor,
+    completion_mask: torch.Tensor,
+    advantages: torch.Tensor,
+    coef_1: torch.Tensor,
+    per_token_loss: torch.Tensor,
+    opsd_loss: Optional[torch.Tensor] = None,
+    combined_loss: Optional[torch.Tensor] = None,
+    opsd_mask: Optional[torch.Tensor] = None,
+    epsilon_low: float = 0.2,
+    epsilon_high: float = 0.2,
+) -> None:
+    if not opsd_debug.should_log_detail(global_step):
+        return
+
+    opsd_debug.log_detail_banner(global_step, "LOSS & GRADIENT SIGNAL")
+
+    with torch.no_grad():
+        m = completion_mask.float()
+        valid_per_sample = m.sum(dim=1).clamp(min=1.0)
+        sample_grpo = (per_token_loss * m).sum(dim=1) / valid_per_sample
+
+        fields: dict[str, Any] = {
+            "grpo_loss_scalar": float(grpo_loss.detach().item()),
+            "grpo_per_sample_mean": float(sample_grpo.mean().item()),
+            "grpo_per_sample_max": float(sample_grpo.max().item()),
+            "completion_mask_tokens_mean": float(valid_per_sample.mean().item()),
+        }
+        fields.update(_tensor_stats(advantages, "advantages"))
+        fields.update(_tensor_stats(per_token_logps, "per_token_logps"))
+        fields.update(_tensor_stats(old_per_token_logps, "old_per_token_logps"))
+        fields.update(_tensor_stats((per_token_logps - old_per_token_logps) * m, "logps_delta_masked"))
+        fields.update(_tensor_stats(coef_1 * m, "coef_1_masked"))
+
+        low_clip = ((coef_1 < 1 - epsilon_low) & (advantages.unsqueeze(1) < 0) & (m > 0)).float().sum()
+        high_clip = ((coef_1 > 1 + epsilon_high) & (advantages.unsqueeze(1) > 0) & (m > 0)).float().sum()
+        fields["clipped_low_tokens"] = int(low_clip.item())
+        fields["clipped_high_tokens"] = int(high_clip.item())
+        fields["grpo_zero_loss_rate"] = float((sample_grpo.abs() < 1e-12).float().mean().item())
+
+        if opsd_loss is not None:
+            fields["opsd_loss_scalar"] = float(opsd_loss.detach().item())
+        if combined_loss is not None:
+            fields["combined_loss_scalar"] = float(combined_loss.detach().item())
+        if opsd_mask is not None:
+            fields["opsd_batch_count"] = int(opsd_mask.sum().item())
+
+        # Weak-signal hints
+        hints: list[str] = []
+        if fields.get("advantages/abs_mean", 1.0) < 1e-6:
+            hints.append("advantages≈0 → GRPO per-token loss≈0")
+        if fields.get("grpo_zero_loss_rate", 0) > 0.9:
+            hints.append("most samples have ~0 GRPO loss")
+        if fields.get("completion_mask_tokens_mean", 0) < 2:
+            hints.append("very few effective completion tokens in mask")
+        if opsd_loss is not None and abs(fields.get("opsd_loss_scalar", 0)) < 1e-8:
+            hints.append("OPSD JSD≈0 → student/teacher nearly identical on completion")
+        fields["weak_signal_hints"] = hints or ["none"]
+
+    opsd_debug.log_detail("loss", "GRPO / OPSD loss breakdown", **fields)
+
+
+def summarize_loss_health(
+    *,
+    grpo_loss: torch.Tensor,
+    per_token_logps: torch.Tensor,
+    completion_mask: torch.Tensor,
+    advantages: torch.Tensor,
+    per_token_loss: torch.Tensor,
+    opsd_loss: Optional[torch.Tensor] = None,
+    combined_loss: Optional[torch.Tensor] = None,
+) -> dict[str, float]:
+    """Lightweight loss/signal summary for every-step health monitoring."""
+    with torch.no_grad():
+        m = completion_mask.float()
+        valid_per_sample = m.sum(dim=1).clamp(min=1.0)
+        sample_grpo = (per_token_loss * m).sum(dim=1) / valid_per_sample
+        adv_flat = advantages.detach().float().reshape(-1)
+        fields: dict[str, float] = {
+            "grpo_loss_scalar": float(grpo_loss.detach().item()),
+            "grpo_zero_loss_rate": float((sample_grpo.abs() < 1e-12).float().mean().item()),
+            "advantages_abs_mean": float(adv_flat.abs().mean().item()) if adv_flat.numel() else 0.0,
+            "completion_mask_tokens_mean": float(valid_per_sample.mean().item()),
+        }
+        if opsd_loss is not None:
+            fields["opsd_loss_scalar"] = float(opsd_loss.detach().item())
+        if combined_loss is not None:
+            fields["combined_loss_scalar"] = float(combined_loss.detach().item())
+        logps_delta = (per_token_logps * m).sum() / m.sum().clamp(min=1.0)
+        fields["logps_delta_mean"] = float(logps_delta.item()) if m.sum() > 0 else 0.0
+    return fields
+
+
+def summarize_batch_data_health(
+    samples: list[dict[str, Any]],
+    *,
+    prompt_mask: Optional[torch.Tensor] = None,
+    pixel_values: Optional[torch.Tensor] = None,
+) -> dict[str, Any]:
+    """Batch-level data I/O sanity for health monitoring."""
+    n = max(len(samples), 1)
+    vf_empty = 0
+    prompt_lens: list[int] = []
+    for sample in samples:
+        vf = (
+            sample.get("visual_fact_hint")
+            or sample.get("visual_fact")
+            or sample.get("visual_facts")
+            or ""
+        )
+        if not str(vf).strip():
+            vf_empty += 1
+        if sample.get("prompt"):
+            prompt_lens.append(len(str(sample["prompt"])))
+
+    out: dict[str, Any] = {
+        "visual_fact_empty_rate": vf_empty / n,
+        "batch_size": len(samples),
+        "prompt_len_mean": float(sum(prompt_lens) / len(prompt_lens)) if prompt_lens else 0.0,
+    }
+
+    if prompt_mask is not None:
+        with torch.no_grad():
+            lengths = prompt_mask.sum(dim=1).float()
+            out["prompt_tokens_mean"] = float(lengths.mean().item())
+            out["prompt_tokens_max"] = float(lengths.max().item())
+
+    if pixel_values is not None and isinstance(pixel_values, torch.Tensor) and pixel_values.numel() > 0:
+        with torch.no_grad():
+            flat = pixel_values.detach().float().reshape(-1)
+            out["pixel_mean"] = float(flat.mean().item())
+            out["pixel_has_nan"] = bool(torch.isnan(flat).any().item())
+            out["pixel_has_inf"] = bool(torch.isinf(flat).any().item())
+
+    return out
+
+
+def log_opsd_jsd_diagnostics(*, global_step: int) -> None:
+    """Emit cached JSD stats recorded during OPSD loss (zero extra model forwards)."""
+    if not opsd_debug.should_log_detail(global_step):
+        return
+
+    capture = _OPSD_JSD_DETAIL_CAPTURE
+    if capture.get("global_step") != global_step:
+        return
+
+    opsd_debug.log_detail_banner(global_step, "OPSD JSD DECOMPOSITION")
+
+    if capture.get("skipped_memory"):
+        opsd_debug.log_detail(
+            "opsd_jsd",
+            "JSD detail skipped (CUDA memory guard)",
+            global_step=global_step,
+            reason=capture.get("skip_reason", ""),
+            min_free_gib=_detail_min_free_gib(),
+            cuda_free_gib=cuda_free_gib(),
+        )
+        _OPSD_JSD_DETAIL_CAPTURE["active"] = False
+        return
+
+    entries = capture.get("entries") or []
+    if not entries:
+        opsd_debug.log_detail(
+            "opsd_jsd",
+            "no JSD detail captured (no OPSD samples on this rank or capture disabled)",
+            global_step=global_step,
+        )
+        _OPSD_JSD_DETAIL_CAPTURE["active"] = False
+        return
+
+    for k, stats in enumerate(entries):
+        opsd_debug.log_detail("opsd_jsd", f"sample[{k}] token-level JSD", global_step=global_step, **stats)
+
+    _OPSD_JSD_DETAIL_CAPTURE["active"] = False

opsd_utils/leakage.py

@@ -0,0 +1,53 @@
+"""Detect OPSD / completion patterns that indicate privileged-information leakage."""
+from __future__ import annotations
+
+import re
+from typing import Any
+
+_LEAKAGE_PHRASES = (
+    "reference answer",
+    "reference reasoning",
+    "according to the reference",
+    "according to the answer",
+    "参考答案",
+)
+
+_ANSWER_IN_COMPLETION = re.compile(r"(?i)answer:\s*.+")
+
+
+def completion_has_leakage_pattern(
+    text: str,
+    gold_answer: str | None = None,
+    *,
+    min_gold_substring_len: int = 4,
+) -> bool:
+    """Return True if completion text suggests privileged-info cheating."""
+    if not text or not text.strip():
+        return False
+    lower = text.lower()
+    for phrase in _LEAKAGE_PHRASES:
+        if phrase in lower:
+            return True
+    if gold_answer:
+        gold = gold_answer.strip()
+        if len(gold) >= min_gold_substring_len and gold.lower() in lower:
+            # Short numeric answers may false-positive; require Answer: prefix context
+            if _ANSWER_IN_COMPLETION.search(text):
+                return False
+            if len(gold) >= 8:
+                return True
+    return False
+
+
+def privileged_suffix_has_gold(suffix: str, sample: dict[str, Any]) -> bool:
+    if not suffix.strip():
+        return False
+    if "[Reference Answer]" in suffix or "[Reference Reasoning]" in suffix:
+        return True
+    answer = (sample.get("answer") or "").strip()
+    hint = (sample.get("hint") or "").strip()
+    if answer and answer in suffix:
+        return True
+    if hint and len(hint) >= 8 and hint in suffix:
+        return True
+    return False

opsd_utils/opsd_loss.py

@@ -0,0 +1,439 @@
+import torch
+import torch.nn.functional as F
+
+from opsd_utils import debug_log as opsd_debug
+from opsd_utils import diagnostics as opsd_diagnostics
+from opsd_utils.deepspeed_utils import deepspeed_requires_single_student_forward
+from opsd_utils.teacher_batching import (
+    align_teacher_prompt_image_tokens,
+    as_batch_num_images_tensor,
+    get_teacher_vision_for_sample,
+    model_inference_device,
+    move_batch_num_images_to_model_device,
+    move_pixel_values_to_model_device,
+    student_batch_num_images_tensor,
+)
+from opsd_utils.vocab_align import align_cross_model_logits
+
+
+def _slice_image_sizes(image_sizes, index: int):
+    """Slice per-sample image_sizes for student path (one image per batch row)."""
+    if image_sizes is None:
+        return None
+    if isinstance(image_sizes, torch.Tensor):
+        if image_sizes.dim() == 0:
+            return image_sizes
+        return image_sizes[index : index + 1]
+    if isinstance(image_sizes, (list, tuple)):
+        return image_sizes[index]
+    return image_sizes
+
+
+def _slice_image_sizes_batch(image_sizes, start: int, end: int):
+    """Slice image_sizes for a micro-batch row range [start, end)."""
+    if image_sizes is None:
+        return None
+    if isinstance(image_sizes, torch.Tensor):
+        if image_sizes.dim() == 0:
+            return image_sizes
+        if image_sizes.shape[0] >= end:
+            return image_sizes[start:end]
+        return image_sizes
+    if isinstance(image_sizes, (list, tuple)):
+        return image_sizes[start:end] if len(image_sizes) >= end else image_sizes
+    return image_sizes
+
+
+def _teacher_image_counts(inputs: dict, batch_size: int) -> list[int]:
+    """Number of teacher images per batch sample (LLaVA-OV stacks images on dim 0)."""
+    counts = inputs.get("teacher_num_images")
+    if counts is None:
+        return [1] * batch_size
+    if isinstance(counts, torch.Tensor):
+        return [int(max(1, c)) for c in counts.detach().cpu().tolist()]
+    return [int(max(1, c)) for c in counts]
+
+
+def slice_teacher_vision_inputs(
+    teacher_pixel_values,
+    teacher_image_sizes,
+    local: int,
+    num_images_per_sample: list[int],
+):
+    """
+    Slice teacher pixel_values / image_sizes for one batch sample.
+    LLaVA-OneVision uses dim-0 = total images across batch (not batch size).
+    """
+    if teacher_pixel_values is None:
+        return None, None
+    start = sum(num_images_per_sample[:local])
+    end = start + num_images_per_sample[local]
+    t_pixel = teacher_pixel_values[start:end]
+    t_sizes = None
+    if teacher_image_sizes is not None and isinstance(teacher_image_sizes, torch.Tensor):
+        t_sizes = teacher_image_sizes[start:end]
+    return t_pixel, t_sizes
+
+
+def generalized_jsd_loss(student_logits, teacher_logits, mask, beta=0.5):
+    """Token-level generalized JSD on completion positions."""
+    # Cross-model OPD: teacher logits already live on the teacher GPU; avoid
+    # copying them onto the student GPU (vocab × seq is multi-hundred MiB per sample).
+    jsd_device = teacher_logits.device
+    if student_logits.device != jsd_device:
+        student_logits = student_logits.to(jsd_device, non_blocking=True)
+    mask = mask.to(device=jsd_device, non_blocking=True)
+
+    comp_dtype = student_logits.dtype
+    if comp_dtype == torch.float32:
+        comp_dtype = torch.bfloat16
+    if student_logits.dtype != comp_dtype:
+        student_logits = student_logits.to(comp_dtype)
+    if teacher_logits.dtype != comp_dtype:
+        teacher_logits = teacher_logits.to(comp_dtype)
+
+    student_logits, teacher_logits = align_cross_model_logits(student_logits, teacher_logits)
+    student_log_probs = F.log_softmax(student_logits, dim=-1)
+    teacher_log_probs = F.log_softmax(teacher_logits, dim=-1)
+    opsd_debug.log(
+        "vocab_align",
+        "generalized_jsd_loss log_softmax on aligned vocab",
+        student_log_prob_shape=tuple(student_log_probs.shape),
+        teacher_log_prob_shape=tuple(teacher_log_probs.shape),
+        student_exp_sum=float(torch.exp(student_log_probs[0, 0]).sum().item()) if student_log_probs.numel() else None,
+        teacher_exp_sum=float(torch.exp(teacher_log_probs[0, 0]).sum().item()) if teacher_log_probs.numel() else None,
+    )
+
+    if beta == 0:
+        jsd = F.kl_div(student_log_probs, teacher_log_probs, reduction="none", log_target=True)
+    elif beta == 1:
+        jsd = F.kl_div(teacher_log_probs, student_log_probs, reduction="none", log_target=True)
+    else:
+        beta_t = torch.tensor(beta, dtype=student_log_probs.dtype, device=student_log_probs.device)
+        mixture_log_probs = torch.logsumexp(
+            torch.stack([student_log_probs + torch.log1p(-beta_t), teacher_log_probs + torch.log(beta_t)]),
+            dim=0,
+        )
+        kl_teacher = F.kl_div(mixture_log_probs, teacher_log_probs, reduction="none", log_target=True)
+        kl_student = F.kl_div(mixture_log_probs, student_log_probs, reduction="none", log_target=True)
+        jsd = beta_t * kl_teacher + (1 - beta_t) * kl_student
+
+    jsd = jsd.sum(dim=-1)
+    jsd = jsd * mask
+    denom = mask.sum().clamp(min=1.0)
+    return jsd.sum() / denom
+
+
+def _teacher_logits_with_oom_retry(
+    model,
+    processor,
+    teacher_prompt_ids,
+    teacher_prompt_mask,
+    completion_ids,
+    completion_mask,
+    t_pixel,
+    t_sizes,
+    logits_to_keep: int,
+    teacher_batch_num_images=None,
+):
+    """Teacher forward with OOM micro-batch halving (decision E). Batch dim is already 1 in OPSD loop."""
+    if processor is not None:
+        teacher_prompt_ids, teacher_prompt_mask = align_teacher_prompt_image_tokens(
+            model,
+            processor,
+            teacher_prompt_ids,
+            teacher_prompt_mask,
+            t_pixel,
+            t_sizes,
+            batch_num_images=teacher_batch_num_images,
+        )
+    teacher_device = model_inference_device(model)
+    teacher_prompt_ids = teacher_prompt_ids.to(teacher_device)
+    teacher_prompt_mask = teacher_prompt_mask.to(teacher_device)
+    completion_ids = completion_ids.to(teacher_device)
+    completion_mask = completion_mask.to(teacher_device)
+    t_pixel = move_pixel_values_to_model_device(model, t_pixel)
+    teacher_batch_num_images = move_batch_num_images_to_model_device(model, teacher_batch_num_images)
+    teacher_input = torch.cat([teacher_prompt_ids, completion_ids], dim=1)
+    teacher_attn = torch.cat([teacher_prompt_mask, completion_mask], dim=1)
+    oom_retries = 0
+    while True:
+        try:
+            with torch.no_grad():
+                return model(
+                    input_ids=teacher_input,
+                    attention_mask=teacher_attn,
+                    pixel_values=t_pixel,
+                    image_sizes=t_sizes,
+                    batch_num_images=teacher_batch_num_images,
+                ).logits[:, -logits_to_keep - 1 : -1, :]
+        except RuntimeError as exc:
+            if "out of memory" not in str(exc).lower():
+                raise
+            oom_retries += 1
+            opsd_debug.log(
+                "teacher_forward_oom",
+                "teacher OPSD forward OOM, clearing cache and retrying",
+                micro_batch_size=teacher_input.shape[0],
+                oom_retries=oom_retries,
+            )
+            if torch.cuda.is_available():
+                torch.cuda.empty_cache()
+            if oom_retries >= 3:
+                raise
+
+
+def slice_student_completion_logits(full_logits: torch.Tensor, logits_to_keep: int) -> torch.Tensor:
+    """Completion-token logits aligned with ``_get_per_token_logps`` / OPSD JSD."""
+    logits = full_logits[:, -logits_to_keep - 1 :, :]
+    logits = logits[:, :-1, :]
+    return logits[:, -logits_to_keep:, :]
+
+
+def compute_vlm_opsd_loss(
+    model,
+    student_prompt_ids,
+    student_prompt_mask,
+    student_pixel_values,
+    student_image_sizes,
+    teacher_prompt_ids,
+    teacher_prompt_mask,
+    teacher_pixel_values,
+    completion_ids,
+    completion_mask,
+    beta=0.5,
+    teacher_image_sizes=None,
+    processor=None,
+    teacher_batch_num_images=None,
+    teacher_model=None,
+    global_idx: int | None = None,
+    capture_jsd_detail: bool = False,
+    tokenizer=None,
+    student_logits=None,
+) -> torch.Tensor:
+    """
+    OPSD / OPD: student vs teacher prompt, shared student completion.
+    When teacher_model is set, cross-model OPD (e.g. frozen 7B teacher); else self-OPSD.
+    """
+    teacher_model = teacher_model if teacher_model is not None else model
+    opsd_debug.log(
+        "opsd_loss",
+        "compute_vlm_opsd_loss enter",
+        beta=beta,
+        student_prompt_shape=tuple(student_prompt_ids.shape),
+        teacher_prompt_shape=tuple(teacher_prompt_ids.shape),
+        completion_shape=tuple(completion_ids.shape),
+        has_teacher_pixel_values=teacher_pixel_values is not None,
+        teacher_pixel_values_shape=(
+            tuple(teacher_pixel_values.shape) if teacher_pixel_values is not None else None
+        ),
+    )
+    student_batch_num_images = student_batch_num_images_tensor(
+        student_pixel_values, student_prompt_ids.shape[0]
+    )
+    if processor is not None and student_pixel_values is not None:
+        student_prompt_ids, student_prompt_mask = align_teacher_prompt_image_tokens(
+            model,
+            processor,
+            student_prompt_ids,
+            student_prompt_mask,
+            student_pixel_values,
+            student_image_sizes,
+            batch_num_images=student_batch_num_images,
+        )
+
+    student_input = torch.cat([student_prompt_ids, completion_ids], dim=1)
+    student_attn = torch.cat([student_prompt_mask, completion_mask], dim=1)
+
+    logits_to_keep = completion_ids.size(1)
+
+    if student_logits is None:
+        with opsd_debug.timed("opsd_loss", "student forward (grad)"):
+            student_logits = model(
+                input_ids=student_input,
+                attention_mask=student_attn,
+                pixel_values=student_pixel_values,
+                image_sizes=student_image_sizes,
+                batch_num_images=student_batch_num_images,
+            ).logits[:, -logits_to_keep - 1 : -1, :]
+    else:
+        opsd_debug.log(
+            "opsd_loss",
+            "reuse GRPO student completion logits (DeepSpeed single-forward)",
+            student_logits_shape=tuple(student_logits.shape),
+        )
+
+    t_pixel = teacher_pixel_values if teacher_pixel_values is not None else student_pixel_values
+    t_sizes = teacher_image_sizes if teacher_image_sizes is not None else student_image_sizes
+    with opsd_debug.timed("opsd_loss", "teacher forward (no grad)"):
+        teacher_logits = _teacher_logits_with_oom_retry(
+            teacher_model,
+            processor,
+            teacher_prompt_ids,
+            teacher_prompt_mask,
+            completion_ids,
+            completion_mask,
+            t_pixel,
+            t_sizes,
+            logits_to_keep,
+            teacher_batch_num_images=teacher_batch_num_images,
+        )
+
+    cross_model = teacher_model is not model
+    if cross_model:
+        opsd_debug.log(
+            "opsd_loss",
+            "cross-model OPD logits",
+            student_vocab=student_logits.size(-1),
+            teacher_vocab=teacher_logits.size(-1),
+        )
+
+    loss = generalized_jsd_loss(student_logits, teacher_logits, completion_mask.float(), beta=beta)
+
+    if capture_jsd_detail and global_idx is not None:
+        opsd_diagnostics.maybe_capture_opsd_jsd_detail(
+            global_idx=global_idx,
+            student_logits=student_logits,
+            teacher_logits=teacher_logits,
+            completion_mask=completion_mask,
+            completion_ids=completion_ids,
+            beta=beta,
+            tokenizer=tokenizer,
+            student_prompt_len=int(student_prompt_mask.sum().item()),
+            teacher_prompt_len=int(teacher_prompt_mask.sum().item()),
+        )
+
+    del teacher_logits
+    if cross_model and torch.cuda.is_available():
+        torch.cuda.empty_cache()
+
+    opsd_debug.log("opsd_loss", "compute_vlm_opsd_loss done", loss=float(loss.detach().item()))
+    return loss
+
+
+def compute_vlm_opsd_loss_masked_batch(
+    model,
+    opsd_indices: list[int],
+    all_indices: list[int],
+    inputs: dict,
+    beta: float = 0.5,
+    processor=None,
+    teacher_model=None,
+    acc_gate: bool = True,
+    pad_to_count: int | None = None,
+    global_step: int | None = None,
+    tokenizer=None,
+    detail_max_samples: int = 2,
+    student_completion_logits=None,
+) -> torch.Tensor:
+    """Compute mean OPSD loss over opsd_indices within a batch.
+
+    Under DDP every rank must run the *same* number of student/teacher
+    forwards, otherwise the per-forward buffer broadcast (and gradient
+    reduction) collectives desync across ranks and NCCL eventually times out.
+    ``pad_to_count`` is the global-max OPSD sample count across ranks; ranks
+    with fewer (or zero) real samples run extra zero-weighted "dummy" forwards
+    on a valid local row so the collective sequence stays aligned.
+    """
+    real_count = len(opsd_indices)
+    target_count = pad_to_count if pad_to_count is not None else real_count
+    if target_count <= 0:
+        opsd_debug.log("opsd_loss", "compute_vlm_opsd_loss_masked_batch skipped (no OPSD samples)")
+        return torch.tensor(0.0, device=inputs["prompt_ids"].device, requires_grad=True)
+
+    opsd_debug.log(
+        "opsd_loss",
+        "compute_vlm_opsd_loss_masked_batch enter",
+        opsd_indices=opsd_indices,
+        all_indices=all_indices,
+        beta=beta,
+        real_count=real_count,
+        target_count=target_count,
+    )
+    capture_jsd_detail = (
+        global_step is not None and opsd_debug.should_log_detail(global_step)
+    )
+    if capture_jsd_detail:
+        opsd_diagnostics.begin_opsd_jsd_detail_capture(
+            global_step,
+            opsd_indices,
+            max_samples=detail_max_samples,
+        )
+    losses = []
+    idx_map = {g: i for i, g in enumerate(all_indices)}
+    batch_size = inputs["prompt_ids"].shape[0]
+    teacher_img_counts = _teacher_image_counts(inputs, batch_size)
+
+    for step_idx in range(target_count):
+        is_real = step_idx < real_count
+        # Dummy iterations reuse the first available row so shapes stay valid;
+        # their contribution is zeroed out below.
+        global_idx = opsd_indices[step_idx] if is_real else all_indices[0]
+        local = idx_map[global_idx]
+        student_sizes = _slice_image_sizes(inputs.get("img_sizes"), local)
+        t_pixel, teacher_sizes = get_teacher_vision_for_sample(
+            inputs, local, teacher_img_counts
+        )
+        if t_pixel is None:
+            t_pixel = inputs["pixel_values"][local : local + 1]
+            teacher_sizes = student_sizes
+        opsd_debug.log(
+            "opsd_loss",
+            "compute sample OPSD loss",
+            global_idx=global_idx,
+            local_idx=local,
+            teacher_num_images=teacher_img_counts[local],
+            student_image_sizes=student_sizes,
+            teacher_image_sizes=teacher_sizes,
+            teacher_pixel_values_shape=tuple(t_pixel.shape) if t_pixel is not None else None,
+        )
+        n_img = teacher_img_counts[local]
+        teacher_batch_num_images = as_batch_num_images_tensor(n_img, t_pixel)
+        if not is_real and deepspeed_requires_single_student_forward():
+            # ZeRO-1/2: avoid extra student forwards (even loss*0 still backprops).
+            losses.append(torch.zeros((), device=inputs["prompt_ids"].device, requires_grad=True))
+            continue
+        precomputed_student_logits = None
+        if student_completion_logits is not None:
+            precomputed_student_logits = student_completion_logits[local : local + 1]
+        with opsd_debug.timed("opsd_loss", f"sample_opsd_loss idx={global_idx}"):
+            loss = compute_vlm_opsd_loss(
+                model,
+                inputs["prompt_ids"][local : local + 1],
+                inputs["prompt_mask"][local : local + 1],
+                inputs["pixel_values"][local : local + 1],
+                student_sizes,
+                inputs["teacher_prompt_ids"][local : local + 1],
+                inputs["teacher_prompt_mask"][local : local + 1],
+                t_pixel,
+                inputs["completion_ids"][local : local + 1],
+                inputs["completion_mask"][local : local + 1],
+                beta=beta,
+                teacher_image_sizes=teacher_sizes,
+                processor=processor,
+                teacher_batch_num_images=teacher_batch_num_images,
+                teacher_model=teacher_model,
+                global_idx=global_idx if is_real else None,
+                capture_jsd_detail=capture_jsd_detail and is_real,
+                tokenizer=tokenizer,
+                student_logits=precomputed_student_logits,
+            )
+            if not is_real:
+                # Keep the autograd graph / DDP collective alive but contribute nothing.
+                loss = loss * 0.0
+            elif acc_gate and "acc_rewards" in inputs:
+                acc_val = float(inputs["acc_rewards"][global_idx].item())
+                loss = loss * max(0.0, 1.0 - acc_val)
+        losses.append(loss)
+
+    # Mean over real samples only; dummy (zero-weighted) forwards keep the
+    # collective sequence aligned across ranks without skewing the loss scale.
+    mean_loss = torch.stack(losses).sum() / max(real_count, 1)
+    opsd_debug.log(
+        "opsd_loss",
+        "compute_vlm_opsd_loss_masked_batch done",
+        mean_loss=float(mean_loss.detach().item()),
+        real_count=real_count,
+        target_count=target_count,
+    )
+    return mean_loss

opsd_utils/privileged/debug_artifacts.py

@@ -0,0 +1,104 @@
+"""Save privileged teacher images to disk on detail steps."""
+from __future__ import annotations
+
+import json
+import os
+from typing import Any, Optional
+
+from PIL import Image
+
+from opsd_utils import debug_log as opsd_debug
+
+_saved_counts: dict[int, int] = {}
+_output_dir: Optional[str] = None
+_cfg: dict[str, Any] = {}
+
+
+def configure(output_dir: Optional[str] = None, privileged_debug_cfg: Optional[dict[str, Any]] = None) -> None:
+    global _output_dir, _cfg
+    _output_dir = output_dir
+    _cfg = dict(privileged_debug_cfg or {})
+    _saved_counts.clear()
+
+
+def _image_subdir() -> str:
+    return _cfg.get("image_subdir", "logs/images")
+
+
+def maybe_save_privileged_images(
+    global_step: Optional[int],
+    sample_idx: int,
+    full_img: Optional[Image.Image],
+    crop_img: Optional[Image.Image],
+    meta: Optional[dict[str, Any]] = None,
+    output_dir: Optional[str] = None,
+    privileged_debug_cfg: Optional[dict[str, Any]] = None,
+) -> Optional[str]:
+    """
+    Save teacher privileged images when should_log_detail(global_step) is true.
+    Returns base path prefix if saved, else None.
+    """
+    if global_step is None:
+        return None
+    if not opsd_debug.should_log_detail(global_step):
+        return None
+
+    cfg = privileged_debug_cfg if privileged_debug_cfg is not None else _cfg
+    if not cfg.get("save_images", True):
+        return None
+
+    max_samples = int(cfg.get("max_samples_per_detail", 2))
+    count = _saved_counts.get(global_step, 0)
+    if count >= max_samples:
+        return None
+
+    base_out = output_dir or _output_dir
+    if not base_out:
+        opsd_debug.log(
+            "privileged_debug",
+            "skip image save (no output_dir)",
+            global_step=global_step,
+            sample_idx=sample_idx,
+        )
+        return None
+
+    subdir = os.path.join(base_out, _image_subdir() if cfg is _cfg else cfg.get("image_subdir", "logs/images"))
+    os.makedirs(subdir, exist_ok=True)
+
+    prefix = os.path.join(subdir, f"step_{int(global_step):06d}_idx_{sample_idx}")
+    saved_paths: list[str] = []
+
+    if full_img is not None:
+        full_path = f"{prefix}_full.png"
+        full_img.save(full_path)
+        saved_paths.append(full_path)
+
+    if crop_img is not None:
+        crop_path = f"{prefix}_crop.png"
+        crop_img.save(crop_path)
+        saved_paths.append(crop_path)
+
+    meta_path = f"{prefix}_meta.json"
+    meta_payload = dict(meta or {})
+    meta_payload.update(
+        {
+            "global_step": global_step,
+            "sample_idx": sample_idx,
+            "saved_paths": saved_paths,
+        }
+    )
+    with open(meta_path, "w", encoding="utf-8") as f:
+        json.dump(meta_payload, f, ensure_ascii=False, indent=2)
+
+    _saved_counts[global_step] = count + 1
+    opsd_debug.log_detail(
+        "privileged_debug",
+        "privileged images saved",
+        global_step=global_step,
+        sample_idx=sample_idx,
+        prefix=prefix,
+        saved_paths=saved_paths,
+        meta_path=meta_path,
+        **{k: v for k, v in (meta or {}).items() if k not in ("full_size", "crop_size")},
+    )
+    return prefix

opsd_utils/privileged/providers.py

@@ -0,0 +1,127 @@
+from typing import Any
+
+from PIL import Image
+
+from data_utils.chart.deplot_pipeline import format_deplot_for_teacher, is_deplot_placeholder
+from data_utils.privileged_schema import parse_visual_fact
+from opsd_utils.privileged.base import PrivilegedContextProvider
+from opsd_utils.privileged.image_utils import heuristic_crop_from_visual_fact, load_rgb
+
+
+DEFAULT_FORMAT_ONLY_HINT = (
+    "Use the following structure in your response:\n"
+    "Goal: ...\nObservation: ...\nReasoning: ...\nAnswer: ..."
+)
+
+
+class FormatOnlyProvider(PrivilegedContextProvider):
+    """Structure hint only — no gold answer or reference reasoning (anti-leakage)."""
+
+    def __init__(self, hint_text: str | None = None):
+        self._hint_text = (hint_text or DEFAULT_FORMAT_ONLY_HINT).strip()
+
+    def build_teacher_suffix(self, sample: dict[str, Any]) -> str:
+        return self._hint_text
+
+
+class TextProvider(PrivilegedContextProvider):
+    def __init__(self, include_gold: bool = True):
+        self.include_gold = include_gold
+
+    def build_teacher_suffix(self, sample: dict[str, Any]) -> str:
+        if not self.include_gold:
+            return ""
+        parts = []
+        hint = (sample.get("hint") or "").strip()
+        answer = (sample.get("answer") or "").strip()
+        if hint:
+            parts.append(f"[Reference Reasoning]\n{hint}")
+        if answer:
+            parts.append(f"[Reference Answer]\n{answer}")
+        return "\n\n".join(parts)
+
+
+class VisualFactsProvider(PrivilegedContextProvider):
+    """B1: raw JSON visual facts; F1+F2 merge hint and deplot sources."""
+
+    def _collect_visual_fact_parts(self, sample: dict[str, Any]) -> list[str]:
+        parts: list[str] = []
+        hint_vf = sample.get("visual_fact_hint")
+        if hint_vf:
+            text = parse_visual_fact(hint_vf)
+            if text:
+                parts.append(f"[Visual Facts - Hint]\n{text}")
+
+        deplot_vf = sample.get("visual_fact_deplot")
+        if deplot_vf and not is_deplot_placeholder(deplot_vf):
+            text = format_deplot_for_teacher(deplot_vf)
+            if text:
+                parts.append(f"[Visual Facts - DePlot]\n{text}")
+
+        primary = sample.get("visual_fact") or sample.get("visual_facts")
+        if primary and not (hint_vf or deplot_vf):
+            text = parse_visual_fact(primary)
+            if text:
+                parts.append(f"[Visual Facts]\n{text}")
+        elif primary and (hint_vf or deplot_vf):
+            text = parse_visual_fact(primary)
+            if text:
+                parts.append(f"[Visual Facts - Combined]\n{text}")
+
+        return parts
+
+    def build_teacher_suffix(self, sample: dict[str, Any]) -> str:
+        parts = self._collect_visual_fact_parts(sample)
+        return "\n\n".join(parts)
+
+
+class CropProvider(PrivilegedContextProvider):
+    """Returns evidence crop as second teacher image (dual-image path uses image_utils)."""
+
+    def build_teacher_suffix(self, sample: dict[str, Any]) -> str:
+        return ""
+
+    def build_teacher_images(self, sample: dict[str, Any], crop_cfg: dict[str, Any] | None = None) -> list[Image.Image]:
+        image = sample.get("image")
+        if image is None:
+            return []
+        full = load_rgb(image)
+        if full is None:
+            return []
+        crop, _, _ = heuristic_crop_from_visual_fact(full, sample, crop_cfg)
+        return [crop]
+
+
+class HybridProvider(PrivilegedContextProvider):
+    def __init__(
+        self,
+        provider_names: list[str],
+        crop_cfg: dict[str, Any] | None = None,
+        *,
+        text_include_gold: bool = True,
+        format_only_hint: str | None = None,
+    ):
+        self._providers: list[PrivilegedContextProvider] = []
+        self._crop_cfg = crop_cfg or {}
+        for name in provider_names:
+            if name == "text":
+                self._providers.append(TextProvider(include_gold=text_include_gold))
+            elif name == "format_only":
+                self._providers.append(FormatOnlyProvider(format_only_hint))
+            elif name == "visual_facts":
+                self._providers.append(VisualFactsProvider())
+            elif name == "crop":
+                self._providers.append(CropProvider())
+
+    def build_teacher_suffix(self, sample: dict[str, Any]) -> str:
+        chunks = [p.build_teacher_suffix(sample) for p in self._providers]
+        chunks = [c for c in chunks if c.strip()]
+        return "\n\n".join(chunks)
+
+    def build_teacher_images(self, sample: dict[str, Any]) -> list[Image.Image]:
+        for p in self._providers:
+            if isinstance(p, CropProvider):
+                imgs = p.build_teacher_images(sample, self._crop_cfg)
+                if imgs:
+                    return imgs
+        return []