[CI][RFC] Replace black-jupyter with ruff-format (#37210)
Co-authored-by: Alison Shao <a.shao@wustl.edu>
This commit is contained in:
co-authored by
Alison Shao
parent
2641e427be
commit
28262c20df
@@ -149,7 +149,7 @@ def test_main(
|
||||
for with_topk in (False, True):
|
||||
if local_rank == 0:
|
||||
print(
|
||||
f'[testing] Running with {"FP8" if isinstance(current_x, tuple) else "BF16"}, {"with" if with_topk else "without"} top-k (async={async_mode}, previous={previous_mode}) ...',
|
||||
f"[testing] Running with {'FP8' if isinstance(current_x, tuple) else 'BF16'}, {'with' if with_topk else 'without'} top-k (async={async_mode}, previous={previous_mode}) ...",
|
||||
flush=True,
|
||||
end="",
|
||||
)
|
||||
@@ -192,9 +192,9 @@ def test_main(
|
||||
|
||||
# Checks
|
||||
recv_gbl_rank_prefix_sum = handle[-4]
|
||||
assert gbl_num_tokens_per_rank[rank].item() == recv_x.size(
|
||||
0
|
||||
), f"{gbl_num_tokens_per_rank[rank].item()} != {recv_x.size(0)}"
|
||||
assert gbl_num_tokens_per_rank[rank].item() == recv_x.size(0), (
|
||||
f"{gbl_num_tokens_per_rank[rank].item()} != {recv_x.size(0)}"
|
||||
)
|
||||
assert (
|
||||
gbl_num_tokens_per_expert.view(num_ranks, -1)[rank].tolist()
|
||||
== recv_num_tokens_per_expert_list
|
||||
@@ -316,10 +316,13 @@ def test_main(
|
||||
tune_args = {"x": current_x, "handle": handle, "config": config}
|
||||
t = bench(lambda: buffer.dispatch(**tune_args))[0]
|
||||
if t < best_time:
|
||||
best_time, best_results = t, (
|
||||
num_sms,
|
||||
nvl_chunk_size,
|
||||
rdma_chunk_size,
|
||||
best_time, best_results = (
|
||||
t,
|
||||
(
|
||||
num_sms,
|
||||
nvl_chunk_size,
|
||||
rdma_chunk_size,
|
||||
),
|
||||
)
|
||||
if local_rank == 0:
|
||||
print(
|
||||
@@ -328,7 +331,7 @@ def test_main(
|
||||
)
|
||||
if local_rank == 0:
|
||||
print(
|
||||
f'[tuning] Best dispatch ({"FP8" if isinstance(current_x, tuple) else "BF16"}): SMs {best_results[0]}, NVL chunk {best_results[1]}, RDMA chunk {best_results[2]}: {rdma_send_bytes / 1e9 / best_time:.2f} GB/s (RDMA), {nvl_recv_bytes / 1e9 / best_time:.2f} GB/s (NVL)',
|
||||
f"[tuning] Best dispatch ({'FP8' if isinstance(current_x, tuple) else 'BF16'}): SMs {best_results[0]}, NVL chunk {best_results[1]}, RDMA chunk {best_results[2]}: {rdma_send_bytes / 1e9 / best_time:.2f} GB/s (RDMA), {nvl_recv_bytes / 1e9 / best_time:.2f} GB/s (NVL)",
|
||||
flush=True,
|
||||
)
|
||||
print("", flush=True)
|
||||
@@ -385,10 +388,13 @@ def test_main(
|
||||
flush=True,
|
||||
)
|
||||
if t < best_time:
|
||||
best_time, best_results = t, (
|
||||
num_sms,
|
||||
nvl_chunk_size,
|
||||
rdma_chunk_size,
|
||||
best_time, best_results = (
|
||||
t,
|
||||
(
|
||||
num_sms,
|
||||
nvl_chunk_size,
|
||||
rdma_chunk_size,
|
||||
),
|
||||
)
|
||||
|
||||
if local_rank == 0:
|
||||
|
||||
@@ -120,7 +120,7 @@ def test_main(
|
||||
for with_topk in (False, True):
|
||||
if local_rank == 0:
|
||||
print(
|
||||
f'[testing] Running with {"FP8" if isinstance(current_x, tuple) else "BF16"}, {"with" if with_topk else "without"} top-k (async={async_mode}, previous={previous_mode}) ...',
|
||||
f"[testing] Running with {'FP8' if isinstance(current_x, tuple) else 'BF16'}, {'with' if with_topk else 'without'} top-k (async={async_mode}, previous={previous_mode}) ...",
|
||||
flush=True,
|
||||
end="",
|
||||
)
|
||||
@@ -162,9 +162,9 @@ def test_main(
|
||||
|
||||
# Checks
|
||||
rank_prefix_matrix = handle[0]
|
||||
assert gbl_num_tokens_per_rank[rank].item() == recv_x.size(
|
||||
0
|
||||
), f"{gbl_num_tokens_per_rank[rank].item()} != {recv_x.size(0)}"
|
||||
assert gbl_num_tokens_per_rank[rank].item() == recv_x.size(0), (
|
||||
f"{gbl_num_tokens_per_rank[rank].item()} != {recv_x.size(0)}"
|
||||
)
|
||||
assert (
|
||||
gbl_num_tokens_per_expert.view(num_ranks, -1)[rank].tolist()
|
||||
== recv_num_tokens_per_expert_list
|
||||
@@ -280,7 +280,7 @@ def test_main(
|
||||
)
|
||||
if local_rank == 0:
|
||||
print(
|
||||
f'[tuning] Best dispatch ({"FP8" if isinstance(current_x, tuple) else "BF16"}): SMs {best_results[0]}, NVL chunk {best_results[1]}, {nvl_recv_bytes / 1e9 / best_time:.2f} GB/s (NVL)',
|
||||
f"[tuning] Best dispatch ({'FP8' if isinstance(current_x, tuple) else 'BF16'}): SMs {best_results[0]}, NVL chunk {best_results[1]}, {nvl_recv_bytes / 1e9 / best_time:.2f} GB/s (NVL)",
|
||||
flush=True,
|
||||
)
|
||||
print("", flush=True)
|
||||
|
||||
@@ -36,9 +36,9 @@ def test_main(
|
||||
|
||||
# NOTES: the integers greater than 256 exceeds the BF16 precision limit
|
||||
rank_offset = 128
|
||||
assert (
|
||||
num_ranks - rank_offset < 257
|
||||
), "Too many ranks (exceeding test precision limit)"
|
||||
assert num_ranks - rank_offset < 257, (
|
||||
"Too many ranks (exceeding test precision limit)"
|
||||
)
|
||||
|
||||
x = torch.ones((num_tokens, hidden), dtype=torch.bfloat16, device="cuda") * (
|
||||
rank - rank_offset
|
||||
@@ -55,9 +55,9 @@ def test_main(
|
||||
|
||||
# Randomly mask some positions
|
||||
for i in range(10):
|
||||
topk_idx[random.randint(0, num_tokens - 1), random.randint(0, num_topk - 1)] = (
|
||||
-1
|
||||
)
|
||||
topk_idx[
|
||||
random.randint(0, num_tokens - 1), random.randint(0, num_topk - 1)
|
||||
] = -1
|
||||
|
||||
# Check dispatch correctness
|
||||
do_check = True
|
||||
@@ -114,9 +114,9 @@ def test_main(
|
||||
assert (
|
||||
num_valid_tokens == (recv_layout_range & int_mask).sum().item()
|
||||
), f"{num_valid_tokens} != {recv_layout_range & int_mask}.sum().item()"
|
||||
assert (
|
||||
num_valid_tokens == (all_topk_idx == expert_id).sum().item()
|
||||
), f"{num_valid_tokens} != {(all_topk_idx == expert_id).sum().item()}"
|
||||
assert num_valid_tokens == (all_topk_idx == expert_id).sum().item(), (
|
||||
f"{num_valid_tokens} != {(all_topk_idx == expert_id).sum().item()}"
|
||||
)
|
||||
|
||||
# Check received data
|
||||
recv_x = recv_x[:num_valid_tokens]
|
||||
@@ -127,9 +127,10 @@ def test_main(
|
||||
recv_x[:, -128:] - recv_src_info.view(-1, 1) % num_tokens
|
||||
).sum().item() == 0
|
||||
for j in range(num_ranks):
|
||||
begin_idx, count = (recv_layout_range[j] >> 32).item(), (
|
||||
recv_layout_range[j] & int_mask
|
||||
).item()
|
||||
begin_idx, count = (
|
||||
(recv_layout_range[j] >> 32).item(),
|
||||
(recv_layout_range[j] & int_mask).item(),
|
||||
)
|
||||
assert (recv_x_amin == j - rank_offset).sum().item() == (
|
||||
all_topk_idx[j] == expert_id
|
||||
).sum().item()
|
||||
@@ -145,9 +146,9 @@ def test_main(
|
||||
# Check combine correctness
|
||||
for zero_copy in (False, True):
|
||||
if zero_copy:
|
||||
buffer.get_next_low_latency_combine_buffer(handle)[
|
||||
:, :, :
|
||||
] = simulated_gemm_x
|
||||
buffer.get_next_low_latency_combine_buffer(handle)[:, :, :] = (
|
||||
simulated_gemm_x
|
||||
)
|
||||
out = torch.empty(
|
||||
(num_tokens, hidden), dtype=torch.bfloat16, device="cuda"
|
||||
)
|
||||
@@ -203,9 +204,9 @@ def test_main(
|
||||
)
|
||||
large_gemm_with_hook(hook) if return_recv_hook else None
|
||||
if zero_copy:
|
||||
buffer.get_next_low_latency_combine_buffer(handle)[
|
||||
:, :, :
|
||||
] = simulated_gemm_x
|
||||
buffer.get_next_low_latency_combine_buffer(handle)[:, :, :] = (
|
||||
simulated_gemm_x
|
||||
)
|
||||
combined_x, event, hook = buffer.low_latency_combine(
|
||||
simulated_gemm_x,
|
||||
topk_idx,
|
||||
|
||||
@@ -238,8 +238,7 @@ class _ElasticScaleUpEndToEndBase(CustomTestCase):
|
||||
join_end = rank_offset + join_tp
|
||||
if join_end > len(visible_devices):
|
||||
raise RuntimeError(
|
||||
f"Scale-up requires {join_end} visible GPUs, got "
|
||||
f"{len(visible_devices)}"
|
||||
f"Scale-up requires {join_end} visible GPUs, got {len(visible_devices)}"
|
||||
)
|
||||
env["CUDA_VISIBLE_DEVICES"] = ",".join(visible_devices[rank_offset:join_end])
|
||||
base_joining_log = os.environ.get(
|
||||
|
||||
@@ -16,7 +16,6 @@ from sglang.test.test_utils import CustomTestCase
|
||||
|
||||
|
||||
class TestFlashinferDispatcher(CustomTestCase):
|
||||
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
server_args = ServerArgs(model_path="dummy")
|
||||
@@ -125,8 +124,7 @@ class TestFlashinferDispatcher(CustomTestCase):
|
||||
self.assertTrue(
|
||||
torch.all(
|
||||
received_hidden_states[
|
||||
expected_source_rank
|
||||
* num_tokens : (expected_source_rank + 1)
|
||||
expected_source_rank * num_tokens : (expected_source_rank + 1)
|
||||
* num_tokens
|
||||
]
|
||||
== 100.0 + expected_source_rank
|
||||
@@ -229,8 +227,7 @@ class TestFlashinferDispatcher(CustomTestCase):
|
||||
self.assertTrue(
|
||||
torch.all(
|
||||
received_hidden_states[
|
||||
expected_source_rank
|
||||
* num_tokens : (expected_source_rank + 1)
|
||||
expected_source_rank * num_tokens : (expected_source_rank + 1)
|
||||
* num_tokens
|
||||
]
|
||||
== 100.0 + expected_source_rank
|
||||
|
||||
Reference in New Issue
Block a user