AsyncGRPOConfig : dtype asyncgrpo training - #6774
Conversation
|
The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update. |
![cursor[bot]](https://avatars.githubusercontent.com/in/1210556?s=60&v=4){kind=small}
There was a problem hiding this comment.
Cursor Bugbot has reviewed your changes using default effort and found 1 potential issue.
❌ Bugbot Autofix is OFF. To automatically fix reported issues with cloud agents, enable autofix in the Cursor dashboard.
Want higher recall? High effort reviews run extra passes and find more bugs. A team admin can switch effort levels in the Cursor dashboard.
Reviewed by Cursor Bugbot for commit 630d840. Configure here.
| logger.warning( | ||
| f"The vLLM server serves in {vllm_dtype} but the weights sent to it are {train_dtype}. Set `dtype` in " | ||
| f"`AsyncGRPOConfig` to '{vllm_dtype}', or start the server with `--dtype {train_dtype}`." | ||
| ) |
There was a problem hiding this comment.
Warning suggests harmful dtype fix
Medium Severity
The new mismatch warning equally recommends setting AsyncGRPOConfig.dtype to the vLLM dtype or aligning the server to the trainer. In the common case (trainer float32, vLLM bfloat16), the first option pushes users onto bfloat16 training—the precision this change and its cited measurements intentionally move away from—so the remediation text works against the feature’s purpose.
Reviewed by Cursor Bugbot for commit 630d840. Configure here.
AmineDiro
changed the title
The two-node Slurm layout (8 trainer ranks + 8 vLLM engines) becomes one h200x8 split 4 + 4, with `gradient_accumulation_steps` doubled to keep the recipe's 128 completions per optimizer step. `resume` is now idempotent: the in-job script reads `global_step` out of the checkpoint and trains STEPS_PER_JOB more, up to TOTAL_STEPS, so a 2400-step run is a chain of jobs on one trackio run. A constant LR is what makes that safe — there is no decay horizon for a per-job `max_steps` to distort. Carries over the four things the Slurm runs proved load-bearing: top_p=1.0 (0.95 collapsed four runs), fp32 master weights under bf16/fp16 autocast rather than bf16 parameters (which went nowhere: +0.031 over 2400 steps), 128 completions per step with token_budget=0, and adam_beta2=0.95. The precision-mismatch warning from PR #6774 fires in this configuration and is expected: mixed precision means fp32 weights going to a bf16 server, which is what every good run did. Code now reaches the job as a `git archive` of the branch, synced to a bucket and mounted read-only. Fetching a branch tarball from GitHub in-job failed three runs in a row with 429/503 — the jobs egress IP is shared and rate-limited. `origin` stays the record; the bucket is only transport. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
1 participant
What does this PR do?
Added dtype to
AsyncGRPOTrainerwith inference dtype detection to raise a warning if mismatch is detectedSee:
https://huggingface.co/papers/2510.26788
https://huggingface.co/spaces/aminediroHF/trainer-generator-bf16-mismatch
Note
Medium Risk
Changes default model load precision and touches GRPO importance-ratio sensitivity; mismatch handling is warn-only but misconfiguration could skew training.
Overview
Adds a
dtypeknob onAsyncGRPOConfig(defaultfloat32) so the policy is loaded viamodel_init_kwargsinstead of a hardcodedtorch.float32inAsyncGRPOTrainer. Adtypeentry inmodel_init_kwargsstill wins.VLLMClient.get_dtype()reads the served weight precision from vLLM’s/server_info.WeightTransferClient.init_weight_transfercompares that to the trainer’s weight dtypes (mode ofdtype_names) and logs a warning with hints to alignAsyncGRPOConfig.dtypeorvllm serve --dtype, aimed at trainer–generator precision mismatch in async GRPO.Reviewed by Cursor Bugbot for commit 630d840. Bugbot is set up for automated code reviews on this repo. Configure here.