EngineX-Hygon/sglang
5
0
Fork 0
Code Issues Pull Requests Actions 7 Projects Releases Wiki Activity

Fix expert distribution dumping causes OOM (#6967)

Browse Source
This commit is contained in:
fzyzcjy
2025-06-09 08:41:14 +08:00
committed by GitHub
parent 9ecb18568b
commit 86fe943bc3
1 changed files with 7 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

7
python/sglang/srt/managers/expert_distribution.py
View File

@@ -703,6 +703,7 @@ class _StatAccumulator(_UtilizationRateAccumulatorMixin):
dtype=torch.int32, dtype=torch.int32,
device=self._server_args.device, device=self._server_args.device,
) )
self._first_dump = True
def append( def append(
self, self,
@@ -727,9 +728,15 @@ class _StatAccumulator(_UtilizationRateAccumulatorMixin):
num_logical_experts=self._expert_location_metadata.num_logical_experts, num_logical_experts=self._expert_location_metadata.num_logical_experts,
physical_to_logical_map=self._expert_location_metadata.physical_to_logical_map, physical_to_logical_map=self._expert_location_metadata.physical_to_logical_map,
) )
if self._first_dump:
self._first_dump = False
torch.cuda.empty_cache()
torch.distributed.all_reduce( torch.distributed.all_reduce(
logical_count_of_buffered_step, op=torch.distributed.ReduceOp.SUM logical_count_of_buffered_step, op=torch.distributed.ReduceOp.SUM
) )
output = dict( output = dict(
rank=self._rank, rank=self._rank,
logical_count=logical_count_of_buffered_step, logical_count=logical_count_of_buffered_step,