Closed FanhaiLu1 closed 4 months ago
This looks great!! Mostly nits, but IIUC this same approach would work for JAX/JetStream right?
Thanks Allen for reviewing it! Yes, the Jax/Jetstream should use same approach in general, we could extract common part of code (For example: Master code could be exactly same for both Jax and Pytorch), so both Jax and Pytorch part can share same code base.
Could we also add a basic unit test for Ray? One way is to take advantage of Ray's "fake" cluster if you look at this file for instance.
Could we also add a basic unit test for Ray? One way is to take advantage of Ray's "fake" cluster if you look at this file for instance.
Looks great with Ray's "fake" cluster test. Can I add test with another PR (I tested current PR e2e manually)? I also plan to add unit test for RayWorkers in next PR.
Looks great with Ray's "fake" cluster test. Can I add test with another PR (I tested current PR e2e manually)? I also plan to add unit test for RayWorkers in next PR.
Sounds good to me! Thanks for patiently addressing all of the comments, LGTM!
This PR enable pytorch engine multiple host on TPU POD slices.
MVP Goal:
The current PR is MVP version of multiple host, two goals in MPV:
Load weight and sharding in multiple host
Compute meaningful decode result
Result validation:
Weight and sharding
With this PR, the weight and sharding on v5e-16 (4 host , total 16 chips) use 50 of memory compared with 8 chips. It worked as expected.
The weight and sharding on v5e-8 (4 host , total 8 chips)
Meaningful Result
With this PR, first two line of the result on v5e-16 (4 host , total 16 chips). The result are meaningful and looking good to human.
first two line of the result on v5e-8 (4 host , total 8 chips)
Caveats
As this is the MVP, it's important to note that there could be some limitations in terms of performance and accuracy right now.