xnn.common.train.trainer

Training loop: batch training, device selection, validation, checkpointing.

Keeps the moving parts explicit rather than hiding them in a framework, so the pipeline is easy to follow and extend.

Multi-GPU / multi-node data parallelism uses native PyTorch DDP and is driven purely by the environment: when the process was spawned by a distributed launcher that sets RANK / LOCAL_RANK / WORLD_SIZE (torchrun --nproc-per-node N -m xnn train ..., Slurm + torchrun, or an accelerate launch configured for multi-GPU), the trainer initializes the process group, shards the data loaders with DistributedSampler, wraps the model in DistributedDataParallel, all-reduces the logged metrics, and writes checkpoints from rank 0 only. A plain python / xnn invocation runs the unchanged single-process pipeline.

Sharded strategies (FSDP, DeepSpeed ZeRO) are deliberately not used: force and stress losses back-propagate through gradients taken with create_graph=True (see ForceStressOutput), a double backward that DDP supports but sharded wrappers do not.

Functions

resolve_device(name)

Resolve a device specification string into a torch.device.

Classes

Trainer(cfg, train_set[, val_set, test_set])

Batch training loop with validation, scheduling and checkpointing.