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.

[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?