Skip to content

Kimi K3.1 async GRPO: Miles adapters, engine samplers, NC sampler pool, run up, cost guard and dashboard - #346

Open
michaellchung wants to merge 44 commits into
mainfrom
k31-async-grpo
Open

michaellchung wants to merge 44 commits into
mainfrom
k31-async-grpo

Conversation

@michaellchung

@michaellchung michaellchung commented Nov 10, 2026 •

Copy link
Copy Markdown
Contributor

Summary

Adds end-to-end Kimi K3.1 async GRPO training to yeto:

  • Miles adapters for the K3.1 lineage: ports island, sampler island, buffer reader with short-group carry-over and loss-masked padding, padding-aware metrics and metadata checks, release plugin that restores the base checkpoint metadata, ledger set-aside when a trainer restarts from the base.
  • Engine samplers (SGLang and vLLM servers): R3 routed experts, top-K log-probs, fp32 lm-head, in-place MXFP4 expert reload served with the base tokenizer, and a bound on how long an agentic round may defer a release (reload_max_defer_s, default 600 s).
  • National Compute sampler pool: elastic join/leave, automatic replacement on eviction and before lease expiry, per-node cost accounting, cost guard and resource ledger, yeto cost.
  • yeto run up: one-config launch with preflight, budget / wall-clock / silence backstops; task dispenser; run health breakers; remote sandbox hosts for agentic tasks with orphan sweeps.
  • Trainer robustness: non-finite reward groups, unconvertible groups, sampler-wait timeout measured on progress, early config refusals, randomized fuzz of the buffer read and train chain.
  • Dashboard views for async GRPO runs; docker/miles-k31 image build; examples/k31 specs and launch scripts; sampler node scripts and scripts/prune_layers.py.
  • Docs: customer guide (docs/KIMI_K31_ASYNC_RL_GUIDE.md), 262k full-parameter layout, NC pool, cost guard, dispenser, breaker, sandbox hosts, Miles RL and CLI reference.
  • Fix: codepack keeps the index mtime on its temporary index copy (a same-size edit within the second after a commit was not seen as dirty).

Validation

  • Kimi K3.1, 64-expert full-parameter, 2x8 B300 trainer + NC MI355X vLLM sampler: 4 healthy optimizer steps (grad_norm 0.012-0.033, KL <= 0.07), a release published every step (LATEST committed, config identical to the base, MXFP4 expert blocks changed), in-place hot reload of v1 on the sampler with a matching policy hash.
  • NC elasticity: eviction fenced and replaced automatically; the other island was unaffected; no duplicate dispatch.
  • Qwen3-0.6B single-GPU rehearsal: 10 steps without error, 6 sampler hot reloads, ledger set-aside on restart, short-group padding and carry-over, stale groups.
  • Qwen3-0.6B trainer + 2 NC MI355X agentic sampler islands: 6 forced hot reloads under always-in-flight rounds (sampler followed the trainer to v10, KL 0.0013-0.0027); hard eviction fenced the island and released its open dispense claims within about a minute; the replacement started on LATEST and resumed the predecessor's position; orphan trial containers swept on same-shard restart and, with the new predecessor sweep, after a cross-shard replacement.
  • Fixes from that run: a replacement sampler sweeps its fenced predecessor's trial containers; a row dispensed twice in one round gets its own prompt id; a recorder session collision loses one trial instead of the island.
  • Tests: full local suite 10161 passed, 0 failed (189 skipped, 7 xfailed); after the reload-bound commits, the sampler / release / agentic / island / policy subset 2395 passed, 0 failed; after the NC fixes, the sampler / agentic / sandbox / recorder / dispense subset 2667 passed, 0 failed.

🤖 Generated with Claude Code

michaellchung and others added 30 commits November 11, 2026 01:19
…guard, run up, sandbox hosts and dashboard

Adds the Kimi K3.1 async GRPO path to yeto: the Miles adapters (ports island,
sampler island, state plugin), the engine sampler with SGLang and vLLM
adapters (R3 routed experts, top-K log-probs, fp32 lm-head), the National
Compute sampler pool with elastic join/leave, the cost guard and resource
ledger, `yeto run up` with budget and silence backstops, sandbox hosts for
agentic tasks, the task dispenser and breaker, and the dashboard views for
async GRPO runs.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…up, checkpoint layer pruning

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
… scripts

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…s and dashboard

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
… sandbox hosts, Miles RL and CLI reference

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…r aside

Restarting the trainer of a run from the base (no cut / release to
continue from) failed in ledger.rebase(0) with a LedgerError after the
whole trainer start-up: the batch ledger on the island volume still held
rounds an earlier launch recorded. yeto run up now adds
--rl-ledger-set-aside-before <now> on a trainer restart (and its
preflight says so and dry-runs that argv); the launcher forwards it, and
the driver moves the ledger to ledger.before-<ts> at a start at rollout 0
when every recorded round predates the cutoff. Rounds recorded by the
same launch (an in-place restart) are still refused.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
A short round is filled with loss-free padding samples (copies of the
shortest real sample, metadata.yeto_padding); Miles' rollout metrics
(rewards, response lengths, truncation) counted them. A buffer reader now
sets --custom-rollout-log-function-path to log_rollout_data, which runs
Miles' default logging over the real samples only (the custom hook off,
no recursion) and adds rollout/padded_samples and padded_fraction.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
… it costs

Padding samples carry no loss, but Miles divides the round's loss by the
full global batch, so the real samples of a padded round train with
(1 - padded_fraction) of their gradient. yeto does not rescale it (the
loss normalization is Miles'); the reader now reports
sample_buffer/padded_fraction and docs/MILES_RL.md explains the effect
and what to change when the fraction stays high.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
The island entry writes miles_args.yeto_rl_ledger_set_aside_before
(--rl-ledger-set-aside-before); the island-args inventory must name every
driver-side key (test_island_args_registry_covers_every_driver_side_key).

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…ver as stale

The short-round top-up (_fill) consumed too-old groups, a fenced
incarnation's groups and quarantined ones unread without counting them,
and a carried group that aged out while it waited was counted stale but
lost its dispense position. They are now counted like the main
selection counts them (stale_groups with stale_positions,
stale_incarnation_groups, quarantined_groups).

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
The real-R2 e2e test takes its profile and bucket from
YETO_E2E_R2_PROFILE / YETO_E2E_R2_BUCKET and skips without them; the
example-placeholder checks compare against that variable instead of a
bucket name; the dual_clipfrac check skips without YETO_TEST_MILES_K31_B8
(no checkout-path default); the cost-guard fixture uses a neutral app name.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
The launch layer imported the Miles adapter statically (k31_checks,
manifest, modal_islands, rl_island_task), which the import-boundary rule
forbids. The registry gains the pins and kimi_k31_model roles and the
launch code loads them (and expert_selection) by role.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…on-RL literals

The neutral schema named sglang (--rl-sampler-engine choices); the flag
serves the Kimi-K3.1 LoRA target check and now lives with the Miles
adapter's flags. The flag literals of the NC agentic setup (pip, tar,
aws s3, the node agent) and of the selection-tool hint in k31_checks are
pinned as non-RL literals.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…'s Miles

pins.MILES_LINEAGES["k31"] gained MILES_K31_FULL_E_COMMIT with the full-form
image pin; the lineage check still listed the commits up to FULL_D.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Covers architecture, prerequisites, run/pool/sampler-spec configuration,
yeto run up/scale/status/down, automatic NC replacement, resume paths,
code pack upgrades, the dashboard, breakers and cost protection, and a
troubleshooting FAQ. Linked from README.md and examples/k31/README.md.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Section 9 covers the pinned Miles K3.1 images, docker/miles-k31 build and
version lock, --rl-image, and git-patch vs release code packs. Section 10
covers TP/PP/EP/ETP/CP, the expert selection file (64 -> 128), the memory
estimator with observed headroom, and recommended 128-expert layouts.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Add docs/K31_FULL_262K_LAYOUT.md (PP3 x TP8 x CP2 x EP16 on 6x8 B300:
measured 524k split 32/28/33, recommended 262k split 32/29/32) and make
section 10 of the guide a summary that links to it. Section 9 now pins the
full-parameter trainer image 2de5d796, names the LoRA image, adds the
vLLM ROCm sampler image digest, and clarifies the default registry.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…dle polls apart

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…e registry

MILES_NEXT_IMAGE is now the K3.1 image under YETO_IMAGE_REGISTRY (default
ghcr.io/agentenv); the check that the default image is not on that registry
predates it. Check the digest pin instead.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
dirty_patch copied .git/index with copyfile, so the copy had a fresh mtime and
git no longer treated entries written in the same second as racily clean. A
same-size edit made within that second (same mtime/ctime seconds, same size)
then looked unchanged: the patch was empty and the pack not dirty, while the
pack digest still moved. test_codepack relaunch hit it about 1 run in 12.
copy2 keeps the mtime, so git re-reads those files.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…amples in the metadata check

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…d, not trained

One NaN reward makes the whole group's centered advantages NaN in Miles
(a non-finite loss, then the driver's nonfinite_grad_norm stop with
nan_skip off by default); an all-NaN group stops the driver at once
(nonfinite_reward). A failed reward call on one sampler must not end the
trainer: the reader passes such groups over like a too-small group and
counts them (sample_buffer/nonfinite_reward_groups). Found by the
trainer read-chain fuzz.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…ole round

The trainer's take raised once wait_s (default 3600 s) had passed since
the take began, even while the samplers kept writing groups: a round that
fills slowly (long agentic trajectories, a sampler that reloads weights
after a release) failed the trainer although nothing was stuck. The wait
now starts again at every group the buffer selects; it still fails a
round whose samplers wrote nothing for wait_s.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
The K3.1 guide's run config template has budget.usd 0 next to the 0
prices it asks the customer to fill in. With a price filled in and the
cap left at 0, the monitor's first spend estimate is a budget breach and
the trainer is stopped minutes after its start. Preflight now names it.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…the start

YETO_RL_XHW_GUARD=action=downweight with std normalization on (Miles'
default) only failed in xhw_round, at the first round the guard tripped:
the trainer stopped mid-run on a setting it could have refused at once.
The reader now checks it when it is built.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
The padding of a short round copied the sample with the fewest tokens;
an empty response (an engine that stopped at once) wins that, and gives
Miles zero-length response rows, where a negative slice x[-0:] is the
whole sequence. Padding now copies the shortest sample with a response
(an all-empty round keeps the old choice).

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
… out of the round

One malformed sample (an R3 routing row mixing pads and expert ids, an
expert id past the model's, another layer count) raised in
miles_sample_dict and stopped the trainer. The round now leaves that
group out (consumed, padded instead, sample_buffer/unconvertible_groups);
a take whose every group is malformed is taken again, and only
MAX_UNCONVERTIBLE_TAKES such takes in a row (a systematic cause) fail
the round.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Real FileSampleBuffer (single index or merged sampler shards), BatchReader,
the Miles reader round (_BufferRound, Miles stubbed), the yeto metadata
hooks over a dir sink, handle_from_metadata, IslandDriver._accept_batch
and BatchLedger. Each seed runs 24 rounds over whole, short, too-small,
oversized, stale, mixed-version, fenced, quarantined, non-finite-reward,
empty-response and malformed-R3 groups, a batch size that changes per
round, releases between rounds, cuts (also with the round taken ahead),
reader resets, cut restores, stalled samplers and restarts from the base;
after every round it checks no exception, no position trained twice or
skipped without a reason, a monotone watermark, trained samples equal to
the metadata's, G * n samples, and the ledger. YETO_FUZZ_SEEDS scales it.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
A Qwen3-0.6B rehearsal (Miles trainer on one H100, vLLM 0.11 engine
sampler on one L4) found that the sampler could not restart its engine
on any release: the trainer's HF export writes tokenizer_config.json in
the transformers-5 form (extra_special_tokens as a list), and the 4.x
tokenizer of the sampler image fails on it ('list' object has no
attribute 'keys'). The sampler retried ten times and left INFRA_FAILED,
so the trainer starved.

RL never changes the tokenizer, so the vLLM server (both lines) now
starts on a release with --tokenizer <base_dir>, and the SGLang server
with --tokenizer-path <base_dir>. An explicit tokenizer flag in
extra_args wins; the cold start on the base is unchanged.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
michaellchung and others added 14 commits November 11, 2026 04:30
…lease

The trainer image saves the HF export with transformers 5, which moves
rope_theta / rope_scaling into rope_parameters, torch_dtype into dtype,
and writes extra_special_tokens of tokenizer_config.json as a list. A
transformers-4 sampler (vLLM 0.11 in the Qwen3-0.6B rehearsal) read no
rope_theta from a release and served it with the default RoPE base:
mean KL was 0.0007 on base groups and 0.43 on the groups of every
release, with no error anywhere.

A full release (the staged writer and the HF round exports) now puts
back every top-level key of the base config.json the export dropped
(the export's own keys are kept; quantization_config and
compression_config are never copied onto a release), and the tokenizer
files are the base's: RL does not change them.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
A streaming agentic round holds the weights while any group's trials run.
With trials of 20+ minutes there is never a gap, so a non-live engine
(vLLM in-place reload) deferred every release after the first one.
PolicyFollower.max_defer_s loads a ready release once hold() has deferred
it that long; rl_sampler_policy_loaded carries forced and deferred_s, and
FollowerCounters.forced_loads counts such loads. A round marked by
generating() is never forced. Trials spanning the switch keep the drain
rule (lost as served_by_other_version, or aborted by the records proxy).

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…t round

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…containers

An NC replacement continues a fenced shard under a new shard id
(continues_shard), so the stale-incarnation sweep keyed on the sampler's
own shard never reached the Harbor trial containers an evicted or
hard-killed node left on the sandbox host. They kept running and held
sandbox room for the rest of the run.

After the handoff, the replacement now removes every container of the
predecessor shard (trial_labels.sweep_shard, SandboxGate.sweep_shard)
once the predecessor left or a pool fence holds its incarnation lease,
and emits rl_sandbox_orphans_removed with the shard. A predecessor that
is still running unfenced is left alone.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
A dispenser can hand out one row at two positions of the same round (an
epoch wrap of a small dataset: positions 74 and 75 both mapped to row7 in
q06-nc-e2e-20261110T2255Z). The agentic recorder sessions and the buffer
group key are built from the prompt id, so both groups opened
'<inc>-r410-row7-0' and the second open stopped the island. The repeat
now gets its dispense position appended (row7-p75).

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
The recorder's "already open" refusal was a plain HarnessConfigError, and
any exception escaping a trial stops the island (INFRA_FAILED). A
duplicate session id is a per-trial fault: the recorder now raises
SessionAlreadyOpen (a HarnessConfigError subclass), and the agentic round
loses that trial as 'session_collision', counts it
(session_collisions_total in the round event) and keeps running.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…ey header, home path or CJK text)

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Resolve --tuning full admission in rl_args and sample_roles: accept both the
SAO streaming preset from main and the K3.1 whole-checkpoint recipe.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
A synced group's sampler islands run as no-sync islands, and the Miles
island entry refuses --rl-single-island-no-sync together with
--sync-preset decoupled. The launch accepted --rl-sync-preset decoupled
with --rl-island-sample-role M:sampler, so the sampler island died at
start-up (exit 2) while the trainer islands kept running. Refuse the
combination in check_synced_group, before any GPU starts.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…ed roster

The async sync mode gives a fixed-roster decoupled run a quorum below the
roster (quorum 1 of 2). A fragment round can then commit without an island
whose local horizon was not reached yet; the syncer permits that
fragment's next round, and the island failed with 'decoupled RL received
conflicting fragment permits' (island 1 of a two-island Modal run, after
outer step 3). The supersede rule (and the drop of a closed round's
permit on its broadcast) applied only to an island pool; it holds for any
roster, because the syncer permits a fragment's next round only after the
previous one committed. With a full-roster quorum nothing changes.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
ps truncates args to the terminal width when stdout is not a tty, which cut
the run name off long worker command lines (seen in CI: the worker was not
found). Pass -ww so the full argv is matched.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
CI has no Modal SDK, so cmd_launch stopped at the SDK preflight before it
reached the live-worker check the test exercises.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant