Adding default XLA/GPU env vars for all JAX-based containers

[terrykong]

Addresses #105 and #109

• --xla_gpu_enable_latency_hiding_scheduler=true: Allows XLA:GPU to move communication collectives to increase overlap with compute kernels.
• --xla_gpu_enable_async_all_gather=true: Allows XLA:GPU to run All Gather NCCL kernels on a separate CUDA stream to allow overlap with compute kernels.
• --xla_gpu_enable_async_reduce_scatter=true: Allows XLA:GPU to run Reduce Scatter NCCL kernels on a separate CUDA stream to allow overlap with compute kernels.
• --xla_gpu_enable_triton_gemm=false: Disallows Trition GeMM kernels; uses CUBLAS GeMM kernels instead. CUBLAS kernels are currently better tuned for GPUs and thus provide better performance.
• CUDA_DEVICE_MAX_CONNECTIONS=1: Use a single queue for GPU work, lowers latency of each stream operation. OK since XLA already orders launches.
• NCCL_IB_SL=1: defines the InfiniBand Service Level (1)

@nluehr @ashors1

[yhtang]

Could someone please edit the PR description to document the rationale behind each flag?

[terrykong]

Could someone please edit the PR description to document the rationale behind each flag?

@nluehr @abhinavgoel95 ?

[abhinavgoel95]

1. --xla_gpu_enable_latency_hiding_scheduler=true: Allows XLA:GPU to move communication collectives to increase overlap with compute kernels.
2. --xla_gpu_enable_async_all_gather=true: Allows XLA:GPU to run All Gather NCCL kernels on a separate CUDA stream to allow overlap with compute kernels.
3. --xla_gpu_enable_async_reduce_scatter=true: Allows XLA:GPU to run Reduce Scatter NCCL kernels on a separate CUDA stream to allow overlap with compute kernels.
4. --xla_gpu_enable_triton_gemm=false: Disallows Trition GeMM kernels, uses CUBLAS GeMM kernels instead.

[nluehr]

@abhinavgoel95 what about the following?
NCCL_PROTO=LL128
NCCL_AVOID_RECORD_STREAMS=1
NCCL_IB_SL=1

[yhtang]

• NCCL_PROTO=LL128:
  doc says:

  The default is LL,LL128,Simple on platforms which support LL128, LL,Simple otherwise.

  Related issue: What is LL128 Protocol? nccl#281

• NCCL_AVOID_RECORD_STREAMS=1: undocumented in https://docs.nvidia.com/deeplearning/nccl/user-guide/docs/env.html.

• NCCL_IB_SL=1
  doc says:

  Defines the InfiniBand Service Level.
  For more information, see the InfiniBand specification Volume 1 (https://www.infinibandta.org/ibta-specifications-download) or vendor documentation.

  and the link in the quoted text above is dead.

[nluehr]

@abhinavgoel95 How much speedup do we see from setting NCCL_PROTO=LL128 in practice?
The combiner threshold may come into play here as for larger collectives, I think LL128 could reduce effective bandwidth by ~6%.

[nluehr]

From my reading of the NCCL code, the only safe way to enable LL128 is to not specify NCCL_PROTO at all.
LL128 is valid only on Volta, Ampere, and Hopper with NVLink. PCIE can cause packet re-ordering that results in silent corruption.
Since we expect users will run our models on a variety of clusters, I think we shouldn't define NCCL_PROTO either in the JAX dockerfile nor in the paxml run scripts. If removing this results in worse performance, we should open a bug with the NCCL team.

[ashors1]

We're also sometimes interested in setting XLA_PYTHON_CLIENT_MEM_FRACTION=0.9 to increase the amount of memory allocated to XLA. Is this something we'd want to consider adding to the base container as well?

[nluehr]

Whether we want to set XLA_PYTHON_CLIENT_MEM_FRACTION depends on the workload (particularly whether any other GPU libraries are being used outside of XLA).

My opinion is that this would be OK to set in the paxml container, but for the JAX container in general we shouldn't set it because there we expect users to extend it in varied workflows.

[yhtang]

NCCL_AVOID_RECORD_STREAMS is NOT documented anywhere. Shall we open an NVBug?

[nluehr]

As best I can tell, NCCL_AVOID_RECORD_STREAMS is a feature of pytorch rather than NCCL. (It seems it's now been renamed to TORCH_NCCL_AVOID_RECORD_STREAMS here).
So I think it's safe to drop it entirely.

[yhtang]

Did we figure out a source for the undocumented NCCL_AVOID_RECORD_STREAMS variable?

[terrykong]

@abhinavgoel95

[abhinavgoel95]

@yhtang @terrykong it is safe to drop NCCL_AVOID_RECORD_STREAMS. We do not need to upstream this. It is a PyTorch specific feature.

[terrykong]

I removed mention of that env var. Can I re-request your review @yhtang ?