Fix Harvey Gemma tool rollouts - #5
Draft
ShubyM wants to merge 12 commits into
Draft
Conversation
- make push-to-cluster (skaffold run): builds the working tree's images with content-addressed inputDigest tags, pushes to gcr.io/$GCP_PROJECT, and deploys the FFT time-slice manifests with built images substituted — including the OPEN_RL_WORKER_IMAGE env var (skaffold resourceSelector), so trainer/sampler pods launched after the deploy run the new server image. The wrapper waits for the gateway rollout and deletes stale worker pods. - Gitignored local.mk (seeded from local.mk.example) for per-developer Makefile settings such as GCP_PROJECT. - IMAGE_TAG for the manual build-images/push-images flow gains a content hash of uncommitted changes so dirty-tree pushes produce fresh tags. - Manifests reference the canonical CI-published images (ghcr.io/gke-labs/open-rl/*) with the deployed tag pinned in one place; gitignored k8s/deploy/local overlay for per-developer manifest overrides. - make clean-workers removes runtime-launched trainer/sampler pods.
# Conflicts: # .gitignore
compute_target_logprobs materialized the full [batch, seq, vocab] logits tensor plus a second fp32 copy from log_softmax, both retained through backward on the GPU holding the lm_head. For large-vocab models (Gemma ~256K) that activation is the dominant term and OOMs the trainer. Run the backbone for hidden states and project through the vocab in chunks, reducing to logit[target] - logsumexp with per-chunk activation checkpointing, so peak head activation is [chunk, vocab] instead of [batch*seq, vocab]. The full-logits fallback (unknown backbone or OPEN_RL_FUSED_LOGPROB=0) now also uses logit - logsumexp to drop the redundant fp32 log_softmax allocation. Env knobs: OPEN_RL_FUSED_LOGPROB (default 1), OPEN_RL_LOGPROB_CHUNK (default 1024). Verified numerically identical (values and gradients) to the original log_softmax(...).gather path, including Gemma final-logit softcapping.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What changed
submittool, theread_documentrename, and duplicated output instructionstop_p, andtop_kthrough the torch FFT and LoRA sampling pathWhy
The live tool schema disagreed with Harvey's prompt and teacher traces: the prompt taught
read, while Open-RL exposedread_document. The environment also added asubmittool that stock Harvey does not define. Separately, the torch sampler discarded Gemma's stop and sampling controls, and the canonical thinking template can pre-open a thought channel in the prompt that must be reconstructed before parsing sampled continuation tokens.Together these issues caused Gemma to claim documents were unavailable or to emit valid tool calls that the rollout parser rejected.
Impact
An exact BF16
google/gemma-4-E4B-itvLLM smoke run on a stock Harvey task now discovers the mounted documents and successfully executes the canonicalreadtool. The remaining incomplete-task behavior in that sample was model search strategy—it searched only for.docxand skipped an.xlsx—rather than environment access.Validation
make test: 87 passed, 5 skippedglobfound the mounted agreement andreadreturned its contentsThe standalone Harvey test file was intentionally omitted from this commit.