Repository navigation
perf: prefer GPU types per job with openrl.gpu_types - #372
Draft
hodaaaaaaaaaa wants to merge 2 commits into
Draft
hodaaaaaaaaaa wants to merge 2 commits into
hodaaaaaaaaaa wants to merge 2 commits into
Conversation
Workers land on the smallest GPU that fits, through the claim's tiers and the tight-fit node-affinity weights. That is the wrong choice for a job that runs several times faster on a bigger device, and a job had no way to say so. accelerator.preferredTypes lists accelerator types, most preferred first. When set, the pod's preferred node-affinity terms are one per type on the node label --accelerator-type-label (OPEN_RL_ACCELERATOR_TYPE_LABEL, default cloud.google.com/gke-accelerator), weighted 100 and 10 less per position, floored at 1, in place of the tight-fit terms. The required affinity is unchanged, so a busy preferred type still overflows to any GPU that fits. Without the field the pod renders as before.
openrl.gpu_types is a comma-separated, ordered list of accelerator types (values of the cloud.google.com/gke-accelerator node label). It is stored on the model's metadata and sent on both the trainer and sampler Workloads as accelerator.preferredTypes. A shared LoRA runtime's Workloads carry the preference of the job that created them. Servers that do not launch workers as pods refuse it.
This branch has not been deployed
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
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.
Problem
The scheduler puts each worker on the smallest GPU that fits. A gemma-4-E2B LoRA trainer asks for about 15 GiB, so it lands on an L4 next to free A100s and trains several times slower. A job has no way to say which GPU type it should prefer.
Change
A new
openrl.gpu_typessetting lists accelerator types in order of preference, as values of the node labelcloud.google.com/gke-accelerator(for examplenvidia-a100-80gb;nvidia-l4; semicolons, since TINKER_TAGS splits on commas). The API server puts it on the job's Workloads asaccelerator.preferredTypes, and the scheduler then ranks nodes by type (weights 100, 90, ...) instead of by size. It is a preference: a worker still lands on any GPU that fits when those are busy.--accelerator-type-label/OPEN_RL_ACCELERATOR_TYPE_LABELnames the node label on other clusters.Results
Measured on GKE with eight GSM8K GRPO jobs launched together (gemma-4-E2B and Qwen3-0.6B, LoRA and full fine-tuning, 7 A100 80GB + 5 L4), with
openrl.gpu_types=nvidia-a100-80gbon the gemma-4-E2B jobs against main, two runs each: