Skip to content

Commit 47c2b28

Browse files
committed
tuner: update binary search in tuner decision script.
Today, this implementation assumes that there are no repeated regions, as every algo-proto choice has a single region. E.g. tree-ll, tree-ll128, tree-ll, tree-ll128 won't work if the binary search happens to have min and max sizes in two distinct regions of same algo-proto choice. This change makes it continue unless the max and min sizes diff by 2x, which we assume is too granular to define two regions. Signed-off-by: Feng Ji <fengnji@amazon.com>
1 parent f155893 commit 47c2b28

File tree

1 file changed

+2
-1
lines changed

1 file changed

+2
-1
lines changed

contrib/python/ofi_nccl/tuner/cli/wrapper.py

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@
22

33
import ctypes
44
import logging
5+
import math
56
import pandas as pd
67
import numpy as np
78
import os
@@ -185,7 +186,7 @@ def analyze_message_range(self, coll_type: NCCLFunc, min_size: int, max_size: in
185186
def bisect(min_size, max_size):
186187
min_decision = self.get_coll_info(coll_type, min_size)
187188
max_decision = self.get_coll_info(coll_type, max_size)
188-
if min_decision == max_decision or max_size - min_size <= 1:
189+
if (min_decision == max_decision and math.log2(max_size) - math.log2(min_size) <= 1) or max_size - min_size <= 1:
189190
return [(min_size, min_decision), (max_size, max_decision)]
190191

191192
mid_size = (min_size + max_size) // 2

0 commit comments

Comments
 (0)