[hotfix] use the original implementation in 8785 (#8994)

This commit is contained in:
Cheng Wan
2025-08-08 21:47:25 -07:00
committed by GitHub
parent fd7e15b76d
commit a47baff12c

View File

@@ -553,6 +553,10 @@ class CommunicateSummableTensorPairFn:
forward_batch.gathered_buffer[: forward_batch.input_ids.shape[0]],
hidden_states,
)
if hidden_states.data_ptr() is global_hidden_states.data_ptr():
hidden_states = torch.empty_like(hidden_states)
if allow_reduce_scatter and forward_batch.dp_padding_mode.is_max_len():
# When using padding, all_reduce is skipped after MLP and MOE and reduce scatter is used here instead.
dp_reduce_scatter_tensor(hidden_states, global_hidden_states)