[perf] UMA-S conserving training on 8x B200: 151.2 ms to 127.9 ms per step, mostly from host synchronizations in the loss

[TarzanZhao]

Summary

Training UMA-S with the repo's perf_check config leaves the GPU idle for about half of every step, for two reasons. One step at 8 ranks issues around 59k aten ops for 10,200 kernels averaging 8.6 µs, so the CPU has little margin to stay ahead of the GPU. And 348 times per step the code reads a value back from the device; each read stops the CPU until the GPU has drained.

This PR removes 276 of those host reads, in the loss, the dataset head wrappers and the graph builder, and cuts the op count of the MoLE layer. Nothing is behind an option, and no config file or dependency changes. On the 8 B200 GPUs I measured, the median step goes from 151.2 ms to 127.9 ms (1.18x). I have not measured other hardware. Issue: #2195.

1. Loss path without host synchronizations. DDPMTLoss counted valid samples with loss[mult_mask].numel() and all-reduced that Python int, which distutils.all_reduce reads back with .item(), once per task and so 15 times a step; forward then branched on isfinite(). The count is now a device scalar, so the all-reduce is enqueued and never waited on, the NaN guard runs unconditionally, and get_output_masks builds the per-dataset masks on the device once per batch instead of one host-to-device copy per task and dataset. 159 host syncs a step go away. There is no switch, and every trainer that uses DDPMTLoss gets the device-side form. Two things go with them: the "Found nans while computing loss" warning, and three asserts in per_structure that each read the device.
2. Dataset head wrappers select on the device. Both DatasetSpecific*Wrapper heads copied batch_full to the host and then fanned every output out to every dataset with new_zeros plus boolean-mask index copies. A shared helper builds one dataset-id tensor on the device and selects with torch.where, which removes up to 40 syncs a step. Same values and gradients, no switch.
3. Graph builder with host-side tables and static sizes. radius_graph_pbc_v2 read device scalars back one image at a time, 131 host syncs a step. It now reads the per-image repetition and atom counts once and builds the unit-cell and cell-index tables on the host in one pass, gives every repeat_interleave its known output size, drops the unique over images present and the edge-count unique whose result was overwritten on the next line, and skips the zero-distance filter in get_pbc_distances for v2 graphs, which exclude those edges already (a new filter_zero_distances argument, default True, so other callers are unaffected); about 115 of the 131 syncs go. Edge index, cell offsets and neighbor counts stay bit-identical to the base on the 40 real batches I checked on CPU, and every otf_graph model on radius_pbc_version 2 — the default — takes this path.
4. MoLE forward over the systems of a batch. MOLE.forward computed x[start:end] @ W_b per system, and with the double backward slice_backward zero-filled a full-size tensor for each slice, 669 slice_backward and 935 zeros a step in the profile. One commit runs all systems in a single bmm over a padded [B, Smax] layout, and the next loops over x.split(sizes) views instead and keeps the padded bmm only where its padding waste is at most 25 % of the input, which brings peak memory back to the baseline. The padded form alone loses on uneven batches (+6.1 GB peak, mean step +2 %), so the two belong together and are measured as a pair.
5. MoLE expert mixing with a contiguous weight gradient. einsum("eoi,be->boi") returned the expert-weight gradient with permuted strides, so DDP copied every one of them into its bucket each step. coefficients @ weights.flatten(1) computes the same product and its gradient is contiguous. On its own it measures inside the run-to-run spread, and DDP still logs "Grad strides do not match bucket view strides" once per rank for some other parameter.

Result

Measured on 8 NVIDIA B200 GPUs in one node, two interleaved runs per version.

	baseline (afba9d8a8)	this PR	change
step time, median of steps 51 to 250	151.2 ms (150.35, 152.03)	127.9 ms (128.61, 127.18)	1.18x, -15.4 %
steps 51 to 250, summed	30.19 s	25.56 s	1.18x
host synchronizations per step	348	72
kernels launched per step	10,200	9,200
peak allocated memory, writing rank	12.87 GB	12.87 GB	unchanged
peak GPU memory per GPU, nvidia-smi	16.0 GB	15.6, 16.0 GB

The sync and kernel counts come from torch.profiler over steps 62 to 64, against the untuned main. Whole-run wall time is 139.2 → 133.2 s: the ~73 s of elastic launch and dataset load before step 0 is untouched. Run-to-run spread is 1.1 % of the median on both arms; inside a run the per-step spread is 4 to 6 % on every arm, from the variable-size batches.

Per commit, cumulative in the order shown, one run each, measured earlier on a different node against the untuned main.

commit	change	median step
	untuned main afba9d8a8	153.7 ms
493b28ed9	loss path without host syncs	147.0 ms (-4.4 %)
dcc5bb7ee	MoLE padded bmm	143.1 ms (the mean rose 2 %, from padding waste)
a50b7cd72	head wrappers select on device	137.8 ms (-3.7 %)
8fe396e23	MoLE mixing as mm	138.4 ms¹
cc0b27a62	MoLE split views, gated bmm	132.5 ms¹
416d21047	graph builder host tables	127.6 ms¹

¹ All three sat on top of the fused AdamW setting, which is now a baseline override for both arms rather than part of this PR. Rows 4 and 5 also sat on a loader change that cost about 4 ms; it was reverted before row 6, whose own effect was 133.7 → 127.6 ms (-4.6 %). Read the column as differences between consecutive rows.

Against the untuned main, without the fused AdamW override, this branch is 1.20x (154.0 → 127.9 ms on the same node and day).

Correctness Verification

I ran the unmodified base and this branch on the same 12 steps with seed 42 on 8 ranks and compared what the repo's own debug_checksums_save_path recorder writes, extended to also save the per-task losses, the model outputs and the batch shapes: eight numeric quantities plus two exact checks, over steps 0 to 4 on all 8 ranks, 40 rank-and-step records. Each difference below is the largest over those steps and ranks, as max |a - b| over the tensor divided by max |a|. The tolerances were set before the runs at three times the spread of three runs of the unmodified code, with a floor of 0.001 %.

recorded	baseline vs this PR	tolerance
parameters, mean of |w| per tensor, after backward	0.0021 %	0.013 %
gradients, mean of |g| per tensor	2.6 %	9.7 %
set of parameter and gradient names	identical	exact
loss	0.0017 %	0.0043 %
per-task losses, 15 tasks	0.010 %	0.023 %
energy per system, all 5 heads	0.0090 %	0.034 %
forces per atom, all 5 heads	0.86 %	1.9 %
stress per system, all 5 heads	1.04 %	2.4 %
node embeddings	0.085 %	0.26 %
batch shapes and dtypes, atoms per system, dataset per system	identical	exact

Nothing exceeds its tolerance, the closest call being the per-task losses at 0.46 of theirs. Two further comparisons with the same tolerances also pass 40 of 40: the untuned main against the baseline with ++optimizer.fused=true, which covers the numerics of that override (closest call 0.34), and the untuned main against this branch (0.35).

Values do move inside fp32 rounding, since the MoLE GEMMs run batched or through a plain mm instead of einsum and the NaN guard runs unconditionally; none of the six commits is bit-identical to the base. Only steps 0 to 4 carry a signal: with lr 0.1 from random init and atomics in the message passing the unmodified code differs from itself by about 1 % by step 6, and job.deterministic=true is refused by index_reduce_cuda in escn_moe.py, so there is no deterministic reference to compare against.

The recording code is one commit on the perf/uma-launch-bound-fixes-verify branch of my fork (704f59c8a), not in this PR, and the baseline record comes from that same commit applied to afba9d8a8.

Details: hardware, model, full command, traces

Hardware and environment. One node, 8 × NVIDIA B200 (sm_100, 183 GB each), driver 580.126.20, CUDA 13.0. Python 3.12.14, torch 2.13.0+cu130 (cuDNN 9.2, NCCL 2.29.7), e3nn 0.6.0, torchtnt 0.2.4, fairchem main at afba9d8a8 with pip install -e packages/fairchem-core[dev] plus pandas pymatgen pyarrow, which the benchmark callback needs to import. I installed nothing else for either arm.

Measurement. 250 steps, seed 42, 8 ranks with job.scheduler.mode=LOCAL, two interleaved runs per version (baseline, this branch, baseline, this branch), each run alone on the node. Both arms share one Hydra override that is not part of this PR, ++optimizer.fused=true, which swaps AdamW's foreach path for its fused kernels: on the untuned main it measures 154.0 → 151.4 ms over four fused runs (1.02x, right at the noise threshold, hence four runs), at a cost of 0.3 to 0.6 GB of peak allocated memory on some ranks. Step time is time.perf_counter() around each optimizer step from the repo's own BenchmarkTrainCallback, median of steps 51 to 250 read from <run_dir>/benchmark_results.pkl; steps 1 to 10 are warm-up. Every rank writes that path and the last one to finish wins, but all ranks agree within 1 ms because DDP holds them in step.

Model and job. UMA-S style eSCN-MD MoLE backbone from backbone/K4L2.yaml with the perf_check overrides: 4 blocks, lmax = mmax = 2, 64 experts (moe_layer_type: pytorch), about 290 M parameters, fp32 with TF32 off, all trained from random init. MLP_EFS_Head per dataset (omol, oc20, omat, odac, omc), regress_stress: True and direct_forces: False, so forces and stress are autograd of the energy and every step runs a second-order backward. 15 tasks from tasks/oc20_omol_conserving_all.yaml. AdamW at lr 0.1, cosine schedule, gradient clipping at 100, EMA 0.999, DDP over 8 ranks. Batches come from MaxAtomDistributedBatchSampler with max_atoms: 350: each rank trains on 1 to 8 systems, 138 to 348 atoms, about 4k to 10k edges built on the GPU each step. The data is the synthetic 5-dataset aselmdb corpus the repo's perf_check fixtures build (fake_dataset.py), 1.1 MB, copied to node-local disk before each run. job.seed=42 seeds Python, numpy and torch on every rank and the sampler uses its own seed: 0, so the batches are identical between runs; values are not bit-reproducible, because of atomics in index_add_ and index_reduce_.

The measured job, both arms. Check out afba9d8a8 for the baseline and this branch for the other arm, and run the same command from each tree. Every flag except the last four is what TrainingBenchmarkRunner._run_single passes:

fairchem -c configs/uma/benchmark/perf_check/training_inner.yaml \
  datasets.data_root_dir=<data_root> job.device_type=CUDA bf16=False \
  max_steps=250 max_epochs=null ++job.seed=42 ++job.run_dir=<run_dir> \
  job.scheduler.mode=LOCAL ++job.scheduler.ranks_per_node=8 \
  runner.callbacks.0.benchmark_results_path=<run_dir>/benchmark_results.pkl \
  ++optimizer.fused=true

max_epochs=null because training_inner.yaml ships max_epochs: 1 and an epoch of this corpus is 6 steps per rank at 8 ranks; ++job.seed and ++job.run_dir because the yaml's job: block does not declare them. Drop the last line to measure the untuned main, which ran 153.87 and 154.09 ms on the same node and day. Nothing is set in the environment beyond node-local TMPDIR, caches and data, HF offline, W&B disabled, CUDA_VISIBLE_DEVICES=0-7 and PYTHONPATH=<tree>/src (the env's editable install points at another clone).

Repo tests. I ran the tests for every touched module on this branch: loss, MoLE, MoE wrappers, eSCN-MD, graph, mlip unit, train runner. The only failures are the cases that parametrize radius_pbc_version=3, which need the optional nvalchemiops package and fail the same way on the base.

Settings tried and rejected on the unmodified code, all within the 1.6 % noise floor: OMP_NUM_THREADS=1 and =16, NCCL_CUMEM_ENABLE=0, TORCH_NCCL_AVOID_RECORD_STREAMS=1, and runner.train_eval_unit.print_every=1000. Not tried because they change the numerics or the job: runner.train_eval_unit.tf32=true, bf16=True, ema_decay, clip_grad_norm, datasets.max_atoms, job.deterministic. Not applicable: no attention in this model, gradient_as_bucket_view already on and broadcast_buffers already off, num_workers: 0 in the dataset config (num_workers=1 fails at startup on this base with TypeError: cannot pickle 'Environment' object from the LMDB handles, and a worker-safe form cost 2.9 %), pin_memory hard-coded on, no torch.compile switch in the trainer, and PYTORCH_CUDA_ALLOC_CONF overwritten by common/utils.py:setup_env_vars.

Scope of the measurement. One node, 8 ranks, fp32 with TF32 off, the synthetic perf_check corpus. Most of the gap is GPU idle time from CPU launch and host synchronization cost, so what the change is worth depends on the host as much as on the B200. Not measured: a single GPU, bf16=True, real datasets, UMA-M, radius_pbc_version=3, and inference through umas_fast_gpu, which does not go through the changed MoLE loop.

Correctness run. On the verify branch, the same command with max_steps=12 +runner.train_eval_unit.debug_checksums_save_path=<dir>; for the baseline record, the same recorder commit on afba9d8a8 with the same command. The comparison against the tolerance file is a script in my workspace rather than on the branch; both are available on request.

Traces (torch.profiler, steps 62 to 64, all 8 ranks; open in https://ui.perfetto.dev; the README next to them lists code, script, hardware, env and capture window). The baseline trace is the untuned main, and the optimized trace holds these same six changes with fused: true set in the optimizer yaml instead of on the command line:

• baseline: https://drive.google.com/open?id=1G70aPMgsTgSC8kDG2VI6vX0MmMhGwVM3
• optimized: https://drive.google.com/open?id=1q6qtvmN47vf1GHPKhVSe4_WrolOektaT

[meta-cla]

Hi @TarzanZhao!

Thank you for your pull request and welcome to our community.

Action Required

In order to merge any pull request (code, docs, etc.), we require contributors to sign our Contributor License Agreement, and we don't seem to have one on file for you.

Process

In order for us to review and merge your suggested changes, please sign at https://code.facebook.com/cla. If you are contributing on behalf of someone else (eg your employer), the individual CLA may not be sufficient and your employer may need to sign the corporate CLA.

Once the CLA is signed, our tooling will perform checks and validations. Afterwards, the pull request will be tagged with CLA signed. The tagging process may take up to 1 hour after signing. Please give it that time before contacting us about it.

If you have received this in error or have any questions, please contact us at cla@meta.com. Thanks!

[meta-cla]

Thank you for signing our Contributor License Agreement. We can now accept your code for this (and any) Meta Open Source project. Thanks!

[TarzanZhao]

Ready for review. This PR moves per-step work off the host in the UMA-S conserving training step: 6 commits across the loss, the dataset head wrappers, the graph builder and the MoLE layer. On 8 B200 the median step goes from 151.2 ms to 127.9 ms, 1.18x.

(Edited: the numbers first posted here, 154 to 130 ms, predate a re-measurement that moved ++optimizer.fused=true into the baseline, since it is a setting rather than something this PR contributes. The description has the current figures.)

cc @rayg1234 (merged the most recent changes to graph/ and uma/), @mlazos (authored and merged the latest uma/graph performance PRs, #2143 and #2154), @misko (last touched configs/uma/benchmark/perf_check and escn_moe.py).

Two small asks: this is my first PR here, so the workflow runs need approval; and the label checker needs patch (or minor) plus enhancement, which I cannot add myself. CLA is signed.

[TarzanZhao]

I have split the first of these six changes out on its own: #2203, the loss path without host synchronizations, two files on its own branch, re-measured and re-verified against today's main. This PR goes back to draft while that one is reviewed.