Py sched integration - #2063
Conversation
|
Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA). View this failed invocation of the CLA check for more information. For the most up to date status, view the checks section at the bottom of the pull request. |
2130e5c to
aa371cb
Compare
- distributed_rl_engine.sync_weights: reinstate the router notify_weights_synced call lost when sync moved to the coordinator path (fires only after a successful round) - distributed_rl_engine_test: rewrite the two router-notification tests against the coordinator contract (mock coordinator; old broadcast API is gone) - run_gsm8k_dist_grpo: drop duplicated --inference_addr argparse flag (rebase artifact; argparse raises on startup) - gke/job.yaml: rename mesh env knobs MESH_FSDP/MESH_TP -> TRAINER_FSDP/TRAINER_TP to match the rebased launcher
aa371cb to
1eaef03
Compare
Upstream generate() now rejects bare prompt items; dict items must carry prompt_id and are normalized to RolloutRequests, so the router receives route_key/request_id/prompt hints on a single unified path.
27bba60 to
19cfa87
Compare
|
Thanks for doing the integration @kfswain . Curious do we have measurement on the performance/throughput before and after? |
| fi | ||
| } | ||
|
|
||
| cleanup() { |
There was a problem hiding this comment.
Looks like some changes in this file is reverting the code to old commits.
| ROLLOUT_CMD+=(--use_lora) | ||
| fi | ||
| print_command "Rollout command" PYTHONUNBUFFERED=1 "${ROLLOUT_CMD[@]}" | ||
| PYTHONUNBUFFERED=1 "${ROLLOUT_CMD[@]}" > "$ROLLOUT_LOG" 2>&1 & |
There was a problem hiding this comment.
This will only launch one rollout worker, right ?
| parser.add_argument("--max_prompt_length", type=int, default=1024) | ||
| parser.add_argument("--max_response_length", type=int, default=1024) | ||
| parser.add_argument("--train_micro_batch_size", type=int, default=1) | ||
| parser.add_argument("--trainer_addr", type=str, default="localhost:20000") |
There was a problem hiding this comment.
These address flags are unnecessary after #2060 (merged)
| if isinstance(actor, str): | ||
| self._actors.append(ActorHandle.from_address(actor)) | ||
| elif isinstance(actor, ActorHandle): | ||
| elif callable(getattr(actor, "submit", None)) and callable( |
There was a problem hiding this comment.
This seems like anti-pattern, allowing arbitrary objects to enter ActorPool.
We see improvements with things like DeepSWE that will have multiple turns and need to balance KV affinity with saturation. I was just doing the integration here and prove it works, and start doing deeper tests as I can get my hands on TPU capacity (gsm8k is much easier to test than DeepSWE b/c capacity reasons) |
4 participants
Description
This PR allows for the py-inference-scheduler to be optionally enabled for tunix. This has been tested on the gsm8k path, and includes a readme for integration.