14. Multi-GPU training (exception track)

The admission-gated exception: distributed / multi-GPU training on Azure ML command jobs against min-zero clusters, logging to the same self-hosted MLflow.

Outcome

When — and only when — a workload genuinely needs distributed or multi-GPU training, it runs on Azure ML command jobs against a min-zero cluster, logging to the same self-hosted MLflow as everything else. This is an admission-gated exception, deliberately off the critical path: it is not part of the baseline and Parts I–II ship without it.

## Design — a narrow, gated exception

  • Why an exception, not the baseline. ACA Jobs are single-node; genuine distributed/multi-GPU training needs a cluster scheduler with GPU topology. Azure ML command jobs provide that without us running a cluster fleet — the cluster scales from zero and back to zero.
  • Admission gate. A workload uses this path only after documented admission (it must actually need multi-GPU); it does not become the default training path.
  • Same identity for the model. The AML job logs runs and registers versions to our self-hosted MLflow, so the produced model is identical in kind to a Ch 03 model: models:/<name>/<version>. Promotion, serving, and batch are unchanged.
  • Min-zero cluster. The compute target idles at zero nodes and scales only for the duration of a job, keeping cost bounded.

Build in projects/ml-platform/

projects/ml-platform/
├── src/train_aml/
│   ├── train_distributed.py    # entrypoint: detects LOCAL_RANK; only rank-0
│   │                           #   writes MLflow runs + results-DB record
│   ├── job.yml                 # AML command job definition: compute, distribution,
│   │                           #   BYO image, self-hosted MLflow tracking URI
│   └── requirements.txt        # same mlflow/sklearn pins + torch for distributed
└── infra/modules/aml/          # (applied only when exception is admitted)
    ├── main.tf                 # AML workspace + min-zero GPU cluster (scale_down PT5M)
    ├── variables.tf
    └── outputs.tf              # workspace_name, cluster_name (referenced in job.yml)

The AML workspace and cluster are not in the Phase-0 foundation — they are only provisioned when the exception is explicitly admitted. Everything downstream (promotion, serving, batch, dashboard) is unchanged: the distributed job registers a version in the same self-hosted MLflow registry by the same URI.

How the pieces connect

Admission gate

Use this path only when a measured workload cannot fit single-node ACA training. Admission is a written justification and explicit approval, not a config flag.

train_distributed.py

The script mirrors train.py; only rank 0 configures MLflow, writes the results row, and registers the model. Every rank trains. The resulting models:/<name>/<version> artifact has the same dataset digest and code-image tag as a chapter 03 version, plus training.backend = aml-distributed.

job.yml

The AML command job references the workload image by digest, the GPU cluster by name, Key Vault-resolved connection values, and distribution.type = PyTorch. Increase instance_count only for genuine multi-node work.

infra/modules/aml/

The module defines an AML workspace and a min-zero compute cluster, but the root module does not instantiate it. Admission work must first supply the external Application Insights resource ID required by the AML workspace, add an explicit root module call, and review the new cost/RBAC surface. Until then, terraform -target=module.aml is not a valid baseline command because no such root instance exists.

This is intentional: checked-in module source documents the exception path; it does not silently expand the deployed platform.

Extensions & boundary

Boundary Course decision
AML pipelines, endpoints, auto-registered model assets Explicitly out of scope; the AML integration stays narrow
Making this the default training path It remains an admission-gated exception

Next: 15 — End-to-end integration walks the whole golden path and summarizes every deferred extension.

Back to top