From b76ed95bcc4ad0d80f18acfaab4280faff2f24a6 Mon Sep 17 00:00:00 2001 From: "Alex Y. Chan" Date: Fri, 21 Aug 2026 18:35:43 +0100 Subject: [PATCH 1/2] [JAX] Fix counting of synced and wrapped collectives in HLO Signed-off-by: Alex Y. Chan --- tests/jax/distributed_test_base.py | 62 ++++++++++++++++++++++++++++-- 1 file changed, 58 insertions(+), 4 deletions(-) diff --git a/tests/jax/distributed_test_base.py b/tests/jax/distributed_test_base.py index 1b59dfb536..673e1712b3 100644 --- a/tests/jax/distributed_test_base.py +++ b/tests/jax/distributed_test_base.py @@ -126,20 +126,74 @@ def get_bytes_per_txt(t): return bytes_count + def get_called_collective_type(line): + """Identify a collective hidden inside an async/fusion wrapper.""" + match = re.search(r"\bcalls=(%[-.\w]+)", line) + if not match: + return None + + computation_name = match.group(1) + computation = re.search( + rf"(?ms)^[ \t]*{re.escape(computation_name)}(?=[ \t(])" + rf".*?^[ \t]*}}[ \t]*$", + target_hlo, + ) + if not computation: + return None + + computation_text = computation.group(0) + has_all_reduce = COLL_AR_KEY in computation_text + has_all_gather = COLL_AG_KEY in computation_text + + if has_all_reduce and not has_all_gather: + return COLL_AR_KEY + if has_all_gather and not has_all_reduce: + return COLL_AG_KEY + + return None + def count_collectives(splitted_hlo): result = generate_collectives_count(0, 0, 0) + for line in splitted_hlo: txt = line.split() + + # strip optional HLO syntax prefix + if txt and txt[0] == "ROOT": + txt = txt[1:] + # Asynchronous collectives are represented by *-start and *-done # instructions, so count only *-start. Synchronous collectives are # represented by a single instruction without either suffix. - is_async_start = len(txt) > 0 and start_symb in txt[0] - is_sync_collective = "collective_backend_config" in line and sync_symb in line + is_async_start = txt and start_symb in txt[0] + is_sync_collective = ( + "collective_backend_config" in line + and sync_symb in line + ) + if is_async_start or is_sync_collective: - if COLL_AR_KEY in txt[0]: + if is_sync_collective: + is_all_reduce = re.search( + r"\ball-reduce\s*\(", line + ) + is_all_gather = re.search( + r"\ball-gather\s*\(", line + ) + else: + called_collective = get_called_collective_type(line) + is_all_reduce = ( + COLL_AR_KEY in txt[0] + or called_collective == COLL_AR_KEY + ) + is_all_gather = ( + COLL_AG_KEY in txt[0] + or called_collective == COLL_AG_KEY + ) + + if is_all_reduce: result[COLL_AR_KEY] += count_bytes(txt) - elif COLL_AG_KEY in txt[0]: + elif is_all_gather: result[COLL_AG_KEY] += count_bytes(txt) else: result[COLL_OTHER_KEY] += count_bytes(txt) From 2f019b1672b4a4e7d93f1e0be1d7a07cc562b16e Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Fri, 21 Aug 2026 17:38:52 +0000 Subject: [PATCH 2/2] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- tests/jax/distributed_test_base.py | 27 ++++++--------------------- 1 file changed, 6 insertions(+), 21 deletions(-) diff --git a/tests/jax/distributed_test_base.py b/tests/jax/distributed_test_base.py index 673e1712b3..5e05947399 100644 --- a/tests/jax/distributed_test_base.py +++ b/tests/jax/distributed_test_base.py @@ -134,8 +134,7 @@ def get_called_collective_type(line): computation_name = match.group(1) computation = re.search( - rf"(?ms)^[ \t]*{re.escape(computation_name)}(?=[ \t(])" - rf".*?^[ \t]*}}[ \t]*$", + rf"(?ms)^[ \t]*{re.escape(computation_name)}(?=[ \t(])" rf".*?^[ \t]*}}[ \t]*$", target_hlo, ) if not computation: @@ -155,7 +154,6 @@ def get_called_collective_type(line): def count_collectives(splitted_hlo): result = generate_collectives_count(0, 0, 0) - for line in splitted_hlo: txt = line.split() @@ -167,29 +165,16 @@ def count_collectives(splitted_hlo): # instructions, so count only *-start. Synchronous collectives are # represented by a single instruction without either suffix. is_async_start = txt and start_symb in txt[0] - is_sync_collective = ( - "collective_backend_config" in line - and sync_symb in line - ) + is_sync_collective = "collective_backend_config" in line and sync_symb in line if is_async_start or is_sync_collective: if is_sync_collective: - is_all_reduce = re.search( - r"\ball-reduce\s*\(", line - ) - is_all_gather = re.search( - r"\ball-gather\s*\(", line - ) + is_all_reduce = re.search(r"\ball-reduce\s*\(", line) + is_all_gather = re.search(r"\ball-gather\s*\(", line) else: called_collective = get_called_collective_type(line) - is_all_reduce = ( - COLL_AR_KEY in txt[0] - or called_collective == COLL_AR_KEY - ) - is_all_gather = ( - COLL_AG_KEY in txt[0] - or called_collective == COLL_AG_KEY - ) + is_all_reduce = COLL_AR_KEY in txt[0] or called_collective == COLL_AR_KEY + is_all_gather = COLL_AG_KEY in txt[0] or called_collective == COLL_AG_KEY if is_all_reduce: result[COLL_AR_KEY] += count_bytes(txt)