feat(experimental): integrate TPU-Raiden weight synchronization - #1907
Draft
khatwanimohit wants to merge 1 commit into
Draft
feat(experimental): integrate TPU-Raiden weight synchronization#1907khatwanimohit wants to merge 1 commit into
khatwanimohit wants to merge 1 commit into
Conversation
…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`.
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.
tunix/experimental/orchestrator(WeightSyncCoordinator,WeightSyncHandler,WeightSyncSource,WeightSyncDestination,WorkerRoundTracker).local_launcher: uses native TPU-RaidenWeightSynchronizerAPI with dynamic host IP resolution.pathways: usesweight_synchronizer_ffi(init_weight_synchronizer,multi_h2d) for Shared Pathways Service (SPS) and GKE multi-slice execution.PeftTrainerV2andLegacyVllmSamplerAdapter.compute_onmonkey-patching).k8s_launcher.sh) and runner node entrypoints for distributed multi-slice deployments.tunix/tests/experimental_weight_sync_test.py.Resolves #<issue_number_goes_here>
Reference
Colab Notebook
Checklist