Version 3.0.1¶
Version 3.0.1 fixes a large time and memory regression in low-precision Euclidean distances, including the default distance used by Contrastive and TripletMargin losses. Public signatures and minimum supported Python, PyTorch, and NumPy versions are unchanged. No migration is required.
Low-precision Euclidean distances¶
In v3.0.0, float16 and BF16 matrices always used the direct torch.cdist
kernel to preserve nearby-point accuracy. Its CUDA backward pass can allocate
memory proportional to queries × references × embedding dimensions.
Large eager Euclidean calls now use float32 matrix evaluation when a numerical screen shows that it is safe. Self comparisons reuse prepared inputs and keep their diagonal exactly zero. Nearby or duplicate points, other norms, small matrices, compiled execution, CUDA graph capture, and reduced float32 matmul precision retain the accurate direct calculation. Global precision settings are unchanged.
The improvement is measured on standard-normal inputs; clustered training embeddings can still take the slower direct path. Values and gradients are checked for random, nearby, duplicate, and mixed-scale float16/BF16 inputs.
Measured impact¶
The fix was measured at commit 32bcfa8, whose numerical code is included
unchanged in this release. Each successful timing configuration averages
900 forward/backward pairs across three fresh processes after warmup.
At BF16 batch 2,048, 128 dimensions, and four samples per class:
| GPU | Loss | v3.0.0 → fix time per pair (ms) | v3.0.0 → fix peak allocation (MiB) |
|---|---|---|---|
| RTX 5090 | Contrastive | 13.3 → 2.23 | 2,118 → 356 |
| RTX 5090 | Batch-hard + Triplet | 16.4 → 1.38 | 2,118 → 151 |
| A100 | Contrastive | 23.4 → 6.32 | 2,118 → 308 |
| A100 | Batch-hard + Triplet | 30.0 → 5.47 | 2,118 → 103 |
These are loss/miner measurements, excluding an encoder, optimizer, and data loading. Peak allocation is total PyTorch tensor allocation, not total device usage. The fix does not beat v2.9.0 in every case: RTX 5090 batch-hard triplet still takes 25.6% longer in this configuration. See the complete comparison and limitations for process ranges, controls, numerical checks, and reproduction instructions.
Broader benchmark coverage¶
The new five-GPU benchmark report compares released v3.0.0 with v2.9.0 on RTX 4090, RTX 5090, L40S, A100, and H100 systems. It covers seven common losses plus batch-hard triplet mining, absolute forward and backward time, total and incremental peak memory, and verified batch limits. CSV tables, SVG scaling charts, raw measurements, and hardware/source provenance are included. Those original release results remain separate from the two-GPU fix measurements.