Distributed training basics in MLOps - Step-by-Step Execution

Process Flow - Distributed training basics

Start: Prepare Data & Model

↓

Split Data Across Nodes

↓

Send Model & Data to Each Node

↓

Each Node Trains Locally

↓

Nodes Share Gradients/Weights

↓

Aggregate Updates on Master Node

↓

Update Global Model

↓

Repeat Until Training Complete

↓

End

The flow shows how data and model are split, trained in parallel on nodes, then updates are combined to improve the global model repeatedly.

Execution Sample

MLOps

for epoch in range(2):
  for node in nodes:
    node.train_one_batch()
  aggregate_gradients()
  update_global_model()

This code simulates two training epochs where each node trains locally, then gradients are aggregated and the global model is updated.

Process Table

Step	Epoch	Node	Action	Local Model State	Global Model State
1	0	Node 1	Train batch	Weights updated locally	Initial weights
2	0	Node 2	Train batch	Weights updated locally	Initial weights
3	0	Node 3	Train batch	Weights updated locally	Initial weights
4	0	Master	Aggregate gradients	N/A	Aggregated weights from nodes
5	0	Master	Update global model	N/A	Global model updated
6	1	Node 1	Train batch	Weights updated locally	Global model updated
7	1	Node 2	Train batch	Weights updated locally	Global model updated
8	1	Node 3	Train batch	Weights updated locally	Global model updated
9	1	Master	Aggregate gradients	N/A	Aggregated weights from nodes
10	1	Master	Update global model	N/A	Global model updated
11	2	Exit	Training complete	N/A	Final global model

💡 Training stops after completing 2 epochs with global model updated.

Status Tracker

Variable	Start	After Step 5	After Step 10	Final
Global Model Weights	Initial weights	Aggregated weights from epoch 0	Aggregated weights from epoch 1	Final global model weights
Node 1 Local Weights	Initial weights	Updated after epoch 0 batch	Updated after epoch 1 batch	N/A
Node 2 Local Weights	Initial weights	Updated after epoch 0 batch	Updated after epoch 1 batch	N/A
Node 3 Local Weights	Initial weights	Updated after epoch 0 batch	Updated after epoch 1 batch	N/A

Key Moments - 3 Insights

Why do nodes train locally before aggregation?

Each node updates its local model with its data batch first (see steps 1-3 and 6-8), so the master can aggregate meaningful gradients from all nodes.

What happens if the global model is not updated after aggregation?

Without updating the global model (steps 5 and 10), nodes would keep training on stale weights, preventing learning progress.

Why do we repeat training for multiple epochs?

Repeating epochs (steps 1-10) allows the model to improve gradually by multiple rounds of local training and global updates.

Visual Quiz - 3 Questions

Look at the execution table, what is the global model state after step 5?

A Aggregated weights from nodes

B Initial weights

C Final global model weights

D Weights updated locally

💡 Hint

Check the 'Global Model State' column at step 5 in the execution table.

At which step does the training complete?

A Step 10

B Step 11

C Step 9

D Step 6

💡 Hint

Look for the 'Exit' action in the 'Node' column in the execution table.

If nodes did not share gradients, how would the global model state change?

A It would update faster

B It would update normally

C It would remain at initial weights

D It would become random

💡 Hint

Refer to steps 4 and 5 where aggregation updates the global model.

Concept Snapshot

Distributed training splits data and model across nodes.
Each node trains locally on its data batch.
Nodes share gradients or weights with a master node.
Master aggregates updates and refreshes the global model.
Process repeats for multiple epochs to improve model accuracy.

Full Transcript

Distributed training basics involve splitting data and model across multiple nodes. Each node trains locally on its assigned data batch, updating its local model weights. After local training, nodes share their gradients or updated weights with a master node. The master node aggregates these updates to form a new global model. This updated global model is then sent back to the nodes for the next training round. This cycle repeats for several epochs until the model converges or training completes. This approach speeds up training by parallelizing work and combining knowledge from all nodes.

Practice

1. What is the main purpose of distributed training in machine learning?

easy

A. To avoid using GPUs during training

B. To split the training workload across multiple machines or GPUs

C. To increase the learning rate automatically

D. To reduce the size of the training dataset

Solution

1. Step 1: Understand distributed training goal

   Distributed training is designed to share the training task among several machines or GPUs to speed up the process.

2. Step 2: Analyze options

   Only To split the training workload across multiple machines or GPUs correctly describes this purpose. Options A, B, and C do not relate to workload distribution.

3. Final Answer:

   To split the training workload across multiple machines or GPUs -> Option B

4. Quick Check:

   Distributed training = workload split [OK]

Hint: Distributed training means sharing work across machines [OK]

Common Mistakes:

• Thinking distributed training reduces dataset size
• Confusing distributed training with hyperparameter tuning
• Believing distributed training avoids GPU use

2. Which of the following is the correct way to initialize a process group for distributed training in PyTorch?

easy

A. torch.distributed.init_process_group(backend='nccl', rank=0, world_size=1)

B. torch.init_process_group(backend='nccl', rank=0, world_size=1)

C. torch.distributed.start_process_group(backend='nccl', rank=0, world_size=1)

D. torch.distributed.init_group(backend='nccl', rank=0, world_size=1)

Solution

1. Step 1: Identify correct function name

   The correct function to initialize communication is torch.distributed.init_process_group.

2. Step 2: Check syntax correctness

   torch.distributed.init_process_group(backend='nccl', rank=0, world_size=1) uses the correct module and function name with proper parameters. Other options use incorrect function names or modules.

3. Final Answer:

   torch.distributed.init_process_group(backend='nccl', rank=0, world_size=1) -> Option A

4. Quick Check:

   Correct init function = torch.distributed.init_process_group(backend='nccl', rank=0, world_size=1) [OK]

Hint: Use torch.distributed.init_process_group to start communication [OK]

Common Mistakes:

• Using wrong function names like start_process_group
• Calling init_process_group from wrong module
• Misspelling function or module names

3. Given the following code snippet for distributed training setup, what is the output of print(rank, world_size)?

import torch.distributed as dist
rank = 2
world_size = 4
print(rank, world_size)

medium

A. 4 2

B. Error: rank and world_size undefined

C. 0 1

D. 2 4

Solution

1. Step 1: Analyze variable assignments

   Variables rank and world_size are assigned values 2 and 4 respectively before the print statement.

2. Step 2: Understand print output

   Printing rank and world_size will output '2 4' exactly as assigned.

3. Final Answer:

   2 4 -> Option D

4. Quick Check:

   Print rank, world_size = 2 4 [OK]

Hint: Print variables as assigned to see output [OK]

Common Mistakes:

• Confusing rank with world_size order
• Assuming variables are undefined
• Expecting automatic values without assignment

4. You wrote this code to initialize distributed training but get an error:

import torch.distributed as dist
dist.init_process_group(backend='nccl', rank=0)

What is missing that causes the error?

medium

A. The rank parameter should be a string

B. The backend parameter is incorrect

C. The world_size parameter is missing

D. The import statement is wrong

Solution

1. Step 1: Check init_process_group parameters

   The function requires both rank and world_size parameters to know the total number of processes.

2. Step 2: Identify missing parameter

   The code misses world_size, which causes the error.

3. Final Answer:

   The world_size parameter is missing -> Option C

4. Quick Check:

   Missing world_size causes error [OK]

Hint: Always provide world_size with rank in init_process_group [OK]

Common Mistakes:

• Omitting world_size parameter
• Using wrong backend names
• Passing rank as string instead of int

5. In a distributed training setup with 4 GPUs, you want each process to know its rank and the total number of processes. Which code snippet correctly sets this up and prints the rank and world size?

hard

A. import torch.distributed as dist
dist.init_process_group(backend='nccl')
rank = dist.get_rank()
world_size = dist.get_world_size()
print(rank, world_size)

B. import torch.distributed as dist
world_size = 4
rank = dist.get_rank()
dist.init_process_group(backend='nccl', rank=rank, world_size=world_size)
print(rank, world_size)

C. import torch.distributed as dist
rank = 0
world_size = 4
dist.init_process_group(backend='nccl', rank=rank, world_size=world_size)
print(rank, world_size)

D. import torch.distributed as dist
rank = dist.get_rank()
world_size = dist.get_world_size()
dist.init_process_group(backend='nccl', rank=rank, world_size=world_size)
print(rank, world_size)

Solution

1. Step 1: Understand correct initialization order

   dist.init_process_group must be called before calling dist.get_rank() or dist.get_world_size() to initialize communication.

2. Step 2: Analyze each option

   import torch.distributed as dist dist.init_process_group(backend='nccl') rank = dist.get_rank() world_size = dist.get_world_size() print(rank, world_size) initializes the process group first, then gets rank and world size, then prints them. Other options either get rank before initialization or pass rank manually without initialization.

3. Final Answer:

   import torch.distributed as dist dist.init_process_group(backend='nccl') rank = dist.get_rank() world_size = dist.get_world_size() print(rank, world_size) -> Option A

4. Quick Check:

   Init first, then get rank/world_size = import torch.distributed as dist dist.init_process_group(backend='nccl') rank = dist.get_rank() world_size = dist.get_world_size() print(rank, world_size) [OK]

Hint: Initialize before getting rank and world size [OK]

Common Mistakes:

• Calling get_rank before init_process_group
• Passing rank manually without init
• Not calling init_process_group at all