[DO NOT MERGE] Run SFT, RL and MOPD on self-hosted Tinker-compatible servers

hosted/opd/backend.py

@@ -239,6 +239,76 @@ def close(self):
         self.service.close()
 
 
+class SelfHosted:
+    """A self-hosted Tinker-compatible server at $TINKER_BASE_URL (see hosted/selfhosted/README.md).
+
+    The servers do not authenticate, but the SDK requires a key that starts with `tml-`. A teacher path of the form
+    `base:<model name>` samples a frozen model that the server serves by name (verl-tinker teachers); any other
+    path is a sampler checkpoint on the same server. Subclasses set what each server accepts.
+    """
+    name = 'self-hosted'
+    ACCEPTS_SEED = True  # create_lora_training_client(seed=...)
+    ACCEPTS_METADATA = True  # user_metadata on the session and on each training client
+    BASE_PREFIX = 'base:'
+
+    def __init__(self, experiment):
+        self.url = os.environ['TINKER_BASE_URL']
+        os.environ.setdefault('TINKER_API_KEY', 'tml-self-hosted')
+        server = json.loads(Path(os.environ['TINKER_SERVER_INFO']).read_text()) if os.environ.get('TINKER_SERVER_INFO') else {}
+        self.account = {'endpoint': self.url, 'server': self.name, 'server_info': server}
+        metadata = {'user_metadata': {'experiment': experiment}} if self.ACCEPTS_METADATA else {}
+        self.service = tinker.ServiceClient(base_url=self.url, api_key=os.environ['TINKER_API_KEY'], **metadata)
+
+    def capabilities(self, model):
+        caps = self.service.get_server_capabilities()
+        assert model in [m.model_name for m in caps.supported_models], f'{model} is not served at {self.url}'
+        return caps.model_dump()
+
+    def new_student(self, model, rank, seed, metadata):
+        return self.service.create_lora_training_client(
+            model, rank=rank, seed=seed if self.ACCEPTS_SEED else None,
+            **({'user_metadata': metadata} if self.ACCEPTS_METADATA else {}))
+
+    def resume_student(self, state_path):
+        return self.service.create_training_client_from_state_with_optimizer(state_path)
+
+    def sampler(self, model=None, path=None):
+        if path and path.startswith(self.BASE_PREFIX):
+            return self.service.create_sampling_client(base_model=path[len(self.BASE_PREFIX):])
+        return self.service.create_sampling_client(**({'model_path': path} if path else {'base_model': model}))
+
+    def sync(self, train, name):
+        path = train.save_weights_for_sampler(name=name).result().path
+        return self.sampler(path=path), path
+
+    def use_tokenizers(self, student, teacher):
+        pass
+
+    def close(self):
+        pass
+
+
+class SkyRL(SelfHosted):
+    """SkyRL's Tinker API server (`python -m skyrl.tinker.api --backend megatron`). LoRA alpha is a server setting, and
+    the seed is ignored. A sampler checkpoint stays servable only while its model is loaded, so the server runs with
+    --session-timeout-sec -1 and teachers are trained on the same server instance that serves them."""
+    name = 'skyrl'
+
+
+class Miles(SelfHosted):
+    """Miles' Tinker gateway (`serve_tinker.py`). It rejects a LoRA seed and user metadata, and caps the rank at the
+    server's --lora-rank."""
+    name = 'miles'
+    ACCEPTS_SEED = False
+    ACCEPTS_METADATA = False
+
+
+class Verl(SelfHosted):
+    """verl-tinker (`python -m verl_tinker.start`). One trainable model per server: teachers are frozen models that the
+    server config serves by name (`base:<name>` paths), and every training run needs its own server."""
+    name = 'verl'
+
+
 def service_config(resources, experiment=None):
     """from_firetitan_config arguments for the lease's trainer and deployment; closing never deletes them."""
     return dict(
@@ -255,4 +325,5 @@ def service_config(resources, experiment=None):
         **({'user_metadata': {'experiment': experiment}} if experiment else {}))
 
 
-BACKENDS = {'tinker': Tinker, 'fireworks': Fireworks, 'fireworks-serverless': FireworksServerless}
+BACKENDS = {'tinker': Tinker, 'fireworks': Fireworks, 'fireworks-serverless': FireworksServerless,
+            'skyrl': SkyRL, 'miles': Miles, 'verl': Verl}

hosted/opd/rl.py

@@ -0,0 +1,167 @@
+"""Reinforcement learning (GRPO) on a Tinker-compatible API: a fresh Qwen3.6-35B-A3B LoRA learns the reasoning_gym
+tasks from the verifier reward alone. Thinking is on.
+
+Every update the student samples `--group-size` responses to each of `--prompts` training problems (split equally
+between `--domains`). The reward is the shared scorer's 0/1 score; each response's advantage is its reward normalized
+within its prompt's group, applied to every response token. The update is one clipped PPO step (ratio 0.8-1.2).
+Groups with no reward variance are kept with zero advantage. The settings are those of the released GRPO teachers
+(hosted/tinker/README.md): 32 prompts x 8 responses, LR 1e-5 for LoRA, Adam betas 0.9/0.98, gradient clip 1.
+
+    python hosted/opd/rl.py --domains caesar_cipher --output /abs/rl-caesar --backend skyrl
+    python hosted/opd/rl.py --output /abs/rl-both --updates 20
+
+Records, evaluation, checkpoints and --policy-lag behave as in opd.py and mopd.py (see common.py).
+"""
+import argparse
+import json
+import statistics
+import time
+
+import tinker
+
+from common import Experiment, add_protocol_arguments, check_protocol_arguments, policy_datum, run
+from shared import recipe, scoring
+
+GRPO = {'prompts': 32, 'group_size': 8, 'clip_low_threshold': 0.8, 'clip_high_threshold': 1.2}
+ADAM = {'beta1': 0.9, 'beta2': 0.98, 'eps': 1e-8, 'weight_decay': 0.0, 'grad_clip_norm': 1.0}
+
+
+def group_advantages(rewards):
+    """Rewards normalized within one prompt's group; a group with one response or no variance gets zeros."""
+    mean = statistics.mean(rewards)
+    std = statistics.stdev(rewards) if len(rewards) > 1 else 0.0
+    return [(r - mean) / (std + 1e-6) for r in rewards]
+
+
+class RL(Experiment):
+    """An Experiment whose rollouts are groups scored by the verifier instead of by a teacher."""
+
+    def start(self):
+        a = self.args
+        self.teachers = {}
+        if a.resume:
+            train = self.backend.resume_student(a.resume)
+        else:
+            train = self.backend.new_student(self.student_model[0], a.rank, a.seed, {'run': self.out.name, 'kind': self.kind})
+        try:
+            (self.out/'model-info.json').write_text(train.get_info().model_dump_json(indent=2))
+        except Exception as error:
+            (self.out/'model-info.json').write_text(json.dumps({'unavailable': str(error)[:300]}))
+        return train, self.backend.sync(train, f'update-{a.start_update:03d}')[0]
+
+    def rollout(self, student, domain, row, seed, t0):
+        """One prompt's group. A multi-response request leaves the seed unset, as in the cookbook."""
+        a = self.args
+        prompt = self.prompt(row)
+        begin = time.monotonic()
+        result = student.sample(tinker.ModelInput.from_ints(prompt), num_samples=a.group_size,
+                                sampling_params=self.params(1.0, seed if a.group_size == 1 else None)).result()
+        sampled = time.monotonic()
+        assert len(result.sequences) == a.group_size
+        samples = []
+        for j, sequence in enumerate(result.sequences):
+            tokens, logprobs = list(sequence.tokens), list(sequence.logprobs)
+            assert len(tokens) == len(logprobs)
+            response = self.tokenizer.decode(tokens, skip_special_tokens=True)
+            samples.append({'sample': j, 'tokens': tokens, 'logprobs': logprobs, 'response': response,
+                            'score': scoring.score(response, row['label']), 'stop_reason': sequence.stop_reason})
+        for s, advantage in zip(samples, group_advantages([s['score'] for s in samples])):
+            s['advantage'] = advantage
+        return {'domain': domain, 'index': row['_index'], 'seed': seed, 'prompt_tokens': prompt, 'samples': samples,
+                'submitted': begin - t0, 'sampled': sampled - t0, 'scored': sampled - t0}
+
+    def update(self, train, batch, prefetch):
+        a = self.args
+        update = batch['update']
+        t0 = time.monotonic()
+        self.log('update_start', update=update + 1, prompts=len(batch['futures']))
+        groups = [f.result() for f in batch['futures']]
+        rollout_seconds = time.monotonic() - t0
+        lag = update - batch['version']
+        assert 0 <= lag <= a.policy_lag, 'Refuse a stale or future-policy batch'
+        upcoming = prefetch()
+        self.save_rollouts(f'update-{update+1:03d}', [{'update': update + 1, 'policy_lag': lag, **g} for g in groups])
+
+        train_start = time.monotonic()
+        flat = [(g, s) for g in groups for s in g['samples']]
+        datums = [policy_datum(g['prompt_tokens'], s['tokens'], s['logprobs'], [s['advantage']] * len(s['tokens']))
+                  for g, s in flat]
+        fb = train.forward_backward(datums, loss_fn='ppo', loss_fn_config={
+            'clip_low_threshold': a.clip_low, 'clip_high_threshold': a.clip_high})
+        opt = train.optim_step(tinker.AdamParams(learning_rate=a.learning_rate, **ADAM))
+        fb, opt = fb.result(), opt.result()
+        train_seconds = time.monotonic() - train_start
+        mismatch = [float(x) - y for (g, s), o in zip(flat, fb.loss_fn_outputs, strict=True) if 'logprobs' in o
+                    for x, y in zip(o['logprobs'].data[len(g['prompt_tokens']) - 1:], s['logprobs'], strict=True)]
+
+        sync_start = time.monotonic()
+        student, sampler_path = self.backend.sync(train, f'update-{update+1:03d}')
+        sync_seconds = time.monotonic() - sync_start
+        self.append('metrics.jsonl', self.rl_metrics(update, groups, t0, rollout_seconds, train_seconds, sync_seconds,
+                                                     mismatch, sampler_path, fb, opt, lag))
+        return student, upcoming
+
+    def rl_metrics(self, update, groups, t0, rollout_seconds, train_seconds, sync_seconds, mismatch, sampler_path,
+                   fb, opt, lag):
+        samples = [s for g in groups for s in g['samples']]
+        sampled = [g['sampled'] for g in groups]
+        response_tokens = sum(len(s['tokens']) for s in samples)
+        prompt_tokens = sum(len(g['prompt_tokens']) * len(g['samples']) for g in groups)
+        by_domain = {}
+        for d in self.domains:
+            gs = [g for g in groups if g['domain'] == d]
+            ss = [s for g in gs for s in g['samples']]
+            by_domain[d] = {'groups': len(gs), 'responses': len(ss), 'reward': statistics.mean(s['score'] for s in ss),
+                            'groups_with_signal': sum(len({s['score'] for s in g['samples']}) > 1 for g in gs),
+                            'mean_response_tokens': statistics.mean(len(s['tokens']) for s in ss)}
+        update_seconds = time.monotonic() - t0
+        metric = {
+            'update': update + 1, 'policy_lag': lag, 'groups': len(groups), 'responses': len(samples),
+            'learning_rate': self.args.learning_rate, 'prompt_tokens': prompt_tokens, 'response_tokens': response_tokens,
+            'max_response_tokens': max(len(s['tokens']) for s in samples),
+            'truncated': sum(s['stop_reason'] == 'length' for s in samples), 'by_domain': by_domain,
+            'reward': statistics.mean(s['score'] for s in samples),
+            'nonzero_advantage_responses': sum(s['advantage'] != 0 for s in samples),
+            'mean_unique_responses_per_group': statistics.mean(len({s['response'] for s in g['samples']}) for g in groups),
+            'timing': {'rollout': rollout_seconds, 'train': train_seconds, 'sync': sync_seconds, 'update': update_seconds,
+                       'first_group': min(sampled), 'median_group': statistics.median(sampled),
+                       'last_group': max(sampled)},
+            'throughput': {'sampled_tokens_per_second': response_tokens / max(sampled),
+                           'trained_tokens_per_second': (prompt_tokens + response_tokens) / train_seconds,
+                           'update_response_tokens_per_second': response_tokens / update_seconds},
+            'train_sampler_logprob_delta_mean': statistics.mean(mismatch) if mismatch else None,
+            'train_sampler_logprob_delta_abs_mean': statistics.mean(abs(x) for x in mismatch) if mismatch else None,
+            'sampler_path': sampler_path, 'forward_backward_metrics': fb.metrics, 'optimizer_metrics': opt.metrics}
+        self.log('update_complete', **{k: v for k, v in metric.items() if k != 'forward_backward_metrics'})
+        return metric
+
+
+def main():
+    parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
+    parser.add_argument('--domains', choices=recipe.DOMAINS, nargs='+', default=list(recipe.DOMAINS))
+    parser.add_argument('--group-size', type=int, default=GRPO['group_size'])
+    parser.add_argument('--clip-low', type=float, default=GRPO['clip_low_threshold'])
+    parser.add_argument('--clip-high', type=float, default=GRPO['clip_high_threshold'])
+    add_protocol_arguments(parser)
+    parser.set_defaults(prompts=GRPO['prompts'])
+    args = parser.parse_args()
+    check_protocol_arguments(parser, args)
+    if args.group_size < 2: parser.error('--group-size must be at least 2 for group-normalized advantages')
+    if args.prompts % len(args.domains): parser.error('--prompts must split evenly over the domains')
+    domains = [d for d in recipe.DOMAINS if d in args.domains]
+    share = args.prompts // len(domains)
+
+    class Run(RL):
+        def record_inputs(self):
+            super().record_inputs()
+            config = json.loads((self.out/'config.json').read_text())
+            config.update(loss='ppo', advantage='group-normalized verifier reward', adam=ADAM,
+                          clip=[args.clip_low, args.clip_high], teachers=None)
+            (self.out/'config.json').write_text(json.dumps(config, indent=2))
+
+    run(lambda: Run(args, 'rl', {d: None for d in domains}),
+        lambda experiment, update: [p for d in domains for p in experiment.train_prompts(d, update, share)])
+
+
+if __name__ == '__main__':
+    main()

hosted/selfhosted/campaign.sh

@@ -0,0 +1,58 @@
+#!/bin/bash
+# Run the rest of one framework's campaign on a running server, one stage at a time, then summarize every run.
+#
+#   CAMPAIGN=/shared/... CLIENT_VENV=... hosted/selfhosted/campaign.sh FRAMEWORK SERVER_DIR [STAGE ...]
+#
+# Stages (default: sft rl mopd), each run through client.sh so it gets its own telemetry window:
+#   sft   both SFT teachers at once (sft.py train, the released-teacher defaults); skipped if they already exist
+#   rl    GRPO on both domains (rl.py defaults)
+#   mopd  MOPD from the two SFT teachers (mopd.py defaults); teachers.json is built from the teachers' teacher.json
+# A stage waits for the previous one and stops the campaign if it fails. Outputs: $CAMPAIGN/FRAMEWORK/out/<run>,
+# $CAMPAIGN/FRAMEWORK/runs/<run>.{log,window.json,summary.json}.
+set -euo pipefail
+: "${CAMPAIGN:?CAMPAIGN must be set}"
+FRAMEWORK=$1; SERVER=$(cd "$2" && pwd); shift 2
+if [[ $# -eq 0 ]]; then STAGES=(sft rl mopd); else STAGES=("$@"); fi
+OUT=$CAMPAIGN/$FRAMEWORK/out; RUNS=$CAMPAIGN/$FRAMEWORK/runs
+SELFHOSTED=$CAMPAIGN/opd/hosted/selfhosted
+mkdir -p "$OUT" "$RUNS"
+log() { echo "$(date -u +%FT%TZ) $*"; }
+
+wait_run() {  # Wait for a client run's window file; fail if it exited non-zero.
+    local name=$1
+    until [[ -f $RUNS/$name.window.json ]]; do sleep 30; done
+    local code; code=$(python3 -c "import json,sys; print(json.load(open(sys.argv[1]))['exit_code'])" "$RUNS/$name.window.json")
+    "$CAMPAIGN/client-venv/bin/python" "$SELFHOSTED/summarize.py" "$RUNS/$name.window.json" "$OUT/$name" \
+        --telemetry-code "$CAMPAIGN/telemetry-scripts" || log "summary failed for $name"
+    [[ $code == 0 ]] || { log "$name exited $code; stopping"; exit 1; }
+    log "$name complete"
+}
+
+start() {  # start NAME -- COMMAND...
+    local name=$1; shift 2
+    [[ -f $RUNS/$name.window.json ]] && { log "$name already ran"; return; }
+    [[ -f $RUNS/$name.log ]] || bash "$SELFHOSTED/client.sh" "$SERVER" "$name" "$FRAMEWORK" -- "$@"
+}
+
+for stage in "${STAGES[@]}"; do
+    case $stage in
+    sft)
+        for d in caesar_cipher simple_geometry; do
+            start "sft-$d" -- python hosted/opd/sft.py train --backend "$FRAMEWORK" --domain "$d" --output "$OUT/sft-$d"
+        done
+        wait_run sft-caesar_cipher; wait_run sft-simple_geometry ;;
+    rl)
+        start rl -- python hosted/opd/rl.py --backend "$FRAMEWORK" --output "$OUT/rl"
+        wait_run rl ;;
+    mopd)
+        python3 - "$OUT" > "$OUT/teachers.json" <<'PY'
+import json, sys
+out = sys.argv[1]
+print(json.dumps({d: json.load(open(f'{out}/sft-{d}/teacher.json'))['sampler_path'] for d in ('caesar_cipher', 'simple_geometry')}, indent=2))
+PY
+        start mopd -- python hosted/opd/mopd.py --backend "$FRAMEWORK" --teachers "$OUT/teachers.json" --output "$OUT/mopd"
+        wait_run mopd ;;
+    *) log "unknown stage $stage"; exit 2 ;;
+    esac
+done
+log "campaign stages done: ${STAGES[*]}"

hosted/selfhosted/client.sh

@@ -0,0 +1,33 @@
+#!/bin/bash
+# Run one client (smoke.py, sft.py, rl.py, opd.py or mopd.py) against a running self-hosted server, as a step of the
+# server's Slurm job on its head node, with the server's endpoint and description in the environment.
+#
+#   CAMPAIGN=/shared/... hosted/selfhosted/client.sh SERVER_DIR NAME BACKEND -- python hosted/opd/sft.py train ...
+#
+# SERVER_DIR is the server's output directory (it holds server.json). The client's console goes to
+# SERVER_DIR/../runs/NAME.log, and NAME.window.json records its UTC start and end, the server job and the exit
+# code, so the server's telemetry can be cut to the run. The step runs in the background and survives the caller.
+# CLIENT_VENV (default $CAMPAIGN/client-venv) selects the client environment: SkyRL pins tinker 0.25.0, whose
+# responses the 0.30 SDK rejects (it requires sample_sequence_ids), so SkyRL clients use a 0.25.0 environment.
+set -euo pipefail
+: "${CAMPAIGN:?CAMPAIGN must be set}"
+SERVER=$(cd "$1" && pwd); NAME=$2; BACKEND=$3; shift 3
+[[ ${1:-} == -- ]] && shift
+RUNS=$(dirname "$SERVER")/runs
+VENV=${CLIENT_VENV:-$CAMPAIGN/client-venv}
+mkdir -p "$RUNS"
+read -r JOB HEAD ENDPOINT < <(python3 -c "import json,sys; s=json.load(open(sys.argv[1])); print(s['slurm_job_id'], s['head'], s['endpoint'])" "$SERVER/server.json")
+[[ -e $RUNS/$NAME.log ]] && { echo "$RUNS/$NAME.log exists; pick a new name" >&2; exit 1; }
+printf -v COMMAND '%q ' "$@"
+nohup srun --overlap --jobid="$JOB" -N1 -w "$HEAD" bash -c "
+    export HF_HOME=$CAMPAIGN/hf HF_HUB_OFFLINE=1 PYTHONDONTWRITEBYTECODE=1 PATH=$VENV/bin:\$PATH
+    export CAMPAIGN_PYDEPS=$VENV/lib/python3.12/site-packages
+    export TINKER_BASE_URL=$ENDPOINT TINKER_SERVER_INFO=$SERVER/server.json
+    cd $CAMPAIGN/opd
+    start=\$(date -u +%Y-%m-%dT%H:%M:%S.%3NZ)
+    $COMMAND; code=\$?
+    python3 -c 'import json,sys; json.dump(dict(name=sys.argv[1], backend=sys.argv[2], server=sys.argv[3], slurm_job_id=sys.argv[4],
+        start_utc=sys.argv[5], end_utc=sys.argv[6], exit_code=int(sys.argv[7]), command=sys.argv[8]), open(sys.argv[9], \"w\"), indent=2)' \
+        $NAME $BACKEND $SERVER $JOB \$start \$(date -u +%Y-%m-%dT%H:%M:%S.%3NZ) \$code '$COMMAND' $RUNS/$NAME.window.json
+    exit \$code" > "$RUNS/$NAME.log" 2>&1 &
+echo "started $NAME on $HEAD (job $JOB): $RUNS/$NAME.log"