| """jevlab.vllm_reader.VllmReader with one unmerged LoRA adapter on top: the layer-2 entry's engine. |
| |
| Everything VllmReader does is unchanged (rows, read policy, label log-probabilities, temperatures); the only difference |
| is that vLLM is started with LoRA enabled and every generate call carries the adapter's LoRARequest. The adapter must be |
| the --vllm-out copy of train/average_adapters.py (k_proj LoRA copied to v_proj on Gemma-4's full-attention layers, where |
| V = k_proj(x)); without it vLLM would serve V without the trained change. LoRA strength lives in the adapter's |
| lora_alpha (average_adapters.py --scale), so the engine has no strength option. |
| |
| python -m decision_index run --engine jevlab.vllm_lora_reader:VllmLoraReader \ |
| --option lora=runs/lora-pilot-10k-vllm --option layout=header --option read_twice=true ... --out runs/NAME |
| |
| max_lora_rank must be one of vLLM's allowed values (1, 8, 16, 32, 64, 128, 256, 320, 512) and at least the adapter's |
| rank (r 64 for one arm, 192 for three averaged arms -> 256). By default it is the smallest allowed value that holds the |
| adapter's rank. |
| """ |
| import hashlib |
| import json |
| import os |
|
|
| from jevlab.vllm_reader import VllmReader |
|
|
| ALLOWED_RANKS = (1, 8, 16, 32, 64, 128, 256, 320, 512) |
|
|
|
|
| def adapter_rank(path): |
| with open(os.path.join(path, "adapter_config.json")) as f: |
| cfg = json.load(f) |
| return int(cfg["r"]), cfg |
|
|
|
|
| def adapter_digest(path): |
| """sha256 of every file of the adapter directory, by name.""" |
| out = {} |
| for name in sorted(os.listdir(path)): |
| full = os.path.join(path, name) |
| if os.path.isfile(full): |
| h = hashlib.sha256() |
| with open(full, "rb") as f: |
| for block in iter(lambda: f.read(1 << 20), b""): |
| h.update(block) |
| out[name] = h.hexdigest() |
| return out |
|
|
|
|
| class _LoraLLM: |
| """vllm.LLM stand-in that adds the adapter's LoRARequest to every generate call; everything else is forwarded.""" |
|
|
| def __init__(self, llm, request): |
| self._llm, self._request = llm, request |
|
|
| def generate(self, prompts, sampling_params=None, **kwargs): |
| kwargs.setdefault("lora_request", self._request) |
| return self._llm.generate(prompts, sampling_params, **kwargs) |
|
|
| def __getattr__(self, name): |
| return getattr(self._llm, name) |
|
|
|
|
| class VllmLoraReader(VllmReader): |
| name = "jevlab-vllm-lora-reader" |
|
|
| def __init__(self, lora, max_lora_rank=0, lora_name="layer2", llm_options=None, backend=None, lora_request=None, **options): |
| lora = os.path.abspath(os.path.expanduser(str(lora))) |
| rank, cfg = adapter_rank(lora) |
| max_lora_rank = int(max_lora_rank) or next(r for r in ALLOWED_RANKS if r >= rank) |
| if max_lora_rank not in ALLOWED_RANKS or max_lora_rank < rank: |
| raise ValueError(f"max_lora_rank {max_lora_rank} must be one of {ALLOWED_RANKS} and >= the adapter rank {rank}") |
| opts = dict(llm_options or {}) |
| opts.update(enable_lora=True, max_lora_rank=max_lora_rank, max_loras=1) |
| super().__init__(llm_options=opts, backend=backend, **options) |
| if lora_request is None: |
| from vllm.lora.request import LoRARequest |
|
|
| lora_request = LoRARequest(lora_name, 1, lora) |
| self.backend.llm = _LoraLLM(self.backend.llm, lora_request) |
| self.lora_path, self.lora_rank, self.lora_alpha = lora, rank, float(cfg.get("lora_alpha", rank)) |
| self.provenance.update(kind="LoRA", adapter={ |
| "path": lora, "rank": rank, "lora_alpha": self.lora_alpha, "target_modules": cfg.get("target_modules"), |
| "max_lora_rank": max_lora_rank, "files_sha256": adapter_digest(lora)}) |
|
|
