BERT Cross Encoder LoRA — GPU Test Status (slack/data-airflow#69908)

gist_update2.md

BERT Cross Encoder LoRA — GPU Test Status

PR: slack/data-airflow#69908 Branch: lora-adapter Script: scripts/test_lora_gpu.py Last updated: 2026-04-29 (rebased onto master, code review fixes, new diagnostic runs)

Current Blocker: vLLM V1 Engine vs BERT Pooling

The ml-base-gpu:latest Docker image has a vLLM version where the V0 engine has been completely removed. AsyncLLMEngine is now an alias for vllm.v1.engine.async_llm.AsyncLLM, and from_engine_args() unconditionally returns V1. The VLLM_USE_V1=0 env var has no effect.

The V1 engine does not support BERT pooling models, so engine initialization fails with:

RuntimeError: Engine core initialization failed. See root cause above. Failed core proc(s): {}

Diagnostic Evidence (2026-04-29)

From scripts/vllm_diag4.py run on ray-cluster-girafarig-12af36f0:

{
  "vllm.engine.async_llm_engine.AsyncLLMEngine": {
    "found": true,
    "actual_module": "vllm.v1.engine.async_llm"
  },
  "vllm.v1.engine.async_llm.AsyncLLM": {
    "found": true,
    "actual_module": "vllm.v1.engine.async_llm"
  }
}

from_engine_args signature returns AsyncLLM (V1) directly — no V0 fallback path exists.

Options to Unblock

1. Wait for vLLM V1 to support BERT pooling — upstream issue
2. Use --image with an older Docker image that has vLLM with V0 engine support (pre-V1-only transition)
3. Use a different inference path — e.g., direct PyTorch inference with PEFT get_peft_model() instead of vLLM's LoRA manager (loses vLLM's batching/scheduling benefits)

What's Been Solved

Issue	Fix
S3 credentials on KubeRay	service_account="moonlet" (default quasar lacks IRSA)
HF Hub model resolution	Patched config.json with custom architecture
Model registration	ModelRegistry.register_model() with class object directly
tensor_parallel_size double assignment	Set on engine_args after create_vllm_args()
Adapter HF lookups	base_model_name_or_path → local model path

Code Review Fixes (2026-04-29)

1. Replaced bare assert with MLPredictionError in BertCrossEncoderLoRAVLLM
2. Fixed validation order: adapter_id <= 0 before adapter_path is None
3. Added truncate_prompt_tokens=-1 to Gemma2EmbeddingLoRAVLLM's PoolingParams
4. Removed dead BertEmbeddingModelWithLoRA class
5. Removed unrelated files from the branch

Cluster Runs

Cluster	Job	Result
ray-cluster-metagross-6277d3a5	script-job-granbull-840dd19d	LocalEntryNotFoundError — HF Hub lookup
ray-cluster-snorunt-1377eb4b	script-job-ampharos-3548ae0a	V1 engine RuntimeError
ray-cluster-politoed-4e1341d2	script-job-sableye-f016ef11	TypeError: multiple values for tensor_parallel_size
ray-cluster-wynaut-ebc77b7d	script-job-qwilfish-a4fd6df3	V1 engine RuntimeError
ray-cluster-misdreavus-8a3714bc	script-job-feebas-67017d02	V1 engine RuntimeError
ray-cluster-forretress-50bada78	script-job-illumise-a87bc7e9	V1 engine RuntimeError
ray-cluster-sealeo-36fb480e	script-job-dustox-13e7ba90	V1 engine RuntimeError (post-rebase)
ray-cluster-armaldo-40915f23	script-job-furret-350014e9	V1 engine RuntimeError (V0 workaround attempt)
ray-cluster-girafarig-12af36f0	script-job-hoppip-97cf7ef7	Diagnostic: confirmed V0 removed

Test Script Architecture

The GPU test script (scripts/test_lora_gpu.py) does the following:

1. Downloads base model and adapter from S3
2. Patches config.json to use custom architecture
3. Registers BertForSequenceClassificationWithLoRA via ModelRegistry
4. Creates AsyncEngineArgs with enable_lora=True
5. Runs inference with LoRARequest pointing to the adapter
6. Validates output shape and score ranges