Update front-page readme with link to XLA flag doc

[terrykong]

Followup for #674 that links to this new document

[instinct79]

Should we delete all xla flags from the front page readme, and link to this document and maybe the scripts for PAX or T5X for benchmark specific flags ?

[nouiz]

Should we delete all xla flags from the front page readme, and link to this document and maybe the scripts for PAX or T5X for benchmark specific flags ?

We need to tell which flag are enabled in the container. But maybe we can just list them and tell to see that page for their description?

[instinct79]

I think we should point to our scripts that enable these flags. Otherwise, there are duplicate documentation. Perhaps, the environment variables can stay here.

[terrykong]

Since the flags in there can change over time, what about instead including instructions to view the flags. Example:

docker run --rm quay.io/skopeo/stable inspect docker://ghcr.io/nvidia/jax:upstream-pax | jq -r '.Env[]' | grep '^XLA_FLAGS='

# which returns

XLA_FLAGS= --xla_gpu_enable_latency_hiding_scheduler=true --xla_gpu_enable_async_all_gather=true --xla_gpu_enable_async_reduce_scatter=true --xla_gpu_enable_triton_gemm=false

[instinct79]

That's a good idea. Let's do that.

[yhtang]

Since the flags in there can change over time, what about instead including instructions to view the flags. Example:

docker run --rm quay.io/skopeo/stable inspect docker://ghcr.io/nvidia/jax:upstream-pax | jq -r '.Env[]' | grep '^XLA_FLAGS='

# which returns

XLA_FLAGS= --xla_gpu_enable_latency_hiding_scheduler=true --xla_gpu_enable_async_all_gather=true --xla_gpu_enable_async_reduce_scatter=true --xla_gpu_enable_triton_gemm=false

I would consider it a less-readable approach as a reader will have to break from the mental flow and do scripting to obtain these flags.

A more user-friendly but also much heavier solution is to set up a proper docs building process (e.g. readthedocs) so that the docs page is automatically populated with the flags set in the containers.

[instinct79]

Maybe provide a sample output from one of our latest containers for now ? The flags are not going to change so rapidly and will atleast provide the users some idea of what is enabled for our benchmarks.

[yhtang]

Maybe provide a sample output from one of our latest containers for now ? The flags are not going to change so rapidly and will atleast provide the users some idea of what is enabled for our benchmarks.

Good point. @terrykong what do you think?

[terrykong]

Done 👍