MNIST works with TP?

tt-train/sources/examples/mnist_mlp/main.cpp

@@ -64,10 +64,14 @@ float evaluate(DataLoader &test_dataloader, Model &model, size_t num_targets) {
     model->eval();
     float num_correct = 0;
     float num_samples = 0;
+    auto *device = &ttml::autograd::ctx().get_device();
     for (const auto &[data, target] : test_dataloader) {
         auto output = (*model)(data);
-        auto output_vec = ttml::core::to_vector(output->get_value());
-        auto target_vec = ttml::core::to_vector(target->get_value());
+        ttml::core::MeshToXTensorVariant<float> composer = ttml::core::VectorMeshToXTensor<float>(device->shape());
+        auto output_xtensor = ttml::core::to_xtensor(output->get_value(), composer)[0];
+        auto target_xtensor = ttml::core::to_xtensor(target->get_value(), composer)[0];
+        auto output_vec = std::vector<float>(output_xtensor.begin(), output_xtensor.end());
+        auto target_vec = std::vector<float>(target_xtensor.begin(), target_xtensor.end());
         for (size_t i = 0; i < output_vec.size(); i += num_targets) {
             auto predicted_class = std::distance(
                 output_vec.begin() + i,
@@ -194,12 +198,25 @@ int main(int argc, char **argv) {
 
     LossAverageMeter loss_meter;
     int training_step = 0;
+
+    auto get_loss_value = [device](const TensorPtr &loss) {
+        ttml::core::MeshToXTensorVariant<float> composer = ttml::core::VectorMeshToXTensor<float>(device->shape());
+        auto loss_xtensors = ttml::core::to_xtensor(loss->get_value(), composer);
+        // sum of loss xtensors
+        float loss_float =
+            std::accumulate(loss_xtensors.begin(), loss_xtensors.end(), 0.0F, [](float acc, auto &xtensor) {
+                return acc + xtensor(0);
+            });
+
+        return loss_float / static_cast<float>(loss_xtensors.size());
+    };
+
     for (size_t epoch = 0; epoch < config.num_epochs; ++epoch) {
         for (const auto &[data, target] : train_dataloader) {
             optimizer.zero_grad();
             auto output = (*model)(data);
             auto loss = ttml::ops::cross_entropy_loss(output, target);
-            auto loss_float = ttml::core::to_vector(loss->get_value())[0];
+            auto loss_float = get_loss_value(loss);
             loss_meter.update(loss_float, config.batch_size);
             if (training_step % config.logging_interval == 0) {
                 fmt::print("Step: {:5d} | Average Loss: {:.4f}\n", training_step, loss_meter.average());

tt-train/sources/ttml/modules/distributed/linear.cpp

@@ -34,7 +34,7 @@ autograd::TensorPtr RowParallelLinear::operator()(autograd::TensorPtr tensor) {
 
     // do not pass bias
     tensor = ops::linear_op(tensor, m_weight, /* bias */ nullptr);
-    tensor = ops::distributed::all_reduce(tensor, tensor->rank() - 1U);
+    tensor = ops::distributed::all_reduce(tensor);
     if (m_bias != nullptr) {
         tensor = ops::add(tensor, m_bias);
     }

tt-train/sources/ttml/ops/distributed/comm_ops.cpp

@@ -31,9 +31,9 @@ autograd::TensorPtr all_gather(const autograd::TensorPtr& tensor, int dim) {
     return out;
 }
 
-autograd::TensorPtr all_reduce(const autograd::TensorPtr& tensor, int dim) {
-    auto out = autograd::create_tensor(
-        ttnn::experimental::all_reduce(tensor->get_value(), ttnn::operations::reduction::ReduceType::Sum, dim));
+autograd::TensorPtr all_reduce(const autograd::TensorPtr& tensor) {
+    auto out = autograd::create_tensor(ttnn::experimental::all_reduce(
+        tensor->get_value(), ttnn::operations::reduction::ReduceType::Sum, 1, std::nullopt, ttnn::ccl::Topology::Ring));
     autograd::GradFunction grad = [tensor, out]() { tensor->set_grad(out->get_grad()); };
     auto links = autograd::get_links(tensor);
     out->set_node(autograd::ctx().add_backward_node(std::move(grad), links));

tt-train/sources/ttml/ops/distributed/comm_ops.hpp

@@ -9,7 +9,7 @@
 namespace ttml::ops::distributed {
 
 autograd::TensorPtr scatter(const autograd::TensorPtr& tensor, int dim);
-autograd::TensorPtr all_reduce(const autograd::TensorPtr& tensor, int dim);
+autograd::TensorPtr all_reduce(const autograd::TensorPtr& tensor);
 autograd::TensorPtr all_gather(const autograd::TensorPtr& tensor, int dim);
 
 }  // namespace ttml::ops::distributed