tuning

tune_batch_size

tune_batch_size(
    trainer,
    module,
    datamodule=None,
    *,
    effective_batch_size,
    pin_path=None,
    num_devices=None,
    batch_arg="batch_size",
    retune=False,
    **scale_kwargs,
)

Pin the effective batch; find the largest fitting device batch, set accumulation.

Runs Tuner.scale_batch_size once to find the largest device batch that fits (found_max), then chooses device_batch = the largest divisor of effective_batch_size / num_devices that is <= found_max and sets trainer.accumulate_grad_batches accordingly. Because device_batch divides the per-device target exactly, the realized effective batch always equals effective_batch_size on any hardware — no drift.

found_max (the raw hardware probe) is written to pin_path; a later run reads it and skips the search, re-deriving device_batch/accumulation for that run's effective_batch_size/num_devices. retune=True forces a fresh search.

Parameters:

Name	Type	Description	Default
effective_batch_size	int	The pinned, hardware-independent quantity device_batch * accumulate_grad_batches * num_devices. Must be divisible by num_devices.	required
pin_path	str, Path, or None	Sidecar YAML storing found_max_device_batch. Defaults to <trainer.default_root_dir>/mushin_batch_pin.yaml. Inside a Hydra --multirun an explicit pin_path is required (the default dir is shared across jobs).	None
num_devices	int or None	Defaults to trainer.num_devices * trainer.num_nodes.	None
batch_arg	str	Attribute on the module (or datamodule) the tuner scales and this helper sets; forwarded to the tuner as batch_arg_name.	'batch_size'
retune	bool	Ignore any existing pin and search again.	False
**scale_kwargs		Forwarded to Tuner.scale_batch_size (e.g. mode, steps_per_trial).	{}

Returns:

Type	Description
BatchPin

tune_learning_rate

tune_learning_rate(
    trainer,
    module,
    datamodule=None,
    *,
    pin_path=None,
    lr_attr="lr",
    retune=False,
    **lr_find_kwargs,
)

Record-and-pin Lightning's LR finder.

Runs Tuner.lr_find once, writes the suggested learning rate to pin_path, and sets module.<lr_attr>. A later call reads the pin and skips the (stochastic) range test; retune=True forces a fresh search. Learning rate is hardware-independent, so there is no device math.

Parameters:

Name	Type	Description	Default
pin_path	str, Path, or None	Sidecar YAML. Defaults to <trainer.default_root_dir>/mushin_lr_pin.yaml.	None
lr_attr	str	Attribute on the module set to the found learning rate.	'lr'
retune	bool	Ignore any existing pin and search again.	False
**lr_find_kwargs		Forwarded to Tuner.lr_find (e.g. min_lr, max_lr, num_training).	{}