DPointNet multi-GPU support#

This page describes DPointNet’s opt-in sample-parallel training support. Start with DPointNet for a working network and DPointNet training for objectives, optimizers and evaluation. Use documentation from the same BMTK revision as your installation.

Each GPU holds a complete copy of the network and processes different input trials. The replicas combine their gradients to make one shared weight update:

Global batch of input trials
           |
    split by sample
      /          \
  GPU 0          GPU 1
full network   full network
local trials   local trials
      \          /
       mean gradient
            |
    one shared update

This increases the number of trials processed per update and can reduce the per-GPU memory needed for a given global batch. It does not divide neurons or connections between GPUs, and it does not pool GPU memory: the complete network and each local batch must still fit on every device. Both execution modes below are single-host engineering previews; omission of run.distribution preserves existing behavior.

Choose an execution mode#

Start with MPI when opting into multi-GPU training. One TensorFlow process per GPU is the recommended approach for the qualified DPointNet workload. Single-process MirroredStrategy remains an explicitly selected alternative. This is the manual’s recommended starting point, not an automatic runtime default: neither mode is enabled merely because several GPUs are visible. Omitting run.distribution preserves the previous execution and batch behavior.

Approach	Explicit selection	When to use it
MPI (recommended)	`strategy="mpi"` plus the MPI GPU-selection launcher	Starting point for supported single-host multi-GPU training
MirroredStrategy (alternative)	`strategy="mirrored"` in one TensorFlow process	An alternative to qualify and benchmark for your workload

In the qualified two-V10032 full-V1 tests, optimized MPI achieved 1.881528x throughput scaling, while optimized single-process MirroredStrategy achieved 1.551438x. Each result used its own matched single-GPU control in a separate qualification run; these are not a simultaneous head-to-head benchmark. They support the MPI recommendation for that workload, not a claim that MPI transport is intrinsically faster than TensorFlow’s collectives or that every network and GPU will scale the same way.

Memory use: MPI versus MirroredStrategy#

MPI does not inherently require more GPU memory than MirroredStrategy. At the same local batch size, both modes keep a complete network, optimizer state and local trial/BPTT state on each GPU. A separate process does not add another full network copy on that same GPU. Runtime allocations, temporary buffers and device placement can differ, so neither mode has a universal memory advantage.

Separate MPI processes can increase host RAM use: each has its own Python and TensorFlow runtime, graphs and host-side network data, plus MPI input and gradient staging buffers. The qualification below did not measure a matched host-RAM comparison, so it does not establish a host-memory overhead figure.

The historical two-V10032 full-V1 qualification runs above used local batch 16/global batch 32, 500 timesteps, FP32/NEST and checkpoint interval 25. Their recorded TensorFlow allocator lifetime peaks were:

Mode	First GPU (GiB)	Second GPU (GiB)
MPI (ranks 0 and 1)	9.82	9.45
MirroredStrategy	10.59	8.97

MPI did not increase the largest per-GPU TensorFlow peak in these runs. These are separate qualification runs, not a controlled head-to-head memory experiment or a universal memory ranking. TensorFlow allocator peaks are not total driver GPU memory or a minimum VRAM requirement: they do not include all CUDA/driver allocations or describe allocator reservations. End-of-run driver snapshots were approximately 16.3–16.5 GiB per GPU; those snapshots were not sampled driver peaks.

For either mode, splitting a fixed global batch can reduce per-GPU trial state and activation memory compared with placing that entire batch on one GPU. It does not halve the network/optimizer memory or pool device memory: the complete network and local batch must still fit on every GPU.

Sample-parallel training#

An explicit run.distribution opts into equal, disjoint sample batches on one host. Both modes hold the complete network on every replica; neither partitions neurons or exchanges spikes between GPUs.

run.batch_size and the sole training condition’s batch_size are per-replica sizes. The optional global_batch_size must equal this size times the selected replica count. There is one optimizer update per global batch. Device ordinals refer to TensorFlow’s visible allocation, not arbitrary physical GPU numbers. Explicit missing, duplicate or remote devices raise.

The initial scope is full FP32, BPTT, one condition, equal full batches and zero-state initialization. Mixed precision, condition parallel/series/ series_accumulate, local learning rules, warmup/cached initial states and optimizer clipping are rejected until separately qualified. No sample microbatch-accumulation option is implied by condition accumulation.

TargetFiringRate is a sample-separable mean and is admitted. Its existing rate_units="Hz" default is unchanged. Explicit rate_units="spikes_per_step" compares each sample/neuron’s time-mean spike count with firing_rate in spikes per simulation step; 0.005 matches the engineering prototype’s probability-space objective at 1-ms steps. This is not the same loss scale as an Hz-space squared error. Select units deliberately. Other objectives must declare and independently qualify data_parallel_reduction="replica_mean"; nonlinear population/batch statistics cannot in general be computed by averaging local scalar losses. Unqualified objectives raise instead of silently changing the objective.

The training iterator produces one global host batch using the input modules’ existing RNG/seed configuration, then splits inputs and targets into contiguous, disjoint sample slices. Each slice is explicitly placed on its replica before the temporal loop. Creating PerReplica values alone does not establish device residence: a peer’s full batch can otherwise be copied repeatedly inside checkpoint/recomputation functions. This preview intentionally uses the host/reference generator instead of the existing per-device LGN iterator, whose historical semantics replicate a parameter batch. Internal Poisson noise retains replica-specific replay-safe streams; changing the replica count is not a bitwise noise-reproduction guarantee.

Build both operators for the intended pool before running. Compilation and logical-CPU tests alone do not qualify two physical GPUs. Record the resolved distribution, precision, acceleration report and actual operator devices.

Every replica refreshes recurrent and trainable input compute shadows after the update. Checkpoints must explicitly include all canonical masters, including frozen named-input weights, and optimizer/noise state; validate restoration by mutating and checking those resources, not only by accepting assert_consumed().

This preview is not a package release, convergence claim, all-GPU qualification, multi-node backend or replacement for network/neuron-partitioned simulation.

One process per GPU (MPI engineering preview)#

run.distribution.strategy="mpi" explicitly selects the recommended single-host execution mode. Each process owns one complete network on its sole visible GPU. Only sample batches are split; neurons, connections and temporal state are not partitioned. Omission and explicit "mirrored" keep their existing behavior. This is a separately qualified opt-in backend, not a promise of near-linear scaling on every workload.

For two processes with local batch 16, use this configuration fragment:

{
  "run": {
    "dtype": "float32",
    "batch_size": 16,
    "distribution": {
      "strategy": "mpi",
      "devices": ["/GPU:0"],
      "global_batch_size": 32
    }
  }
}

Launch inside a two-GPU allocation with an MPI runtime and mpi4py from the selected environment:

mpirun -n 2 python -m bmtk.utils.dpointnet_mpi run_dpointnet.py config.json

mpirun is the traditional spelling in BMTK examples. mpiexec is also valid when provided by the selected MPI installation; most installations expose both as equivalent launchers. Use the launcher compatible with the MPI library used by mpi4py rather than mixing Open MPI and MPICH installations. The required extra step is python -m bmtk.utils.dpointnet_mpi, which assigns each process its GPU before TensorFlow import, regardless of launcher spelling.

The optional package extra is bmtk[dpointnet-mpi]; installing the extra does not select or qualify an MPI runtime, TensorFlow/CUDA ABI or GPU family. Keep existing immutable environments unchanged. The launcher selects rank ordinals within inherited CUDA_VISIBLE_DEVICES before TensorFlow import. Never replace a Slurm allocation such as 2,3 with physical 0,1. Every process then calls its assigned GPU /GPU:0. The inherited list must contain exactly one distinct GPU per process. Multi-host jobs are rejected. The launcher’s explicit --cpu option is for CPU qualification.

Runnable configuration example#

The source checkout includes examples/dpointnet_mpi/config.train.json and run_training.py. They reuse the checked-in 300-neuron SONATA network and cell/synapse models in examples/dpointnet_all2all; no network download or regeneration is required. Run these commands from the checkout root with an MPI runtime, TensorFlow and the optional dpointnet-mpi extra installed:

# CPU smoke: two processes, local batch 2, global batch 4.
mpirun -n 2 python -m bmtk.utils.dpointnet_mpi --cpu \
  examples/dpointnet_mpi/run_training.py --output /tmp/dpointnet-mpi

# Ordinary CPU control, without the new distribution configuration.
CUDA_VISIBLE_DEVICES="" python examples/dpointnet_mpi/run_training.py \
  --serial --output /tmp/dpointnet-serial

Use separate, fresh output directories. The example selects FP32, NEST hard-forward/soft-surrogate dynamics, 64 timesteps, checkpoint interval 16, zero initial state, seeded Poisson inputs and a target-rate objective in Hz. It performs two epochs of two Adam updates and logs Completed 4 optimizer updates. once. The copied resolved configuration and run.log appear in the output directory. Chief callbacks write callbacks/losses.csv, callbacks/performance.csv and SONATA edge files under callbacks/trained_weights. These are best-epoch weight exports, not a complete optimizer/noise resume checkpoint. The serial control retains batch 2, whereas MPI combines batch 4; these commands are smoke tests, not a matched gradient or speed comparison. Four updates do not establish convergence or biological accuracy.

For the GPU version, build both native operators for the selected TensorFlow/CUDA environment as described in DPointNet performance and GPU setup. Inside an allocation exposing exactly two qualified GPUs, omit --cpu:

mpirun -n 2 python -m bmtk.utils.dpointnet_mpi \
  examples/dpointnet_mpi/run_training.py --output /path/to/new/results

A single-node Slurm allocation can run one MPI launcher which starts both TensorFlow processes:

#!/usr/bin/env bash
#SBATCH --nodes=1
#SBATCH --ntasks=1
#SBATCH --cpus-per-task=8
#SBATCH --gres=gpu:2
#SBATCH --mem=16G
#SBATCH --time=00:05:00
set -euo pipefail
# Activate the qualified environment and cd to the source checkout first.
export TF_NUM_INTRAOP_THREADS=4 TF_NUM_INTEROP_THREADS=4
export OMP_NUM_THREADS=4 OPENBLAS_NUM_THREADS=4
mpirun -n 2 python -m bmtk.utils.dpointnet_mpi \
  examples/dpointnet_mpi/run_training.py --output /path/to/new/results

Partition/account, MPI process-manager integration and GPU/host-memory sizing are site-specific; these resource sizes cover the small example, not V1. Both child processes must inherit the full two-GPU allocation list. A launch mode which restricts each process to one token before this launcher runs is not supported. Do not override the scheduler’s GPU assignment.

Normal BMTK configuration and output ownership#

Every rank executes the same normal entry point:

from bmtk.simulator import dpointnet

config = dpointnet.Config.from_json("config.json")
config.build_env()
network = dpointnet.RNN.from_config(config)
network.run()

Under the MPI launcher, DPointNet’s I/O adapter obtains rank, size and barriers from BMTK’s shared communicator. Rank zero alone creates/removes the output directory, opens the shared log and copies the resolved configuration. All ranks synchronize before using those outputs. overwrite_output_dir retains its existing meaning; it is false in the example. Even with overwrite disabled, log/config initialization is not an append/resume API. Repeated build_env() is a no-op; explicit build_env(force=True) reinitializes the environment collectively. Nonchief informational messages are suppressed by default and do not open or truncate the shared log. Rank-local fatal errors reach the launcher without waiting at an I/O error barrier, so a chief or worker failure aborts the job instead of waiting forever for a different collective.

bmtk.utils.io.ioutils.set_world_comm(comm) may configure the shared communicator before Config.build_env() and model construction; I/O and training use that same communicator. A duplicated communicator is tested. Changing communicators after setup/training starts is unsupported, and the launcher’s failure handling still aborts the entire launched MPI job. Without the launcher, serial I/O remains independent of optional mpi4py. An MPI configuration without the launcher, or a non-MPI configuration under the launcher, raises before environment output setup. Plain mpiexec -n 2 python run_training.py is therefore not a substitute for the early-device-selection launcher.

Support summary#

Surface	MPI backend scope
Parallelism	One host; one complete network per process; equal sample batches
Precision and learning	FP32, BPTT, one condition, zero state, no optimizer clipping
Objective	Qualified sample-separable means, including TargetFiringRate
Communication	Host MPI Scatter and gradient mean; not NCCL or CUDA-aware MPI
Outputs and callbacks	Chief-owned environment and training outputs; collective control flow
Restore	Explicit complete checkpoint plus collective synchronization
Not admitted	Multi-node, mixed precision, local rules, combined training/inference

Implementation and validation#

The chief alone generates each global host batch. Numeric sample and target leaves are scattered together into contiguous local slices, preserving Poisson counts, input-generator seeds and sample identities. Each process computes a local sample-mean gradient through the existing BPTT/checkpoint path. Nonchief receive buffers are reused by leaf shape and dtype; chief in-place Scatter uses its original sample slice without a redundant receive copy. Eligible contiguous, writable FP32 GPU inputs use a zero-copy CPU DLPack view before the required local GPU upload; other dtypes retain ordinary TensorFlow conversion. Uploads complete before these host buffers can be reused or changed. This avoids repeated host snapshots without retaining caller-owned mutable storage on GPU. The same completed-upload boundary is used for reusable reduced-gradient buffers, preserving earlier GPU gradient snapshots across later updates. With one process, finite canonical tensor gradients remain on their original device; no host gradient download, mean reduction or re-upload is needed. Canonical dense FP32 gradients from all trainable surfaces are packed into one host MPI sum reduction per update and divided by process count before local Adam or another admitted optimizer runs. This is not a per-timestep collective, CUDA-aware MPI or neuron exchange. Globally averaged scalar metrics use a separate reduction.

The eager TrainingEngine.train_step boundary coordinates compiled local gradient and optimizer functions; MPI communication is outside TensorFlow graphs. The high-level rnn.train() path uses it automatically. The private traced _distributed_train_step is not an MPI entry point. Optimizer constraints and recurrent/named-input shadow refresh remain in the ordinary update lifecycle. Internal Poisson streams fold in the process rank and retain checkpoint replay; changing process count does not reproduce identical noise.

Startup broadcasts chief model variables, all canonical recurrent and named-input masters (including frozen inputs), optimizer slots and noise variables. After explicitly restoring a chief checkpoint, every process must collectively call rnn.data_parallel.synchronize(rnn.training_engine) before another update. Checkpoints must still enumerate all intended masters: this does not add an automatic generic resume API or repair an incomplete checkpoint.

Input generation and training callbacks run only on the chief. Training loss metrics and early-stop decisions are synchronized so all ranks take the same number of updates. Custom callbacks must respect this ownership: mutating chief weights/optimizer/noise without a collective synchronization is not supported. Uncaught worker errors are printed and abort the MPI job rather than leaving peers blocked. Tracebacks and pending standard output/error are flushed before aborting, since MPI abort bypasses Python’s normal shutdown. Run all ranks through the launcher and maintain identical collective control flow.

The initial FP32/BPTT/one-condition/zero-state/equal-full-batch/no-clipping and sample-separable-loss restrictions above still apply. Dense canonical gradients and numeric or Boolean sample leaves are required. Mixed precision, stateful online population-rate updates, local rules and multi-node execution need separate qualification. Combined MPI training/inference in rnn.run() is rejected to prevent duplicate inference files; run inference separately.

Contributor tests and measured scope#

The DPointNet MPI smoke CI workflow installs the optional MPI dependency and runs real CPU subprocesses in Keras 2 and Keras 3 environments, in addition to mocked-communicator unit tests. From the checkout root:

BMTK_DPOINTNET_DISABLE_FUSED_CUDA=1 BMTK_RUN_MPI_TESTS=1 \
  python -m pytest -q tests/simulator/dpointnet/test_mpi_integration.py
mpirun -n 2 python -m bmtk.utils.dpointnet_mpi --cpu \
  tests/simulator/dpointnet/process_parallel_probe.py --output /tmp/mpi-proof.json

The integration tests cover ordinary serial I/O, one/two MPI ranks, configured communicators, overwrite/no-overwrite, repeated/forced environment setup, chief directory/config-copy errors, worker fatal logging and the complete JSON-to-training example with exactly four optimizer updates and finite loss/weight outputs. Error tests have external time limits and reject timeout as success. The numerical oracle independently compares gradients/Adam updates, constraints, state, explicit restore, callback/input ownership and Poisson replay. CPU smoke tests do not qualify physical GPU placement, native kernels or speed.

The pre-integration V10032 qualification used 66,658 neurons, 24,450,554 recurrent edges, 500 timesteps, FP32/NEST, checkpoint interval 25 and recurrent-only Adam. Twenty synchronized full-output/diagnostic-write samples after three excluded warmups measured 3.686884 s for the efficient single-GPU batch-16 reference and 3.919032 s for two processes at local 16/global 32: 1.881528x throughput. This is not a universal speed guarantee, single-GPU batch-32 comparison, convergence result or a new benchmark of subsequent I/O/documentation changes. The earlier GTX1080Ti 1.953x prototype used a different, counts-only loop and must not be treated as a matched hardware comparison.

Alternative: single-process MirroredStrategy#

Choose run.distribution.strategy="mirrored" explicitly to run both GPUs inside one TensorFlow process. This is an alternative, not the recommended starting point for the qualified workload above. Both it and MPI remain opt-in. Do not launch this configuration once per MPI rank.

This configuration fragment, not a standalone experiment, selects two visible GPU ordinals:

{
  "run": {
    "dtype": "float32",
    "batch_size": 16,
    "distribution": {
      "strategy": "mirrored",
      "devices": ["/GPU:0", "/GPU:1"],
      "global_batch_size": 32
    }
  },
  "training": {
    "training_approach": "single",
    "initial_state": {"module": "zero_state"}
  }
}

Run the ordinary Python entry point once, without mpirun or mpiexec:

python run_dpointnet.py config.mirrored.json

If devices is omitted, MirroredStrategy selects its local devices. The shared FP32/BPTT/loss/batch/initial-state restrictions above still apply. It uses TensorFlow’s gradient collectives: canonical gradients are restored locally, and the local scalar mean is divided by replica count before differentiation. The optimizer’s sum reduction then produces the global mean gradient. Training and validation metrics are reduced before the single host callback.

Native CSR metadata is owned per selected device, and native admission checks the selected homogeneous GPU architectures against both operator manifests. The fixed-queue CSR routing arrays are also resident on each worker. Immutable cell tensors used by timestep kernels (including scalar timestep, threshold and surrogate/gradient gains) are materialized as worker-local eager captures before temporal/checkpoint graphs are traced. A local identity inside a checkpoint graph still fetches its source at every chunk boundary. The rollout frame is nested and thread-local, not a global cell-value cache. Assignable coefficient variables remain live resources; replacing immutable tensors requires ordinary TensorFlow retracing. Localize scalar dependencies as well as large arrays: a four-byte peer fetch on every timestep can serialize otherwise independent GPU streams. The live type-identity dispatch predicate is copied to the host once per rollout, rather than read back from the GPU at every forward/recomputed/VJP step. Coefficient identity is still rechecked and incompatible live values retain the existing fallback or explicit-required error.