Register tensors with symmetric memory for qwen (#18643)
This commit is contained in:
@@ -324,7 +324,11 @@ class Qwen2MoeSparseMoeBlock(nn.Module):
|
||||
final_hidden_states = self._forward_router_experts(hidden_states)
|
||||
|
||||
if shared_output is not None:
|
||||
final_hidden_states = final_hidden_states + shared_output
|
||||
# In-place add is required to keep final_hidden_states in the
|
||||
# symmetric memory pool (when --enable-symm-mem is used).
|
||||
# An out-of-place add would allocate a new tensor outside symm
|
||||
# memory, breaking subsequent symmetric collective operations.
|
||||
final_hidden_states += shared_output
|
||||
if self.tp_size > 1 and not use_reduce_scatter:
|
||||
final_hidden_states = tensor_model_parallel_all_reduce(final_hidden_states)
|
||||
|
||||
|
||||
Reference in New Issue
Block a user