Distributed training and training time: MLA-C02 study notes and practice questions
What to know about Distributed training and training time for the MLA-C02 exam, from the ML Model and Foundation Model (FM) Development domain (24% of the exam guide), then 1 free question to test it. Part of the MLA-C02 study guide.
These questions ask how to make training finish sooner, or fit a model that doesn’t fit in one GPU’s memory.
Two kinds of parallelism
Data parallelism
Model parallelism
Splits
The data: each GPU holds a full model copy and processes different batches
The model: layers or tensors spread across GPUs
Solves
Training takes too long on large datasets
The model is too large for one GPU’s memory
Tools
SageMaker distributed data parallel, PyTorch DDP, FSDP
SageMaker model parallel library, tensor and pipeline parallelism, sharded data parallelism
Fully sharded data parallel (FSDP) and sharded data parallelism split model states across GPUs, combining both ideas.
Saving memory and time
Mixed precision (FP16 or BF16) speeds training and halves memory for activations.
Activation checkpointing recomputes some activations instead of storing them, trading compute for memory.
Gradient accumulation simulates larger batches on limited memory.
Larger batches may need a scaled learning rate.
Hardware
GPU instances (P and G families) for deep learning; AWS Trainium (Trn instances) for lower-cost training of supported models.
Use instances with fast networking (EFA) for multi-node training.
CPU instances often suffice for XGBoost and classic ML.
Data loading bottlenecks
If GPUs sit idle (low GPU use in Debugger or CloudWatch), the input pipeline is too slow: use Fast File mode or FSx for Lustre, more data loader workers, larger files and efficient formats.
Resilience
SageMaker HyperPod keeps large clusters running, replacing faulty nodes and resuming from checkpoints.
Common traps
Data parallelism for a model that doesn’t fit in memory. Each GPU still needs the whole model.
Adding GPUs while they’re waiting on data.
Bigger batches without adjusting the learning rate.
A ski resort group is training a 13-billion-parameter language model on SageMaker AI for resort-guest support. The model does not fit in memory on the GPU instances the team has budgeted for, and the training job fails with out-of-memory errors at the first batch.
The team wants to keep the same model architecture and the same training data, and to use the budgeted instance type without making the model smaller.
Which TWO changes meet these requirements? (Choose two.)
Choose two answers.
AA. Turn on data parallelism, so that each GPU holds a full copy of the model and processes a different slice of the data.
BB. Turn on model parallelism, so that the layers or parameters of the model are split across GPUs.
CC. Increase the batch size so that more examples fit in each training step.
DD. Turn on activation checkpointing, so that intermediate activations are recomputed during the backward pass instead of being kept in memory.
EE. Switch the training data from text to compressed archives, which reduces the GPU memory the model needs.
Show the answer and explanation
Answers: B, D
A. Wrong: Data parallelism still gives each GPU a full copy of the model, so the out-of-memory error remains. It speeds up training when the model fits but does not solve this problem.
C. Wrong: A larger batch needs more memory for activations, so it makes the out-of-memory error worse. The fix needs to reduce the memory the model uses.
D. Right: Correct. Activation checkpointing trades extra computation for lower memory use, because it stores fewer intermediate activations. Combined with model parallelism, it helps the large model fit without changing its architecture. Learn more: Activation checkpointing (opens in a new tab)
E. Wrong: Compressing the input files does not change the memory the model parameters and activations need on the GPU. The problem is the model, not the storage format.
Explanation
Model parallelism splits the parameters across GPUs, and activation checkpointing reduces the memory used for intermediate results, so the same model can train on the budgeted instances. Data parallelism and larger batches keep a full model copy or add memory, and the input format does not change GPU memory.
Want more on Distributed training and training time? The 5 full-length practice sets have 5 more Distributed training and training time questions, 325 in all, none of them in the free exam. One payment of $14.99, no subscription. See the practice sets →