Add offline symbolic and parallel weight-sync plan computation. - #1066
Open
copybara-service[bot] wants to merge 1 commit into
Open
copybara-service[bot] wants to merge 1 commit into
copybara-service[bot] wants to merge 1 commit into
Conversation
Previously, `ReshardPlanner` computed the entire resharding transfer schedule sequentially on the centralized controller at step 0 after all workers registered their dynamic `ip:port` endpoints. For large topologies and multi-hundred-billion-parameter models, this centralized calculation delayed step 0 execution. This change introduces two complementary capabilities to eliminate step 0 planning overhead: - Offline symbolic planning and parallel worker loading: `ReshardPlanner.compute_offline_schedule` and `save_offline_plan` compute the N-D resharding schedule ahead of time using deterministic `raiden_symbolic://` endpoint placeholders and serialize per-worker `ControlRequest` protobufs (`<unit>.pb`) alongside `offline_schedule.json`. At runtime, `load_offline_worker_plan`, `load_offline_worker_plans_parallel`, and `load_offline_schedule` load the precomputed plans in parallel and bind symbolic placeholders to live `ip:port` endpoints in `O(unique_plans)`. Configurable via `offline_plan_path` / `RAIDEN_OFFLINE_PLAN_PATH` and `save_offline_plan_path` / `RAIDEN_SAVE_OFFLINE_PLAN_PATH`. - Parallel per-worker schedule computation: `compute_transfer_schedule_from_metadata` now accepts `target_src_units` and `parallel_worker_planning`, allowing each source worker's schedule slice to be computed independently or across a parallel worker pool (`compute_schedules_in_parallel_workers`) and merged via `merge_worker_schedules`. Configurable via `parallel_worker_planning` / `RAIDEN_PARALLEL_WORKER_PLANNING` or `RAIDEN_PLANNING_MODE=parallel_worker`. PiperOrigin-RevId: 988717008
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.
Add offline symbolic and parallel weight-sync plan computation.
Previously,
ReshardPlannercomputed the entire resharding transfer schedule sequentially on the centralized controller at step 0 after all workers registered their dynamicip:portendpoints. For large topologies and multi-hundred-billion-parameter models, this centralized calculation delayed step 0 execution.This change introduces two complementary capabilities to eliminate step 0 planning overhead:
ReshardPlanner.compute_offline_scheduleandsave_offline_plancompute the N-D resharding schedule ahead of time using deterministicraiden_symbolic://endpoint placeholders and serialize per-workerControlRequestprotobufs (<unit>.pb) alongsideoffline_schedule.json. At runtime,load_offline_worker_plan,load_offline_worker_plans_parallel, andload_offline_scheduleload the precomputed plans in parallel and bind symbolic placeholders to liveip:portendpoints inO(unique_plans). Configurable viaoffline_plan_path/RAIDEN_OFFLINE_PLAN_PATHandsave_offline_plan_path/RAIDEN_SAVE_OFFLINE_PLAN_PATH.compute_transfer_schedule_from_metadatanow acceptstarget_src_unitsandparallel_worker_planning, allowing each source worker's schedule slice to be computed independently or across a parallel worker pool (compute_schedules_in_parallel_workers) and merged viamerge_worker_schedules. Configurable viaparallel_worker_planning/RAIDEN_PARALLEL_WORKER_PLANNINGorRAIDEN_PLANNING_MODE=parallel_worker.