Refactor readme

[FanhaiLu1]

This RP refactor readme to make it more clear. Here is are several principle for the refactor:

1. Keep readme clean and simple
2. Trade Jax and Pytorch inference stack equally
3. Separate engine implementation README from JetStream README

[JoeZijunZhou]

I will let @gangji decide how we showcase the engines. If we need to make this change, shall we add the maxtext user guide to /docs/tutorials/online-inference-with-maxtext-engine.md. Similarly, we could also add /docs/tutorials/online-inference-with-pytorch-engine.md here.

[FanhaiLu1]

I will let @gangji decide how we showcase the engines. If we need to make this change, shall we add the maxtext user guide to /docs/tutorials/online-inference-with-maxtext-engine.md. Similarly, we could also add /docs/tutorials/online-inference-with-pytorch-engine.md here.

Move maxtext to '/docs/tutorials/' sound reasonable. Ideally, I feel the user guide can be owned by maxtext.
For pytorch path, we have the readme in Jetstream-Pytorch repo, JetStream can refer it.

[gangji]

In general, having a docs folder is a great idea. I am also a fan of keeping tests/docs close to the code. Also, having multiple source of truths might cause maintenance overhead. One solution is link to the https://cloud.google.com/tpu/docs/tutorials/LLM/jetstream. Open to suggestions. Thanks!

[qihqi]

Agree with Fanhai's change.

In principle:
Jetstream is a library that both maxtext and jetstream-pytorch depends.
Jetstream's docs should focus on "How to implement the Engine API and use jetstream itself"; i.e. if huggingface want to also implement the Engine API and serve their models what steps they follow etc. it should not focus how to use any particular impl of the engine, instead, it can link to those as reference examples.

[FanhaiLu1]

In general, having a docs folder is a great idea. I am also a fan of keeping tests/docs close to the code. Also, having multiple source of truths might cause maintenance overhead. One solution is link to the https://cloud.google.com/tpu/docs/tutorials/LLM/jetstream. Open to suggestions. Thanks!

I feel we need two docs: 1: official doc for release version, this is the "https://cloud.google.com/tpu/docs/tutorials/LLM/jetstream" 2: Latest doc with current git code, this one we could use the git readme, the latest change also been updated to readme, but not to official doc as it's not release yet.

In current PR, I keep both git readme and the office doc. The overhead maintenance is minor, anyway engineer need to update the readme with change together if the change either brings new feature or current readme doesn't reflect the code (i.e. Zijun recent dns server address to ip change)

[gangji]

cc @vipannalla.

[vipannalla]

Agree with others, we need separate doc folders. JetStream main README should focus on engine APIs and how to make it easy for anyone to write their own engine. We should just have links to Jetstream-Pytorch and MaxText as reference implemtations.

[vipannalla]

How about "Currently, there are two reference engine implementations available -- one for Jax models and another for Pytorch models."

[FanhaiLu1]

Thanks, changed.

[vipannalla]

This file should really be in MaxText repo, similar to pytorch one.

[FanhaiLu1]

This file should really be in MaxText repo, similar to pytorch one.

We can do it later (after separation maxtext training vs inference ), there is no place I can put in maxtext repo right now.

[vipannalla]

Sounds good.