Skip to content

Add offline symbolic and parallel weight-sync plan computation. - #1066

Open
copybara-service[bot] wants to merge 1 commit into
mainfrom
test_988717008
Open

copybara-service[bot] wants to merge 1 commit into
mainfrom
test_988717008

Conversation

@copybara-service

Copy link
Copy Markdown

Add offline symbolic and parallel weight-sync plan computation.

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.

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
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant