搭建一个 RL 微调平台，训练 JEV 式决策模型

{"billing": 0.974, "technical": 0.011, "account": 0.010, "shipping": 0.005}

pip install mlx mlx-lm

import json
import mlx.core as mx
import mlx.nn as nn
from mlx_lm import load
from mlx_lm.tuner import linear_to_lora_layers

model, tok = load("mlx-community/Qwen3-4B-4bit")   # 4-bit so it fits in 24 GB
model.freeze()                                        # base weights stay frozen
linear_to_lora_layers(model, 8, {"rank": 16, "scale": 8.0, "dropout": 0.0})

SYSTEM = "Judge how well the Option answers the Question, given the State."

def render(state, question, options):
    if not isinstance(state, str):
        state = json.dumps(state, ensure_ascii=False)
    prefix = ("<|im_start|>system\n" + SYSTEM + "<|im_end|>\n"
              "<|im_start|>user\n<State>: " + state + "\n<Question>: " + question + "\n")
    branches = [f"<Option>: {o}<|im_end|>" for o in options]
    return prefix, branches

class Gero(nn.Module):
    def __init__(self, backbone, hidden=2560):
        super().__init__()
        self.backbone = backbone
        self.scorer = nn.Linear(hidden, 1, bias=False)
        # Hidden states have a norm around 300, so 1e-3 gives scores spread
        # around 0.3: close to uniform at the start, but not so tiny that
        # gradients die.
        self.scorer.weight = mx.random.normal(self.scorer.weight.shape) * 1e-3

m = Gero(model)

class Record:
    """Prefix pass: remember this layer's keys and values."""
    def __init__(self):
        self.offset = 0
    def update_and_fetch(self, k, v):
        self.k, self.v = k, v
        self.offset += k.shape[2]
        return k, v

class SharePrefix:
    """Branch pass: every branch sees the same prefix keys and values."""
    def __init__(self, rec):
        self.pk, self.pv, self.offset = rec.k, rec.v, rec.offset
    def update_and_fetch(self, k, v):
        B = k.shape[0]                      # number of branches
        pk = mx.broadcast_to(self.pk, (B,) + self.pk.shape[1:])
        pv = mx.broadcast_to(self.pv, (B,) + self.pv.shape[1:])
        return mx.concatenate([pk, k], 2), mx.concatenate([pv, v], 2)

def run(backbone, ids, mask, caches):
    t = backbone.model
    h = t.embed_tokens(ids)
    for layer, c in zip(t.layers, caches):
        h = layer(h, mask, c)
    return t.norm(h)

def branch_logits(m, pre_ids, opt_ids):
    recs = [Record() for _ in m.backbone.model.layers]
    run(m.backbone, mx.array([pre_ids]), "causal", recs)      # prefix, once

    P, Lb = len(pre_ids), max(len(o) for o in opt_ids)
    ids = mx.array([o + [0] * (Lb - len(o)) for o in opt_ids])  # pad branches
    r = mx.arange(Lb)
    mask = mx.concatenate([mx.ones((Lb, P), dtype=mx.bool_),   # see all of the prefix
                           r[:, None] >= r[None, :]], axis=1)  # causal inside the branch
    h = run(m.backbone, ids, mask, [SharePrefix(x) for x in recs])

    last = mx.array([len(o) - 1 for o in opt_ids])
    return m.scorer(h[mx.arange(len(opt_ids)), last].astype(mx.float32)).squeeze(-1)

            prefix tokens        its own tokens
branch A:   [1 1 1 1 1 1 1 1]    [causal]
branch B:   [1 1 1 1 1 1 1 1]    [causal]      <- but never A's tokens

pre, brs = render("Charged twice for my invoice.", "Which team?", ["billing", "technical", "shipping"])
enc = lambda s: tok.encode(s, add_special_tokens=False)
a = mx.softmax(branch_logits(m, enc(pre), [enc(b) for b in brs]))
b = mx.softmax(branch_logits(m, enc(pre), [enc(b) for b in reversed(brs)]))
print(a, b[::-1])   # should be identical

import random

NOUL_OPTS = [("no", "yes"), ("false", "true"), ("incorrect", "correct")]

def make_yes_no(R, i):
    n = R.randint(2, 8)
    shelves = R.sample(range(1, 100), n)
    items = [f"{R.choice(['QX','FJ','RT','ZP'])}-{R.randint(100,999)}{R.choice('ABCDEFG')}" for _ in range(n)]
    ask = R.randrange(n)

    truth = (i % 2 == 0)                         # exactly 50/50 on the fact
    # the wrong item is taken from ANOTHER shelf in the same state,
    # so "is this code in the text?" never gives away the answer
    probe = items[ask] if truth else items[(ask + 1) % n]
    negated = (i // 2) % 2 == 1                  # exactly 50/50 on negation

    if negated:
        q = f"Is it false that shelf {shelves[ask]} holds {probe}?"
    else:
        q = f"Does shelf {shelves[ask]} hold {probe}?"
    state = "Manifest: " + "; ".join(f"shelf {s} holds {it}" for s, it in zip(shelves, items)) + "."
    no, yes = R.choice(NOUL_OPTS)
    return dict(state=state, question=q, options=[no, yes], target=int(truth != negated))

from collections import Counter

def check_leak(train, holdout):
    """Train and holdout must not share a state, not just an index."""
    seen = {json.dumps(e["state"], sort_keys=True) for e in train}
    leaked = [e for e in holdout if json.dumps(e["state"], sort_keys=True) in seen]
    assert not leaked, f"{len(leaked)} holdout states also appear in train"

def check_position(data):
    """The right answer should not prefer one position."""
    for k in {e_k for e_k in (len(e["options"]) for e in data) if e_k <= 6}:
        rows = [e for e in data if len(e["options"]) == k]
        top = Counter(e["target"] for e in rows).most_common(1)[0][1] / len(rows)
        assert top < 1 / k + 0.10, f"k={k}: one slot holds {top:.0%} of answers"

def check_majority_text(data):
    """If always picking the same option TEXT scores well, the item can be solved without reading."""
    gold = Counter(e["options"][e["target"]] for e in data)
    worst = gold.most_common(1)[0][1] / len(data)
    assert worst < 0.6, f"answer text '{gold.most_common(1)[0][0]}' is right {worst:.0%} of the time"

def check_string_presence(data):
    """'Pick whichever option appears in the state' must not work."""
    solved = 0
    for e in data:
        hits = [i for i, o in enumerate(e["options"]) if o in json.dumps(e["state"])]
        solved += hits == [e["target"]]
    assert solved / len(data) < 0.05, f"string presence solves {solved/len(data):.0%}"

def stage2_loss(model, batch):
    total = 0.0
    for e in batch:
        logits = branch_logits(model, e["pre_ids"], e["opt_ids"])
        logp = logits - mx.logsumexp(logits)
        target = mx.array(e["soft"])            # one-hot for normal items
        total = total - (target * logp).sum()
    return total / len(batch)

import mlx.optimizers as optim

opt_scorer = optim.AdamW(learning_rate=5e-4)
opt_lora = optim.AdamW(learning_rate=3e-5)
step_fn = nn.value_and_grad(m, stage2_loss)

for step in range(steps):
    batch = rng.sample(train, 8)
    loss, g = step_fn(m, batch)
    g, _ = optim.clip_grad_norm(g, 1.0)
    p = m.trainable_parameters()
    m.update(opt_scorer.apply_gradients({"scorer": g["scorer"]}, {"scorer": p["scorer"]}))
    m.update(opt_lora.apply_gradients({"backbone": g["backbone"]}, {"backbone": p["backbone"]}))
    mx.eval(m.parameters(), opt_scorer.state, opt_lora.state)

truth = mx.array([0.174, 0.261, 0.005, 0.001, 0.141, 0.418])   # true answer odds

def naive_loss(z):
    p = mx.softmax(z)
    c = mx.clip(p, 1e-6, 1 - 1e-6)
    # choose option a with prob p(a), score how well p(a) predicts "a is right"
    per_choice = truth * mx.log(c) + (1 - truth) * mx.log1p(-c)
    return -(p * per_choice).sum()

truth  [0.174, 0.261, 0.005, 0.001, 0.141, 0.418]
naive  [0.200, 0.181, 0.208, 0.208, 0.202, 0.001]

def rl_loss(logits, soft, G, rng, lam=1.0):
    logp = logits - mx.logsumexp(logits)
    k = len(soft)
    outcomes = rng.choices(range(k), weights=soft, k=G)   # G sampled outcomes

    top = int(mx.argmax(logp).item())
    c = mx.clip(mx.exp(logp[top]), 1e-6, 1 - 1e-6)

    full = mx.stack([logp[y] for y in outcomes]).mean()
    hit_rate = sum(1 for y in outcomes if y == top) / G
    calib = hit_rate * mx.log(c) + (1 - hit_rate) * mx.log1p(-c)
    return -(full + lam * calib)

truth  [0.174, 0.261, 0.005, 0.001, 0.141, 0.418]
ours   [0.174, 0.261, 0.005, 0.001, 0.141, 0.418]

def murphy(rows, bins=10):
    """rows: list of (confidence, correct) pairs."""
    n = len(rows)
    base = sum(o for _, o in rows) / n
    groups = [[] for _ in range(bins)]
    for c, o in rows:
        groups[min(bins - 1, int(c * bins))].append((c, o))
    rel = res = 0.0
    for g in groups:
        if not g:
            continue
        conf = sum(c for c, _ in g) / len(g)
        acc = sum(o for _, o in g) / len(g)
        rel += len(g) / n * (conf - acc) ** 2
        res += len(g) / n * (acc - base) ** 2
    return rel, res

def selective(rows, thr=0.9):
    hi = [(c, o) for c, o in rows if c >= thr]
    return dict(coverage=len(hi) / len(rows),
                accuracy=sum(o for _, o in hi) / len(hi),
                promised=sum(c for c, _ in hi) / len(hi))

correct = soft[pred]          # not: 1.0 if pred == argmax(soft) else 0.0

from mlx_lm.utils import dequantize_model
from mlx.utils import tree_flatten, tree_unflatten

fused = [(n, mod.fuse(dequantize=True)) for n, mod in m.backbone.named_modules() if hasattr(mod, "fuse")]
m.backbone.update_modules(tree_unflatten(fused))
bb = dequantize_model(m.backbone)

weights = {k: v.astype(mx.bfloat16) for k, v in tree_flatten(bb.parameters())}
weights.pop("lm_head.weight", None)               # tied embeddings, no LM head needed
weights["score.weight"] = m.scorer.weight.astype(mx.bfloat16)
mx.save_safetensors("gero-4b/model.safetensors", weights, metadata={"format": "pt"})
# plus config.json with architectures=["Qwen3ForSequenceClassification"], num_labels=1