[Bugfix] fix quickreduce acc error in cudagraph mode (#29508)
Signed-off-by: Haoyang Li <lihaoyang0109@gmail.com>
This commit is contained in:
haoyangli0109
@@ -28,16 +28,25 @@ __global__ __quickreduce_launch_bounds_two_shot__ static void allreduce_prototyp
|
||||
int rank,
|
||||
uint8_t** dbuffer_list,
|
||||
uint32_t data_offset,
|
||||
uint32_t flag_color,
|
||||
uint32_t* d_flag_counters,
|
||||
int64_t data_size_per_phase) {
|
||||
int block = blockIdx.x;
|
||||
int grid = gridDim.x;
|
||||
|
||||
// Read this block's counter from device memory and bump it here in the
|
||||
// kernel. Keeping the value in device memory (instead of a host scalar
|
||||
// baked into the launch) lets every CUDA-graph replay see a fresh color.
|
||||
uint32_t flag_color = d_flag_counters[blockIdx.x];
|
||||
|
||||
while (block < num_blocks) {
|
||||
AllReduceKernel::run(A, B, N, block, rank, dbuffer_list, data_offset, flag_color, data_size_per_phase);
|
||||
block += grid;
|
||||
flag_color++;
|
||||
}
|
||||
// The whole block ends up with the same value, so a single writer suffices.
|
||||
if (threadIdx.x == 0 && threadIdx.y == 0) {
|
||||
d_flag_counters[blockIdx.x] = flag_color;
|
||||
}
|
||||
}
|
||||
|
||||
#define TWOSHOT_DISPATCH(__codec) \
|
||||
@@ -57,7 +66,7 @@ __global__ __quickreduce_launch_bounds_two_shot__ static void allreduce_prototyp
|
||||
rank, \
|
||||
dbuffer_list, \
|
||||
data_offset, \
|
||||
flag_color, \
|
||||
d_flag_counters, \
|
||||
this->kMaxProblemSize); \
|
||||
} else if (world_size == 4) { \
|
||||
using LineCodec = __codec<T, 4>; \
|
||||
@@ -75,7 +84,7 @@ __global__ __quickreduce_launch_bounds_two_shot__ static void allreduce_prototyp
|
||||
rank, \
|
||||
dbuffer_list, \
|
||||
data_offset, \
|
||||
flag_color, \
|
||||
d_flag_counters, \
|
||||
this->kMaxProblemSize); \
|
||||
} else if (world_size == 8) { \
|
||||
using LineCodec = __codec<T, 8>; \
|
||||
@@ -93,7 +102,7 @@ __global__ __quickreduce_launch_bounds_two_shot__ static void allreduce_prototyp
|
||||
rank, \
|
||||
dbuffer_list, \
|
||||
data_offset, \
|
||||
flag_color, \
|
||||
d_flag_counters, \
|
||||
this->kMaxProblemSize); \
|
||||
}
|
||||
|
||||
@@ -112,7 +121,7 @@ struct DeviceComms {
|
||||
static int constexpr kMaxWorldSize = 8;
|
||||
|
||||
bool initialized = false;
|
||||
uint32_t flag_color = 1;
|
||||
uint32_t* d_flag_counters = nullptr;
|
||||
int world_size;
|
||||
int rank;
|
||||
|
||||
@@ -145,6 +154,14 @@ struct DeviceComms {
|
||||
// Clear the flags buffer.
|
||||
HIP_CHECK(hipMemset(dbuffer, 0, flags_buffer_size));
|
||||
|
||||
// A per-block color counter that the kernel advances itself. Seed it with
|
||||
// 1 rather than 0 so it never matches the freshly zeroed flags buffer.
|
||||
HIP_CHECK(hipMalloc(&d_flag_counters, kMaxNumBlocks * sizeof(uint32_t)));
|
||||
{
|
||||
std::vector<uint32_t> init_color(kMaxNumBlocks, 1u);
|
||||
HIP_CHECK(hipMemcpy(d_flag_counters, init_color.data(), kMaxNumBlocks * sizeof(uint32_t), hipMemcpyHostToDevice));
|
||||
}
|
||||
|
||||
// Device-side list of IPC buffers.
|
||||
buffer_list.resize(world_size);
|
||||
HIP_CHECK(hipMalloc(&dbuffer_list, world_size * sizeof(uint8_t*)));
|
||||
@@ -169,6 +186,12 @@ struct DeviceComms {
|
||||
}
|
||||
|
||||
void destroy() {
|
||||
// This buffer is created before `initialized` becomes true, so release it
|
||||
// on its own check to keep a half-finished init from leaking it.
|
||||
if (d_flag_counters) {
|
||||
HIP_CHECK(hipFree(d_flag_counters));
|
||||
d_flag_counters = nullptr;
|
||||
}
|
||||
if (initialized) {
|
||||
for (int i = 0; i < world_size; i++) {
|
||||
if (i != rank) {
|
||||
@@ -229,8 +252,7 @@ struct DeviceComms {
|
||||
break;
|
||||
}
|
||||
HIP_CHECK(cudaGetLastError());
|
||||
// Rotate the flag color.
|
||||
flag_color += divceil(N, grid);
|
||||
// The color now advances on-device inside the kernel; no host-side bump.
|
||||
}
|
||||
};
|
||||
|
||||
|
||||
@@ -15,6 +15,7 @@ from sglang.srt.distributed.communication_op import ( # noqa
|
||||
tensor_model_parallel_all_reduce,
|
||||
)
|
||||
from sglang.srt.distributed.device_communicators.quick_all_reduce import (
|
||||
QuickAllReduce,
|
||||
qr_rocm_arch_available,
|
||||
)
|
||||
from sglang.srt.distributed.parallel_state import (
|
||||
@@ -258,6 +259,117 @@ def qr_variable_input(rank, world_size):
|
||||
num += 1
|
||||
|
||||
|
||||
def qr_graph_replay(rank, world_size, quant_mode="FP", num_replays=10):
|
||||
"""Capture ONE CUDA graph with a single quick-reduce and replay it many
|
||||
times with changing input. Every rank contributes the same value v in a
|
||||
round, so the true all-reduce sum is world_size * v; the FP regime is
|
||||
lossless, so the comparison is bit-exact.
|
||||
|
||||
The pre-fix kernel bakes the per-block flag color into the graph launch and
|
||||
reuses it on every replay -- the waiting peer is satisfied by the previous
|
||||
round's residual flag and reads stale data, giving wrong results on some
|
||||
replays. The fixed kernel advances the color on-device each replay.
|
||||
"""
|
||||
os.environ["ROCM_QUICK_REDUCE_QUANTIZATION"] = quant_mode
|
||||
os.environ["ROCM_QUICK_REDUCE_CAST_BF16_TO_FP16"] = "0"
|
||||
device = torch.device(f"cuda:{rank}")
|
||||
torch.cuda.set_device(device)
|
||||
|
||||
# gloo (CPU) group: QuickAllReduce must attach to a non-NCCL group; it is
|
||||
# used only for the one-time IPC-handle exchange.
|
||||
dist.init_process_group(
|
||||
backend="gloo",
|
||||
init_method="tcp://127.0.0.1:29500",
|
||||
rank=rank,
|
||||
world_size=world_size,
|
||||
)
|
||||
qr = QuickAllReduce(group=dist.group.WORLD, device=device)
|
||||
assert not qr.disabled, (
|
||||
"quick-reduce unavailable on this arch/env "
|
||||
"(needs ROCm MI300 gfx94/gfx95, even GPU count, same node, "
|
||||
"and a non-NONE ROCM_QUICK_REDUCE_QUANTIZATION)."
|
||||
)
|
||||
|
||||
N = 1 << 21 # 4 MB fp16, above the QR size threshold for the direct path
|
||||
inp = torch.empty(N, dtype=torch.float16, device=device)
|
||||
out = torch.empty(N, dtype=torch.float16, device=device)
|
||||
|
||||
# Warmup, then capture a graph with EXACTLY ONE quick-reduce.
|
||||
inp.fill_(1.0)
|
||||
qr.quick_all_reduce(inp, out=out)
|
||||
torch.cuda.synchronize()
|
||||
dist.barrier()
|
||||
|
||||
graph = torch.cuda.CUDAGraph()
|
||||
with torch.cuda.graph(graph):
|
||||
qr.quick_all_reduce(inp, out=out)
|
||||
torch.cuda.synchronize()
|
||||
dist.barrier()
|
||||
|
||||
try:
|
||||
for v in range(1, num_replays + 1):
|
||||
inp.fill_(float(v)) # in-place: same value on every rank
|
||||
dist.barrier()
|
||||
graph.replay()
|
||||
torch.cuda.synchronize()
|
||||
dist.barrier()
|
||||
expected = float(v * world_size)
|
||||
assert torch.all(out.float() == expected), (
|
||||
f"[rank {rank}] round {v}: got {out.float().flatten()[0].item()}, "
|
||||
f"expected {expected} (stale-flag corruption across replays)"
|
||||
)
|
||||
finally:
|
||||
dist.destroy_process_group()
|
||||
|
||||
|
||||
class TestQuickreduceGraphReplay(CustomTestCase):
|
||||
"""Regression test for the QuickReduce CUDA-graph stale-flag bug.
|
||||
|
||||
Unlike test_graph_allreduce (which captures a fresh graph each iteration
|
||||
and replays it once), this captures a single graph and replays it many
|
||||
times -- the exact scenario the on-device flag-color fix addresses.
|
||||
"""
|
||||
|
||||
TP_SIZES = [4, 8]
|
||||
|
||||
@unittest.skipIf(
|
||||
not qr_rocm_arch_available(),
|
||||
"Only test Quick AllReduce on ROCm architectures >= gfx94*",
|
||||
)
|
||||
def test_quick_allreduce_graph_replay(self):
|
||||
for tp_size in self.TP_SIZES:
|
||||
world_size = tp_size
|
||||
if world_size > torch.cuda.device_count():
|
||||
continue
|
||||
|
||||
multiprocessing.set_start_method("spawn", force=True)
|
||||
timeout = 120
|
||||
processes = []
|
||||
for rank in range(tp_size):
|
||||
p = multiprocessing.Process(
|
||||
target=qr_graph_replay, args=(rank, tp_size)
|
||||
)
|
||||
p.start()
|
||||
processes.append((rank, p))
|
||||
for rank, p in processes:
|
||||
p.join(timeout=timeout)
|
||||
if p.is_alive():
|
||||
for r, proc in processes:
|
||||
if proc.is_alive():
|
||||
proc.terminate()
|
||||
proc.join()
|
||||
raise RuntimeError(
|
||||
f"QuickReduce graph-replay hang detected after {timeout}s!"
|
||||
)
|
||||
for rank, p in processes:
|
||||
self.assertEqual(
|
||||
p.exitcode,
|
||||
0,
|
||||
f"QuickReduce graph-replay (tp={tp_size}, rank={rank}) "
|
||||
f"produced wrong results -- stale-flag bug not fixed.",
|
||||
)
|
||||
|
||||
|
||||
class TestQuickreduceVariableInput(CustomTestCase):
|
||||
"""
|
||||
When the tensor parallelism is set to 4 or 8, frequent changes
|
||||
|
||||
Reference in New Issue
Block a user