Trainer and Module

gpulayer.pytorch gives you a small training API: put the model and its training logic in a Module, and let a Trainer run the loop. The Trainer also decides where the loop runs.

import torch
from torch import nn
import gpulayer.pytorch as glt


class Classifier(glt.Module):
    def __init__(self, lr: float = 1e-3):
        super().__init__()
        self.lr = lr
        self.net = nn.Sequential(nn.Flatten(), nn.Linear(28 * 28, 128), nn.ReLU(), nn.Linear(128, 10))

    def forward(self, x):
        return self.net(x)

    def training_step(self, batch, batch_idx):
        x, y = batch
        loss = nn.functional.cross_entropy(self(x), y)
        self.log("train_loss", loss)
        return loss

    def validation_step(self, batch, batch_idx):
        x, y = batch
        logits = self(x)
        self.log("val_loss", nn.functional.cross_entropy(logits, y))
        self.log("val_acc", (logits.argmax(1) == y).float().mean())

    def configure_optimizers(self):
        return torch.optim.Adam(self.parameters(), lr=self.lr)


if __name__ == "__main__":
    model = Classifier()
    trainer = glt.Trainer(max_epochs=10, accelerator="auto")
    trainer.fit(model, train_loader, val_loader)
    print(trainer.callback_metrics)

Module

glt.Module is an nn.Module with these hooks:

self.log(name, value) records a scalar (a float or a tensor). The Trainer averages each logged name over the steps of an epoch and across all processes, then stores the result in trainer.callback_metrics and trainer.history. If training_step logs nothing, the Trainer logs the returned loss as train_loss.

Batches arrive on the right device already. The Trainer moves tensors, and lists, tuples and dicts of tensors, before calling your step.

Trainer arguments

Unsupported combinations fail with an error instead of quietly doing something else. Examples: precision="16-mixed" on MPS, accelerator="cuda" on a machine without CUDA, devices=2 with MPS, or accelerator="auto" with target="remote" (the remote device is unknown until the job is placed, so name it).

Callbacks

Callbacks hook into the training loop. Two are built in:

from gpulayer.pytorch import EarlyStopping, ModelCheckpoint, Trainer

best = ModelCheckpoint(monitor="val_loss")                 # saves checkpoints/best.ckpt
trainer = Trainer(max_epochs=50, callbacks=[EarlyStopping(monitor="val_loss", patience=3), best])
trainer.fit(model, train_loader, val_loader)
model = MyModel.load_from_checkpoint(best.best_model_path)

• EarlyStopping(monitor, mode="min", patience=3, min_delta=0.0) stops when the metric has not improved for patience epochs.
• ModelCheckpoint(monitor, mode="min", filename="best.ckpt") saves a checkpoint each time the metric improves. best_model_path points at it, also when training ran on workers (the file comes back with the job's outputs).
• The monitored name must be logged with self.log(...); otherwise the error lists the names that were logged.

Write your own by subclassing Callback and overriding any of on_fit_start, on_train_epoch_start, on_train_batch_end(trainer, model, loss, batch_idx), on_train_epoch_end (after validation; trainer.callback_metrics holds the epoch's values) and on_fit_end. Set trainer.should_stop = True to end training after the current epoch; trainer.save_checkpoint("name.ckpt") writes a checkpoint.

In distributed runs callbacks run on every rank. Epoch metrics are averaged across ranks first, and if any rank sets should_stop, every rank stops after the same epoch. When fit launches training on workers, the callbacks run there; in the launching process only on_fit_end runs.

Where fit runs

When fit submits a job, it re-runs your script on the workers with the same command-line arguments. On the workers, the same fit call runs the training loop. When the job finishes, the launching process downloads the outputs and loads the final weights into your local model. Code after fit therefore sees a trained model in both cases.

Because the script runs again on the workers, code at module level and after fit runs there too. Guard work that should only happen once, on your machine:

trainer.fit(model, train_loader, val_loader)
if glt.is_launcher_process():
    torch.save(model.state_dict(), "final.pt")

After training

Checkpoints and resuming

With checkpointing on, global rank 0 writes checkpoints/last.ckpt and metrics.json after every epoch. Inside a job they go to $GPULAYER_OUTPUT_DIR and come back as artifacts (gpulayer_runs/<job-id>/...); in-process they go to gpulayer_runs/<directory>-<script>/, for example gpulayer_runs/my_project-train/ for my_project/train.py.

model = Classifier.load_from_checkpoint("gpulayer_runs/my_project-train/checkpoints/last.ckpt")
trainer = glt.Trainer(max_epochs=20)
trainer.fit(model, train_loader, val_loader,
            ckpt_path="gpulayer_runs/my_project-train/checkpoints/last.ckpt")

A checkpoint holds state_dict, optimizer, epoch, global_step and history. ckpt_path restores the optimizer and continues from the next epoch.

Validation only

metrics = trainer.validate(model, val_loader)   # {"val_loss": ..., "val_acc": ...}

validate runs in the current process.

What the Trainer does not do yet

Logger integrations, learning-rate monitoring and strategies other than DDP (such as FSDP) are not implemented. If you need them, keep your existing training code and run it with launch().