Enable Pipeline Parallelism on jax worker

[Chenyaaang]

Description

The implementation for Pipeline Parallelism are splitted into the following small PRs.

• Util functions for PP on Jax (Create Utils functions for PP #1042)
• worker changes (The current PR)
• runner changes
• platform changes
• torchax changes
• Jax changes
• Single host changes (this will be in vLLM)
• Multi host changes (Ray)

This PR is to modify Jax worker to support PP.

• The worker __init__ takes in the current worker's IP and its previous worker's IP, to start transfer server and connection later.
• During execute_model, for the PP workers who not in the first rank, they need to receive intermediate tensor from the previous worker. For the PP workers who is not the last rank, they need to send intermediate tensor to their next worker.
• The profiler should profile every PP worker, so add subfolders under the parent profile_dir to save profiles for each worker.

Tests

E2E test has verified the whole PP implementation works properly.

Checklist

Before submitting this PR, please make sure:

• I have performed a self-review of my code.
• I have necessary comments in my code, particularly in hard-to-understand areas.
• I have made or will make corresponding changes to any relevant documentation.

[github-actions]

Description

Start with a short description of what the PR does and how this is a change from
the past.

The rest of the description includes relevant details and context, examples:

• why is this change being made,
• the problem being solved and any relevant context,
• why this is a good solution,
• some information about the specific implementation,
• shortcomings of the solution and possible future improvements.

If the change fixes a bug or a Github issue, please include a link, e.g.,:
FIXES: b/123456
FIXES: #123456

Tests

Please describe how you tested this change, and include any instructions and/or
commands to reproduce.

Checklist

Before submitting this PR, please make sure:

• I have performed a self-review of my code.
• I have necessary comments in my code, particularly in hard-to-understand areas.
• I have made or will make corresponding changes to any relevant documentation.

[mrjunwan-lang]

do we need report errors if the settings are not in these 2 ways?

[Chenyaaang]

I use v6e8 as an example to provide 2 ways to subslice the chips. I was thinking if the customer is using other chips, they should replace with their own topology. Do you have any better idea to interpret this?

[yixinshi]

Can the topology be passed as config variables (and make the default as one of v6e's supported topology) or at least a parameter of init_device() so people would find the needed changes more easily? And please move lines 136-141 to line 152.

[Chenyaaang]

There are 6 vars in total:

• TPU_PROCESS_ADDRESSES: all workers are the same
• TPU_PROCESS_PORT: different in each worker
• CLOUD_TPU_TASK_ID: different in each worker
• TPU_PROCESS_BOUNDS: all workers are the same
• TPU_CHIPS_PER_PROCESS_BOUNDS: all workers are the same
• TPU_VISIBLE_CHIPS: different in each worker

It might be difficult to pass from configs. I made it as parameters of init_device() method.

[Chenyaaang]

@sixiang-google Hi Xiang, can you help me take a look at init_device() will my change to local_devices affect disagg serving? Thanks

[yixinshi]

Good work! A general comment: shall we have more specific PR title here for each PR?