Distributed training and hyperparameter search

Use distributed execution after the same model, dataset, and overrides complete on one GPU. It adds process coordination and cluster resource requirements while retaining the same dataset and checkpoint compatibility requirements.

Choose an execution mode¶

--gpus gives the number of PyTorch DDP processes for the standard train command. For Ray commands, --ray-gpus gives the number of GPUs assigned to each training worker or trial. Select the visible CUDA devices through the execution environment.

Single-node DDP¶

This example uses four visible GPUs on one machine:

uv run python3 mlpf/pipeline.py \
  --spec-file particleflow_spec.yaml \
  --model-name pyg-cld-v1 \
  --production-name cld \
  --data-dir /path/to/cld/tfds \
  --prefix cld-ddp- \
  train \
  --gpus 4 \
  --sampler_mode interleaved-shards \
  --dtype bfloat16 \
  --num_steps 100000

Each rank gets its own data-loader process and GPU. The effective event mix depends on dataset-group batch sizes, gpu_batch_multiplier, world size, and sampler behavior. Record all of these values alongside the nominal batch size.

Choose the event order and group step shares using Dataset mixing and loss monitoring. These settings apply to both single-GPU and DDP training. For groups containing several physics samples, use --sampler_mode interleaved-shards.

All ranks must see the same files and experiment directory. Typical failures are mismatched visible GPUs, NCCL networking, insufficient file-descriptor limits, and one rank running out of memory.

Ray Train¶

Run a small local Ray check before connecting to a cluster:

uv run python3 mlpf/pipeline.py \
  --spec-file particleflow_spec.yaml \
  --model-name pyg-cld-v1 \
  --production-name cld \
  --data-dir data/tfds/tensorflow_datasets/cld \
  --experiments-dir experiments \
  --prefix cld-ray-local- \
  ray-train \
  --data_config 1 \
  --ray-local \
  --ray-cpus 2 \
  --ray-gpus 0 \
  --model.attention.attention_type math \
  --num_steps 2 \
  --ntrain 10 \
  --nvalid 10 \
  --ntest 10

For a GPU worker, use --ray-gpus 1 and settings appropriate to that GPU. --ray-local creates the local Ray environment; cluster execution uses the Ray environment supplied by the site. Ray keeps its own run and checkpoint metadata under the configured experiment storage.

Ray Tune hyperparameter search¶

Searches use mlpf/raytune/search_space.py. Review that file before starting: a search can launch many full training trials. The selected model’s raytune.local_dir must be a writable absolute storage path in a custom specification.

After setting that path, a bounded local check is:

uv run python3 mlpf/pipeline.py \
  --spec-file /path/to/reviewed-spec.yaml \
  --model-name pyg-cld-v1 \
  --production-name cld \
  --data-dir /path/to/cld/tfds \
  ray-hpo \
  --name cld-search \
  --ray-local \
  --ray-cpus 2 \
  --ray-gpus 0 \
  --model.attention.attention_type math \
  --raytune-num-samples 2 \
  --num_steps 2 \
  --ntrain 10 \
  --nvalid 10

An existing compatible Tune directory is resumed automatically, including unfinished and errored trials. Use a new search name when the search space or scientific question changes.

Reproducibility checklist¶

Record the repository commit, complete resolved configuration, dataset versions and partitions, checkpoint, visible hardware, world size, worker resources, sampler mode, precision, and random seeds. Compare runs at equal effective data exposure and validation cadence. TensorBoard is written locally; Comet logging is optional through --comet or --comet-offline.