From b05591a1e8df570d01f86ec45bf8a8c000c9f21c Mon Sep 17 00:00:00 2001 From: Terry Kong Date: Fri, 20 Sep 2024 15:10:40 -0700 Subject: [PATCH] introduce threshold bytes Signed-off-by: Terry Kong --- rosetta/rosetta/projects/pax/xla_flags/gpt-126m.env | 4 +++- rosetta/rosetta/projects/pax/xla_flags/gpt-5b.env | 4 +++- 2 files changed, 6 insertions(+), 2 deletions(-) diff --git a/rosetta/rosetta/projects/pax/xla_flags/gpt-126m.env b/rosetta/rosetta/projects/pax/xla_flags/gpt-126m.env index e169bb1fb..e5b97b466 100644 --- a/rosetta/rosetta/projects/pax/xla_flags/gpt-126m.env +++ b/rosetta/rosetta/projects/pax/xla_flags/gpt-126m.env @@ -1,12 +1,14 @@ set -x +THRESHOLD_BYTES=33554432 export XLA_FLAGS="\ --xla_gpu_enable_latency_hiding_scheduler=true \ --xla_allow_excess_precision \ --xla_gpu_enable_highest_priority_async_stream=true \ --xla_gpu_enable_triton_softmax_fusion=false \ - --xla_gpu_all_reduce_combine_threshold_bytes=33554432 \ + --xla_gpu_all_reduce_combine_threshold_bytes=${THRESHOLD_BYTES} \ --xla_gpu_graph_level=0 \ --xla_gpu_enable_cudnn_fmha=false \ " export XLA_PYTHON_CLIENT_MEM_FRACTION=0.8 +unset THRESHOLD_BYTES set +x diff --git a/rosetta/rosetta/projects/pax/xla_flags/gpt-5b.env b/rosetta/rosetta/projects/pax/xla_flags/gpt-5b.env index c43dfc80d..26c819143 100644 --- a/rosetta/rosetta/projects/pax/xla_flags/gpt-5b.env +++ b/rosetta/rosetta/projects/pax/xla_flags/gpt-5b.env @@ -1,11 +1,13 @@ set -x +THRESHOLD_BYTES=51200 export XLA_FLAGS="\ --xla_gpu_enable_latency_hiding_scheduler=true \ --xla_allow_excess_precision \ --xla_gpu_enable_highest_priority_async_stream=true \ --xla_gpu_enable_triton_softmax_fusion=false \ - --xla_gpu_all_reduce_combine_threshold_bytes=51200 \ + --xla_gpu_all_reduce_combine_threshold_bytes=${THRESHOLD_BYTES} \ --xla_gpu_graph_level=0 \ " export XLA_PYTHON_CLIENT_MEM_FRACTION=0.8 +unset THRESHOLD_BYTES set +x