[qb Compiler] backend="torch" combined mxq_compile() fails at device-allocation graph coloring for Qwen2.5-VL (OLMOCR2) — "some operators are missing during coloring"

qb Compiler
1

Product & Environment

  • Product: qbcompiler
  • Compiler Version: 1.2.0
  • OS: Ubuntu 24.04.1 LTS, kernel 6.8.0-106-generic

Description

Summary

Compiling a Qwen2.5-VL-based VLM (allenai/olmOCR-2-7B-1025, Qwen2_5_VLForConditionalGeneration) with backend="torch" and the combined mxq_compile() entry point (parse + quantize in one call, device="gpu", config_preset="multimodal") fails deterministically — every run, same stack trace — right after FX tracing/parsing completes successfully, inside qbcompiler.model_dict.parser.device_alloc.DeviceAllocator._color_graph():

Exception: some operators are missing during coloring.

Parsing itself is not the problem: the exact same model, with the exact same set of workarounds applied, was already successfully parsed (mblt_compile(), export-only, no quantize) into a valid .mblt in an earlier session — 92.91% op coverage, no fatal errors. The failure only appears when the combined mxq_compile() path is used, because that path additionally runs allocate_graph_to_devices() / graph coloring, which mblt_compile() alone never reaches.
I’m filing this because I specifically need GPU-based compilation (not CPU) to work for this model, and this appears to be a compiler-internal bug or an undocumented constraint, not something fixable from user code.

Environment

Component Version
qbcompiler 1.2.0
Python 3.10.21
torch 2.7.1+cu126
torchvision 0.22.1+cu126
transformers 4.57.6
accelerate 1.15.0
einops 0.8.2
psutil 7.2.2
OS Ubuntu 24.04.1 LTS, kernel 6.8.0-106-generic
GPU 1x NVIDIA H200 NVL
Driver 570.211.01 (CUDA 12.8)
GPU partitioning MIG enabled, 2x 1g.18gb instances exposed via CUDA_VISIBLE_DEVICES (both slices together, to fit the 16.6 GB BF16 checkpoint)
Target device aries-rb
Model: allenai/olmOCR-2-7B-1025 — Qwen2_5_VLForConditionalGeneration, 28 decoder layers, hidden_size 3584, ~16.6 GB in BF16 across 4 safetensors shards.

What I’m trying to do

Compile this model end-to-end (parse + quantize) in a single mxq_compile() call with device="gpu", backend="torch", config_preset="multimodal". Since hf_config-based auto-loading only builds a text-only chat internally (no way to include an image, confirmed by reading qbcompiler/compiler/compiler.py’s _parsing_task()), I load the model and processor myself, build a real chat with an image, capture the real forward-pass kwargs via qbcompiler’s own InputCaptureCtxManager/DefaultInputsCaptureContainer, and pass that as feed_dict= directly to mxq_compile(model=<loaded model>, feed_dict=..., ...).
To get FX tracing to succeed at all for this architecture, several workarounds are applied from the compile script itself (no changes to installed qbcompiler/transformers files):

  1. Monkeypatch Qwen2_5_VLForConditionalGeneration.forward at the class level to the inspect.unwrap()‘d real function (works around a torch.fx._patch_function "varnames is too small" CodeType construction bug caused by transformers’ functools.wraps-based forward decorator).
  2. Replace Qwen2_5_VLModel.get_image_features() with a version using qbcompiler’s own VisionModelForQwen2_5_VL wrapper (works around torch.split()/torch.cat() proxy issues and data-dependent fancy indexing in the vision tower’s rotary/window-index computation).
  3. Reimplement repreprocess_pixel_values() (qbcompiler ships an einops-based version, but einops.rearrange() doesn’t recognize qbcompiler’s MbltProxy during tracing) using plain .reshape()/.permute().
  4. Source-patch Qwen2_5_VLModel.forward to split a 2-positional-arg .to(device, dtype) call into two single-arg .to() calls (qbcompiler’s IR builder only supports single-arg .to()).
  5. Source-patch the same forward to replace masked_scatter() (entirely unimplemented in this qbcompiler version — flat NotImplementedError) with a torch.cat splice, using the real, contiguous image-token span precomputed from the captured input_ids.
  6. Replace Qwen2_5_VLModel.get_placeholder_mask() with a version that compares against a torch.tensor scalar instead of a bare Python int (works around make_tensor_cmp() crashing on a plain int).
  7. Swap in qbcompiler’s CachedQwen2_5_VLTextRotaryEmbedding (primed via set_rope()) and PatchedQwen2_5_VLSdpaAttention for the language model’s rotary embedding and every decoder layer’s self-attention (works around the same data-dependent-indexing issue as #2, in the text decoder’s M-RoPE computation).
    All 7 workarounds succeed — FX tracing and parsing complete cleanly (✔ Parsing deep learning model... Done!), reaching a Summary by Operator Type report with HFPatchedFunction: 60 instances (the only meaningfully “uncovered” op type, expected from the autowrapped helper functions above) before the fatal error below.

Full reproduction script

import sys
import types
import inspect
import traceback
import numpy as np
import torch
from PIL import Image
from transformers import AutoProcessor, DynamicCache, Qwen2_5_VLForConditionalGeneration
from qbcompiler import mxq_compile
from qbcompiler.configs import LlmConfig
from qbcompiler.model_dict.parser.backend.hf.util import (
    DefaultInputsCaptureContainer,
    InputCaptureCtxManager,
)
from qbcompiler.model_dict.parser.backend.torch.util import wrap_tensor
from qbcompiler.model_dict.parser.backend.torch.object_wrapper import set_attention_mask
MODEL_PATH = "/path/to/olmOCR2"   # allenai/olmOCR-2-7B-1025, downloaded locally
SAVE_PATH = "/path/to/olmOCR2_trial.mxq"
model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
    MODEL_PATH, torch_dtype=torch.bfloat16, trust_remote_code=True, device_map="auto"
).eval()
processor = AutoProcessor.from_pretrained(MODEL_PATH, trust_remote_code=True)
device = model.device
# Placeholder document page image. Size must be a multiple of 112px on each
# side (VisionModelForQwen2_5_VL.set_grid_thw()'s llm_grid_h/w % vit_win==0
# constraint, vit_win = window_size(112)//spatial_merge_size(2)//patch_size(14) = 4).
main_image = Image.new("RGB", (336, 448), color=(255, 255, 255))
messages = [{"role": "user", "content": [
    {"type": "text", "text": "Attached is one page of a document that you must process. "
                              "Just return the plain text representation of this document as if you were reading it naturally.\n"},
    {"type": "image"},
]}]
text = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
inputs = processor(text=[text], images=[main_image], padding=True, return_tensors="pt")
inputs = {k: v.to(device) for k, v in inputs.items()}
inputs_container = DefaultInputsCaptureContainer()
with InputCaptureCtxManager(model, 1, inputs_container):
    try:
        model.generate(**inputs, max_new_tokens=1, do_sample=False)
    except RuntimeError as e:
        if "max num call limit reached" not in str(e):
            raise
feed_dict_raw = inputs_container.captured_kwargs[-1]
DROP_KEYS = {"inputs_embeds", "pixel_values_videos", "video_grid_thw",
             "second_per_grid_ts", "return_dict", "use_cache"}
fd_inputs = {}
for k, v in feed_dict_raw.items():
    if k in DROP_KEYS or v is None:
        continue
    fd_inputs[k] = wrap_tensor(k, v.to(model.device)) if isinstance(v, torch.Tensor) else v
fd_inputs["input_ids"].src_shape[-1].set_dynamic(True)
fd_inputs["attention_mask"].src_shape[-1].set_dynamic(True)
set_attention_mask(fd_inputs["attention_mask"], "causal_mask")
fd_inputs["position_ids"].src_shape[-1].set_dynamic(True)
fd_inputs["cache_position"].src_shape[0].set_dynamic(True)
fd_inputs.setdefault("past_key_values", DynamicCache())
fd_inputs["logits_to_keep"] = 1
# --- Fix 1: class-level forward unwrap (torch.fx _patch_function bug) ---
ModelClass = type(model)
ModelClass.forward = inspect.unwrap(ModelClass.forward)
# --- Fix 2/3: vision tower wrapper + torch-native repreprocess ---
from transformers.models.qwen2_5_vl.modeling_qwen2_5_vl import Qwen2_5_VLModel
from qbcompiler.model_dict.parser.backend.fx_hf_extensions.transformers.models.qwen2_5_vl import (
    VisionModelForQwen2_5_VL, CachedQwen2_5_VLTextRotaryEmbedding, PatchedQwen2_5_VLSdpaAttention,
)
def repreprocess_pixel_values(pixel_values, grid_thw, patch_size=14, temporal_patch_size=2, merge_size=2, channels=3):
    gt, raw_gh, raw_gw = grid_thw
    Mh = Mw = merge_size; pt = temporal_patch_size; c = channels; ph = pw = patch_size
    gh, gw = raw_gh // Mh, raw_gw // Mw
    x = pixel_values.reshape(gt, gh, gw, Mh, Mw, c, pt, ph, pw)
    x = x.permute(0, 6, 5, 1, 2, 7, 3, 4, 8)
    return x.reshape(gt, pt * c, gh * gw * ph, Mh * Mw * pw)
vlm_model = model.model
real_grid_thw = feed_dict_raw["image_grid_thw"]
gt, gh, gw = (int(v) for v in real_grid_thw[0])
visual_wrapper = VisionModelForQwen2_5_VL(vlm_model)
visual_wrapper.set_grid_thw(real_grid_thw)
visual_wrapper = visual_wrapper.to(next(vlm_model.visual.parameters()).device)
vlm_model.visual_wrapper = visual_wrapper
def patched_get_image_features(self, pixel_values, image_grid_thw=None):
    images = repreprocess_pixel_values(pixel_values, (gt, gh, gw))
    images = images.to(next(self.visual_wrapper.parameters()).dtype)
    image_embeds = self.visual_wrapper(images)
    return (image_embeds.reshape(-1, image_embeds.shape[-1]),)
Qwen2_5_VLModel.get_image_features = patched_get_image_features
# --- Fix 4/5: source-patch forward for .to(device,dtype) split + masked_scatter splice ---
import textwrap, re
real_model_forward = inspect.unwrap(Qwen2_5_VLModel.forward)
src = textwrap.dedent(inspect.getsource(real_model_forward))
src_lines = src.splitlines()
def_idx = next(i for i, line in enumerate(src_lines) if line.lstrip().startswith("def "))
src = "\n".join(src_lines[def_idx:])
for name in ("image_embeds", "video_embeds"):
    src = src.replace(
        f"{name} = torch.cat({name}, dim=0).to(inputs_embeds.device, inputs_embeds.dtype)",
        f"{name} = torch.cat({name}, dim=0).to(inputs_embeds.dtype).to(inputs_embeds.device)",
    )
real_input_ids_row = feed_dict_raw["input_ids"][0].to("cpu")
image_token_id = model.config.image_token_id
image_positions = (real_input_ids_row == image_token_id).nonzero(as_tuple=True)[0]
img_start, img_end = int(image_positions[0]), int(image_positions[0]) + int(image_positions.numel())
pattern = re.compile(
    r"image_mask, _ = self\.get_placeholder_mask\(\s*"
    r"input_ids, inputs_embeds=inputs_embeds, image_features=image_embeds\s*"
    r"\)\s*"
    r"inputs_embeds = inputs_embeds\.masked_scatter\(image_mask, image_embeds\)"
)
new_block = (
    f"inputs_embeds = torch.cat([inputs_embeds[:, :{img_start}, :], "
    f"image_embeds.reshape(1, -1, inputs_embeds.shape[-1]), inputs_embeds[:, {img_end}:, :]], dim=1)"
)
src, _ = pattern.subn(new_block, src)
ns = {}
exec(compile(src, "<patched forward>", "exec"), real_model_forward.__globals__, ns)
Qwen2_5_VLModel.forward = ns["forward"]
# --- Fix 6: scalar-comparison placeholder mask ---
def patched_get_placeholder_mask(self, input_ids, inputs_embeds, image_features=None, video_features=None):
    tok = torch.tensor(self.config.image_token_id, dtype=input_ids.dtype, device=input_ids.device)
    mask = (input_ids == tok).unsqueeze(-1).expand_as(inputs_embeds).to(inputs_embeds.device)
    return mask, torch.zeros_like(mask)
Qwen2_5_VLModel.get_placeholder_mask = patched_get_placeholder_mask
# --- Fix 7: cached M-RoPE + matching attention patch ---
text_model = model.model.language_model
target_device = next(text_model.rotary_emb.buffers()).device
text_model.rotary_emb = text_model.rotary_emb.to("cpu")
cached_rotary = CachedQwen2_5_VLTextRotaryEmbedding(text_model.rotary_emb)
cached_rotary.set_rope(feed_dict_raw["position_ids"].to("cpu"))
text_model.rotary_emb = cached_rotary.to(target_device)
for layer in text_model.layers:
    layer.self_attn = PatchedQwen2_5_VLSdpaAttention(layer.self_attn)
# --- The actual call that fails ---
llm_config = LlmConfig(
    apply=True,
    attributes=LlmConfig.Attributes(
        maxSequenceLength=256, maxCacheLength=256, maxDataLength=256,
        calibration=LlmConfig.Attributes.Calibration(randomSeqLength=256, useFullSeqLength=True),
        runtime=LlmConfig.Attributes.Runtime(batchSize=1, dynamicRope=False, dynamicMask=False),
    ),
)
mxq_compile(
    model=model,
    target_device="aries-rb",
    backend="torch",
    device="gpu",
    config_preset="multimodal",
    feed_dict=fd_inputs,
    output_meta={"type": "list", "keys": [0]},
    llm_config=llm_config,
    use_random_calib=True,
    save_path=SAVE_PATH,
)

Run with (both MIG slices exposed, needed just to load the 16.6 GB checkpoint):

export CUDA_VISIBLE_DEVICES=<mig-uuid-0>,<mig-uuid-1>
python compile_olmocr2.py

Exact error (reproduced twice, byte-identical stack trace both times)

Error occurred in parsing task!
Traceback (most recent call last):
  File ".../qbcompiler/compiler/compiler.py", line 1116, in compile
    result_path = _parsing_task(**_kw)
  File ".../qbcompiler/compiler/compiler.py", line 791, in _parsing_task
    parser.parse(
  File ".../qbcompiler/model_dict/parser/parser.py", line 680, in parse
    self._post_graph_transform(
  File ".../qbcompiler/model_dict/parser/parser.py", line 475, in _post_graph_transform
    md = self.allocate_graph_to_devices(
  File ".../qbcompiler/model_dict/parser/parser.py", line 371, in allocate_graph_to_devices
    ).allocate(sg_main, activation_tracer)
  File ".../qbcompiler/model_dict/parser/device_alloc.py", line 291, in allocate
    self._color_graph()
  File ".../qbcompiler/model_dict/parser/device_alloc.py", line 992, in _color_graph
    raise Exception("some operators are missing during coloring.")
Exception: some operators are missing during coloring.
✔ Parsing deep learning model... Done!
COMPILE FAILED: Exception some operators are missing during coloring.

(Full outer traceback shows this is reached via mxq_compile() → frontend.py:923 model_dict.compile() → frontend.py:535 super().compile() → the same _parsing_task chain above.)
Two non-fatal warnings appear during parsing, immediately before this, and may be related (same op type, fx_patched_fn0, is involved in both):

[ERROR] ViTRopeUpdate failed during transform on op `fx_patched_fn0`: IndexError('list index out of range'). This rule will be skipped from now on.
Traceback (most recent call last):
  File ".../qbcompiler/model_dict/parser/graph_transform.py", line 1170, in process
    res = rule.transform(op, sg, wd)
  File ".../qbcompiler/model_dict/parser/transform_operator/operator_folding/hf_patched.py", line 1074, in transform
    return self._transform0(op, sg, wd)
  File ".../qbcompiler/model_dict/parser/transform_operator/operator_folding/hf_patched.py", line 1084, in _transform0
    sg.activations[op.inputs[0]].name.split("model_blocks_")[1].split("_")[0]
IndexError: list index out of range

This fires twice (same op, same error) and each time qbcompiler logs it and disables that transform rule going forward — it doesn’t crash the run by itself, but it may be leaving the graph in a state that later causes the coloring failure.

Key observation: this is specific to the combined compile/quantize path, not parsing

The exact same model with the exact same 7 workarounds, called instead with the export-only mblt_compile() (same feed_dict, same backend="torch", same device="gpu", same config_preset="multimodal", just producing a .mblt instead of a .mxq) succeeds completely — 92.91% op-coverage report, valid .mblt written to disk, no exceptions. The only difference in the failing run is calling mxq_compile() instead, which internally reaches allocate_graph_to_devices() → _color_graph() — a stage mblt_compile() never exercises.
This tells us:

  • FX tracing, IR construction, and the parser’s own op-coverage checks are all fine for this model + our workarounds.
  • The failure is specifically in the compiler’s device-allocation/graph-coloring stage, which is required for quantization but not for a plain .mblt export.

Questions

  1. What does “some operators are missing during coloring” mean exactly — which operators, and why would they be present for mblt_compile()'s IR construction but missing for mxq_compile()'s device-allocation coloring pass on the identical graph?
  2. Is _color_graph() failing because of the custom/patched ops we introduce (fx_patched_fn0 / HFPatchedFunction, from the workarounds needed to trace this architecture at all), or is this a general gap in backend="torch" + config_preset="multimodal" support for the combined compile path?
  3. Is there a known-working reference example of compiling a Qwen2_5_VLForConditionalGeneration model (or any VLM) end-to-end (not parse-only) with backend="torch" in qbcompiler 1.2.0? The tutorials/model zoo I’ve found don’t appear to cover a Qwen2.5-VL-family torch-backend compile with a custom feed_dict (bypassing hf_config, which doesn’t support image inputs at all — confirmed by reading _parsing_task()'s source, it hardcodes a text-only chat).
  4. Separately (not the focus of this thread, but related): is device="gpu" for the quantize stage of a 7B-class model expected to work at all under a MIG-partitioned H200? We’ve also seen Failed to set device errors from the native FB quantizer on the two-step mxq_compile_V2(model=<persisted .mblt>) path for this same model on GPU (both single-MIG and dual-MIG), and a separate NVML CUDACachingAllocator assertion crash quantizing a different (text-only, 4B) model on GPU under a single MIG slice. Happy to open that as a separate thread if it’s unrelated to the coloring issue above.
    Happy to share the full compile log, the persisted .mblt from the successful parse-only run, or run additional diagnostics — this is currently the only thing blocking us from compiling OLMOCR2 for aries-rb at all with GPU-accelerated compilation.

Logs & Attachments

2

Hello Uday,

  1. What does “some operators are missing during coloring” mean exactly — which operators, and why would they be present for mblt_compile()'s IR construction but missing for mxq_compile()'s device-allocation coloring pass on the identical graph?

    → It seems the model needs special treatment in parsing. I will share the sample parsing code for Qwen2.5-VL for now. In the future we are planning to support OLMOCR2 model natively.

    qwen2_5_vl_using_mblt_compile.py (4.5 KB)

  2. Is _color_graph() failing because of the custom/patched ops we introduce (fx_patched_fn0 / HFPatchedFunction, from the workarounds needed to trace this architecture at all), or is this a general gap in backend="torch" + config_preset="multimodal" support for the combined compile path?

    → Please check answer above.

  3. Is there a known-working reference example of compiling a Qwen2_5_VLForConditionalGeneration model (or any VLM) end-to-end (not parse-only) with backend="torch" in qbcompiler 1.2.0? The tutorials/model zoo I’ve found don’t appear to cover a Qwen2.5-VL-family torch-backend compile with a custom feed_dict (bypassing hf_config, which doesn’t support image inputs at all — confirmed by reading _parsing_task()'s source, it hardcodes a text-only chat).

    → Please check the parsing code provided.

  4. Separately (not the focus of this thread, but related): is device="gpu" for the quantize stage of a 7B-class model expected to work at all under a MIG-partitioned H200? We’ve also seen Failed to set device errors from the native FB quantizer on the two-step mxq_compile_V2(model=<persisted .mblt>) path for this same model on GPU (both single-MIG and dual-MIG), and a separate NVML CUDACachingAllocator assertion crash quantizing a different (text-only, 4B) model on GPU under a single MIG slice. Happy to open that as a separate thread if it’s unrelated to the coloring issue above.
    Happy to share the full compile log, the persisted .mblt from the successful parse-only run, or run additional diagnostics — this is currently the only thing blocking us from compiling OLMOCR2 for aries-rb at all with GPU-accelerated compilation.

    → It depends on how much memory you partitioned for H200 MIG. For 7B model, ~30GB VRAM should be enough at the moment. In the near future, we are planning to provide memory efficient way of compiling a model.