diff --git a/CHANGELOG.md b/CHANGELOG.md index 1b49609..7372d53 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,6 +2,18 @@ ## Unreleased +### Chunked soft RankIC + +- `RankICLoss` computes soft-rank columns in blocks above `chunk_size` (default + 256) and uses non-reentrant activation checkpointing in training. Pairwise + temporary matrices grow as N × chunk_size rather than N², and backward + recomputes blocks instead of retaining every sigmoid activation. The pairwise + arithmetic remains quadratic and recomputation can cost more runtime. +- `chunk_size=None` retains the dense reference. The same soft-rank objective + supports both-input and higher-order gradients; reduction rounding can differ + across blocks/devices. Inputs must be aligned one-dimensional cross-sections. + This is memory scaling support, not predictive-quality or speedup evidence. + ### Purged date validation - `TrainConfig.purge_dates` removes complete observed signal dates immediately diff --git a/docs/backtest_assumptions.md b/docs/backtest_assumptions.md index 4affb02..be13715 100644 --- a/docs/backtest_assumptions.md +++ b/docs/backtest_assumptions.md @@ -76,7 +76,12 @@ sample/date counts and numeric date-ID ranges so the chosen split can be inspect The default `batch_mode="auto"` uses full daily cross-sections for `ICLoss` and `RankICLoss`, with one optimizer step per date and equally weighted daily losses. `batch_size` is ignored for daily batching: memory depends on the largest daily -cross-section, and soft RankIC still requires quadratic pairwise rank memory. +cross-section. Soft RankIC uses `chunk_size=256` by default above that many +stocks, with N × chunk_size pairwise matrices and activation checkpointing during +training. `RankICLoss(chunk_size=None)` retains the dense reference. Pairwise work +remains quadratic and backward recomputation adds runtime; reduction rounding +can differ across blocks/devices. This does not bound model activations or total +training memory, and no GPU performance estimate is provided here. Use `batch_mode="date"` for a custom daily objective. Ordinary regression uses flat sample batches and sample-weighted epoch loss; it retains the final partial batch. Explicit `val_split_mode="random"` splits whole dates in daily mode and diff --git a/src/mlquant/models/losses.py b/src/mlquant/models/losses.py index 7a6a4a1..710a7f1 100644 --- a/src/mlquant/models/losses.py +++ b/src/mlquant/models/losses.py @@ -24,6 +24,7 @@ import torch import torch.nn as nn +from torch.utils.checkpoint import checkpoint class AdjMSELoss(nn.Module): @@ -88,18 +89,43 @@ class RankICLoss(nn.Module): Spearman = Pearson on ranks. We use the *soft-rank* trick (a differentiable approximation that sorts the inputs through a softmax-temperature lens) so the loss is gradient-friendly. + + Above ``chunk_size`` stocks, compute columns in blocks and checkpoint their + activations. Pairwise work remains quadratic; temporary matrices scale as + N * chunk_size. ``chunk_size=None`` restores the dense reference path. """ - def __init__(self, temperature: float = 1.0) -> None: + def __init__(self, temperature: float = 1.0, *, chunk_size: int | None = 256) -> None: super().__init__() if not math.isfinite(temperature) or temperature <= 0: raise ValueError("temperature must be finite and positive") + if chunk_size is not None and (type(chunk_size) is not int or chunk_size < 1): + raise ValueError("chunk_size must be a positive integer or None") self.temperature = temperature + self.chunk_size = chunk_size def forward(self, pred: torch.Tensor, target: torch.Tensor) -> torch.Tensor: + if pred.shape != target.shape or not pred.numel(): + raise ValueError("correlation inputs must have the same nonempty shape") + if pred.ndim != 1: + raise ValueError("RankIC inputs must be one-dimensional cross-sections") return -_pearson(self._soft_rank(pred), self._soft_rank(target)) + @staticmethod + def _rank_columns(full, columns, temperature): + return torch.sigmoid((columns.unsqueeze(0) - full.unsqueeze(1)) / temperature).sum(0) + def _soft_rank(self, x: torch.Tensor) -> torch.Tensor: - # Differentiable rank: softmax over pairwise differences. - diffs = (x.unsqueeze(0) - x.unsqueeze(1)) / self.temperature - return torch.sigmoid(diffs).sum(dim=0) + if self.chunk_size is None or x.numel() <= self.chunk_size: + return self._rank_columns(x, x, self.temperature) + ranks = [] + for columns in x.split(self.chunk_size): + if torch.is_grad_enabled() and x.requires_grad: + # Non-reentrant checkpointing preserves autograd.grad and higher + # derivatives while recomputing one block during backward. + rank = checkpoint(self._rank_columns, x, columns, self.temperature, + use_reentrant=False, preserve_rng_state=False) + else: + rank = self._rank_columns(x, columns, self.temperature) + ranks.append(rank) + return torch.cat(ranks) diff --git a/tests/test_chunked_rankic.py b/tests/test_chunked_rankic.py new file mode 100644 index 0000000..6ce35b3 --- /dev/null +++ b/tests/test_chunked_rankic.py @@ -0,0 +1,89 @@ +"""Chunked soft rank keeps its objective/derivatives with bounded activations.""" +from unittest.mock import patch + +import pytest +import torch + +from mlquant.models.losses import RankICLoss + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.float64]) +@pytest.mark.parametrize("temperature", [.05, 1., 7.]) +@pytest.mark.parametrize("chunk", [1, 4, 9]) +def test_chunked_loss_and_both_gradients_match_dense_with_ties_and_strides(dtype, temperature, chunk): + values = torch.tensor([.1, 9., -.7, 9., .1, 9., .9, 9., 1.2, 9., -.4, 9., .5, 9., -.2, 9., 2.1, 9.], dtype=dtype) + pred = values[::2].requires_grad_() + target = torch.tensor([.6, -.1, .6, .3, -.7, .2, .9, -.4, .8], dtype=dtype, requires_grad=True) + assert not pred.is_contiguous() + dense = RankICLoss(temperature, chunk_size=None)(pred, target) + actual = RankICLoss(temperature, chunk_size=chunk)(pred, target) + expected_grads = torch.autograd.grad(dense, (pred, target), retain_graph=True) + actual_grads = torch.autograd.grad(actual, (pred, target)) + tolerance = 2e-6 if dtype == torch.float32 else 1e-12 + torch.testing.assert_close(actual, dense, atol=tolerance, rtol=tolerance) + for actual_grad, expected_grad in zip(actual_grads, expected_grads): + torch.testing.assert_close(actual_grad, expected_grad, atol=tolerance, rtol=tolerance) + + +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32, torch.float64]) +@pytest.mark.parametrize("constant", ["prediction", "target", "both"]) +def test_constant_chunked_batches_keep_zero_gradients(dtype, constant): + pred = torch.ones(17, dtype=dtype) if constant != "target" else torch.arange(17, dtype=dtype) + target = torch.ones(17, dtype=dtype) if constant != "prediction" else torch.arange(17, dtype=dtype) + pred.requires_grad_() + target.requires_grad_() + loss = RankICLoss(chunk_size=4)(pred, target) + loss.backward() + assert loss.item() == 0 + assert pred.grad.eq(0).all() and target.grad.eq(0).all() + + +def test_checkpointed_rankic_supports_first_and_second_derivatives(): + pred = torch.tensor([.4, -.2, .7, 1.1], dtype=torch.float64, requires_grad=True) + target = torch.tensor([-.1, .3, .6, .8], dtype=torch.float64, requires_grad=True) + loss = RankICLoss(.8, chunk_size=2) + assert torch.autograd.gradcheck(loss, (pred, target)) + assert torch.autograd.gradgradcheck(loss, (pred, target)) + + +def test_large_cross_section_never_builds_a_dense_pairwise_matrix_or_saves_quadratic_activations(): + n, chunk = 513, 32 + pred = torch.linspace(-2, 2, n, requires_grad=True) + target = torch.sin(torch.arange(n).float()) + saved, shapes = {}, [] + real_sigmoid = torch.sigmoid + + def keep(tensor): + storage = tensor.untyped_storage() + saved[storage.data_ptr()] = storage.nbytes() + return tensor + + def observe(values): + shapes.append(tuple(values.shape)) + return real_sigmoid(values) + + with patch("torch.sigmoid", side_effect=observe): + with torch.autograd.graph.saved_tensors_hooks(keep, lambda tensor: tensor): + loss = RankICLoss(chunk_size=chunk)(pred, target) + retained_bytes = sum(saved.values()) + loss.backward() + assert all(rows == n and columns <= chunk for rows, columns in shapes) + assert len(shapes) > 2 * ((n + chunk - 1) // chunk) # Backward recomputes blocks. + assert retained_bytes < 64 * n * pred.element_size() + assert pred.grad.isfinite().all() + saved.clear() + with torch.autograd.graph.saved_tensors_hooks(keep, lambda tensor: tensor): + dense = RankICLoss(chunk_size=None)(pred, target) + assert sum(saved.values()) >= n * n * pred.element_size() + torch.testing.assert_close(loss, dense, atol=2e-6, rtol=2e-6) + + +@pytest.mark.parametrize("bad", [0, -1, True, 1.5]) +def test_invalid_chunk_size_is_rejected(bad): + with pytest.raises(ValueError, match="positive integer or None"): + RankICLoss(chunk_size=bad) + + +def test_rankic_rejects_multidimensional_cross_sections(): + with pytest.raises(ValueError, match="one-dimensional"): + RankICLoss()(torch.ones(2, 3), torch.ones(2, 3))