Multi-GPU training basics in TensorFlow - Deep Dive

Overview - Multi-GPU training basics

What is it?

Multi-GPU training means using more than one graphics card to teach a machine learning model faster. Instead of one GPU doing all the work, the task is split across several GPUs working together. This helps handle bigger models or larger amounts of data in less time. It is like having many helpers instead of one to finish a big job quickly.

Why it matters

Training big machine learning models on a single GPU can take a very long time, sometimes days or weeks. Multi-GPU training speeds this up, making it possible to experiment and improve models faster. Without it, progress in AI would be slower, and some complex models would be too slow or impossible to train efficiently. This technology helps bring powerful AI to real-world applications sooner.

Where it fits

Before learning multi-GPU training, you should understand basic machine learning model training on a single GPU or CPU, and how TensorFlow works. After this, you can explore advanced distributed training techniques, like multi-node training across many machines, or optimization methods to improve training speed and accuracy.

Mental Model

Core Idea

Multi-GPU training splits the work of teaching a model across several GPUs to speed up learning by processing data in parallel.

Think of it like...

Imagine you have a huge pile of homework to grade. If you do it alone, it takes a long time. But if you share the pile with friends, each grading a part, the whole pile gets done much faster.

┌───────────────┐
│ Training Data │
└──────┬────────┘
       │ Split batches
       ▼
┌───────────┐  ┌───────────┐  ┌───────────┐
│  GPU 1    │  │  GPU 2    │  │  GPU 3    │
│ Processes │  │ Processes │  │ Processes │
│ part of   │  │ part of   │  │ part of   │
│ data      │  │ data      │  │ data      │
└────┬──────┘  └────┬──────┘  └────┬──────┘
     │             │             │
     └─────┬───────┴─────┬───────┘
           ▼             ▼
       Combine gradients and update model
           │
           ▼
     Updated Model Parameters

Build-Up - 7 Steps

1

Foundation

What is GPU and why use it

Concept: Introduce what a GPU is and why it helps in machine learning.

A GPU (Graphics Processing Unit) is a special chip designed to handle many tasks at once, originally for graphics. Machine learning uses GPUs because they can do many calculations in parallel, making training faster than using a CPU (Central Processing Unit).

Result

Learners understand that GPUs speed up training by doing many calculations simultaneously.

Knowing why GPUs are faster than CPUs helps you appreciate why multi-GPU setups can speed up training even more.

2

Foundation

Single-GPU training basics

Concept: Explain how training works on one GPU before adding complexity.

Training a model means showing it many examples and adjusting it to make better predictions. On one GPU, data is processed in batches. The GPU calculates how wrong the model is (loss), then updates the model to improve it. This repeats many times.

Result

Learners see the step-by-step process of training on a single GPU.

Understanding single-GPU training is essential before learning how to split this work across multiple GPUs.

3

Intermediate

Data parallelism concept

🤔Before reading on: Do you think multi-GPU training splits the model or the data? Commit to your answer.

Concept: Introduce data parallelism, the main way multi-GPU training works.

In data parallelism, the model is copied on each GPU. The training data is split into parts, and each GPU processes its part independently. After processing, the GPUs share what they learned to update the model together.

Result

Learners understand that multi-GPU training usually splits data, not the model.

Knowing that data parallelism copies the model but splits data clarifies how GPUs work together without confusion.

4

Intermediate

TensorFlow MirroredStrategy usage

🤔Before reading on: Do you think MirroredStrategy requires manual GPU management or automates it? Commit to your answer.

Concept: Show how TensorFlow helps manage multi-GPU training with MirroredStrategy.

TensorFlow's MirroredStrategy automatically copies the model to each GPU and handles splitting data and combining results. You wrap your training code inside this strategy, and TensorFlow does the hard work of syncing updates.

Result

Learners see how to write simple code to use multiple GPUs without complex manual setup.

Understanding MirroredStrategy reduces the barrier to using multi-GPU training in real projects.

5

Intermediate

Gradient aggregation and synchronization

🤔Before reading on: Do you think GPUs update the model independently or synchronize after each batch? Commit to your answer.

Concept: Explain how GPUs share learning by combining gradients before updating the model.

Each GPU calculates gradients (how to improve the model) from its data part. These gradients are sent to a central place where they are averaged (aggregated). The model is updated with this combined information, ensuring all GPUs stay in sync.

Result

Learners understand the importance of synchronizing updates to keep the model consistent.

Knowing gradient synchronization prevents confusion about why models trained on multiple GPUs stay the same.

6

Advanced

Handling batch size and learning rate

🤔Before reading on: Should you keep the same learning rate when increasing batch size in multi-GPU training? Commit to your answer.

Concept: Discuss how increasing batch size affects training and how to adjust learning rate.

When using multiple GPUs, the effective batch size grows because each GPU processes a batch. Larger batch sizes can require increasing the learning rate to keep training stable and efficient. This adjustment helps the model learn well without slowing down or diverging.

Result

Learners grasp how to tune hyperparameters for multi-GPU training success.

Understanding batch size and learning rate interplay is key to avoiding slow or unstable training.

7

Expert

Performance bottlenecks and communication overhead

🤔Before reading on: Do you think adding more GPUs always speeds up training linearly? Commit to your answer.

Concept: Reveal the limits of multi-GPU scaling due to communication costs.

While more GPUs can speed training, they need to communicate gradients frequently. This communication takes time and can become a bottleneck, especially with many GPUs or slow connections. Efficient algorithms and hardware help reduce this overhead but cannot eliminate it entirely.

Result

Learners realize that multi-GPU training has practical limits and trade-offs.

Knowing communication overhead explains why adding GPUs has diminishing returns and guides better system design.

Under the Hood

Underneath, multi-GPU training duplicates the model on each GPU's memory. Each GPU processes a slice of the input data independently, computing gradients locally. These gradients are then sent through a communication layer (like NCCL or MPI) to a central coordinator or all GPUs collectively, where they are averaged. The averaged gradients update the model parameters identically on all GPUs, keeping them synchronized. This cycle repeats every training step.

Why designed this way?

This design balances speed and simplicity. Copying the model avoids complex splitting of model parts, which can be hard to manage. Data parallelism leverages GPUs' strength in parallel computation. Synchronizing gradients ensures model consistency. Alternatives like model parallelism exist but are more complex and less common for many tasks.

┌───────────────┐
│ Input Dataset │
└──────┬────────┘
       │ Split into batches
       ▼
┌───────────────┐   ┌───────────────┐
│   GPU 1       │   │   GPU 2       │
│ Model Copy 1  │   │ Model Copy 2  │
│ Compute Grad  │   │ Compute Grad  │
└──────┬────────┘   └──────┬────────┘
       │                 │
       └───────┬─────────┘
               ▼
      Gradient Aggregation
               │
               ▼
       Update Model Copies
               │
               ▼
       Next Training Step

Myth Busters - 4 Common Misconceptions

Quick: Does multi-GPU training always make your model more accurate? Commit to yes or no.

Common Belief: Using more GPUs always improves model accuracy because training is faster.

Reality: Multi-GPU training speeds up training but does not inherently improve accuracy. Accuracy depends on data, model, and training settings, not just speed.

Why it matters: Believing speed equals accuracy can lead to ignoring proper tuning and evaluation, resulting in wasted resources and poor models.

Quick: Do you think multi-GPU training splits the model across GPUs by default? Commit to yes or no.

Common Belief: Multi-GPU training splits the model itself across GPUs to share the workload.

Reality: Most multi-GPU training uses data parallelism, where each GPU has a full copy of the model. The data, not the model, is split.

Why it matters: Misunderstanding this can cause confusion about memory use and how to debug training issues.

Quick: Is it true that adding more GPUs always speeds up training proportionally? Commit to yes or no.

Common Belief: Doubling GPUs always halves training time perfectly.

Reality: Adding GPUs speeds training but with diminishing returns due to communication overhead and synchronization delays.

Why it matters: Expecting perfect scaling can lead to frustration and poor hardware investment decisions.

Quick: Do you think you must manually manage GPU communication in TensorFlow multi-GPU training? Commit to yes or no.

Common Belief: You have to write complex code to handle data splitting and gradient syncing manually.

Reality: TensorFlow's MirroredStrategy automates these tasks, making multi-GPU training easier to implement.

Why it matters: Believing manual management is required can discourage beginners from trying multi-GPU training.

Expert Zone

1

Gradient accumulation can be combined with multi-GPU training to simulate even larger batch sizes without increasing memory use per GPU.

2

Mixed precision training often pairs with multi-GPU setups to speed up computation and reduce memory, but requires careful loss scaling to avoid errors.

3

The choice of communication backend (NCCL, MPI, or gRPC) significantly affects multi-GPU training speed and stability, especially on different hardware setups.

When NOT to use

Multi-GPU training is not ideal for very small models or datasets where overhead outweighs speed gains. In such cases, single-GPU or CPU training is simpler and more efficient. For extremely large models that don't fit on one GPU, model parallelism or specialized distributed training frameworks are better alternatives.

Production Patterns

In production, multi-GPU training is often combined with automated hyperparameter tuning and checkpointing to manage long training runs. Cloud platforms provide managed multi-GPU instances with optimized communication. Also, mixed precision and gradient checkpointing are used to maximize GPU memory and speed.

Connections

Distributed Computing

Multi-GPU training is a form of distributed computing focused on parallelizing machine learning tasks across GPUs.

Understanding distributed computing principles helps grasp how data and tasks are split and synchronized in multi-GPU training.

Parallel Processing in CPUs

Both multi-GPU training and CPU parallel processing split work to speed up tasks, but GPUs are specialized for many small calculations simultaneously.

Knowing CPU parallelism concepts clarifies why GPUs are preferred for training deep learning models.

Teamwork in Project Management

Multi-GPU training is like a team dividing work and sharing progress to finish a project faster.

Recognizing this connection helps understand synchronization and communication overhead as coordination challenges in teams.

Common Pitfalls

#1 Ignoring batch size increase when using multiple GPUs.

Wrong approach:

Using the same batch size per GPU without adjusting learning rate or total batch size.
strategy = tf.distribute.MirroredStrategy()
with strategy.scope():
model.compile(optimizer='adam', loss='sparse_categorical_crossentropy')
model.fit(dataset.batch(32), epochs=5)

Correct approach:

Adjust batch size and learning rate to account for multiple GPUs.
strategy = tf.distribute.MirroredStrategy()
with strategy.scope():
model.compile(optimizer=tf.keras.optimizers.Adam(learning_rate=0.001 * strategy.num_replicas_in_sync),
loss='sparse_categorical_crossentropy')
model.fit(dataset.batch(32 * strategy.num_replicas_in_sync), epochs=5)

Root cause:

Not realizing that effective batch size grows with number of GPUs, requiring hyperparameter tuning.

#2 Manually splitting data and updating model outside TensorFlow strategy.

Wrong approach:

Splitting dataset manually and updating model without MirroredStrategy.
for gpu_id in range(num_gpus):
data_part = dataset.shard(num_gpus, gpu_id)
model.train_on_batch(data_part)
# No synchronization of gradients

Correct approach:

Use MirroredStrategy to handle data splitting and synchronization automatically.
strategy = tf.distribute.MirroredStrategy()
with strategy.scope():
model.compile(...)
model.fit(dataset.batch(...))

Root cause:

Underestimating complexity of synchronization and gradient aggregation.

#3 Expecting linear speedup with every added GPU.

Wrong approach:

Adding more GPUs and assuming training time halves each time without measuring.
# Add 4 GPUs expecting 4x speedup blindly

Correct approach:

Measure performance and consider communication overhead.
# Use profiling tools to find bottlenecks and optimize accordingly

Root cause:

Ignoring communication costs and synchronization delays in multi-GPU setups.

Key Takeaways

Multi-GPU training speeds up machine learning by splitting data across GPUs that each have a full model copy.

TensorFlow's MirroredStrategy automates the complex parts of multi-GPU training, making it accessible to beginners.

Synchronizing gradients after each batch keeps all model copies consistent across GPUs.

Increasing batch size in multi-GPU training requires adjusting learning rates to maintain training quality.

Adding more GPUs improves speed but has limits due to communication overhead and synchronization delays.

Practice

1. What is the main purpose of using tf.distribute.MirroredStrategy() in TensorFlow?

easy

A. To train a model on multiple GPUs simultaneously

B. To reduce the size of the dataset

C. To convert a model to run on CPU only

D. To save the model in a special format

Solution

1. Step 1: Understand the role of MirroredStrategy

   tf.distribute.MirroredStrategy() is designed to distribute training across multiple GPUs to speed up the process.

2. Step 2: Compare options with the purpose

   Options A, B, and C do not relate to multi-GPU training. Only training a model on multiple GPUs simultaneously correctly states the purpose.

3. Final Answer:

   To train a model on multiple GPUs simultaneously -> Option A

4. Quick Check:

   Multi-GPU training = MirroredStrategy usage [OK]

Hint: MirroredStrategy means training on many GPUs at once [OK]

Common Mistakes:

• Thinking it reduces dataset size
• Confusing with CPU-only training
• Assuming it changes model saving format

2. Which of the following is the correct way to start using multi-GPU training in TensorFlow?

easy

A. model = build_model()
strategy = tf.distribute.MirroredStrategy()
model.compile()

B. strategy = tf.distribute.MirroredStrategy()
with strategy.scope():
model = build_model()

C. strategy = tf.distribute.MirroredStrategy()
model = build_model()
model.fit()

D. model = build_model()
model.fit()
strategy = tf.distribute.MirroredStrategy()

Solution

1. Step 1: Identify correct usage of strategy.scope()

   Model building and compiling must happen inside strategy.scope() to enable multi-GPU support.

2. Step 2: Check each option for correct order

   Only strategy = tf.distribute.MirroredStrategy() with strategy.scope(): model = build_model() builds the model inside the strategy scope. Others build or compile outside, which disables multi-GPU benefits.

3. Final Answer:

   strategy = tf.distribute.MirroredStrategy() with strategy.scope(): model = build_model() -> Option B

4. Quick Check:

   Build model inside strategy.scope() = strategy = tf.distribute.MirroredStrategy() with strategy.scope(): model = build_model() [OK]

Hint: Always build and compile inside strategy.scope() [OK]

Common Mistakes:

• Building model before strategy.scope()
• Compiling model outside strategy.scope()
• Calling fit before defining strategy

3. Given the code below, what will be the output shape of the predictions if the input batch size is 64 and there are 2 GPUs?

strategy = tf.distribute.MirroredStrategy()
with strategy.scope():
    model = tf.keras.Sequential([
        tf.keras.layers.Dense(10, activation='softmax')
    ])
    model.compile(optimizer='adam', loss='sparse_categorical_crossentropy')

import numpy as np
inputs = np.random.random((64, 20))
predictions = model.predict(inputs)
print(predictions.shape)

medium

A. (10, 64)

B. (32, 10)

C. (128, 10)

D. (64, 10)

Solution

1. Step 1: Understand batch size and MirroredStrategy behavior

   MirroredStrategy splits the batch across GPUs internally but the output predictions shape matches the input batch size.

2. Step 2: Check model output shape

   The model outputs 10 classes per input. Input batch size is 64, so output shape is (64, 10).

3. Final Answer:

   (64, 10) -> Option D

4. Quick Check:

   Output shape = (batch_size, classes) = (64, 10) [OK]

Hint: Output shape matches input batch size, not multiplied by GPUs [OK]

Common Mistakes:

• Multiplying batch size by number of GPUs
• Confusing output shape dimensions
• Assuming output shape is (10, 64)

4. You wrote this code but get an error: RuntimeError: tf.distribute.MirroredStrategy() must be used with model creation inside its scope. What is the likely mistake?

strategy = tf.distribute.MirroredStrategy()
model = build_model()
with strategy.scope():
    model.compile(optimizer='adam', loss='mse')

medium

A. Model was compiled outside the strategy.scope()

B. MirroredStrategy cannot be used with compile()

C. Model was built outside the strategy.scope()

D. The optimizer 'adam' is not supported

Solution

1. Step 1: Analyze error message and code

   The error says model creation must be inside strategy.scope(). Here, model is built before entering the scope.

2. Step 2: Identify correct fix

   Building the model inside the scope fixes the error. Compiling inside scope is correct.

3. Final Answer:

   Model was built outside the strategy.scope() -> Option C

4. Quick Check:

   Build model inside scope to avoid RuntimeError [OK]

Hint: Build model inside strategy.scope() to fix this error [OK]

Common Mistakes:

• Building model before strategy.scope()
• Thinking compile() causes error
• Blaming optimizer choice

5. You want to train a large model on 4 GPUs using TensorFlow. Which of the following steps correctly applies multi-GPU training and ensures the model can handle larger batch sizes efficiently?

hard

A. Use tf.distribute.MirroredStrategy(), build and compile the model inside its scope, then increase batch size in model.fit()

B. Build the model first, then use tf.distribute.MirroredStrategy() only during model.fit() with default batch size

C. Use tf.distribute.MirroredStrategy() but keep batch size small to avoid memory issues

D. Train the model separately on each GPU without using MirroredStrategy and average weights manually

Solution

1. Step 1: Apply MirroredStrategy correctly

   Use tf.distribute.MirroredStrategy() and build/compile the model inside its scope to enable multi-GPU training.

2. Step 2: Adjust batch size for efficiency

   Increasing batch size in model.fit() leverages multiple GPUs to handle larger data efficiently.

3. Final Answer:

   Use tf.distribute.MirroredStrategy(), build and compile the model inside its scope, then increase batch size in model.fit() -> Option A

4. Quick Check:

   Build inside scope + increase batch size = efficient multi-GPU training [OK]

Hint: Build inside scope and increase batch size for multi-GPU speed [OK]

Common Mistakes:

• Building model outside strategy scope
• Not increasing batch size to use GPUs fully
• Manually averaging weights instead of using MirroredStrategy