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

2025-05-27 12:22:59 -07:00
parent 673ff668f7
commit fa6723f08f
2 changed files with 5 additions and 13 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_mode(),
+                input_data_scatter_mode=ScatterMode.model_input_output(),
                positions=positions,
                forward_batch=forward_batch,
                hidden_states=hidden_states,