tp world_size fix

Signed-off-by: sahil suneja <[email protected]>

fms_extras/modules/attention.py

@@ -348,8 +348,8 @@ def forward(
         # if use_cache=True, we return the hidden_state as well as the kv cache.
         # We only reduce the output, and keep the cache thread-local
         if use_cache:
-            out = reduce_from_tensor_model_parallel_region(out_par[0])
+            out = reduce_from_tensor_model_parallel_region(out_par[0], self.world_size)
             return out, out_par[1]
         else:
-            out = reduce_from_tensor_model_parallel_region(out_par)
+            out = reduce_from_tensor_model_parallel_region(out_par, self.world_size)
             return out