08. 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 (docs/08), deliberately off the critical path: it is not part of the baseline and Phases 0–5 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

This path is used only when a workload demonstrably cannot fit single-node training on an ACA Job. The gate is a written justification + explicit approval — not a config flag. Most training stays on ACA (train.py, Ch 03).

train_distributed.py

Mirrors train.py almost exactly; the structural difference is a single _local_rank() guard: only rank 0 calls configure_mlflow, record_run, and mlflow.start_run. All ranks train, but only rank 0 writes lineage and registers the version. The registered version is indistinguishable from a Ch 03 version — same models:/<name>/<version> URI, same dataset digest and code image tag.

The training.backend = aml-distributed tag on the MLflow run is the only marker that this path was used; serving and batch never see it.

job.yml

An AML command job referencing: - Our container image by digest (BYO — not a curated AML environment). - The GPU compute cluster by name (azureml:<cluster-name>). - Secrets (mlflow-tracking-uri, results-pghost) resolved from Key Vault at submit time — never in the YAML. - distribution.type = PyTorch (extend instance_count for genuine multi-node).

Submitted with az ml job create -f job.yml --workspace-name <name>.

infra/modules/aml/

azurerm_machine_learning_workspace wired to the shared storage, Key Vault, and App Insights (all from the foundation). azurerm_machine_learning_compute_cluster with min_node_count = 0 and scale_down_nodes_after_idle_duration = PT5M — zero cost at idle, scales per job. The GPU VM size and max_nodes are set at admission time, not hardcoded.

The module is not referenced from root main.tf — it is applied as a standalone module only if/when the exception is triggered:

terraform apply -target=module.aml -var-file infra/environments/dev.tfvars

## Extensions & boundary

Deferred / boundary Contract
AML pipelines, endpoints, auto-registered model assets docs/08 — explicitly out of scope; AML is narrow
Making this the default training path docs/00 invariant 9 — it stays an exception

Next: 09 — Broker upgrade — the other conditional track, for fan-out that outgrows the results-DB continuation rule.

Back to top