From 9436ff2f6e8009cd507a4b55a5a00e922a8cdcbb Mon Sep 17 00:00:00 2001 From: Bubbliiiing <47347516+bubbliiiing@users.noreply.github.com> Date: Wed, 9 Jul 2025 13:49:28 +0800 Subject: [PATCH] update discrete_sampler (#249) --- videox_fun/utils/discrete_sampler.py | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/videox_fun/utils/discrete_sampler.py b/videox_fun/utils/discrete_sampler.py index 84d23d2..f9948d6 100644 --- a/videox_fun/utils/discrete_sampler.py +++ b/videox_fun/utils/discrete_sampler.py @@ -17,12 +17,15 @@ class DiscreteSampling: if world_size % i != 0 or num_idx % (world_size // i) != 0: i += 1 else: - self.group_num = world_size // i + if i >= sp_size: + self.group_num = world_size // i + elif sp_size > world_size: + self.group_num = 1 + else: + self.group_num = world_size // sp_size break assert self.group_num > 0 assert world_size % self.group_num == 0 - if self.group_num >= sp_size: - self.group_num = self.group_num // sp_size # the number of rank in one group self.group_width = world_size // self.group_num self.sigma_interval = self.num_idx // self.group_num