nvmmonkey/opt125m-ltr-ranker

OPT-125M Output-Length Ranker (Learning-to-Rank scheduling for LLM serving)

A small output-length ranker used to schedule LLM inference requests shortest-predicted-first (SJF), the mechanism behind the Learning-to-Rank (LTR) scheduler of Fu et al. (NeurIPS 2024) and the FDU latency study (Kumar et al.). Given a chat prompt it emits one scalar score — higher score ⇒ shorter expected output ⇒ run this request earlier — which is mapped to a vLLM request priority. It is the B1 ranker in the capstone KVCache-Coordinated Latency Optimization.

• Architecture: facebook/opt-125m backbone (AutoModel, fp32) + masked-mean pooling over the prompt tokens + a Linear(768 → 1) scoring head.
• Objective: ListMLE listwise ranking loss (list size 16), seed 0, 10 epochs.
• Labels: real output lengths from LMSYS-Chat-1M (the FDU paper's methodology).
• Target serve model: Llama-3.1-8B-Instruct (the ranker itself is model-agnostic — it ranks prompts, not tokens).

Training

Base	facebook/opt-125m
Loss	ListMLE (listwise), list size 16
Epochs	10 (resumed 5→10, continuous curve), seed 0
Data	LMSYS-Chat-1M prompt → output-length lists
ListMLE loss	25.67 → 23.04 → 19.98 → 18.02 → 16.22 → 15.01 → 13.84 → 12.92 → 12.51 → 11.67
Ranking quality	Kendall τ ≈ 0.71 on a held-out LMSYS split (score vs true length)

Monotone loss with no overfit rebound at 10 epochs (the paper's balance point).

The target-sampled output-length labels are published separately: nvmmonkey/llama31-8b-output-lengths (LMSYS prompts withheld per license).

Usage

The scoring head is custom, so load the state dict onto the architecture shipped here (modeling_opt_ranker.py):

import torch, importlib.util
from huggingface_hub import snapshot_download
from safetensors.torch import load_file
from transformers import AutoTokenizer

repo = snapshot_download("<your-hf-user>/opt125m-ltr-ranker")

spec = importlib.util.spec_from_file_location("m", f"{repo}/modeling_opt_ranker.py")
m = importlib.util.module_from_spec(spec); spec.loader.exec_module(m)

ranker = m.build_ranker("facebook/opt-125m")          # OPT backbone + linear score head
ranker.load_state_dict(load_file(f"{repo}/model.safetensors"))
ranker.eval()

tok = AutoTokenizer.from_pretrained(repo)
enc = tok(["Explain quantum computing in one sentence.",
           "Write a 2000-word essay on the French Revolution."],
          return_tensors="pt", padding=True, truncation=True, max_length=512)

with torch.no_grad():
    scores = ranker(enc.input_ids, enc.attention_mask)  # higher => shorter output => schedule first
print(scores)   # the short prompt should score higher than the long one

To serve with it, stamp each request's vLLM priority from the score (see the code repo's ltr/scheduler/priority.py and serving/serve_b1_ltr.sh, launched with --scheduling-policy priority).

Provenance, license & citation

• Base model facebook/opt-125m is under the OPT license (research / non-commercial); this fine-tune inherits it.
• Training labels come from LMSYS-Chat-1M, under the LMSYS-Chat-1M dataset license — non-commercial, agree to its terms before use.
• Method reproduces prior work — please cite it:
  • Y. Fu et al., "Efficient LLM scheduling by learning to rank," NeurIPS 2024. — github.com/hao-ai-lab/vllm-ltr
  • A. Saravana Kumar et al., "An empirical study on latency reduction techniques for large language models," FDU, 2026.

Trained for the FDU CSCI 6806 capstone (Guoliang Liu, Wenhui Kang, Junpeng Huang).