Hands-on Practical: Simulating Distributed SGD

Simulating the behavior of synchronous and asynchronous Stochastic Gradient Descent (SGD) helps understand their fundamental differences and trade-offs in a controlled environment. While a distributed system involves complex networking and hardware considerations, a simulation allows isolation of algorithmic behaviors.

Our goal is not to build a production-ready distributed framework but to gain intuition about how distributing gradient computations and updates affects the optimization process. We'll use a simple problem, like linear regression, and simulate multiple "workers" contributing to the learning process.

Simulation Setup

We'll consider a standard supervised learning problem: minimizing a loss function $L(\theta)$ defined over a dataset $D = \{ (x_i, y_i) \}_{i=1}^N$. In a distributed setting, we typically partition the data $D$ across $K$ workers. Let $D_k$ be the subset of data assigned to worker $k$.

1. The Problem: We can use linear regression as a simple, convex example. The loss for a single data point $(x_i, y_i)$ with parameters $\theta$ is $l(x_i, y_i; \theta) = \frac{1}{2} (x_i^T \theta - y_i)^2$. The total loss is $L(\theta) = \frac{1}{N} \sum_{i=1}^N l(x_i, y_i; \theta)$. The gradient for a single point is $\nabla l(x_i, y_i; \theta) = (x_i^T \theta - y_i) x_i$.

2. Simulating Workers and Data: We can represent the workers and their data partitions in Python. We'll generate some synthetic data and split it.

import numpy as np

# Generate synthetic data for linear regression
N = 1000  # Total data points
d = 10    # Feature dimension
X = np.random.rand(N, d)
true_theta = np.random.rand(d, 1)
y = X @ true_theta + 0.1 * np.random.randn(N, 1) # Add some noise

# Simulation parameters
num_workers = 4
learning_rate = 0.01
num_iterations = 100

# Partition data among workers (simple split)
indices = np.arange(N)
np.random.shuffle(indices)
split_indices = np.array_split(indices, num_workers)
worker_data = [(X[split_indices[k]], y[split_indices[k]]) for k in range(num_workers)]

# Initialize parameters (shared across simulations)
theta_sync = np.zeros((d, 1))
theta_async = np.zeros((d, 1))

# Store loss history
loss_history_sync = []
loss_history_async = []

# Define loss and gradient functions
def compute_loss(X_data, y_data, theta_current):
    m = len(y_data)
    if m == 0:
        return 0.0
    predictions = X_data @ theta_current
    loss = (1/(2*m)) * np.sum((predictions - y_data)**2)
    return loss

def compute_gradient(X_batch, y_batch, theta_current):
    m = len(y_batch)
    if m == 0:
        return np.zeros_like(theta_current)
    predictions = X_batch @ theta_current
    gradient = (1/m) * X_batch.T @ (predictions - y_batch)
    return gradient

# Global loss calculation for tracking
def calculate_global_loss(theta_current):
    return compute_loss(X, y, theta_current)

Simulating Synchronous SGD

In synchronous SGD, all workers perform computation in lockstep.

1. Each worker computes the gradient on its local data subset (or a mini-batch thereof).
2. A central parameter server (or an All-Reduce operation) aggregates these gradients.
3. The aggregated gradient is used to update the global model parameters.
4. The updated parameters are broadcast back to all workers.
5. A new iteration begins only after all workers have received the updated parameters.

Let's simulate this:

# --- Synchronous SGD Simulation ---
theta_sync = np.zeros((d, 1)) # Re-initialize for clarity
loss_history_sync = [calculate_global_loss(theta_sync)]

print("Running Synchronous SGD Simulation...")
for iteration in range(num_iterations):
    gradients = []
    # 1. Workers compute gradients (using their full local data subset here for simplicity)
    for k in range(num_workers):
        X_k, y_k = worker_data[k]
        grad_k = compute_gradient(X_k, y_k, theta_sync)
        gradients.append(grad_k)

    # 2. Aggregate gradients (simple averaging simulates parameter server/All-Reduce)
    aggregated_gradient = np.mean(gradients, axis=0)

    # 3. Update global parameters
    theta_sync = theta_sync - learning_rate * aggregated_gradient

    # 4. (Implicit) Parameters are "broadcast" (theta_sync is now ready for next iter)

    # Track global loss
    current_loss = calculate_global_loss(theta_sync)
    loss_history_sync.append(current_loss)

    if (iteration + 1) % 10 == 0:
         print(f"Sync Iteration {iteration+1}/{num_iterations}, Loss: {current_loss:.4f}")

print("Synchronous SGD Simulation Complete.")

The main aspect is the explicit aggregation step where all gradients from the current iteration are combined before the update.

Simulating Asynchronous SGD

In asynchronous SGD, workers operate independently without waiting for others.

1. Each worker computes a gradient on its local data (or mini-batch).
2. It sends the gradient (or the parameter update) to the parameter server immediately.
3. The parameter server updates the global parameters using the received gradient. This update might be based on slightly old parameters if other workers have updated them since the current worker fetched its parameters.
4. The worker fetches the latest parameters from the server (which might already include updates from other workers) and starts the next computation.

Simulating the true parallelism and network delays is complex. We can approximate the core behavior by allowing updates to happen based on potentially "stale" parameters. A simple way is to have each worker compute its gradient based on the parameters it last saw, and the central parameter is updated immediately upon receiving a gradient.

# --- Asynchronous SGD Simulation ---
theta_async = np.zeros((d, 1)) # Re-initialize
loss_history_async = [calculate_global_loss(theta_async)]
# Keep track of parameters each worker 'thinks' it has (simulates staleness)
worker_local_thetas = [theta_async.copy() for _ in range(num_workers)]

print("\nRunning Asynchronous SGD Simulation...")
# We need to simulate the asynchronous nature.
# Instead of rigid iterations, we can simulate a total number of gradient computations.
total_gradients_to_compute = num_iterations * num_workers

gradients_computed = 0
while gradients_computed < total_gradients_to_compute:
    # Randomly pick a worker to finish its computation 'next'
    worker_id = np.random.randint(num_workers)

    # Worker fetches the 'current' parameters (could be slightly stale in reality)
    # For simulation simplicity, let's assume it fetches the latest global theta
    # A more complex simulation could introduce delays here.
    current_theta_for_worker = theta_async.copy()

    # Worker computes gradient based on its data and the parameters it has
    X_k, y_k = worker_data[worker_id]
    grad_k = compute_gradient(X_k, y_k, current_theta_for_worker)

    # Parameter server immediately applies the update
    theta_async = theta_async - learning_rate * grad_k
    gradients_computed += 1

    # Track global loss periodically (e.g., every 'num_workers' computations ~ one effective epoch)
    if gradients_computed % num_workers == 0:
        current_loss = calculate_global_loss(theta_async)
        loss_history_async.append(current_loss)
        effective_iteration = gradients_computed // num_workers
        if (effective_iteration) % 10 == 0:
             print(f"Async Effective Iteration {effective_iteration}/{num_iterations}, Loss: {current_loss:.4f}")

print("Asynchronous SGD Simulation Complete.")

# Ensure loss history lengths match for plotting (pad async if needed)
while len(loss_history_async) < len(loss_history_sync):
    loss_history_async.append(loss_history_async[-1])
loss_history_async = loss_history_async[:len(loss_history_sync)]

In this simulation, the critical difference is that theta_async is updated immediately by each worker's gradient, without waiting for others. This models the core idea that updates are applied as they arrive, potentially based on different versions of the parameters.

Analysis and Visualization

Now, let's compare the convergence behavior. We expect synchronous SGD to have a smoother convergence path, while asynchronous SGD might converge faster in terms of wall-clock time (simulated here by the number of gradient computations) but potentially exhibit more erratic behavior or slower convergence per gradient step due to stale updates.

Comparison of global loss convergence for simulated Synchronous and Asynchronous SGD over iterations (or effective iterations for Async). Note the potentially faster initial drop but potentially noisier convergence of the asynchronous method.

Discussion Points:

• Convergence Speed: Did asynchronous SGD show faster initial progress in terms of gradient computations? How did its final convergence compare to synchronous SGD?
• Stability: Was the convergence curve for asynchronous SGD noisier? This often happens because gradients are computed using stale parameters.
• Staleness: Our async simulation is basic. How would introducing explicit delays (simulating network latency or slower workers) affect the results? Staleness can significantly degrade performance if not managed.
• Implementation Complexity: Notice that even in simulation, the logic for asynchronous updates requires careful handling of parameter versions, which hints at the greater complexity in real systems.
• Problem Dependence: The impact of asynchrony heavily depends on the problem (e.g., convexity, gradient variance) and hyperparameters (learning rate).

Further Exploration

• Mini-Batches: Modify the simulation to use mini-batches within each worker instead of their full local dataset for each gradient computation.
• Simulate Staleness: Introduce a delay mechanism. When a worker computes a gradient, base it on parameters from $\tau$ steps ago. How does increasing $\tau$ affect convergence?
• Vary Number of Workers: Rerun the simulations with different num_workers. How does this impact the performance gap and staleness effects?
• Different Learning Rates: Compare the methods using different learning rates. Asynchronous methods might sometimes tolerate or even benefit from slightly higher learning rates, but can also become more unstable.

This practical simulation provides a tangible understanding of the core mechanics and trade-offs involved in synchronous versus asynchronous distributed SGD. It highlights why synchronous methods are often easier to reason about but potentially slower, while asynchronous methods offer potential speedups at the cost of complexity and potential stability issues arising from stale gradients.