Cache NaFlex patch interpolation and expose variable-patch capabilities

[chenghuichen]

Motivation

While fixing variable patch-size integration in mlfoundations/open_clip#1195, I noticed a few generic improvements that belong in timm. In particular, every non-base patch-size forward rebuilt the resize basis and ran torch.linalg.pinv, even though training usually cycles through a small fixed set of patch sizes.

The existing PatchEmbedResamplerFixedOrigSize already had most of the required caching logic, but PatchEmbedInterpolator did not use it and the cached matrices were persistent buffers. Downstream integrations also had no reliable way to discover whether a NaFlex model supports variable patches, while eval transforms always flattened patches before the model could infer their spatial size.

Changes

Cached interpolation and prewarming

PatchEmbedInterpolator now shares a PatchEmbedResamplerFixedOrigSize between linear and convolutional weight resampling. Each matrix is computed once per target patch size and device, then kept in a runtime-only cache. Cache misses produce regular float32 tensors even under inference mode or autocast; device or dtype conversion clears stale entries. The cache adds no parameters or persistent buffers, so checkpoint keys remain unchanged; the standalone resample_patch_embed function and both linear patch layouts continue to behave as before.

The cache can be populated before torch.compile or another captured execution path:

model = timm.create_model(
    'naflexvit_base_patch16_gap',
    enable_patch_interpolator=True,
).to(device)
model.prewarm_patch_interpolator([(8, 8), (12, 12), (16, 16)])
model = torch.compile(model)

Prewarming should happen after moving the model to its execution device. The base patch size is a no-op. Each unique non-base patch size retains one float32 resampling matrix; there is intentionally no eviction policy because NaFlex training normally uses a small configured set of patch sizes.

Capability discovery and evaluation layout

NaFlexVit.supports_patch_interpolation reports the resolved capability: a linear projection with an enabled interpolator and no incompatible input normalization.

create_transform and transforms_imagenet_eval now accept patchify_flatten=False, returning spatial patches such as [N, Ph, Pw, C]. The default remains True, so existing transform calls retain their flattened output.

Benchmark

Both benchmarks ran on an NVIDIA A10 with PyTorch 2.11.0+cu128.

Patch-resampling cache

benchmark_patch_interpolator_cache.py compares a float32 channels-last projection (embed_dim=768, in_chans=3, base patch 16x16). Values are median latency per call.

Target patch	Previous uncached	Cache miss	Cache hit	Hit speedup	Cache size
8x8	2.319 ms	2.337 ms	0.428 ms	5.41x	64 KiB
12x12	5.986 ms	5.998 ms	0.857 ms	6.98x	144 KiB
32x32	12.973 ms	12.994 ms	3.442 ms	3.77x	1024 KiB

Cache misses match the previous path, while repeated sizes avoid resize-matrix construction and torch.linalg.pinv.

Synthetic OpenCLIP training step

benchmark_open_clip_patch_interpolation_e2e.py compares forced misses and cache hits over a full single-GPU OpenCLIP training step. The run used naflex_ViT-B-16, BF16 autocast, batch size 32, sequence length 196, and 32x32 patches; ten pairs were measured after warmup.

Mode	Median step	IQR	Samples/s
Previous behavior / forced miss	150.732 ms	0.264 ms	212.30
Cache hit	139.797 ms	0.943 ms	228.90

This synthetic workload shows a 7.25% latency reduction and 7.82% throughput increase with a 1 MiB cache. End-to-end gains will vary with workload and patch-size distribution.

Tests

46 targeted tests passed, covering cache reuse and dtype, gradients and checkpoint state, prewarming with torch.compile, capability reporting, non-base patch forwards, and spatial eval patchification.

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.