fix communicator for non-dp lm head (#6662)

2025-05-27 02:31:12 -07:00
parent b18416fbf8
commit a3d7f4b673
2 changed files with 13 additions and 5 deletions
@@ -451,7 +451,7 @@ class Qwen2MoeModel(nn.Module):
            hidden_states, residual = model_forward_maybe_tbo(
                layers=self.layers,
                enable_tbo=True,
-                input_data_scatter_mode=ScatterMode.model_input_output(),
+                input_data_scatter_mode=ScatterMode.model_input_mode(),
                positions=positions,
                forward_batch=forward_batch,
                hidden_states=hidden_states,