Support TP in attention for two batch overlap (#6634)

This commit is contained in:
fzyzcjy
2025-05-27 11:28:12 +08:00
committed by GitHub
parent ebd1ed49d4
commit 32cd707002
4 changed files with 104 additions and 8 deletions

View File

@@ -1613,6 +1613,9 @@ class DeepseekV2Model(nn.Module):
forward_batch=forward_batch,
hidden_states=hidden_states,
residual=residual,
input_data_scatter_mode=self.layers[
normal_num_layers - 1
].layer_scatter_modes.layer_output_mode,
zero_allocator=zero_allocator,
)