Fix world-size-one aliasing in MLP batch sync (#34997)

Co-authored-by: wangwenchen0407 <wangwenchen@meta.com>
This commit is contained in:
Lianmin Zheng
2026-08-17 00:04:09 -07:00
committed by GitHub
co-authored by wangwenchen0407
parent f7a404e9c3
commit 12a455a910
@@ -137,10 +137,14 @@ class MLPSyncBatchInfo:
local_info_tensor = self._get_local_tensor(device=device)
fallback_tensor = self._get_fallback_tensor(device=device)
info_width = local_info_tensor.numel()
# Inactive max_world_size slots must decode as IDLE.
global_info_tensor = fallback_tensor.expand(
self.dp_size, self.tp_size * self.cp_size, info_width
).contiguous()
# Inactive max_world_size slots must decode as IDLE. repeat() (not
# expand().contiguous()) so the buffer never aliases fallback_tensor:
# at world size 1 the expanded view is already contiguous, contiguous()
# is a no-op, and the masked fallback writes below would then read and
# write the same storage.
global_info_tensor = fallback_tensor.repeat(
self.dp_size, self.tp_size * self.cp_size, 1
)
if use_all_reduce:
# Admission can expose different WORLD sizes; use fixed global slots.