main.py

import torch
from torch import nn
from torchvision import datasets, transforms
from torch.utils.data import DataLoader, Subset
from torch.optim import Optimizer, SGD
from mpi4py import MPI
import time
import numpy as np

torch.manual_seed(42)  # For reproducibility

CONFIG = {
    "epochs": 20,
    "batch_size": 64,
    "learning_rate": 0.01,
    "momentum": 0.9,
    "data_dir": "./data",  # change on Mahti
}

# device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
device = "cpu"


class FeedForwardNN(nn.Module):
    """
    A simple feedforward neural network for FashionMNIST classification.
    The network consists of three fully connected layers with ReLU activations.
    The input layer expects 784 features (28x28 images flattened),
    """
    def __init__(self):
        super(FeedForwardNN, self).__init__()
        self.flatten = nn.Flatten()
        self.layers = nn.Sequential(
            nn.Linear(784, 128),
            nn.ReLU(),
            nn.Linear(128, 64),
            nn.ReLU(),
            nn.Linear(64, 10),  # 10 classes for FashionMNIST
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = self.flatten(x)
        x = self.layers(x)
        return x


def load_data(rank: int, size: int) -> tuple[DataLoader, DataLoader]:
    """
    Load FashionMNIST dataset and distribute it across ranks.
    Args:
        rank (int): The rank of the current process.
        size (int): Total number of processes.
    Returns:
        tuple: DataLoader for training and testing datasets.
    """

    try:
        transform = transforms.Compose(
            [transforms.ToTensor(), transforms.Normalize((0.286,), (0.353,))]
        )

        train_dataset = datasets.FashionMNIST(
            root=CONFIG["data_dir"],
            train=True,
            download=True,
            transform=transform,
        )

        test_dataset = datasets.FashionMNIST(
            root=CONFIG["data_dir"],
            train=False,
            download=True,
            transform=transform,
        )
    except Exception as e:
        raise RuntimeError(f"Error loading dataset: {e}.")

    # Split dataset into subsets for each rank
    total_samples = len(train_dataset)
    sample_per_rank = total_samples // size
    remainder = total_samples % size
    indices = list(range(total_samples))

    # Distribute samples among ranks
    start_index = rank * sample_per_rank + min(rank, remainder)
    extra = 1 if rank < remainder else 0
    local_samples = sample_per_rank + extra
    end_index = start_index + local_samples
    local_indices = indices[start_index:end_index]
    print(f"Rank {rank} has {local_samples} samples. indices {start_index}:{end_index}")

    # Create DataLoader for local dataset
    local_dataset = Subset(train_dataset, local_indices)
    train_loader = DataLoader(
        local_dataset, batch_size=CONFIG["batch_size"], shuffle=True
    )
    test_loader = DataLoader(
        test_dataset, batch_size=CONFIG["batch_size"], shuffle=False
    )
    return train_loader, test_loader


def synchronize_gradients(model: nn.Module, size: int, comm: MPI.Comm) -> None:
    """Synchronize gradients across all processes."""

    for param in model.parameters():
        if param.grad is None:
            continue
        # Gather gradients from all ranks
        grad = param.grad.data.cpu().numpy()
        avg_grad = np.zeros_like(grad)
        try:
            comm.Allreduce(grad, avg_grad, op=MPI.SUM)
            avg_grad /= size  # Average the gradients
            param.grad.data = torch.tensor(avg_grad, device=device)

        except Exception as e:
            raise RuntimeError(f"Error synchronizing gradients: {e}.")


def train(
    model: nn.Module,
    train_loader: DataLoader,
    test_loadr: DataLoader,
    optimizer: Optimizer,
    criterion: nn.Module,
    rank: int,
    size: int,
    comm: MPI.Comm,
) -> None:
    """Train the model with distributed data parallelism."""
       
    model.train()
    start_time = MPI.Wtime()
    for epoch in range(CONFIG["epochs"]):
        epoch_start_time = time.time()
        total_loss = 0.0

        for batch_indx, (data, target) in enumerate(train_loader):
            data, target = data.to(device), target.to(device)

            optimizer.zero_grad()

            output = model(data)  # Forward pass

            loss = criterion(output, target)  # Compute loss

            loss.backward()  # Backward pass

            synchronize_gradients(model, size, comm)  # Synchronize gradients

            optimizer.step()  # Update weights

            total_loss += loss.item()

        avg_loss = total_loss / len(train_loader)  # Average local loss for the epoch

        epoch_duration = time.time() - epoch_start_time

        all_avg_loss = (
            comm.allreduce(avg_loss, op=MPI.SUM) / size
        )  # Average loss across all ranks

        local_paremeter_norm = torch.norm(
            torch.cat([p.view(-1) for p in model.parameters()])
        )
        
        test_accuracy = evaluate(model=model, test_loader=test_loadr, rank=rank, size=size, comm=comm)
        
        # Synchronize all processes
        comm.Barrier()
        info = f"{rank:^5} | {epoch:^7} | {test_accuracy:^12.4f}% |{avg_loss:^10.4f} | {all_avg_loss:^11.4f} | {epoch_duration:^10.2f}"

        all_info = comm.gather(info, root=0)

        if rank == 0 and epoch % 10 == 0:
            print(
                f"{'Rank':^5} | {'Epoch':^7} | {'Test Accuracy':^6} | {'Local Loss':^10} | {'Total Loss':^11} | {'Time (s)':^10}"
            )
            print("-" * 70)
            for row in all_info:
                print(row)
            print("")

    end_time = MPI.Wtime()
    local_training_time = end_time - start_time
    max_training_time = comm.allreduce(local_training_time, op=MPI.MAX)

    if rank == 0:
        print(
            f"Maximum training time across all ranks: {max_training_time:.2f} seconds"
        )
        


        with open("mpi_perf.csv", "a") as f:
            f.write(f"{size},{max_training_time:.4f}\n")
        
        print(f"Logged performance: {size} procs → {max_training_time:.2f} s")

        
        torch.save(model.state_dict(), "fashion_mnist_model.pth")
        print("Model saved as 'fashion_mnist_model.pth'.")


def evaluate(
    model: nn.Module, test_loader: DataLoader, rank: int, size: int, comm: MPI.Comm
) -> None:
    """Evaluate the model on the test dataset."""
    
    model.eval()
    correct = 0
    total = 0

    with torch.no_grad():
        for data, target in test_loader:
            data, target = data.to(device), target.to(device)

            output = model(data)
            pred = output.argmax(
                dim=1, keepdim=True
            )  # Get the index of the max log-probability
            correct += pred.eq(target.view_as(pred)).sum().item()
            total += target.size(0)

    accuracy = 100 * correct / total
    # print(f"Rank {rank} - Test Accuracy: {accuracy:.2f}%")
    return accuracy


def main():
    # Initialize MPI
    comm = MPI.COMM_WORLD
    rank = comm.Get_rank()
    size = comm.Get_size()

    train_loader, test_loadr = load_data(rank=rank, size=size)

    # Initialize model, loss function, and optimizer
    model = FeedForwardNN()
    criterion = nn.CrossEntropyLoss()
    optimizer = SGD(
        model.parameters(), lr=CONFIG["learning_rate"], momentum=CONFIG["momentum"]
    )

    model.to(device)
    print(f"Rank {rank} using device: {device}")

    # Train the model
    train(
        model=model,
        train_loader=train_loader,
        test_loadr=test_loadr,
        optimizer=optimizer,
        criterion=criterion,
        rank=rank,
        size=size,
        comm=comm,
    )

if __name__ == "__main__":
    main()