aybchan · Nov 4, 2026 · Nov 10, 2026
diff --git a/.github/container/pytest-xdist.sh b/.github/container/pytest-xdist.sh
@@ -34,7 +34,7 @@ for worker in $(seq 0 $((WORKERS_PER_GPU-1))); do
   done
 done
 export PYTEST_ADDOPTS
-export XLA_PYTHON_CLIENT_MEM_FRACTION=$(printf '0.%02d' $((75/WORKERS_PER_GPU)))
+unset XLA_PYTHON_CLIENT_MEM_FRACTION
 mkfifo ${FIFO}
 # Handle errors explicitly below here
 set +e

diff --git a/.github/container/test-maxtext.sh b/.github/container/test-maxtext.sh
@@ -15,7 +15,6 @@ usage() {
     echo ""
     echo "  OPTIONS                    DESCRIPTION"
     echo "  -a, --additional-args      Additional args to pass to maxtext.trainers.pre_train.train. Can be passed many times."
-    echo "  --mem-fraction             Specify the percentage of memory to preallocate for XLA. Example: 0.90, 0.85, 0.65". Default to 0.90, contradicting JAX default of 0.75.
     echo "  --model-name               Specify the model names to run [Preferred]. If you specify model name then you do not need to specify decoder-block. Currently supported ootb models:
                                        gemma-2b, gemma-7b, gpt3-175b, gpt3-22b, gpt3-52k, gpt3-6b, llama2-13b, llama2-70b, llama2-7b, llama3-70b, llama3-8b, mistral-7b, mixtral-8x7b"
     echo "  --decoder-block            Specify decoder block to run. Example: llama2, default. Use this option only to define a custom model. This is not preferred, only used in CI"
@@ -50,15 +49,14 @@ usage() {
     exit $1
 }
 
-args=$(getopt -o a:b:s:o:n:h --long additional-args:,mem-fraction:,model-name:,decoder-block:,attn-type:,remat-policy:,batch-per-gpu:,dtype:,quantization:,steps:,help,multiprocess,output:,data-parallel:,fsdp:,tensor-parallel:,tensor-sequence-parallel:,pipeline-parallel:,nodes: -- "$@")
+args=$(getopt -o a:b:s:o:n:h --long additional-args:,model-name:,decoder-block:,attn-type:,remat-policy:,batch-per-gpu:,dtype:,quantization:,steps:,help,multiprocess,output:,data-parallel:,fsdp:,tensor-parallel:,tensor-sequence-parallel:,pipeline-parallel:,nodes: -- "$@")
 if [[ $? -ne 0 ]]; then
     exit $1
 fi
 
 # Default arguments
 HARDWARE='gpu'
 OUTPUT=$(mktemp -d)
-MEM_FRACTION=0.90
 
 MODEL="gpt3-52k"
 DECODER_BLOCK=""
@@ -84,10 +82,6 @@ while [ : ]; do
         ADDITIONAL_ARGS="$ADDITIONAL_ARGS $2"
         shift 2
         ;;
-    --mem-fraction)
-        MEM_FRACTION="$2"
-        shift 2
-        ;;
     --model-name)
         MODEL="$2"
         shift 2
@@ -232,7 +226,6 @@ else
 fi
 
 print_var ADDITIONAL_ARGS
-print_var MEM_FRACTION
 print_var MODEL
 print_var DECODER_BLOCK
 print_var ATTN_TYPE
@@ -260,7 +253,7 @@ pushd ${MAXTEXT_DIR}
 ## Launch
 
 export NVTE_FUSED_ATTN=${ENABLE_FUSED_ATTN}
-export XLA_PYTHON_CLIENT_MEM_FRACTION=${MEM_FRACTION}
+unset XLA_PYTHON_CLIENT_MEM_FRACTION
 
 local_arch=$(local_cuda_arch)
 if [[ "${local_arch}" == "9.0" ]]; then

diff --git a/docs/GPU_performance.md b/docs/GPU_performance.md
@@ -29,7 +29,7 @@ Please note that some of these flags are experimental. All combinations of flags
 - Use CUDA Graphs (command buffer) → see [CUDA graphs](#cuda-graphs)
 
 **OOMs or fragmentation**
-- Set XLA_PYTHON_CLIENT_MEM_FRACTION → see [Memory management](#flags-to-manage-memory-used-in-jaxxla)
+- Set XLA_CLIENT_MEM_FRACTION → see [Memory management](#flags-to-manage-memory-used-in-jaxxla)
 - Tune xla_gpu_memory_limit_slop_factor → see [Memory management](#flags-to-manage-memory-used-in-jaxxla)
 
 **Multi-device single-process hangs**
@@ -41,7 +41,12 @@ Please note that some of these flags are experimental. All combinations of flags
 
 ## Flags to manage memory used in JAX/XLA
 
-- XLA_PYTHON_CLIENT_MEM_FRACTION is a XLA environment variable that allocates a fraction of GPU memory for JAX/XLA.
+Launch scripts unset the deprecated `XLA_PYTHON_CLIENT_MEM_FRACTION`. JAX uses its
+default memory fraction (normally 0.75) unless `XLA_CLIENT_MEM_FRACTION` is set
+externally. For multiple pytest workers sharing a GPU, set `XLA_CLIENT_MEM_FRACTION`
+to divide the available memory among workers.
+
+- XLA_CLIENT_MEM_FRACTION is a XLA environment variable that allocates a fraction of GPU memory for JAX/XLA.
 --  Ideally, should be 1, but in practice less because some memory is used by NVIDIA Libraries, and the JAX framework.
 --  We typically set it to 0.9 or 0.8. At 0.9, XLA gets 90% of GPU memory.
 
@@ -60,7 +65,7 @@ The following environment variable restricts CUDA queues to 1 and is useful when
 See [NCCL Environment Variables](https://docs.nvidia.com/deeplearning/nccl/user-guide/docs/env.html) for more details.
 - NCCL_PROTO: SIMPLE,LL,LL128
 
-The following variable accelerates all-reduce collective on NVLink4/H100. It requires additional GPU memory and may need one to reduce `XLA_PYTHON_CLIENT_MEM_FRACTION` to avoid OOMs if enabled.
+The following variable accelerates all-reduce collective on NVLink4/H100. It requires additional GPU memory and may need one to reduce `XLA_CLIENT_MEM_FRACTION` to avoid OOMs if enabled.
 - NCCL_NVLS_ENABLE:1 
 
 
@@ -171,7 +176,7 @@ The following flag removes extra copies introduced by DUS (dynamic update slice)
 Enable user-buffers in NCCL for zero-copy collectives and send/recv. Needs NCCL_NVLS_ENABLE=1 for AG, AR, RS.
 - --xla_gpu_enable_nccl_user_buffers=true
 
-When user-buffers is enabled, a separate memory pool is created for user-buffer registered memory. Environment variable `XLA_PYTHON_CLIENT_COLLECTIVE_MEM_SIZE_MB` can be used to configure this memory pool. It may also be necessary to reduce `XLA_PYTHON_CLIENT_MEM_FRACTION` to ensure there is enough memory for the user buffer pool.
+When user-buffers is enabled, a separate memory pool is created for user-buffer registered memory. Environment variable `XLA_PYTHON_CLIENT_COLLECTIVE_MEM_SIZE_MB` can be used to configure this memory pool. It may also be necessary to reduce `XLA_CLIENT_MEM_FRACTION` to ensure there is enough memory for the user buffer pool.
 - `XLA_PYTHON_CLIENT_COLLECTIVE_MEM_SIZE_MB=0` (default value) - The user buffer pool will start empty, but will grow during execution as more collective memory is required. This setting can result in extra fragmentation and inefficient memory use.
 - `XLA_PYTHON_CLIENT_COLLECTIVE_MEM_SIZE_MB=<amount of MiB to preallocate>` - The user buffer pool will preallocate this amount of memory at the begining. The number should be high enough to cover peak collective memory usage.
 

diff --git a/docs/frameworks/maxtext/example_slurm.sub b/docs/frameworks/maxtext/example_slurm.sub
@@ -46,7 +46,7 @@ CHECKPOINT_DIR=/opt/maxtext/workspace/llama-checkpoint
 MOUNTS="--container-mounts=$BASE_WORKSPACE_DIR:$WORKSPACE_DIR,$BASE_VOCAB_PATH:$GPT_VOCAB_PATH,$BASE_TFDS_DATA_DIR:/$TFDS_DATA_DIR,$BASE_CHECKPOINT_DIR:$CHECKPOINT_DIR"
 
 # export some necessary flags
-export XLA_PYTHON_CLIENT_MEM_FRACTION=0.90
+unset XLA_PYTHON_CLIENT_MEM_FRACTION
 export CUDA_DEVICE_MAX_CONNECTIONS=1
 export NVTE_FUSED_ATTN=1
 export NCCL_IB_SL=1

diff --git a/docs/frameworks/maxtext/nvfp4/nvfp4_example.sh b/docs/frameworks/maxtext/nvfp4/nvfp4_example.sh
@@ -12,7 +12,7 @@ RS_MULTIPLE=4
 AR_MULTIPLE=4
 AG_MULTIPLE=4
 
-export XLA_PYTHON_CLIENT_MEM_FRACTION=0.9
+unset XLA_PYTHON_CLIENT_MEM_FRACTION
 export NVTE_FUSED_ATTN=1
 
 export XLA_FLAGS="--xla_gpu_enable_latency_hiding_scheduler=true
@@ -32,7 +32,6 @@ export XLA_FLAGS="--xla_gpu_enable_latency_hiding_scheduler=true
                 --xla_disable_hlo_passes=rematerialization"
 
 echo "XLA_FLAGS = ${XLA_FLAGS}"
-echo "XLA_PYTHON_CLIENT_MEM_FRACTION = ${XLA_PYTHON_CLIENT_MEM_FRACTION}"
 
 RUN_SETTINGS="-m maxtext.trainers.pre_train.train maxtext/configs/base.yml run_name=debug_run base_output_directory=./debug_logs hardware=gpu dataset_type=synthetic  model_name=llama3-8b remat_policy='minimal_with_context_and_quantization' scan_layers=False attention='cudnn_flash_te' steps=50 dtype=bfloat16 max_target_length=8192 per_device_batch_size=4 ici_data_parallelism=${ici_DP} dcn_data_parallelism=${dcn_DP} ici_fsdp_parallelism=${ici_FSDP} dcn_fsdp_parallelism=${dcn_FSDP} profiler=nsys enable_checkpointing=false override_model_config=True gradient_accumulation_steps=1 quantization=te_nvfp4_no_rht max_segments_per_seq=32"
 

diff --git a/jax-inference-offloading/examples/decoupled_asynchronous/decoupled_async.sh b/jax-inference-offloading/examples/decoupled_asynchronous/decoupled_async.sh
@@ -269,7 +269,7 @@ export MAX_COMPLETED_PROMPTS
 export DISPATCH_DELAY
 export VLLM_ENFORCE_EAGER
 export VLLM_GPU_MEMORY_UTILIZATION
-export XLA_PYTHON_CLIENT_MEM_FRACTION=0.95
+unset XLA_PYTHON_CLIENT_MEM_FRACTION
 export XLA_FLAGS="--xla_gpu_enable_latency_hiding_scheduler=true
                   --xla_gpu_enable_command_buffer=FUSION,CUBLAS,CUDNN,CUSTOM_CALL
                   --xla_gpu_collective_permute_combine_threshold_bytes=8589934592

diff --git a/jax-inference-offloading/examples/decoupled_synchronous/decoupled_sync.sh b/jax-inference-offloading/examples/decoupled_synchronous/decoupled_sync.sh
@@ -230,7 +230,7 @@ export PARAM_MAPPING_PATH
 export NUM_ITERATIONS
 export VLLM_ENFORCE_EAGER
 export VLLM_GPU_MEMORY_UTILIZATION
-export XLA_PYTHON_CLIENT_MEM_FRACTION=0.95
+unset XLA_PYTHON_CLIENT_MEM_FRACTION
 export XLA_FLAGS="--xla_gpu_enable_latency_hiding_scheduler=true
                   --xla_gpu_enable_command_buffer=FUSION,CUBLAS,CUDNN,CUSTOM_CALL
                   --xla_gpu_collective_permute_combine_threshold_bytes=8589934592

diff --git a/jax-inference-offloading/examples/single_controller/run_single_controller.sh b/jax-inference-offloading/examples/single_controller/run_single_controller.sh
@@ -246,7 +246,7 @@ export SC_JAX_RESULTS_TOPIC="inference/results/sc_forwarded"
 export SC_SYNC_TOPIC="sync/weights_ready"
 export SC_VLLM_EVENT_TOPIC="sc/vllm/events"
 
-export XLA_PYTHON_CLIENT_MEM_FRACTION=0.95
+unset XLA_PYTHON_CLIENT_MEM_FRACTION
 export XLA_FLAGS="--xla_gpu_enable_latency_hiding_scheduler=true
                   --xla_gpu_enable_command_buffer=FUSION,CUBLAS,CUDNN,CUSTOM_CALL
                   --xla_gpu_collective_permute_combine_threshold_bytes=8589934592