Skip to content

constants

constants

Constants and enums for training configuration.

PrecisionType

Bases: StrEnum

PyTorch Lightning precision types.

See: https://lightning.ai/docs/pytorch/stable/common/trainer.html#precision

get_model_dtype 

get_model_dtype()

Get the dtype to convert model parameters to for this precision type.

For mixed precision types, converts to the lower precision dtype to avoid dtype mismatch errors during inference (e.g., bfloat16 input vs float32 bias).

Returns:

Type Description
dtype

torch.dtype to convert model to

Raises:

Type Description
NotImplementedError

For INT8 precision (requires specialized quantization)
Source code in src/versatil/training/constants.py 
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
def get_model_dtype(self) -> torch.dtype:
    """Get the dtype to convert model parameters to for this precision type.

    For mixed precision types, converts to the lower precision dtype to avoid
    dtype mismatch errors during inference (e.g., bfloat16 input vs float32 bias).

    Returns:
        torch.dtype to convert model to

    Raises:
        NotImplementedError: For INT8 precision (requires specialized quantization)
    """
    if self == PrecisionType.INT8:
        raise NotImplementedError(
            "INT8 precision requires post-training quantization. "
            "Use the versatil.quantization module or the post_training_compress endpoint."
        )
    dtype_map = {
        PrecisionType.FP32: torch.float32,
        PrecisionType.FP16_MIXED: torch.float16,
        PrecisionType.BF16_MIXED: torch.bfloat16,
        PrecisionType.FP16_TRUE: torch.float16,
        PrecisionType.BF16_TRUE: torch.bfloat16,
        PrecisionType.FP64: torch.float64,
    }
    return dtype_map[self]

is_mixed 

is_mixed()

Check if this precision type autocasts compute around float32 weights.

Returns:

Type Description
bool

True for the mixed half-precision types, where trainable parameters
bool

must stay in float32 storage so optimizer updates are not rounded
bool

away by the low-precision dtype.
Source code in src/versatil/training/constants.py 
68
69
70
71
72
73
74
75
76
77
78
79
def is_mixed(self) -> bool:
    """Check if this precision type autocasts compute around float32 weights.

    Returns:
        True for the mixed half-precision types, where trainable parameters
        must stay in float32 storage so optimizer updates are not rounded
        away by the low-precision dtype.
    """
    return self in (
        PrecisionType.FP16_MIXED,
        PrecisionType.BF16_MIXED,
    )

autocast 

autocast(device_type)

Return an autocast context matching this precision.

Enabled only for the mixed half-precision types, where forward passes outside the Lightning training loop must reproduce the training-time autocast over mixed float32/low-precision parameters. For all other precisions the returned context is a no-op.

Parameters:

Name Type Description Default
device_type str

Device type string for torch.autocast (e.g. "cuda" or "cpu").

 required
Source code in src/versatil/training/constants.py 
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
def autocast(self, device_type: str) -> torch.autocast:
    """Return an autocast context matching this precision.

    Enabled only for the mixed half-precision types, where forward passes
    outside the Lightning training loop must reproduce the training-time
    autocast over mixed float32/low-precision parameters. For all other
    precisions the returned context is a no-op.

    Args:
        device_type: Device type string for ``torch.autocast`` (e.g.
            ``"cuda"`` or ``"cpu"``).
    """
    return torch.autocast(
        device_type=device_type,
        dtype=self.get_model_dtype() if self.is_mixed() else None,
        enabled=self.is_mixed(),
    )

should_convert_model 

should_convert_model()

Check if model should be converted to a specific dtype for this precision.

Returns:

Type Description
bool

True if model should be converted (for mixed/true half precision types)
Source code in src/versatil/training/constants.py 
 99
100
101
102
103
104
105
106
107
108
109
110
def should_convert_model(self) -> bool:
    """Check if model should be converted to a specific dtype for this precision.

    Returns:
        True if model should be converted (for mixed/true half precision types)
    """
    return self in (
        PrecisionType.FP16_MIXED,
        PrecisionType.BF16_MIXED,
        PrecisionType.FP16_TRUE,
        PrecisionType.BF16_TRUE,
    )

CheckpointFilename

Bases: StrEnum

Standard filenames within a training checkpoint directory.

CheckpointKey

Bases: StrEnum

Keys used within checkpoint state dicts and normalizer params.

CompileMode

Bases: StrEnum

torch.compile optimization modes.

See: https://pytorch.org/docs/stable/generated/torch.compile.html

Float32MatmulPrecision

Bases: StrEnum

Float32 matrix multiplication precision for Tensor Cores.

Controls the precision of float32 matrix multiplications on GPUs with Tensor Cores.

See: https://pytorch.org/docs/stable/generated/torch.set_float32_matmul_precision.html