6.1 Manage Distributed Training for Large and Deep Learning Models

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.

Instance count versus process count is the distinction the exam loves:

1. resources.instance_count (YAML) or JobResourceConfiguration.instance_count (SDK) = how many virtual machines (nodes).
2. 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.
3. 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.