Skip to content

feat(experimental): integrate TPU-Raiden weight synchronization - #1907

Draft
khatwanimohit wants to merge 1 commit into
mainfrom
mohit/raiden-integration
Draft

feat(experimental): integrate TPU-Raiden weight synchronization#1907
khatwanimohit wants to merge 1 commit into
mainfrom
mohit/raiden-integration

Conversation

@khatwanimohit

Copy link
Copy Markdown
Collaborator
  • Implement transport-neutral weight synchronization subsystem in tunix/experimental/orchestrator (WeightSyncCoordinator, WeightSyncHandler, WeightSyncSource, WeightSyncDestination, WorkerRoundTracker).
  • Add fail-safe 3-phase quiesce-and-commit barrier (prepare -> transfer/h2d -> commit) with prefix cache invalidation and rollback on failure.
  • Implement dual backend support:
    • local_launcher: uses native TPU-Raiden WeightSynchronizer API with dynamic host IP resolution.
    • pathways: uses weight_synchronizer_ffi (init_weight_synchronizer, multi_h2d) for Shared Pathways Service (SPS) and GKE multi-slice execution.
  • Ensure deterministic parameter manifest indexing via alphabetical ordering across PeftTrainerV2 and LegacyVllmSamplerAdapter.
  • Upgrade dependencies to JAX 0.11.0 and Flax 0.12.8 with native runtime support (zero compute_on monkey-patching).
  • Update GKE Kubernetes JobSet launcher (k8s_launcher.sh) and runner node entrypoints for distributed multi-slice deployments.
  • Add comprehensive unit and integration tests in tunix/tests/experimental_weight_sync_test.py.

Resolves #<issue_number_goes_here>

It's a good idea to open an issue first for discussion.

Reference

Colab Notebook

Checklist

  • I have added all the necessary unit tests for my change.
  • I have verified that my change does not break existing code and all unit tests pass.
  • I have added all appropriate doc-strings/documentation.
  • My PR is based on the latest changes of the main branch (if unsure, rebase the code).
  • I have signed the Contributor License Agreement.
  • I have followed Contribution Guidelines.

…dual backends

- Implement transport-neutral weight synchronization subsystem in `tunix/experimental/orchestrator` (`WeightSyncCoordinator`, `WeightSyncHandler`, `WeightSyncSource`, `WeightSyncDestination`, `WorkerRoundTracker`).
- Add fail-safe 3-phase quiesce-and-commit barrier (prepare -> transfer/h2d -> commit) with prefix cache invalidation and rollback on failure.
- Implement dual backend support:
  * `local_launcher`: uses native TPU-Raiden `WeightSynchronizer` API with dynamic host IP resolution.
  * `pathways`: uses `weight_synchronizer_ffi` (`init_weight_synchronizer`, `multi_h2d`) for Shared Pathways Service (SPS) and GKE multi-slice execution.
- Ensure deterministic parameter manifest indexing via alphabetical ordering across `PeftTrainerV2` and `LegacyVllmSamplerAdapter`.
- Upgrade dependencies to JAX 0.11.0 and Flax 0.12.8 with native runtime support (zero `compute_on` monkey-patching).
- Update GKE Kubernetes JobSet launcher (`k8s_launcher.sh`) and runner node entrypoints for distributed multi-slice deployments.
- Add comprehensive unit and integration tests in `tunix/tests/experimental_weight_sync_test.py`.
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