Study guide
6.1 Manage Distributed Training for Large and Deep Learning Models
Microsoft AI-300: Operationalizing Machine Learning and Generative AI Solutions Exam
Track your progress as you read
Finished sections get a check mark.
1 Introduction & Exam Overview
1.1 Current Exam Facts 1.2 Skills Measured, Format, and How This Guide Is Organized 1.3 Candidate Profile, Azure ML v2, Foundry, and DevOps Tooling
2 Chapter 2: Azure Machine Learning Workspaces, Datastores, Compute, and IAM
2.1 Create and Manage Azure Machine Learning Workspaces 2.2 Create and Manage Datastores 2.3 Create and Manage Compute Targets 2.4 Configure Identity and Access Management for Workspaces
3 Chapter 3: Azure Machine Learning Data Assets, Environments, Components, and Registries
3.1 Create and Manage Data Assets 3.2 Create and Manage Environments 3.3 Create and Manage Components 3.4 Share Assets Across Workspaces with Registries
4 Chapter 4: Infrastructure as Code, GitHub, and Network Isolation for Azure ML
4.1 GitHub Integration and Git Source Control for Machine Learning Projects 4.2 Deploy Workspaces and Resources with Bicep and Azure CLI 4.3 Automate Resource Provisioning with GitHub Actions 4.4 Restrict Network Access to Machine Learning Workspaces
5 Chapter 5: MLflow Tracking, AutoML, Training Jobs, and Hyperparameter Sweeps
5.1 Configure Experiment Tracking with MLflow 5.2 Use Automated Machine Learning to Explore Optimal Models 5.3 Notebooks, Command Jobs, and Training Scripts 5.4 Automate Hyperparameter Tuning
6 Chapter 6: Distributed Training, Training Pipelines, and Job Comparison
6.1 Manage Distributed Training for Large and Deep Learning Models 6.2 Implement Training Pipelines 6.3 Compare Model Performance Across Jobs
7 Chapter 7: Model Registration, Responsible AI, and Lifecycle Management
7.1 Register MLflow Models and Package Feature Retrieval Specs 7.2 Evaluate Models with Responsible AI Principles 7.3 Manage Model Lifecycle and Archiving
8 Chapter 8: Online Endpoints, Batch Endpoints, Rollout, and Rollback
8.1 Deploy Models to Managed Online Endpoints 8.2 Deploy Models to Batch Endpoints 8.3 Test, Troubleshoot, Progressive Rollout, and Safe Rollback
9 Chapter 9: Data Drift, Production Metrics, Retraining, and Alerts
9.1 Detect and Analyze Data Drift 9.2 Monitor Production Metrics and Trigger Retraining or Alerts
10 Chapter 10: Microsoft Foundry Resources, IAM, Networking, and IaC
10.1 Create and Configure Foundry Resources and Project Environments 10.2 Managed Identities and RBAC for Foundry 10.3 Foundry Network Security and Private Networking 10.4 Deploy Foundry Infrastructure with Bicep and Azure CLI
11 Chapter 11: Deploying and Managing Foundation Models
11.1 Serverless API Endpoints and Managed Compute for Foundation Models 11.2 Select Foundation Models for Specific Use Cases 11.3 Model Versioning and Production Deployment Strategies 11.4 Provisioned Throughput Units for High-Volume Workloads
12 Chapter 12: Prompt Design, Variants, and Git Version Control
12.1 Design and Develop Prompts 12.2 Create Prompt Variants and Compare Performance 12.3 Version Control for Prompts with Git
13 Chapter 13: Evaluation Datasets, Quality Metrics, Safety, and Automated Workflows
13.1 Test Datasets and Data Mapping for Evaluation 13.2 Groundedness, Relevance, Coherence, and Fluency 13.3 Risk and Safety Evaluations 13.4 Built-in and Custom Automated Evaluation Workflows
14 Chapter 14: Continuous Monitoring, Performance, Cost, Tracing, and Logging
14.1 Continuous Monitoring in Microsoft Foundry 14.2 Latency, Throughput, Token Use, and Cost Metrics 14.3 Logging, Distributed Tracing, and Production Debugging
15 Chapter 15: Retrieval Tuning, Embeddings, Hybrid Search, and RAG Evaluation
15.1 Tune Similarity Thresholds, Chunk Sizes, and Retrieval Strategies 15.2 Select and Fine-Tune Embedding Models 15.3 Hybrid Search: Semantic and Keyword Retrieval 15.4 RAG Relevance Metrics and A/B Testing
16 Chapter 16: Fine-Tuning Methods, Synthetic Data, and Promotion to Production
16.1 Advanced Fine-Tuning Methods (LoRA, QLoRA, Full) 16.2 Create and Manage Synthetic Data for Fine-Tuning 16.3 Monitor, Optimize, and Promote Fine-Tuned Models
6.1 Manage Distributed Training for Large and Deep Learning Models
Key Takeaways Explain This Topic AI
SDK/CLI v2 puts distribution on the command job: type pytorch (process_count_per_instance, default 1), type tensorflow (worker_count, optional parameter_server_count), or type mpi (process_count_per_instance is required).
resources.instance_count is the node count; process_count_per_instance is processes per node and should equal GPUs per node for PyTorch DistributedDataParallel. Omit it and Azure Machine Learning launches one process per node.
GPU collectives use NCCL. Multi-node linear scaling wants RDMA/InfiniBand SKUs (the r in Standard_NC24rs_v3). Standard_NC24s_v3 has the same GPU count without InfiniBand.
DeepSpeed and Fully Sharded Data Parallel are libraries in your environment and training script, not Azure Machine Learning resource or job types. Serverless can run distributed jobs when instance_count is set.
Do not distribute small tabular fits. MPI needs a communication-enabled cluster with node-to-node connectivity; isolated-network clusters that block inter-node sockets fail MPI.
Last updated: August 2026
Manage Distributed Training for Large and Deep Learning Models
Quick Answer: Distributed training is a command job with a distribution block. PyTorch uses type: pytorch and process_count_per_instance (usually one process per GPU). TensorFlow uses type: tensorflow and worker_count . Message Passing Interface (MPI) uses type: mpi with a required process_count_per_instance . resources.instance_count is the node count. Serverless works if that count is set. DeepSpeed and Fully Sharded Data Parallel (FSDP) stay in your training code — they are not Azure Machine Learning resource types. Skip distribution for small tabular jobs.
Exam AI-300 Domain 2 asks you to manage distributed training for large and deep learning models . Chapter 5 already covered single-node command jobs, notebooks, automated machine learning, and hyperparameter sweeps. This section is the step you take when one GPU, or one CPU node, cannot hold the model or cannot finish in an acceptable wall-clock time.
Data parallelism versus model parallelism
Azure Machine Learning documents two families of distributed training. Microsoft's own GPU guide is blunt: more than 90 percent of the time, use distributed data parallelism .
Data parallelism copies the full model onto every worker. Each worker trains on a shard of the batch, then workers synchronize gradients (or parameters) so they stay on one consistent model. The entire model must fit in one worker's memory — for GPU jobs, that means one GPU. PyTorch DistributedDataParallel (DDP) and TensorFlow tf.distribute.Strategy implement this pattern.
Model parallelism (also called network parallelism) splits layers or tensors across workers so a model that does not fit on one GPU can still train. Workers exchange activations and shared parameters. FSDP , tensor parallel, and pipeline parallel live here. They are harder to operate and debug.
Azure Machine Learning does not pick the parallelism strategy for you. It launches processes, injects rendezvous environment variables, and lets your script call torch.distributed.init_process_group or the TensorFlow equivalent. If the script is a single-process sklearn.fit , extra nodes sit idle.
The distribution block on a command job
SDK/CLI v2 puts distribution on the command job , not on the compute cluster definition. The YAML schema lists four distribution objects: PyTorchConfiguration , TensorFlowConfiguration , MpiConfiguration , and RayConfiguration . The exam cares about the first three.
- Distribution type Key knobs Who it is for Default process count
- pytorch process_count_per_instance Native torch.distributed / DDP 1 (one process per node if omitted)
tensorflow worker_count , optional parameter_server_count TensorFlow 2.x tf.distribute ; parameter servers are legacy TensorFlow 1.x worker_count defaults to resources.instance_count
mpi process_count_per_instance ( required ) Horovod, many DeepSpeed launchers, custom MPI none — you must set it
ray Ray cluster fields Ray jobs (less common on this exam) not applicable
Instance count versus process count is the distinction the exam loves:
resources.instance_count (YAML) or JobResourceConfiguration.instance_count (SDK) = how many virtual machines (nodes) .
distribution.process_count_per_instance = how many worker processes on each node . For GPU DDP this almost always equals GPUs per node . A node with four GPUs should use process_count_per_instance: 4 .
World size is instance_count × process_count_per_instance . Two Standard_NC24rs_v3 nodes (four GPUs each) with process_count_per_instance: 4 yield eight ranks .
If you omit process_count_per_instance on a PyTorch job, Azure Machine Learning launches one process per node . A four-GPU node then wastes three GPUs. You do not wrap the command in torch.distributed.launch or torchrun . Azure Machine Learning is the launcher. The command remains python train.py ... . $schema: https://azuremlschemas.azureedge.net/latest/commandJob.schema.json
- type: command
- code: src
- command: python train.py --data-dir ${{inputs.cifar}} --epochs ${{inputs.epochs}}
- environment: azureml:AzureML-acpt-pytorch-2.8-cuda12.6@latest
- compute: azureml:gpu-cluster
distribution:
- type: pytorch
- process_count_per_instance: 4
resources:
instance_count: 2
SDK v2 uses command(..., distribution={"type": "PyTorch", "process_count_per_instance": 4}) . For serverless, omit compute and set job.resources = ResourceConfiguration(instance_type="Standard_NC24rs_v3", instance_count=2) .
What Azure Machine Learning injects
For PyTorch , the service sets MASTER_ADDR , MASTER_PORT , WORLD_SIZE , and NODE_RANK on each node, and process-level RANK and LOCAL_RANK . Initialize with torch.distributed.init_process_group(backend='nccl', init_method='env://') . PyTorch's own guidance prefers DDP over DataParallel and over the multiprocessing package for both single-node and multi-node work, which is why Azure Machine Learning examples focus on DDP.
NVIDIA Collective Communications Library (NCCL) is the GPU backend you should use. gloo is a CPU or Ethernet fallback. Using mpi as a PyTorch backend is uncommon when the job already has distribution.type: mpi .
LOCAL_RANK is 0 through GPUs-per-node minus one. Data download, checkpoint writes, and MLflow logging usually run only on local_rank == 0 so you do not stampede storage.
RANK is the global rank, 0 through world size minus one.
For TensorFlow , Azure Machine Learning writes TF_CONFIG before your script starts ( cluster.worker host list and task.index ). Read it with os.environ['TF_CONFIG'] if you must. For TensorFlow 2.x MultiWorkerMirroredStrategy you typically do not parse it yourself. parameter_server_count is for legacy TensorFlow 1.x parameter-server strategy — do not pick it for a 2026 Keras job unless the scenario names that pattern. You can set job.distribution = TensorFlowDistribution(worker_count=2) instead of a dict.
For MPI , OpenMPI (bundled in many curated GPU environments) launches process_count_per_instance ranks per node. Your script uses MPI.COMM_WORLD or a framework that sits on MPI, such as Horovod. MPI ranks must talk to each other . That is why MPI jobs fail on an isolated-network cluster that blocks node-to-node sockets. Use a communication-enabled Azure Machine Learning compute cluster — one that allows inter-node traffic. Isolated clusters and some locked-down no-public-IP designs without extra node-to-node rules are the wrong target for MPI.
GPU clusters, NCCL, and RDMA
Intra-node GPU collectives already use NVLink or PCIe. Multi-node collectives ride the network. Ethernet works, but you lose the linear scaling story (two nodes should approach half the wall-clock of one if communication is cheap). Azure GPU SKUs with Remote Direct Memory Access (RDMA) and InfiniBand (Single Root I/O Virtualization) keep NCCL close to linear.
The SKU hint is the letter r in the size name: Standard_NC24rs_v3 is InfiniBand-enabled; Standard_NC24s_v3 is the same core and GPU count without InfiniBand. Standard_NC24r (no s_v3 ) is an older RDMA SKU without the SR-IOV InfiniBand hardware — do not treat it as current InfiniBand. When you create AmlCompute on a current RDMA SKU, the OS image already includes the Mellanox OpenFabrics Enterprise Distribution (OFED) driver.
Pick GPU families (NC, ND, some H-series) for deep learning. A CPU D-series cluster with distribution.type: pytorch will not magically grow a GPU.
DeepSpeed and FSDP are user code
DeepSpeed is a library you install in the environment and launch from the training script (DeepSpeed launcher, PyTorch distribution, or MPI). Azure Machine Learning supports DeepSpeed on top of PyTorch distribution or MPI. Curated Azure Container for PyTorch images include DeepSpeed, ONNX Runtime, Microsoft Collective Communication Library (MSSCCL), and PyTorch. There is no type: deepspeed compute resource and no workspace asset named DeepSpeed. Autotuning a ds JSON is a DeepSpeed feature you run as a job, not a control-plane object.
FSDP is the same story: torch.distributed.fsdp in your code, still a PyTorch distribution job. Shard size, wrapping policy, and mixed precision are training-script concerns.
Serverless can be distributed
Omit compute (or set a pipeline default_compute: azureml:serverless ) and set resources.instance_count plus resources.instance_type . Serverless is valid for command jobs including distributed training , plus sweeps, AutoML, and parallel jobs. Azure Machine Learning creates a short-lived communicating group, then tears it down. You still consume Azure Machine Learning compute quota . Insufficient GPU quota still fails or queues the job. Dedicated ( Standard ) versus Spot ( job_tier: Spot ) is a queue setting, not a substitute for InfiniBand.
When not to distribute
Do not add a distribution block because the word "scale" appeared in the stem. Counterexamples:
A 200,000-row tabular gradient-boosted tree on Standard_DS12_v2 — communication setup exceeds the training time.
A model that already fits one GPU and finishes in minutes.
A debugging run: start single-process ( instance_count: 1 , no distribution or process_count_per_instance: 1 ) so logs are readable.
Distribution is for large and deep learning models — the skill-measured phrase — not for every training script.
Exam scenario
A computer-vision team trains a segmentation network on two Standard_NC24s_v3 nodes (four GPUs each). The YAML sets distribution.type: pytorch but omits process_count_per_instance , and the SKU has no r . The job starts two processes total, NCCL all-reduces over Ethernet, and epoch time barely improves versus one node. The MLOps fix: set process_count_per_instance: 4 , move the cluster (or serverless instance_type ) to Standard_NC24rs_v3 or an ND InfiniBand size so eight ranks communicate over InfiniBand, and keep DeepSpeed ZeRO as a change inside train.py plus the environment — not a new Azure resource. Confirm the cluster is communication-enabled so MPI fallback or Horovod would also work.
Common trap
Treating DeepSpeed, Horovod, or FSDP as an Azure Machine Learning compute type or job type . The job type stays command . The distribution type is pytorch , tensorflow , or mpi . Another trap: setting instance_count: 4 on a CPU cluster for a scikit-learn fit that is not MPI-aware — you pay for four nodes and train on one. A third trap: putting torchrun in the command string; Azure Machine Learning already launched the processes. A fourth: expecting an isolated-network cluster to run MPI without node-to-node connectivity.