Add JetStream to MaxText container

[DwarKapex]

This is the first step to bring JetStream for MaxText.
The next step will be adding perf tests for couple JetStream + MaxText

[kocchop]

LGTM, please rebase it on top of #1070

[kocchop]

In addition to the pip install -r requirements.txt we also need to add pip install . from the JetStream repo. The google-jetstream that is installed as part of maxtext requirements is not the latest. The above step would uninstall it and install the most recent compiled one from the repo.

[yhtang]

In addition to the pip install -r requirements.txt we also need to add pip install . from the JetStream repo. The google-jetstream that is installed as part of maxtext requirements is not the latest. The above step would uninstall it and install the most recent compiled one from the repo.

Do you mean to add something like

echo "-e ${SRC_PATH_JETSTREAM}" >> /opt/pip-tools.d/requirements-jetstream.in

?

[yhtang]

See comment above

[olupton]

Can we test multi-GPU functionality?

[DwarKapex]

I did not find JetStream unit-test examples on GitHub for multi-GPU. Even more: afaik, it does not support GPUs so far. But it will start GPU support in the nearest future.