Performance

Measure the complete workload: source reading and preprocessing, tensorization, model computation, and prediction writing. Schema shape and requested outputs can matter as much as model width. Compare quality and throughput under the same data, hardware, precision, and optimization budget.

Control Work First

Repeated branches reserve their configured lengths. Nested lengths multiply coordinate counts. Choose lengths and overflow rules that retain the context the task needs; reductions determine how much of each branch reaches its parent. These choices can change model quality and are not interchangeable performance switches.

Padding and structurally skipped coordinates are absent from attention. Eligible encoder work packs present coordinates when the layers have no customizations or active stochastic dropout. This is automatic; it does not require changing the input schema. Null and learned-mask coordinates remain present model inputs.

Export only needed outputs. Root or branch embeddings, repeated leaf outputs, and Category top-k values enlarge the prediction payload. Selecting a larger batch can improve utilization but also changes updates per epoch; account for that when comparing training runs.

Feed The Device

model.batch_size controls encoded model batches. Loader settings belong on the data module; tune them after measuring whether input preparation leaves the device idle:

datamodule = rf.ArrowDataModule(
    model=model,
    train=train_table,
    validate=validation_table,
    num_workers=4,
    persistent_workers=True,
    pin_memory=True,
    prefetch_factor=2,
)

Workers and prefetching consume host memory. Arrow sources and retained data remain on CPU while encoded tensors move to the model’s device. Source readers are replayed by consuming ranks/workers before disjoint row selection, so adding workers does not guarantee less source I/O. Keep Arrow/Polars transformations columnar instead of repeatedly converting observations to Python rows.

The throughput callback estimates rate using completed batches multiplied by configured batch size. It is useful for comparisons with the same batch geometry, not exact row accounting for partial batches. Prediction throughput is retained by the callback and sent directly to the attached logger; it is not a trainer.callback_metrics result.

Use The AdamW Factory

import relflow as rf

model.optimizer = rf.adamw(learning_rate=1e-3)

The factory defaults to fused Torch AdamW. Biases, one-dimensional parameters, and names containing "norm" receive no weight decay by default. Use fused=False for an unsupported device or fused=None for Torch’s selection. When resuming, keep optimizer settings consistent with the saved run.

Compile Encoder And Pool Regions

model.compile() updates the existing model and returns it, so it can follow a constructor or be called programmatically before training or inference.

model.compile()
model.compile(encoders=True, pools=False)
model.compile(encoders=False, pools=True)
model.compile(encoders=False, pools=False)

Each call replaces the previous selection. The last call above restores eager execution.

Switch	Selected region
encoders=True	Branch sequence encoder stacks.
pools=True	Learned attention reductions in branches and decoders.

Coordinate encoders, routing, tensorfield operations, mean pooling, losses, metrics, and optimizer execution keep their ordinary paths. Each selected region uses torch.compile(fullgraph=True). Custom compute implementations remain eager, as do customized regions and regions with nested hooks, including hooks added after compilation.

model.compile(backend="inductor", dynamic=True, options={})

backend="inductor", dynamic=None, and options=None are the defaults. dynamic=None lets Torch decide when to generalize shapes; True requests dynamic shapes and False specializes them. Inductor defaults preserve eager random draws and disable CUDA graphs and automatic tensor padding. Backend-specific options override those defaults. Floating-point results can still differ.

Compilation happens lazily on real inputs. Measure initial compilation and steady-state execution separately, covering the batch sizes, branch occupancy, training mode, and inference mode used by the application. Compilation is not a guaranteed speedup: first establish a useful result on representative work. Compiler errors in selected regions propagate rather than silently selecting an eager fallback.

Schema mutations, checkpoint rebuilds, and model copies do not preserve compilation. Reapply it after loading or changing the model. Compilation is not stored in checkpoints.