Metadata-Version: 2.4
Name: memeff
Version: 0.3.1
Summary: Memory-efficient contrastive losses (CLIP, LiT, Qwen3-Embedding InfoNCE) with Triton kernels
Author: Mikhail Kindulov
License: Apache-2.0
Project-URL: Homepage, https://github.com/b0nce/MemoryEfficientCLIP
Project-URL: Repository, https://github.com/b0nce/MemoryEfficientCLIP
Keywords: clip,contrastive-learning,triton,deep-learning,loss
Classifier: Development Status :: 4 - Beta
Classifier: Intended Audience :: Science/Research
Classifier: License :: OSI Approved :: Apache Software License
Classifier: Programming Language :: Python :: 3
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Classifier: Environment :: GPU :: NVIDIA CUDA
Requires-Python: >=3.9
Description-Content-Type: text/markdown
License-File: LICENSE
Requires-Dist: torch>=2.0
Dynamic: license-file

# memeff — Memory-Efficient Contrastive Losses

Triton-kernel implementations of CLIP, LiT, and Qwen3-Embedding (InfoNCE) contrastive losses that never materialize the B×B similarity matrix, on a single GPU or sharded across a DDP ring. Batch sizes of 300k+ work in practice. Numerically validated against dense autograd references on A100 (sm80), H100 (sm90), and B200 (sm100).

## Installation

```bash
pip install memeff
```

Requires PyTorch ≥ 2.0, Triton ≥ 3.0 (ships with the Linux torch wheels), and a CUDA GPU. Inputs may be fp32, fp16, or bf16 — denominators and gradients always accumulate in fp32.

## Quickstart

```python
import torch
from memeff import MemoryEfficientCLIPLoss

clip_loss = MemoryEfficientCLIPLoss(temperature=0.07)

batch_size, dim = 2 ** 15, 1152
image_features = torch.randn(batch_size, dim, device="cuda", requires_grad=True)
text_features = torch.randn(batch_size, dim, device="cuda", requires_grad=True)

loss = clip_loss(image_features, text_features)
loss.backward()
```

All losses come in a single-GPU and a `Distributed*` (multi-GPU DDP) variant:

| Loss | Single GPU | Multi-GPU DDP |
|---|---|---|
| CLIP (bidirectional) | `MemoryEfficientCLIPLoss` | `DistributedMemoryEfficientCLIPLoss` |
| LiT (locked image tower) | `MemoryEfficientLiTLoss` | `DistributedMemoryEfficientLiTLoss` |
| Qwen3 InfoNCE (both towers) | `MemoryEfficientQwen3Loss` | `DistributedMemoryEfficientQwen3Loss` |
| Qwen3 InfoNCE (locked docs) | `MemoryEfficientLiTQwen3Loss` | `DistributedMemoryEfficientLiTQwen3Loss` |
| Matryoshka CLIP (fused) | `MemoryEfficientMatryoshkaCLIPLoss` | `DistributedMemoryEfficientMatryoshkaCLIPLoss` |
| Matryoshka LiT (fused) | `MemoryEfficientMatryoshkaLiTLoss` | `DistributedMemoryEfficientMatryoshkaLiTLoss` |
| Matryoshka Qwen3 (fused, both towers) | `MemoryEfficientMatryoshkaQwen3Loss` | `DistributedMemoryEfficientMatryoshkaQwen3Loss` |
| Matryoshka Qwen3 (fused, locked docs) | `MemoryEfficientMatryoshkaLiTQwen3Loss` | `DistributedMemoryEfficientMatryoshkaLiTQwen3Loss` |
| Matryoshka, any loss above (eager) | `MatryoshkaLoss(base_loss, dims)` | same wrapper |

## Options (all modules)

- `temperature` — softmax temperature (default 0.07).
- `normalized_inputs=True` — skip the internal L2 normalization.
- `stable=True` (default) — rescale gradients by `sqrt(batch / temperature)` instead of `1 / temperature` to avoid fp32 underflow at very large batches (loss value unchanged; see details below). Pass `stable=False` for the textbook `1 / temperature` gradient scale.
- `tau_plus > 0` — the debiased contrastive loss of [Chuang et al., 2020](https://arxiv.org/abs/2007.00224). `forward` also accepts a per-call override: a float or a **(batch,) tensor of per-row priors** — useful when some samples are known to have approximate copies in the dataset. Costs no extra kernels.
- `label_smoothing > 0` — smoothed softmax targets with the `F.cross_entropy` convention (`(1-eps)` on the positive plus `eps/C` uniform). Costs only O(batch·dim) eager math.

All options compose with each other and with the distributed variants.

<details>
<summary><b>LiT loss (locked image tower)</b></summary>

```python
from memeff import MemoryEfficientLiTLoss

lit_loss = MemoryEfficientLiTLoss(temperature=0.07)
# text first; the image tower is locked and receives no gradient.
loss = lit_loss(text_features, image_features)
```

</details>

<details>
<summary><b>Qwen3 loss (embedding-model InfoNCE)</b></summary>

The improved InfoNCE objective of the Qwen3-Embedding report: an asymmetric query→document row softmax with a false-negative mask (any negative whose similarity exceeds `s(q_i, d_i) + margin` is dropped as a presumed unlabeled positive), row-specific hard negatives, and optional q-q / d-d in-batch negatives.

```python
from memeff import MemoryEfficientQwen3Loss

qwen3_loss = MemoryEfficientQwen3Loss(
    temperature=0.05,
    margin=0.1,               # margin >= 2 disables the false-negative mask
    use_qq_negatives=True,    # queries repel other queries      (optional)
    use_dd_negatives=True,    # documents repel other documents  (optional)
)

# Row-specific hard negatives (batch, K, dim): each query competes only against its own K.
loss = qwen3_loss(query_features, doc_features, hard_negatives)  # hard_negatives optional
```

For a locked document tower (precomputed corpus embeddings) use `MemoryEfficientLiTQwen3Loss`: only the queries receive gradients, and there is no d-d option (with locked documents its repulsion gradient has nowhere to land).

</details>

<details>
<summary><b>Distributed (multi-GPU DDP) usage</b></summary>

One process per GPU (e.g. `torchrun`). Each rank passes **only its shard** of the global batch; `forward` returns this rank's contribution to the loss, and `backward` fills the shard's full gradient of the *global* loss. Hard negatives (Qwen3) stay on their query's rank; per-row `tau_plus` tensors cover the local shard.

```python
import torch.distributed as dist
from memeff import DistributedMemoryEfficientCLIPLoss

dist.init_process_group("nccl")
torch.cuda.set_device(dist.get_rank())

clip_loss = DistributedMemoryEfficientCLIPLoss(temperature=0.07)
partial_loss = clip_loss(image_shard, text_shard)
partial_loss.backward()

global_loss = partial_loss.detach().clone()
dist.all_reduce(global_loss)   # logging only
```

Peak memory: the distributed CLIP and Qwen3 losses keep the assembled travelling tower(s) for the backward — O(global_batch × dim) per rank. The distributed LiT losses re-stream the ring in backward instead, staying at O(local_batch × dim). The full B×B similarity matrix is never materialized anywhere; communication is O(batch) against O(batch²) compute.

</details>

<details>
<summary><b>Matryoshka representation learning (MRL)</b></summary>

Nested-prefix training ([Kusupati et al., 2022](https://arxiv.org/abs/2205.13147)): each prefix `x[..., :m]` is re-normalized and trained with its own full contrastive loss, so at inference the embedding can be truncated to any of the trained dims.

Two implementations:

```python
from memeff import (MatryoshkaLoss, MemoryEfficientQwen3Loss,
                    MemoryEfficientMatryoshkaQwen3Loss)

# 1) Eager wrapper: works with EVERY memeff loss (incl. the DDP variants),
#    any kernel-legal dims. K separate loss passes.
loss_fn = MatryoshkaLoss(MemoryEfficientQwen3Loss(temperature=0.05),
                         dims=(64, 128, 256, 384))

# 2) Fused (CLIP/LiT/Qwen3 families): all K dims in one pass over the
#    similarity blocks. Extra state is O(K * batch) scalar tables -- no
#    per-dim feature copies, no per-dim gradient buffers.
loss_fn = MemoryEfficientMatryoshkaQwen3Loss(
    dims=(64, 128, 256, 384),   # strictly increasing, multiples of 64,
    weights=None,               # ending exactly at d_model (else ValueError)
    temperature=0.05, stable=True, tau_plus=1e-4)
```

The fused `Distributed*Matryoshka*` variants keep the one-pass property across the DDP ring: the towers travel **once** for all K dims (the wrapper re-runs the whole ring per dim), each hop feeds the fused denominator kernel, and the traveling blocks' prefix-norm tables are recomputed on arrival instead of communicated. Backward is one rectangular launch of the same tile kernels over local rows × global columns, plus the plain losses' reduce-scatter for column gradients and one O(K·batch) all-reduce for the re-normalization correction sums. The distributed MRL LiT losses keep the never-assemble contract: peak memory O(local_batch × dim), no gradient communication.

The fused kernels exploit two facts: raw prefix dots are cumulative across feature chunks, and prefix re-normalization is a per-row scalar — so the forward snapshots every dim's denominator in one sweep, and the backward telescopes a per-pair coefficient tile through two chunk walks. All options (`margin`, `stable`, `tau_plus` incl. per-row, `label_smoothing`, hard negatives, q-q/d-d) compose per dim.

Honest numbers (A100-PCIE-40GB, bf16, B=65536, D=384, dims 64/128/256/384, fwd+bwd): plain loss 507 ms, fused MRL 981 ms, eager wrapper 1158 ms. The backward picks between two kernels per ladder: single-walk prefix emission (each boundary's exp2/mask sweep runs once, extra prefix matmuls on cache-hot chunks — wins at small `d_model`) and a telescoping two-walk (minimal matmuls — wins at large `d_model`; at D=1024 with a dense 7-dim ladder it is 1.35x faster than prefix emission and 1.43x faster than the wrapper). Memory: the fused loss adds only O(K·batch) scalars to the loss state (the bench peaks are dominated by the fp32 gradient buffers, identical asymptotics to the non-MRL losses).

Against the naive dense implementation (materialize the B×B similarity per dim, torch autograd; same loss semantics, same GPU/config, loss-only peaks):

| batch | naive ms / GiB | fused ms / GiB | wrapper ms / GiB |
|---|---|---|---|
| 4,096 | 9.4 / 0.48 | **6.9** / 0.10 | 9.4 / 0.07 |
| 8,192 | 30.1 / 1.82 | **14.6** / 0.18 | 16.2 / 0.13 |
| 16,384 | 90.0 / 7.11 | 63.8 / 0.35 | **63.3** / 0.24 |
| 32,768 | 354.2 / 28.21 | **262.3** / 0.68 | 275.1 / 0.46 |
| 65,536 | OOM (>40 GiB) | **981** / 1.33 | 1146 / 0.89 |
| 131,072 | OOM (>40 GiB) | **4196** / 2.65 | 4708 / 1.77 |

Validated against the per-dim dense reference across the full feature matrix in fp32 (≤1e-4, mostly ≤1e-6) and bf16 (≤6e-3): `python test_mrl_qwen3_loss.py`.

</details>

<details>
<summary><b>Stable gradient rescaling (<code>stable=True</code>)</b></summary>

`stable=True` (the default) rescales the gradient by `sqrt(batch / temperature)` instead of `1 / temperature`: the textbook `1 / (batch * temperature)` factor can nullify small values even in fp32, which matters at large batch sizes (300k+ works fine in practice). The loss value is unchanged, only the gradient scale differs, so the learning rate becomes batch-size dependent — use `lr / sqrt(batch * temperature)` to mimic `stable=False`, though at large batches standard values like 1e-4 tend to work well without that correction. The old `StableMemoryEfficientCLIPLoss` / `StableMemoryEfficientLiTLoss` classes remain as deprecated aliases.

</details>

<details>
<summary><b>Debiased contrastive loss (<code>tau_plus</code>)</b></summary>

`tau_plus > 0` switches to the debiased objective of [Chuang et al., 2020](https://arxiv.org/abs/2007.00224): with probability `tau_plus` an in-batch "negative" is actually an unlabeled positive, so each softmax denominator's negative sum `sum_neg` is replaced by `N * g` with

```
g = max((sum_neg / N - tau_plus * pos) / (1 - tau_plus), e^(-1/temperature))
```

where `N` is the nominal negative count and `pos` the positive exponential (the paper's estimator with M = 1). The clamp keeps the estimate at its theoretical minimum; rows where it fires push no gradient into their negatives.

```python
clip_loss = MemoryEfficientCLIPLoss(temperature=0.07, tau_plus=0.1)

# per-row priors: rows with known approximate copies get a higher prior
tau_row = duplicate_rates            # (batch,) tensor, values in [0, 1)
loss = clip_loss(image_features, text_features, tau_plus=tau_row)
```

Each row's denominator is debiased independently, so per-row priors slot straight into the estimator; for CLIP, sample i's prior applies to both its row and its column softmax. On the Qwen3 losses debiasing composes with the false-negative mask: masked entries contribute zero to the negative mean but keep their slot in the nominal count `N = (B-1)(1 + qq + dd) + K`.

Debiasing reuses every kernel untouched — the debiased denominator is a per-row transform of quantities the kernels already produce, and the gradient change rides the per-row divisor vector plus the eager positive-pair seed (`debias_denominators` in `_common.py`). Distributed: the CLIP column transform all-gathers the positive exponentials (and the prior vector, if per-row) — O(batch); the row-only losses debias with no communication.

</details>

<details>
<summary><b>Label smoothing</b></summary>

`label_smoothing=eps` smooths the softmax targets with the `F.cross_entropy` convention: `(1-eps)` on the positive plus `eps/C` uniform over the C candidates. For the Qwen3 losses the candidate set is the nominal one (positive + all nominal negatives); the false-negative mask does not reshape the target — masked entries are presumed positives, and a sliver of attraction toward them is the point of smoothing.

```python
clip_loss = MemoryEfficientCLIPLoss(temperature=0.07, label_smoothing=0.1)
```

Because the smoothed targets still sum to 1, the log-denominator cancels in the difference between the smoothed and unsmoothed loss, which collapses to batch-level scalars (the diagonal trace and dot products of summed embeddings). The correction is therefore a plain differentiable eager term added outside the kernels — O(batch·dim) math, one O(dim) all-reduce in DDP — and composes mechanically with `stable` (the term's gradient is rescaled to match) and `tau_plus`.

Note: `label_smoothing` and `tau_plus` push in related directions (both address "some negatives aren't really negatives") — smoothing adds uniform attraction to all negatives, debiasing removes estimated false-negative repulsion. They compose, but don't stack both at full strength blindly; if the motivation is known duplicates, per-row `tau_plus` is the more targeted tool.

</details>

<details>
<summary><b>Implementation details</b></summary>

**Code layout.** Shared constants, distributed helpers, the debias transform, the smoothing terms, and the masked Qwen3 kernel pair with its ring assembly live in `memeff/_common.py`. The single-GPU CLIP and LiT losses reuse the kernels of their distributed counterparts with the whole batch as one block, so there is exactly one implementation of each kernel.

**CLIP.** Two kernels in `distributed_clip_loss.py`: `clip_denom_kernel` accumulates row and column sum-exp block-wise without materializing the similarity matrix; `clip_grad_both_kernel` recomputes each block once in backward and emits both gradient directions from that pass. The loss value is formed directly in the log2 domain (no exp/log round trip), with a fixed maximum trick of `1/temperature` (numerically stable enough for the CLIP task).

For small enough `d_model` every single-GPU backward (CLIP, LiT, and the Qwen3 family) uses a FlashAttention-shaped kernel (`clip_fa_grad_kernel` in `clip_loss.py`, `qwen3_fa_grad_kernel` in `_common.py`): each program owns a block of gradient rows, streams the opposite tower, and accumulates in a `(rows, next_pow2(d_model))` fp32 register tile written out exactly once — no gradient atomics at all. The row block shrinks as `d_model` grows to keep the accumulator at 128 KB (64 rows up to 512 lanes, 32 rows up to 1024); each direction recomputes its similarities (the FlashAttention-2 dQ vs dK/dV trade), and the Qwen3 self-similarity blocks (q-q / d-d) fold both softmax roles of a pair into one streamed pass. Measured end-to-end (bf16, fwd+bwd) vs the atomic tile-grid kernels:

- **A100 (sm80)**, cap `d_model <= 1024`: CLIP 3.2x at B=65536/D=512 (640 -> 201 ms), 2.8x at D=384 (B=131072: 2071 -> 703 ms), 1.3x/1.6x at B=32768 D=768/1024; Qwen3 2.8x at B=65536/D=384 (514 -> 184 ms) and 3.5x with q-q + d-d negatives (1540 -> 436 ms); Qwen3-LiT 3.1-3.6x.
- **H100 (sm90)**, cap 1024, wider 64-column j-blocks above 512 lanes (fit in the larger shared memory): CLIP 2.6x at D=512, 1.3x/1.6x at 768/1024; Qwen3 2.0x at D=384 (2.4x with q-q + d-d), 1.8x at D=1024 with q-q + d-d; Qwen3-LiT 2.5-2.6x.
- **B200 (sm100)**, cap `d_model <= 512`: 3.3x for CLIP at D=512, 2.8x/3.7x for Qwen3 at D=384 — but above 512 lanes the register-bound FA shapes cannot feed Blackwell's faster tensor cores and the atomic kernels win, so the cap is lower.

Above the per-arch cap everything falls back to the atomic tile-grid kernels (at D=2048 the register budget forces blocks too thin to beat them on every measured arch). Verified on triton 3.1 (torch 2.5, cu124) and triton 3.3 (torch 2.7, cu128).

The fused matryoshka backwards use the same FA shape with one launch per prefix boundary (`mrl_fa_grad_kernel` in `mrl_qwen3_loss.py`): each launch sizes its dots by `next_pow2(m_k)`, folds `w_k` into the gradient scale, and read-modify-writes its boundary's gradient rows and rho row (the re-normalization correction sums) — no atomics anywhere, and the per-boundary restreams still beat the single-pass tile kernels because the boundary sweeps dominate. A fused single-launch variant that kept a prefix-masked B tile in registers ran 23x *slower* (the masked tile round-trips through shared memory at every boundary dot) — the boundary-launch split is the fast shape. Same per-arch caps and block tables as the plain FA kernels (the sweeps reproduced the plain winners on all three arches); fp32 inputs stage twice the shared memory and get their own narrower blocks. Measured vs the better tile kernel per case (bf16, fwd+bwd, B=65536/D=384 ladder 128/256/384 and B=32768/D=1024 ladder 256/512/1024): A100 2.0x base Qwen3, 2.3x with q-q + d-d, 1.5x CLIP, 1.4x LiT, and 1.3-1.5x at D=1024; H100 2.0x/2.4x/2.0x/2.4x, 1.2-1.5x at 1024; B200 2.5x/3.1x/2.4x/3.0x at D=384 (above its 512 cap the prefix tile kernel stays ahead, matching the plain-FA finding).

**LiT.** `lit_denom_kernel` / `lit_grad_kernel`: row-only sum-exp and a single output GEMM — the image tower is locked, so no image gradient and no column denominator.

**Distributed CLIP.** The global batch is sharded one contiguous slice per rank. Only the column (text) tower travels, rotating around a SigLIP-style ring of batched point-to-point transfers, each block prefetched one hop ahead so its transfer overlaps the current block's matmul. Row denominators complete locally; column denominators are one O(batch) all-reduce. In backward the local-row gradient is final and the column gradient reaches its owner via reduce-scatter. Backward tile sizes are selected per GPU architecture (A100 / H100 / B200, with a safe default elsewhere).

**Distributed LiT.** Same ring, specialized: each rank keeps its text rows home and streams the locked image features past them twice (denominator lap, gradient lap). No gradient communication, no assembled towers — peak memory stays O(local_batch × dim).

**Qwen3.** The softmax is row-only, so every negative group ((Q,D), (Q,Q), (D,D)) is another additive contribution to the same per-row fp32 denominator, and one masked kernel pair covers all passes (`qwen3_denom_kernel` / `qwen3_grad_kernel`, with an optional global diagonal exclusion). The mask recomputes identically in backward, so no B×B state is stored. The B×K hard-negative block is handled eagerly in fp32. The DDP variant runs on the same ring as the distributed CLIP loss: row denominators live entirely on the row's home rank, d-d rides the travelling document blocks for free, q-q makes the query tower travel too (doubling ring payload, one extra reduce-scatter), hard negatives never leave their rank.

**Distributed matryoshka.** `distributed_mrl_clip_loss.py` / `distributed_mrl_qwen3_loss.py` contain **no new kernels**: the single-GPU MRL tile kernels already take rectangular grids, separate row/column table strides, and row/column offsets, so the fused DDP variants are pure orchestration — the ring travels once for all K dims, each hop launches the fused denom kernel with that block's locally recomputed prefix-norm table, and the backward is one rectangular launch per similarity block with the plain losses' reduce-scatter plus an O(K·batch) rho all-reduce. The FA-shaped backward stays single-GPU (it wants the whole opposite tower streamed inside one launch).

</details>

<details>
<summary><b>Tests</b></summary>

Both test scripts compare losses and all gradients against dense autograd references (they need a CUDA GPU) and cover the `stable`, `tau_plus` (scalar and per-row), and `label_smoothing` variants and their combinations:

```bash
python test_clip_lit_loss.py                        # CLIP + LiT, single GPU
python test_qwen3_loss.py                           # Qwen3, single GPU
python test_mrl_qwen3_loss.py                       # matryoshka (wrapper + fused)
python test_mrl_clip_loss.py
torchrun --nproc-per-node=2 test_clip_lit_loss.py   # distributed variants
torchrun --nproc-per-node=2 test_qwen3_loss.py
torchrun --nproc-per-node=2 test_mrl_qwen3_loss.py  # distributed fused matryoshka
torchrun --nproc-per-node=2 test_mrl_clip_loss.py
```

The kernels run with ieee fp32 matmuls in the tests (`MEMEFF_INPUT_PRECISION=ieee`) so the comparison is not drowned in tensor-core rounding noise; production runs default to tf32.

</details>

## Citation

```
@misc{memory-efficient-clip-loss,
  author = {Mikhail Kindulov},
  title = {memeff: Memory Efficient CLIP Loss},
  year = {2025},
  publisher = {GitHub},
  url = {https://github.com/b0nce/MemoryEfficientCLIP}
}
```

## License

Apache License 2.0
