doubleblindsubmission/Monet

1# /Monet/sitecustomize.py2# -----------------------------------------------------------------------3# This patch is only for RL codes (vllm==0.8.5), not for SFT and inference.4# So use `MONET_RL_PATCH=1` only in the RL script.5# -----------------------------------------------------------------------6 7import os, sys, importlib, inspect8 9print(f"[sitecustomize] imported from {__file__}", file=sys.stderr)10 11def patch_qwen_monet():12    # Import official and Monet implementations13    import transformers.models.qwen2_5_vl.modeling_qwen2_5_vl as q_official14    import monet_models.transformers.monet_modeling_qwen2_5_vl as q_monet15    print("[Monet RL patch] replacing...")16    off_cls = q_official.Qwen2_5_VLForConditionalGeneration17    mon_cls = q_monet.Qwen2_5_VLForConditionalGeneration18    q_official.QWEN2_5_VL_ATTENTION_CLASSES["flash_attention_2"] = q_monet.QWEN2_5_VL_ATTENTION_CLASSES["flash_attention_2"]19    q_official.QWEN2_5_VL_ATTENTION_CLASSES["sdpa"] = q_monet.QWEN2_5_VL_ATTENTION_CLASSES["sdpa"]20 21    # Debug: check signatures before patch22    try:23        print(24            "[Monet RL patch] official forward sig before:",25            inspect.signature(off_cls.forward),26            file=sys.stderr,27        )28        print(29            "[Monet RL patch] monet    forward sig:",30            inspect.signature(mon_cls.forward),31            file=sys.stderr,32        )33    except Exception:34        pass35 36    # In-place monkey patch: copy methods from Monet class to official class37    off_cls.forward = mon_cls.forward38    q_official.Qwen2_5_VLModel.forward = q_monet.Qwen2_5_VLModel.forward39    q_official.Qwen2_5_VLFlashAttention2.forward = q_monet.Qwen2_5_VLFlashAttention2.forward40    q_official.Qwen2_5_VLSdpaAttention.forward = q_monet.Qwen2_5_VLSdpaAttention.forward41    q_official.Qwen2_5_VLDecoderLayer.forward = q_monet.Qwen2_5_VLDecoderLayer.forward42    #off_cls.__init__ = mon_cls.__init__43 44    # Debug: check signature after patch45    try:46        print(47            "[Monet RL patch] official forward sig after:",48            inspect.signature(off_cls.forward),49            file=sys.stderr,50        )51        print(52            "[Monet RL patch] q_official.Qwen2_5_VLModel.forward sig after:",53            inspect.signature(q_official.Qwen2_5_VLModel.forward),54            file=sys.stderr,55        )56        print(57            "[Monet RL patch] q_official.Qwen2_5_VLFlashAttention2.forward sig after:",58            inspect.signature(q_official.Qwen2_5_VLFlashAttention2.forward),59            file=sys.stderr,60        )61    except Exception:62        pass63 64    print(65        "[Monet RL patch] Patched methods of Qwen2_5_VLForConditionalGeneration in-place",66        file=sys.stderr,67    )68 69 70def patch():71    print("[sitecustomize] patch() called", file=sys.stderr)72 73    os.environ["VLLM_USE_V1"] = "1"74    os.environ["VLLM_NO_USAGE_STATS"] = "1"75 76    workspace = os.path.abspath(".")77    old_path = os.environ.get("PYTHONPATH", "")78    os.environ["PYTHONPATH"] = f"{workspace}:{old_path}" if old_path else workspace79    os.environ["LATENT_START_ID"] = "151666"80    os.environ["LATENT_END_ID"] = "151667"81    os.environ["AVT_LATENT_HOOK_BIN"] = "1"82    sys.modules["vllm.v1.worker.gpu_model_runner"] = importlib.import_module(83        "monet_models.vllm.monet_gpu_model_runner"84    )85    patch_qwen_monet()86    '''sys.modules[87        "transformers.models.qwen2_5_vl.modeling_qwen2_5_vl"88    ] = importlib.import_module(89        "monet_models.transformers.monet_modeling_qwen2_5_vl"90    )'''91 92    print("[Monet RL patch] vllm & transformers patched", file=sys.stderr)93 94patch()95 96 97