[aw471]

Summary

Compiled code no longer chunks its inputs automatically, so one graph serves any batch size. Before, a graph recorded the number of chunks, and a TransformerBlock with a dynamic batch size failed with ConstraintViolationError above 65,535 batch entries. Without torch.compile, chunking is unchanged.

Description of the change

To bound peak memory on CUDA, three places split a large input into chunks and run them one after another:

• split_size in sdm/_memory.py, which the numerical processors ClipSigma, RobustScale and RankGaussian use,
• TransformerBlock.forward, unless an integer batch_size_limit is passed,
• CellEmbedding.forward with batch_size_limit="auto".

The number of chunks depends on the size of the input and on the free memory of the device. torch.compile unrolls the loop over the chunks, so a graph holds a fixed number of them and guards on the sizes that give this number. For a batch size that is marked dynamic, the guard contradicts the mark:

- 0/0: 2 <= query.size()[0] <= 65535  # device=query.device,  # sdm/nn/attention.py:752 in forward
torch.fx.experimental.symbolic_shapes.ConstraintViolationError: Constraints violated (L['query'].size()[0])!

The three places now take one chunk while Dynamo traces them, as they already do while gradients are enabled. An integer batch_size_limit is still honoured in compiled code.

A TransformerBlock without a limit chunked at 65,535 batch entries for the attention kernel. It no longer has to in compiled code: the attention operator of the previous change in this stack slices such a batch when it runs.

The check is torch.compiler.is_dynamo_compiling(). It is False in ordinary Python and True only in code that Dynamo traces. torch.compiler.is_compiling() reads a process-wide flag instead, so it would also switch chunking off for uncompiled calls in other threads while a compile runs.

Cost. Compiled code takes one chunk, so the chunking no longer bounds its peak memory. A caller with inputs that needed several chunks splits them or uses uncompiled predict. For queries of up to 5,000 rows, a graph compiled for one query size has the same speed and peak memory as on main: at 5,000 rows it takes 77.6 ms (4 columns, fp32) and 70.5 ms (50 columns, bf16) on this stack, against 78.2 ms and 71.2 ms on main (small Kumo Tabular model, RTX PRO 6000, PyTorch 2.14.1, 3,000 context rows, 4 estimators). Larger queries were not measured for this change alone.

Reproduction (CPU, PyTorch 2.14.1):

import sys

import torch
from torch._dynamo.utils import counters

from sdm.nn import TransformerBlock

module = TransformerBlock(channels=8, num_query_heads=2, mlp=torch.nn.Identity()).eval()
compiled = torch.compile(module, fullgraph=True, dynamic=False, backend=sys.argv[1])
with torch.inference_mode():
    for call, batch in enumerate((3, 70_000, 9, 140_000, 2)):
        query = torch.randn(batch, 4, 8)
        torch._dynamo.mark_dynamic(query, 0)
        print(f"BEGIN call={call} batch={batch}", flush=True)
        try:
            out = compiled(query=query, batch_size_limit="auto")
            expected = module(query=query, batch_size_limit="auto")
            result = f"matches_eager={torch.allclose(out, expected, atol=1e-6)}"
        except Exception as e:
            result = f"FAILED {type(e).__name__}"
        print(f"END call={call} graphs={counters['stats']['unique_graphs']} {result}", flush=True)

Validation

Ran the script above with TORCH_LOGS="recompiles" and the eager and inductor backends, before and after this change (BEGIN/END lines are printed by the script). Both backends give the same result.

Before:

BEGIN call=0 batch=3
END call=0 graphs=1 matches_eager=True
BEGIN call=1 batch=70000
[__recompiles]     - 0/0: 2 <= query.size()[0] <= 65535
END call=1 graphs=2 FAILED ConstraintViolationError
BEGIN call=2 batch=9
END call=2 graphs=2 matches_eager=True
BEGIN call=3 batch=140000
[__recompiles]     - 0/0: 2 <= query.size()[0] <= 65535
END call=3 graphs=3 FAILED ConstraintViolationError
BEGIN call=4 batch=2
END call=4 graphs=3 matches_eager=True

After, no recompile line in the log:

BEGIN call=0 batch=3
END call=0 graphs=1 matches_eager=True
BEGIN call=1 batch=70000
END call=1 graphs=1 matches_eager=True
BEGIN call=2 batch=9
END call=2 graphs=1 matches_eager=True
BEGIN call=3 batch=140000
END call=3 graphs=1 matches_eager=True
BEGIN call=4 batch=2
END call=4 graphs=1 matches_eager=True

• test_transformer_block_compile_dynamic_batch fails without the change (ConstraintViolationError) and passes with it.
• FULL_TEST=1 pytest test/nn: 149 passed.
• CUDA (RTX PRO 6000, PyTorch 2.14.1, inductor backend): the script gives the same result on CUDA, ConstraintViolationError at 70,000 and 140,000 before and 1 graph after. FULL_TEST=1 pytest -m cuda test/nn: 125 passed.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.