Batch and multi-GPU runs

Solve many variants of one scene in a single batch, spread independent simulations over several GPUs, or split one large simulation across GPUs.

Maxwell has three ways to use more than one simulation's worth of GPU work. Pick by the shape of your problem:

ProblemUseStatus
Many variants of one scene that differ only in source amplitudes or material densities, for example a design sweep or an optimizer populationSimulation.batchExperimental
Many independent simulations, possibly with different grids, on one or more GPUsrun_many with MultiGPUExecutionSupported
One simulation too large for one GPUSimulation.fdtd(..., parallel=FDTDParallelConfig(...))Engineering preview

Every class named here is described on the Batch and execution API page.

Batch variants of one scene#

Warning

Batch execution is experimental in 0.5.1. It supports FDTD scenes on CUDA with only PointDipole sources and only PointMonitor, PlaneMonitor and FluxMonitor monitors. Scenes with ports, thin wires, networks, lumped elements or circuits are rejected when the batch is prepared.

A batch compiles the scene once and steps several variants side by side. Start from an ordinary simulation, then describe what changes with a leading variant dimension:

import torch
import witwin.maxwell as mw
 
scene = mw.Scene(
    domain=mw.Domain(bounds=((-0.2, 0.2),) * 3),
    grid=mw.GridSpec.uniform(0.025),
    boundary=mw.BoundarySpec.pml(num_layers=6),
    device="cuda",
)
scene.add_source(mw.PointDipole(name="drive", position=(-0.1, 0.0, 0.0), polarization="Ez",
                                source_time=mw.GaussianPulse(frequency=1e9, fwidth=0.5e9)))
scene.add_monitor(mw.PointMonitor(name="probe", position=(0.1, 0.0, 0.0), fields=("Ez",)))
sim = mw.Simulation.fdtd(scene, frequencies=(1e9,))
 
scales = mw.BatchParameters(source_scales={"drive": torch.tensor([1.0, 2.0, 3.0], device="cuda")})
batch = mw.Simulation.batch(
    sim,
    parameters=scales,
    execution=mw.BatchExecution(devices=("cuda:0",), batch_size="auto"),
)
print(batch.plan().to_markdown())       # dry check: no GPU work yet
 
results = batch.run()                     # prepare, run and close
results.results()                         # raises the first failed variant, if any
ez = results.stack(lambda r: r.monitor("probe")["Ez"])   # shape (3,)

results[i] is an ordinary Result for variant i. In this example the three probe values scale exactly with the source amplitude.

  • source_scales maps a named source to a tensor of shape (V,) that multiplies its amplitude.
  • region_densities maps a MaterialRegion name to densities of shape (V, nx, ny, nz).
  • BatchExecution(batch_size=...) sets how many variants run at once; "auto" fits it to free memory. List more devices to spread the variants over them.
  • outputs=mw.BatchOutput(monitors=("probe",)) keeps only the monitors you need, and result_device="cpu" moves results off the GPU.

Reuse a prepared batch#

Preparing compiles the scene and sets up memory. To run the same scene with new values, keep the prepared batch open:

with batch.prepare() as prepared:
    first = prepared.run()
    second = prepared.run(parameters=mw.BatchParameters(
        source_scales={"drive": torch.tensor([4.0, 5.0, 6.0], device="cuda")}))

New parameters must have the same names, shapes and dtypes. To handle variants as they finish instead of all at the end, use prepared.iter_results() and call item.raise_for_error() on each BatchItem.

Gradients through a batch#

Give the batch tensors requires_grad=True, set shutoff=0.0 on the simulation, and build a loss from the results:

density = torch.full((2, 4, 4, 4), 0.5, device="cuda", requires_grad=True)
batch = mw.Simulation.batch(
    sim_with_region,                       # its scene holds MaterialRegion "rho" with a (4, 4, 4) density
    parameters=mw.BatchParameters(region_densities={"rho": density}),
    execution=mw.BatchExecution(devices=("cuda:0",), batch_size=2),
)
results = batch.run()
results.results()
loss = sum(r.monitor("probe")["Ez"].abs().square().sum() for r in results)
(grad,) = torch.autograd.grad(loss, density)    # shape (2, 4, 4, 4)

With two variants on a small grid, this forward and backward pass takes under a minute on a recent GPU.

Independent simulations on several GPUs#

run_many runs a list of simulations and returns their results in the same order. Each simulation needs its own Scene object; the scenes may differ in anything.

# make_scene(f) is your own function that builds and returns a new Scene
sims = [mw.Simulation.fdtd(make_scene(f), frequencies=(f,)) for f in (0.9e9, 1.0e9, 1.1e9)]
execution = mw.MultiGPUExecution.ensemble(devices=("cuda:0", "cuda:1"))
seq = mw.run_many(sims, execution=execution)
 
for i, entry in enumerate(seq.entries):
    print(i, entry)                        # a Result or a DistributedFailure
print(seq.failures)                        # failed entries only
first = seq[0]                             # raises if simulation 0 failed

By default each GPU runs one task at a time, tasks go to devices with enough free memory, and one failure does not stop the others (fail_fast=False). seq.records gives the device and timing of every task. A single device such as ("cuda:0",) also works, which is useful for queuing.

run_many rejects SceneModule inputs, simulations with trainable tensors and simulations with a parallel config.

One simulation across several GPUs#

For a grid too large for one GPU, split it along x:

sim = mw.Simulation.fdtd(
    scene,
    frequencies=[1e9],
    parallel=mw.FDTDParallelConfig(devices=("cuda:0", "cuda:1")),
)
result = sim.run()

Each GPU holds a slab of the grid and exchanges boundary layers with its neighbours every step, by direct GPU-to-GPU copies or NCCL (transport="auto" picks one). Point, plane, flux and mode monitor data are assembled on result_device; set gather_fields=True to also collect the full electric field. CUDA graphs and live plotting are off in this mode. See FDTDParallelConfig.

Warning

Joint multi-GPU solves are an engineering preview. In 0.5.1 they are qualified on two GPUs with a real-valued subset of features: no periodic, Bloch or symmetry boundaries along x, no nonlinear or fully anisotropic media, and no advanced source families. A split pays off only for large grids; a 257³ vacuum case ran about 2× faster on two GPUs, while a 129 × 65 × 65 case ran slower than on one.

Related#