diff --git a/.github/container/Dockerfile.jax b/.github/container/Dockerfile.jax index 1f5aabed7..7cc25bb6f 100644 --- a/.github/container/Dockerfile.jax +++ b/.github/container/Dockerfile.jax @@ -97,7 +97,6 @@ ENV BUILD_DATE=${BUILD_DATE} # The following environment variables tune performance ENV XLA_FLAGS="" ENV XLA_FLAGS="${XLA_FLAGS} --xla_gpu_enable_latency_hiding_scheduler=true" -ENV NCCL_NVLS_ENABLE=0 COPY --from=builder ${BUILD_PATH_JAXLIB} ${BUILD_PATH_JAXLIB} COPY --from=builder ${SRC_PATH_JAX} ${SRC_PATH_JAX} diff --git a/README.md b/README.md index ec125d273..a111994f3 100644 --- a/README.md +++ b/README.md @@ -218,10 +218,6 @@ The [JAX image](https://github.com/NVIDIA/JAX-Toolbox/pkgs/container/jax) is emb | --------- | ----- | ----------- | | `--xla_gpu_enable_latency_hiding_scheduler` | `true` | allows XLA to move communication collectives to increase overlap with compute kernels | -| Environment Variable | Value | Explanation | -| -------------------- | ----- | ----------- | -| `NCCL_NVLS_ENABLE` | `0` | Disables NVLink SHARP ([1](https://docs.nvidia.com/deeplearning/nccl/user-guide/docs/env.html#nccl-nvls-enable)). Future releases will re-enable this feature. | - There are various other XLA flags users can set to improve performance. For a detailed explanation of these flags, please refer to the [GPU performance](./rosetta/docs/GPU_performance.md) doc. XLA flags can also be tuned per workload. For example, each script includes a directory [xla_flags](./rosetta/rosetta/projects/maxtext/xla_flags). For a list of previously used XLA flags that are no longer needed, please also refer to the [GPU performance](./rosetta/docs/GPU_performance.md#previously-used-xla-flags) page.