Revert the changes on NCCL symmetric memory (#10210)
Co-authored-by: Yineng Zhang <me@zhyncs.com>
This commit is contained in:
co-authored by
Yineng Zhang
parent
d3ee70985f
commit
d352c29aa0
@@ -13,14 +13,10 @@ from sglang.srt.distributed import (
|
||||
divide,
|
||||
get_tensor_model_parallel_rank,
|
||||
get_tensor_model_parallel_world_size,
|
||||
parallel_state,
|
||||
split_tensor_along_last_dim,
|
||||
tensor_model_parallel_all_gather,
|
||||
tensor_model_parallel_all_reduce,
|
||||
)
|
||||
from sglang.srt.distributed.device_communicators.pynccl_allocator import (
|
||||
use_symmetric_memory,
|
||||
)
|
||||
from sglang.srt.layers.parameter import (
|
||||
BasevLLMParameter,
|
||||
BlockQuantScaleParameter,
|
||||
@@ -1315,9 +1311,7 @@ class RowParallelLinear(LinearBase):
|
||||
# Only fuse bias add into GEMM for rank 0 (this ensures that
|
||||
# bias will not get added more than once in TP>1 case)
|
||||
bias_ = None if (self.tp_rank > 0 or self.skip_bias_add) else self.bias
|
||||
with use_symmetric_memory(parallel_state.get_tp_group()) as sm:
|
||||
output_parallel = self.quant_method.apply(self, input_parallel, bias=bias_)
|
||||
sm.tag(output_parallel)
|
||||
output_parallel = self.quant_method.apply(self, input_parallel, bias=bias_)
|
||||
|
||||
if self.reduce_results and self.tp_size > 1 and not skip_all_reduce:
|
||||
output = tensor_model_parallel_all_reduce(output_parallel)
|
||||
|
||||
Reference in New Issue
Block a user