diff --git a/docs/_static/architecture.png b/docs/_static/architecture.png new file mode 120000 index 00000000..ebcee6a4 --- /dev/null +++ b/docs/_static/architecture.png @@ -0,0 +1 @@ +../assets/architecture.png \ No newline at end of file diff --git a/docs/_static/attn_hyena.jpg b/docs/_static/attn_hyena.jpg new file mode 100644 index 00000000..448b2dd4 Binary files /dev/null and b/docs/_static/attn_hyena.jpg differ diff --git a/docs/_static/cuda_b2b_causal_conv1d.png b/docs/_static/cuda_b2b_causal_conv1d.png new file mode 100644 index 00000000..732681c2 Binary files /dev/null and b/docs/_static/cuda_b2b_causal_conv1d.png differ diff --git a/docs/_static/cuda_fft_conv2d.png b/docs/_static/cuda_fft_conv2d.png new file mode 100644 index 00000000..b02a1b16 Binary files /dev/null and b/docs/_static/cuda_fft_conv2d.png differ diff --git a/docs/_static/cuda_implicit_filter.png b/docs/_static/cuda_implicit_filter.png new file mode 100644 index 00000000..c189cdb3 Binary files /dev/null and b/docs/_static/cuda_implicit_filter.png differ diff --git a/docs/_static/hyena_hyena.jpg b/docs/_static/hyena_hyena.jpg new file mode 100644 index 00000000..6129ed5e Binary files /dev/null and b/docs/_static/hyena_hyena.jpg differ diff --git a/docs/_static/mamba1_hyena.jpg b/docs/_static/mamba1_hyena.jpg new file mode 100644 index 00000000..c586b6f2 Binary files /dev/null and b/docs/_static/mamba1_hyena.jpg differ diff --git a/docs/_static/mamba2_hyena.jpg b/docs/_static/mamba2_hyena.jpg new file mode 100644 index 00000000..fafc49d2 Binary files /dev/null and b/docs/_static/mamba2_hyena.jpg differ diff --git a/docs/_static/operator_comparison.png b/docs/_static/operator_comparison.png new file mode 120000 index 00000000..2ddc7bac --- /dev/null +++ b/docs/_static/operator_comparison.png @@ -0,0 +1 @@ +../assets/operator_comparison.png \ No newline at end of file diff --git a/docs/_static/throughput_scaling.png b/docs/_static/throughput_scaling.png new file mode 120000 index 00000000..25240089 --- /dev/null +++ b/docs/_static/throughput_scaling.png @@ -0,0 +1 @@ +../assets/throughput_scaling.png \ No newline at end of file diff --git a/docs/api_reference/index.rst b/docs/api_reference/index.rst index 09bf3965..d0c3e619 100644 --- a/docs/api_reference/index.rst +++ b/docs/api_reference/index.rst @@ -9,13 +9,15 @@ Organised bottom-up: low-level convolution primitives first, then the mixer modules that compose them, then full networks, then the parallel, core utility, and experiments layers. -See `ops/README.md <../ops/README.html>`_ for the math motivation behind -the FFT-based ops, and ``docs-tracker.md`` at the repo root for the -documentation coverage plan. +Start with the :doc:`Ops Primer <../ops/README>` for the math motivation +behind the FFT-based ops (the convolution theorem, the linear/circular +flavours, and a decision tree for picking a function). ``docs-tracker.md`` +at the repo root tracks the documentation coverage plan. .. toctree:: :maxdepth: 2 + Ops Primer <../ops/README> ops modules networks diff --git a/docs/api_reference/modules.rst b/docs/api_reference/modules.rst index 6893a01c..9b43aa07 100644 --- a/docs/api_reference/modules.rst +++ b/docs/api_reference/modules.rst @@ -12,7 +12,7 @@ residual blocks they rely on. Mixers ------ -Sequence/spatial mixers — Hyena, Mamba, attention variants. +Sequence/spatial mixers: Hyena, Mamba, attention variants. .. autosummary:: :toctree: generated/ diff --git a/docs/api_reference/networks.rst b/docs/api_reference/networks.rst index 4761b18b..401c7a19 100644 --- a/docs/api_reference/networks.rst +++ b/docs/api_reference/networks.rst @@ -6,7 +6,8 @@ Networks ======== End-to-end classification and general-purpose networks composing the -modules above, plus the UNet-ConvNeXt baselines used in benchmark comparisons. +modules above, plus the UNet-ConvNeXt baselines used in benchmark +comparisons. Classification & general-purpose -------------------------------- diff --git a/docs/api_reference/ops.rst b/docs/api_reference/ops.rst index 1ea8ab3a..bd0a90d2 100644 --- a/docs/api_reference/ops.rst +++ b/docs/api_reference/ops.rst @@ -9,6 +9,8 @@ Low-level convolution primitives. Pure-PyTorch reference implementations double as the spec the CUDA kernels must match; the ``subquadratic_ops_torch`` wrappers are the production path on GPU. +.. _ops-fftconv-fp32: + FFT convolutions (reference fp32) --------------------------------- @@ -27,8 +29,10 @@ Use these for correctness and as the spec for the CUDA kernels below. ~ops.fftconv.fftconv3d_fp32_bhl ~ops.fftconv.causal_fftconv1d_fp32_bhl +.. _ops-fftconv-custom: + FFT convolutions (CUDA-accelerated) ------------------------------------ +------------------------------------ Drop-in wrappers around the ``subquadratic_ops_torch`` fused CUDA kernels. 2D non-causal and 1D causal long-conv variants share the same API as the @@ -45,8 +49,10 @@ fp32 reference ops above. ~ops.fftconv_custom.causal_fftconv1d_bhl ~ops.fftconv_custom.causal_fftconv1d_bhl_w_reshape +.. _ops-causal-conv1d: + Direct 1D causal convolutions (CUDA-accelerated) ------------------------------------------------- +------------------------------------------------- Non-FFT CUDA kernels for short and fused 1D causal convolutions. Useful for small kernel sizes (where FFT overhead dominates) and as building @@ -59,8 +65,10 @@ blocks for fused Hyena variants. ~ops.causal_conv1d_custom.causal_conv1d ~ops.causal_conv1d_custom.b2b_causal_conv1d +.. _ops-circular-fftconv: + Circular FFT convolutions -------------------------- +-------------------------- Periodic-boundary FFT convolutions for global mixing without zero padding. @@ -72,8 +80,10 @@ Periodic-boundary FFT convolutions for global mixing without zero padding. ~ops.circular_fftconv.circular_fftconv2d_fp32_bhl ~ops.circular_fftconv.circular_fftconv3d_fp32_bhl +.. _ops-chunking: + Chunking utilities ------------------- +------------------- Helpers to bound the FFT working-set memory by processing along the sequence axis in chunks. @@ -87,10 +97,12 @@ sequence axis in chunks. ~ops.fftconv_chunked.set_default_chunk_size ~ops.fftconv_chunked.get_default_chunk_size +.. _ops-mixed-fftconv: + Mixed boundary-condition FFT convolutions ----------------------------------------- -FFT convolutions with per-axis boundary conditions — periodic on some +FFT convolutions with per-axis boundary conditions: periodic on some spatial axes, zero-padded on others. See :doc:`../ops/mixed_boundary_conditions` for the per-axis algorithm and the ``fft_padding`` API. diff --git a/docs/architecture.md b/docs/architecture.md index c0730f8d..d088b3b4 100644 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -30,42 +30,54 @@ research team can swap any one of them without touching the others. ## What each layer owns -- **nvsubquadratic** (this library) — the PyTorch-native API. +- **nvsubquadratic** (this library): the PyTorch-native API. Sequence/spatial mixers (Hyena, Mamba, attention variants), learned kernels, residual blocks, networks, and the datamodule/wrapper scaffolding consumed by `experiments`. All public surface lives in the {doc}`api_reference/index`. -- **subquadratic-ops** (separate repo) — the fused CUDA kernels. Causal +- **subquadratic-ops** (separate repo): the fused CUDA kernels. Causal Conv1D for short kernels (2–256), FFT-based Causal Conv1D for long kernels (up to 8K–16M), B2B Causal Conv1D for striped Hyena architectures, plus the 1D/2D FFT primitives. nvSubquadratic delegates here via {mod}`subquadratic_ops_torch` and the published docs are at . -- **megatron-core** — Megatron's distributed-training primitives +- **megatron-core**: Megatron's distributed-training primitives (tensor / pipeline / context parallelism). nvSubquadratic uses it via {mod}`nvsubquadratic.parallel.utils`'s `init_parallel_state` and the context-parallel `DistributedDepthwiseConvNd` wrappers. -This layering keeps API ergonomics in nvSubquadratic, kernel -optimisation in subquadratic-ops, and distributed bookkeeping in -megatron-core. Practically: if a kernel is slow, fix it in -subquadratic-ops; if an interface is awkward, fix it here. +API ergonomics live in nvSubquadratic, kernel optimisation in +subquadratic-ops, and distributed bookkeeping in megatron-core. In +practice: if a kernel is slow, fix it in subquadratic-ops; if an +interface is awkward, fix it here. + +## The HyenaND operator + +The operator that gives this stack its name is the +`Short Conv → Gate → Long Conv → Gate` block you see throughout the network +code. {doc}`how_hyenand_works` builds it up from attention, with the full +diagram and a worked trace. Its fused FFT long-convolution path lives in +subquadratic-ops; see {doc}`ops/README`. ## Naming conventions Two conventions show up everywhere in the ops and module code. Both are documented in detail in `docs/ops/README.md`; the short version: -- **`BHL` vs `BLH`** — memory layout. `BHL` is channels-first +- **`BHL` vs `BLH`**: memory layout. `BHL` is channels-first (`[B, H, *spatial]`, matches `torch.nn.ConvNd`); `BLH` is channels-last (`[B, *spatial, H]`, common in transformer code). The FFT runs faster on contiguous spatial axes, so BHL is the fast path. -- **`_w_reshape`** — wrappers that accept BLH input, internally reshape +- **`_w_reshape`**: wrappers that accept BLH input, internally reshape to BHL, run the fast path, and reshape back. Recommended entry point for channels-last callers. -- **`_chunked`** — processes channels in groups to cap peak FFT memory. +- **`_chunked`**: processes channels in groups to cap peak FFT memory. +- **`fp32` vs `fp16`**: internal compute precision. fp16 ops require + power-of-2 spatial dims (cuFFT constraint) and use dual mean-centering + for numerical stability. See the + [FP16 circular FFT convolution report](https://github.com/NVIDIA-BioNeMo/nvSubquadratic/blob/main/reports/fp16_fft_convolution/REPORT.md). So `causal_fftconv1d_fp32_bhl_w_reshape` is a causal 1D FFT conv that accepts channels-last input, runs the fp32 channels-first kernel under @@ -77,27 +89,10 @@ Hyena, attention, CKConv, and Mamba all expose the same `(query, key, value)` mixer signature. The dispatch lives in {class}`nvsubquadratic.modules.sequence_mixer.QKVSequenceMixer`: configure it with a `LazyConfig` over any of the mixers and the rest of -the model code is unchanged. Switching architectures is a one-line -config diff. - -## The lazy-instantiation system - -Every config file under `examples/` is a tree of -{class}`nvsubquadratic.lazy_config.LazyConfig` specs. A spec is a -deferred constructor call: `LazyConfig(SomeClass)(arg=...)` returns a -proxy that holds the target and the arguments, but doesn't instantiate -yet. {func}`nvsubquadratic.lazy_config.instantiate` walks the tree, -resolves nested specs depth-first, evaluates arithmetic strings (for -`L_cache = "max(H, W)"`-style expressions), and constructs the live -object graph. - -This is why most module constructors take `*_cfg: LazyConfig` rather -than concrete instances — the user-facing config file decides what to -instantiate; the module never imports a specific norm/conv/scheduler -class directly. +the model code is unchanged, so switching architectures is a one-line +config change. ## Further reading -- {doc}`api_reference/index` — the curated API. -- `docs/ops/README.md` — math primer for the FFT-based convolution ops. -- `docs-tracker.md` (repo root) — current docstring coverage status. +- {doc}`api_reference/index`: the curated API. +- `docs/ops/README.md`: math primer for the FFT-based convolution ops. diff --git a/docs/assets/architecture.png b/docs/assets/architecture.png new file mode 100644 index 00000000..1f010048 Binary files /dev/null and b/docs/assets/architecture.png differ diff --git a/docs/assets/operator_comparison.png b/docs/assets/operator_comparison.png new file mode 100644 index 00000000..e0f29c50 Binary files /dev/null and b/docs/assets/operator_comparison.png differ diff --git a/docs/assets/throughput_scaling.png b/docs/assets/throughput_scaling.png new file mode 100644 index 00000000..50738db6 Binary files /dev/null and b/docs/assets/throughput_scaling.png differ diff --git a/docs/benchmarks.md b/docs/benchmarks.md index b0c0c817..ecf26d66 100644 --- a/docs/benchmarks.md +++ b/docs/benchmarks.md @@ -1,21 +1,281 @@ # Benchmarks -Throughput numbers and FLOP scaling. The -tables below are included verbatim from the +Throughput numbers, FLOP scaling, FP16 op-level results, and a worked +ImageNet-training optimization case study. The raw measurement tables and the +scripts that reproduce them live in the [`benchmarks/README.md`](https://github.com/NVIDIA-BioNeMo/nvSubquadratic/blob/main/benchmarks/README.md) -single-source — edits should land there, not here. +single source. When a number changes, update it there and mirror it here. ## FLOP scaling -![FLOP scaling for Hyena / attention / CKConv mixers](_static/flop_scaling.png) +FLOPs are the hardware-independent floor on cost: before any kernel, cache, +or bandwidth effect, they show why a *global* operator has to be subquadratic +to scale at all. The plot below sweeps the input resolution of the +ViT-5-Small backbone (7×7 up to 112×112 patches) and counts per-sample FLOPs +for each mixer. + +![FLOP scaling for the ViT-5-Small attention / Hyena / Hyena+FiLM variants](_static/flop_scaling.png) + +Attention's token-mixing cost grows as $O(L^2)$ in the patch count $L$, so +doubling the grid roughly quadruples its FLOPs; the Hyena variants evaluate +the same global receptive field through FFT convolutions in $O(L \log L)$, so +their curve stays close to linear and the gap widens at every resolution +step. There is a crossover: at small grids the constant factors dominate and +attention is cheaper, but past a modest resolution the asymptotics take over +and the FLOP gap compounds. FLOPs are only the theoretical floor, though. The +sections that follow show how much of that advantage survives as real +wall-clock time, and the kernels that make it survive. See [`benchmarks/compare_flops.py`](https://github.com/NVIDIA-BioNeMo/nvSubquadratic/blob/main/benchmarks/compare_flops.py) for the script that produced the plot. -## ViT-5-Small throughput +## Throughput scaling + +The FLOP advantage above translates into wall-clock time. This is +forward-pass time versus sequence length for `flash-attention`, the +official `mamba_chunk_scan_combined` Mamba2 kernel, and `nSubQ`: + +![Forward time vs. sequence length for HyenaND, Attention, and Mamba](_static/throughput_scaling.png) + +HyenaND reaches million-token sequences at 265 ms (1M tokens), while attention +takes ~90 s (a **339×** gap) and the Mamba2 kernel runs out of memory. + +## CUDA kernels (`nSubQ`) + +The throughput curve above is only reachable with kernels that respect the +GPU memory hierarchy. `nSubQ` is the IO-aware CUDA library behind HyenaND: a +suite of fused FFT-convolution and causal-convolution kernels written with +`CuTe` and `cuFFTDx`, benchmarked here on an NVIDIA RTX PRO 6000 Blackwell +Server Edition. + +### Fusion + +Two hardware trends shape the design. GPU matrix-multiply throughput has grown +faster than off-chip (HBM) bandwidth, which favours attention: its dominant +operations ($\mathbf{QK}^\top$ and $\mathbf{PV}$) are GEMMs, and +FlashAttention-style kernels already keep the $L \times L$ score matrix off +HBM. FFT-based operators have the opposite profile: their fast path is bound +by batched-FFT efficiency, memory layout, and bandwidth rather than +tensor-core GEMM throughput. + +A naive FFT convolution is especially HBM-hungry. A single pass materialises +the padded input, the activation spectrum, the filter spectrum, the +inverse-FFT buffer, and the cropped output (at least five HBM round-trips), +and linear-convolution padding inflates the spatial footprint by roughly +$2^D$ in $D$ dimensions before complex storage is even counted. That +constant factor can hide the $O(L \log L)$ advantage entirely for 2D images +and 3D volumes. `nSubQ` closes the gap by fusing the forward FFT, spectral +modulation, inverse FFT, and output segmentation into a single kernel that +keeps every intermediate in on-chip shared memory and registers. + +### Measured kernel speedups + +All figures are module-level, against the equivalent unfused PyTorch path: + +| Kernel | What it fuses | Speedup | Memory | +| ---------------------------- | ------------------------------------------ | --------: | --------------------: | +| Implicit filter generation | SIREN modal-filter synthesis | **>40×** | **32× less HBM** | +| Causal FFTConv1D (`L ≤ 16K`) | RFFT + spectral product + IRFFT + chunking | **~6×** | ~4× less | +| Causal FFTConv1D (`L > 16K`) | 3-kernel Cooley–Tukey (C2C/IC2C + product) | 2–4× | 2–4× less | +| `b2b causal-conv1d` | proj-conv + gate + mixer-conv + gate | **>7.5×** | lower HBM, ½ CP comms | +| 2D `fft-conv2d` (forward) | C2C + complex product + IC2C core | **>5×** | **>2× less** | + +The `causal-conv1d` kernels additionally support filters up to 256 taps +(channel-first) / 128 (channel-last), wider than existing optimised packages. + +### Implicit filter generation + +Profiling the Evo2 long-implicit (LI) layer showed that **30 % of its runtime +was kernel synthesis**, not the convolution itself. A dedicated fused kernel +accelerates filter generation by over **40×** at the module level and cuts HBM +overhead by **32×** (proportional to the modal-filter order). The memory this +frees is what lets a 1B-parameter Evo2 model reach 16M-token context. + +![Implicit filter generation: fused kernel vs PyTorch baseline](_static/cuda_implicit_filter.png) + +### 1D long convolutions + +For sequence lengths up to 16K, the point where HyenaND's asymptotics +overtake fused SDPA, a single kernel fuses the input and filter RFFTs, the +complex product, the IRFFT, and chunking for a **6×** speedup and **4×** lower +memory. Beyond 16K a three-kernel Cooley–Tukey decomposition (recasting the +1D FFT as a 2D FFT) holds a 2–4× edge. For the short projection/mixer +convolutions in the SE and MR layers, the fused `b2b causal-conv1d` kernel +folds projection conv, pre-gate, mixer conv, and post-gate into one +operation: **>7.5×** faster, and with half the context-parallel +communication points. + +![Fused b2b causal-conv1d for the SE / MR layers](_static/cuda_b2b_causal_conv1d.png) + +### 2D FFT convolution + +For vision workloads the `fft-conv2d` kernels run an RFFT along the first +axis, a fused C2C-FFT + complex-multiply + inverse-C2C core along the second, +and a closing IRFFT, keeping the frequency-domain intermediates on chip. +Against an unfused baseline this is over **5×** faster with **>2×** lower +memory on the forward pass. + +![Forward-pass 2D FFT convolution: fused vs unfused](_static/cuda_fft_conv2d.png) + +## ViT-5-Small ImageNet training: an optimization case study + +ImageNet-1k pretraining is where ViT-5-Small's accuracy is validated before the +architecture is used on downstream work, so we needed to train it, and train it +often. An 800-epoch run is the unit of experiment, and its wall-clock cost +sets the pace of the research loop: every recipe change, ablation, or +hyperparameter sweep pays that cost again. When this work began a single run +took roughly **37 hours** on 8 × H100 with the first GPU data pipeline (and +longer still on the original CPU pipeline). The optimizations documented here +bring a full run down to about **12 hours**. Faster runs mean more ideas +tested per GPU-week and a lower cost per result, which is why we invested in the +*training and data pipeline* and not only the model. + +Training ViT-5-Small on ImageNet (22 M params, 224 × 224, batch 256, +8 × H100 SXM 80 GB, BF16) also makes a clean worked example of systematic +pipeline optimization. Every distributed vision job is bounded by model +compute, data augmentation, and storage I/O independently, and fixing one +merely exposes the next. The work below drove end-to-end throughput from +roughly 2.5 it/s to 12.6 it/s (**5×**) by finding and eliminating each +bottleneck in turn. The sections that follow walk each phase: what the profiler +showed, why the fix worked, and what it exposed next. + +Full profiling tables and configuration history live in +[`examples/vit5_imagenet/OPTIMIZATION_TRACKER.md`](../examples/vit5_imagenet/OPTIMIZATION_TRACKER.md). + +______________________________________________________________________ + +### Phase 1: unblock the compiler + +The original model could not run `torch.compile(max-autotune)` at all; it +crashed with CUDA Graph errors. The root cause was a per-forward RoPE cache +built from a Python `dict`, which forced a graph break on every step and +prevented the tracer from ever capturing a static graph. Replacing it with a +precomputed `register_buffer` fixed the graph break, unlocked `max-autotune`, +and as a side effect also enabled CUDA Graph replay between steps. + +Three supporting changes completed the model cleanup: + +- SDPA backend selection left to PyTorch, so it auto-picks CuDNN on H100 + instead of being locked to FlashAttention. +- Redundant `.to(bfloat16)` / `.to(in_dtype)` casts around SDPA removed; + autocast owns precision. +- Manual float32-upcast RMSNorm swapped for QuACK's fused Triton kernel. + +None of these changes touch weights or hyperparameters. + +*Single-GPU model throughput (synthetic data, no I/O):* + +| Configuration | ms/step | img/s | MFU | +| ---------------------------------- | -------: | --------: | --------: | +| Eager — original | 159.2 | 1,608 | 4.6% | +| `torch.compile` (default) | 46.0 | 5,560 | 15.9% | +| `torch.compile` (max-autotune) | *crash* | — | — | +| **After model optimizations** | | | | +| Eager — optimized | 111.1 | 2,305 | 6.6% | +| `torch.compile` (default) | 33.2 | 7,716 | 22.0% | +| **`torch.compile` (max-autotune)** | **32.0** | **8,003** | **22.9%** | + +Theoretical peak on H100 SXM (989 TFLOPS BF16) is ~34,800 img/s (100% MFU). +At 22.9% MFU the remaining gap is compute-bound GEMM and FFT kernel efficiency; +Python and framework overhead are gone. + +______________________________________________________________________ + +### Phase 2: replace the CPU data pipeline + +With compute at 32 ms/step, CPU data loading immediately became visible. +Single-GPU profiling showed torchvision CPU decode and augmentation costing +**105 ms/step**, more than three times the compute budget. Replacing it with +NVIDIA DALI (GPU-pipelined JPEG decode, crop, flip) dropped the data component +to ~42 ms. The compute side was unchanged; the full step shrank because the +pipeline was no longer stalled on the CPU. + +Augmentation transforms were also ported to GPU-friendly forms for +`torch.compile` compatibility: device-cached normalisation tensors via +`register_buffer`, `torch.where`-based blending instead of boolean-index +scatter, vectorised colour jitter via `argsort`. One attempted optimisation, +DALI `fn.transpose` for CHW output, was **reverted** after profiling showed it +added ~47 ms due to an explicit memory copy. Layout conversions that look free +in PyTorch carry a real cost inside the DALI pipeline graph. + +______________________________________________________________________ + +### Phase 3: eliminate I/O variance + +The jump from v2 (6.3 it/s) to `optimized_plus` (12.1 it/s) is the biggest +single step in the campaign: nearly a **2× gain** from a change that has +nothing to do with the model. The mechanism is DDP synchronisation stalls +caused by cross-rank I/O variance on a shared network file system. + +In DDP, every rank must enter `allreduce` together. When DALI fetch times vary +from 0.6 ms to 200+ ms across ranks (the norm on a shared NFS under load), +the fast ranks park at the barrier waiting for the slow ones. Those stalls +appear inflated inside the backward pass in per-phase profiling, making it easy +to misattribute them to compute. Staging ImageNet on each node's local NVMe +(`/scratch`) gave every rank consistent 1–5 ms fetch times and eliminated the +stalls entirely. + +The v2 augmentation changes shaved ~30 ms off per-step time, but were invisible +at DDP scale as long as the I/O stall dominated. Optimisations that show +clearly in single-GPU profiling may be masked by a different bottleneck in +multi-GPU runs. + +*Multi-GPU DDP training throughput (8 × H100, end-to-end):* + +| Version | Dataloader | Storage | it/s | ms/step | Speedup | +| ------------------------ | ------------------------------ | -------------- | -------: | ------: | -------: | +| CPU baseline | torchvision | Network FS | ~2.5 | ~400 | 1.0× | +| v1 (DALI) | DALI | Network FS | 5.3 | 189 | 2.1× | +| v2 (DALI + compiled aug) | DALI | Network FS | 6.3 | 159 | 2.5× | +| **optimized_plus** | **DALI + compiled aug** | **Local NVMe** | **12.1** | **83** | **4.8×** | +| **fused** | **DALI (all aug in pipeline)** | **Local NVMe** | **12.6** | **79** | **5.0×** | + +______________________________________________________________________ + +### Phase 4: fuse augmentations into the DALI pipeline + +After NVMe staging, the step breakdown exposed one remaining serial cost: ~25 ms +of GPU augmentation running in `on_before_batch_transfer`, between DALI fetch +and the forward pass, outside the DALI pipeline's internal overlap window. + +Moving ThreeAugment, ColorJitter, uint8→float, and normalisation **entirely into +the DALI pipeline** (using `enable_conditionals=True` for the per-sample +branching) reduced that component from ~6 ms (v2 measured) to **0.35 ms** +(Mixup + NHWC permute only). The "data (ms)" column in the table below goes +*up* slightly because it now includes augmentation; what matters is that the +**full step goes down** because DALI's internal scheduler overlaps augmentation +with I/O and the previous step's compute. + +*Fused vs. optimised DALI, DDP × 8, H100 SXM, NVMe:* + +| Pipeline | Model | Full step (ms) | Agg throughput (img/s) | Speedup | +| -------------- | --------- | -------------: | ---------------------: | -------: | +| DALI-optimised | Small | 80.6 | 25,401 | — | +| **DALI-fused** | **Small** | **70.4** | **29,103** | **+15%** | +| DALI-optimised | Base | 131.1 | 15,622 | — | +| **DALI-fused** | **Base** | **120.8** | **16,950** | **+9%** | + +______________________________________________________________________ + +### Final state: fully compute-bound + +*Step breakdown after all optimisations (fused DALI, NVMe, compiled):* + +| Component | Time (ms) | % of step | +| ---------------------- | --------: | --------: | +| DALI fetch | ~2 | 3% | +| Mixup/CutMix + permute | 0.35 | 0.5% | +| Forward + backward | ~66 | 94% | +| Optimizer (FusedLAMB) | ~2 | 3% | + +Data loading (decode, augmentation, I/O, and transfer) now accounts for less +than 4% of wall-clock time. The pipeline is fully compute-bound. +Further throughput gains require model-level work: DDP allreduce efficiency, +larger batch sizes, or architectural changes to the mixer itself. + +## Op-level results -```{include} ../benchmarks/README.md ---- -start-after: '# ViT-5-Small Throughput Benchmarks' ---- -``` +- [FP16 circular FFT convolution](https://github.com/NVIDIA-BioNeMo/nvSubquadratic/blob/main/reports/fp16_fft_convolution/REPORT.md). + The full investigation report: the dual-mean-centering derivation, the + accuracy and throughput results against the FP32 reference, and why the + FP16 path was kept opt-in rather than made the default. diff --git a/docs/conf.py b/docs/conf.py index 8f02a59e..a52961c3 100644 --- a/docs/conf.py +++ b/docs/conf.py @@ -92,6 +92,8 @@ def _read_version(): "megatron", "megatron.core", "omegaconf", + "cleanfid", + "diffusers", "pytorch_lightning", "lightning", "matplotlib", @@ -100,6 +102,7 @@ def _read_version(): "h5py", "scipy", "the_well", + "torch_fidelity", "torchmetrics", "torchvision", "timm", diff --git a/docs/examples/index.md b/docs/examples/index.md index 09bc46ef..4d98c176 100644 --- a/docs/examples/index.md +++ b/docs/examples/index.md @@ -11,6 +11,13 @@ The active experimental roadmap (priorities, owners, status) lives at ## Classification +### MNIST / SMNIST + +[`examples/mnist_classification/`](https://github.com/NVIDIA-BioNeMo/nvSubquadratic/tree/main/examples/mnist_classification) +covers MNIST with both attention and Hyena baselines, plus a small CCNN +backbone. [`examples/smnist_classification/`](https://github.com/NVIDIA-BioNeMo/nvSubquadratic/tree/main/examples/smnist_classification) +covers sequential MNIST (1D input). + ### ImageNet [`examples/imagenet_classification/`](https://github.com/NVIDIA-BioNeMo/nvSubquadratic/tree/main/examples/imagenet_classification) @@ -19,12 +26,31 @@ without augmentation, plus tiny variants for laptop sanity checks). Representative entry points: `ccnn_7_512_hyena.py`, `ccnn_7_512_attention.py`. -### TinyImageNet — ViT-5 +### TinyImageNet: ViT-5 [`examples/vit5_imagenet/`](https://github.com/NVIDIA-BioNeMo/nvSubquadratic/tree/main/examples/vit5_imagenet) is the ViT-5 baseline suite (v1–v5) with its own [`TRACKER.md`](https://github.com/NVIDIA-BioNeMo/nvSubquadratic/blob/main/examples/vit5_imagenet/TRACKER.md). +### UCF101 + +[`examples/ucf101_classification/`](https://github.com/NVIDIA-BioNeMo/nvSubquadratic/tree/main/examples/ucf101_classification) +covers video classification with both sequence- and clip-mode datamodules. + +## Diffusion + +### MNIST + +[`examples/mnist_diffusion/`](https://github.com/NVIDIA-BioNeMo/nvSubquadratic/tree/main/examples/mnist_diffusion) +is a small DDPM/JiT diffusion sanity-check. + +### ImageNet + +[`examples/imagenet_diffusion/`](https://github.com/NVIDIA-BioNeMo/nvSubquadratic/tree/main/examples/imagenet_diffusion) +is the full ImageNet diffusion setup. See its +[README](https://github.com/NVIDIA-BioNeMo/nvSubquadratic/blob/main/examples/imagenet_diffusion/README.md) +for the JiT vs Hyena-vs-attention comparison. + ## Spatial recall [`examples/spatial_recall_1d/`](https://github.com/NVIDIA-BioNeMo/nvSubquadratic/tree/main/examples/spatial_recall_1d), @@ -49,6 +75,6 @@ holds the throughput-comparison configs used to produce the numbers in ### The Well [`examples/well/`](https://github.com/NVIDIA-BioNeMo/nvSubquadratic/tree/main/examples/well) -covers The Well PDE benchmark — see its +covers The Well PDE benchmark; see its [README](https://github.com/NVIDIA-BioNeMo/nvSubquadratic/blob/main/examples/well/README.md) for sub-datasets and baselines. diff --git a/docs/getting_started.md b/docs/getting_started.md index c04e5d6f..1c82bb3d 100644 --- a/docs/getting_started.md +++ b/docs/getting_started.md @@ -6,7 +6,7 @@ Apptainer, conda, venv) see the project [README](https://github.com/NVIDIA-BioNe ## Requirements -- CUDA-compatible NVIDIA GPU (Ampere or newer) +- CUDA-compatible NVIDIA GPU - CUDA Toolkit 12.0 or higher - Python 3.11 or higher @@ -43,52 +43,120 @@ the [project README](https://github.com/NVIDIA-BioNeMo/nvSubquadratic/blob/main/ ## Hello, Hyena -A minimal forward pass through a 2D Hyena mixer: +A minimal forward pass through a real 2D Hyena mixer. Everything is +wired with {doc}`LazyConfig `: each `LazyConfig(Cls)(...)` +records the class and its arguments without constructing anything, and a +single `instantiate(...)` call at the end builds the whole tree. This is +exactly how the {doc}`experiments ` configs +assemble their networks, except there the scalar fields are filled by +`"${net.hidden_dim}"`-style interpolation instead of the concrete +integers used below. ```python import torch from nvsubquadratic.lazy_config import LazyConfig, instantiate from nvsubquadratic.modules.hyena_nd import Hyena -from nvsubquadratic.modules.kernels_nd import ( - SIRENKernelND, - SIRENPositionalEmbeddingND, -) -from nvsubquadratic.ops.fftconv import fftconv2d_fp32_bhl +from nvsubquadratic.modules.ckconv_nd import CKConvND +from nvsubquadratic.modules.kernels_nd import SIRENKernelND +from nvsubquadratic.modules.rms_norm_channel_first import RMSNormChannelFirst +from nvsubquadratic.utils.qk_norm import L2Norm device = torch.device("cuda") B, H, X, Y = 2, 64, 32, 32 -x = torch.randn(B, H, X, Y, device=device) - -# A SIREN-parameterised long-range 2D kernel. -kernel_cfg = LazyConfig(SIRENKernelND)( - out_dim=H, - data_dim=2, - mlp_hidden_dim=64, - num_layers=3, - embedding_dim=32, - omega_0=10.0, - L_cache=max(X, Y), - use_bias=True, -) -# Wire a Hyena mixer that consumes the kernel via a global FFT conv. -mixer_cfg = LazyConfig(Hyena)( - global_conv_cfg=LazyConfig(lambda: None)(), # replaced below - short_conv_cfg=LazyConfig(torch.nn.Identity)(), - gate_nonlinear_cfg=LazyConfig(torch.nn.SiLU)(), - pixelhyena_norm_cfg=LazyConfig(torch.nn.Identity)(), - qk_norm_cfg=None, +hyena_cfg = LazyConfig(Hyena)( + # The long-range global convolution. CKConvND owns a SIREN-parameterised + # kernel and applies it as an FFT conv — the kernel is *generated* by the + # SIREN MLP, never random. + global_conv_cfg=LazyConfig(CKConvND)( + data_dim=2, + hidden_dim=H, + kernel_cfg=LazyConfig(SIRENKernelND)( + data_dim=2, + out_dim=H, + mlp_hidden_dim=32, + num_layers=3, + embedding_dim=32, + omega_0=10.0, + hidden_omega_0=1.0, + L_cache=max(X, Y), + use_bias=True, + ), + mask_cfg=LazyConfig(torch.nn.Identity)(), + grid_type="double", # linear (non-circular) convolution + fft_padding="zero", + fft_backend="torch_fft", # portable; "subq_ops" uses the fused 2D CUDA kernel + is_causal=False, + ), + # Depthwise short conv on the concatenated [Q; K; V] (3 * H channels). + short_conv_cfg=LazyConfig(torch.nn.Conv2d)( + in_channels=3 * H, + out_channels=3 * H, + kernel_size=3, + groups=3 * H, + padding=1, + bias=False, + ), + gate_nonlinear_cfg=LazyConfig(torch.nn.SiLU)(), # first gate σ + gate_nonlinear_2_cfg=LazyConfig(torch.nn.Sigmoid)(), # second gate σ₂ + pixelhyena_norm_cfg=LazyConfig(RMSNormChannelFirst)( + dim=H, eps=1e-6, use_quack=False + ), + output_norm_cfg=LazyConfig(RMSNormChannelFirst)(dim=H, eps=1e-6, use_quack=False), + qk_norm_cfg=LazyConfig(L2Norm)(dim=1), # L2 QK-norm on the channel axis ) -# For a self-contained example, skip the LazyConfig dance and call the -# op directly: -kernel = torch.randn(1, H, X, Y, device=device) -y = fftconv2d_fp32_bhl(x, kernel) +# Build the whole module tree in one call. +hyena = instantiate(hyena_cfg).to(device) + +# Hyena consumes channels-last Q, K, V tensors [B, *spatial, C]. In a full +# model these come from a linear projection W·x (see QKVSequenceMixer); here +# we feed random activations to exercise the forward. +q = torch.randn(B, X, Y, H, device=device) +k = torch.randn(B, X, Y, H, device=device) +v = torch.randn(B, X, Y, H, device=device) + +y = hyena(q, k, v) +print(y.shape) # torch.Size([2, 32, 32, 64]) -> [B, X, Y, C] +``` + +In a real network you rarely hold `Q`, `K`, `V` yourself: the +`QKVSequenceMixer` (see +[`mixer_defaults.py`](https://github.com/NVIDIA-BioNeMo/nvSubquadratic/blob/main/examples/spatial_recall_v2/mixer_defaults.py)) +projects a single activation `x` into the three tensors and forwards them +to this `Hyena`. That same factory is what the spatial-recall experiments +instantiate. + +## Going lower: the FFT conv op directly + +The Hyena above ultimately routes its long-range mixing through one of the +FFT-convolution ops. When you only need the convolution itself, with no +gating, kernel generation, or `nn.Module`, you can call the op +directly. Here `kernel` is supplied explicitly (any 2D filter; a SIREN +kernel would normally produce it): + +```python +import torch + +from nvsubquadratic.ops.fftconv import fftconv2d_fp32_bhl + +device = torch.device("cuda") + +B, H, X, Y = 2, 64, 32, 32 +x = torch.randn(B, H, X, Y, device=device) # channels-first [B, C, X, Y] +kernel = torch.randn(1, H, X, Y, device=device) # per-channel filter [1, C, K_x, K_y] + +y = fftconv2d_fp32_bhl(x, kernel) # "same"-size circular-free conv print(y.shape) # torch.Size([2, 64, 32, 32]) ``` +The op casts to fp32 internally for numerical stability and returns the +result in `x`'s original dtype. Note the layout difference: the ops work +**channels-first** `[B, C, *spatial]` (the `_bhl` suffix), whereas the +`Hyena` module's public interface is **channels-last** `[B, *spatial, C]`. + The lower-level FFT ops in {doc}`nvsubquadratic.ops ` are deliberately function-only so higher-level mixers can compose them freely. The @@ -99,8 +167,9 @@ Lightning-driven training pipelines. ## Next steps -- {doc}`architecture` — the three-layer nvSubquadratic / subquadratic-ops - / megatron-core story and the naming conventions used throughout the +- {doc}`architecture`: the three layers (nvSubquadratic, subquadratic-ops, + megatron-core) and the naming conventions used throughout the library. -- {doc}`examples/index` — end-to-end training recipes per dataset. -- {doc}`api_reference/index` — the full curated API surface. +- [`examples/`](https://github.com/NVIDIA-BioNeMo/nvSubquadratic/tree/main/examples): + end-to-end training recipes per dataset. +- {doc}`api_reference/index`: the full curated API surface. diff --git a/docs/glossary.md b/docs/glossary.md new file mode 100644 index 00000000..6b956917 --- /dev/null +++ b/docs/glossary.md @@ -0,0 +1,114 @@ +# Glossary + +Quick definitions for the terms that show up across these docs and the +code. For the narrative that ties them together, see +{doc}`how_hyenand_works`. + +```{glossary} +HyenaND + The library's flagship operator: a `Short Conv → Gate → Long Conv → Gate` + sandwich that achieves a global, data-dependent receptive field in + $O(N \log N)$ on native 1D / 2D / 3D grids. See + {doc}`how_hyenand_works`. + +Subquadratic + Scaling better than $O(N^2)$ in the number of tokens $N$. HyenaND is + $O(N \log N)$; attention is $O(N^2)$. + +Receptive field + The span of input a single output position can depend on. A *global* + receptive field means any output can see the entire input. This is + attention's defining property, which HyenaND reproduces via long + convolutions. + +Data-dependence + Whether the operator's mixing weights are computed from the input + (data-dependent) or fixed after training. Attention is data-dependent + through the attention matrix $A(x)$; HyenaND is data-dependent through + {term}`gating`. + +Gating + Element-wise multiplication of a signal by a data-derived mask + (e.g. $q \odot \mathrm{SiLU}(k)$). HyenaND interleaves gates with its + long convolution to make the effective operator depend on the input + without ever materialising an $N \times N$ matrix. + +Implicit filter + A convolution kernel produced by evaluating a small network + ({term}`SIREN`) on grid coordinates, rather than stored as one learnable + weight per tap. It is compact, and because it is a continuous function of + position it is samplable on a grid of any size or aspect ratio without + retraining. Contrast with an *explicit* filter, whose taps are + learned parameters (as in a classical CNN). + +SIREN + A sinusoidal-activation MLP ($f_\theta$) used to parametrise implicit + filters. Its frequency is controlled by an $\omega_0$ hyperparameter + that scales with grid resolution and dimensionality. See + {doc}`reports` for the dimensional-scaling rule. Implemented in + {mod}`nvsubquadratic.modules.kernels_nd`. + +FiLM + Feature-wise Linear Modulation. Conditions the synthesised kernel on a + control variable $z(\mathbf{x})$ pooled from the input's + {term}`register tokens`, making the kernel input-dependent. Implemented + in {mod}`nvsubquadratic.modules.film`. + +Register tokens + Auxiliary tokens carried alongside the data tokens whose pooled state + feeds the {term}`FiLM` conditioning of the Hyena kernel. + +FFT convolution (FFTConv) + Computing a convolution as an element-wise product in the frequency + domain, $y = \mathcal{F}^{-1}(\mathcal{F}(x) \odot \mathcal{F}(K))$. + Each FFT is $O(N \log N)$ and the total cost is independent of kernel + size, which is what makes a global kernel affordable. Implemented in + {doc}`ops/README`. + +Convolution theorem + The identity that convolution in the spatial domain equals + element-wise multiplication in the frequency domain. The mathematical + basis for {term}`FFT convolution (FFTConv)`. + +Toeplitz matrix + The matrix form of a 1D convolution: each row is a shifted copy of the + filter. Convolving a signal with a filter is the same as multiplying + by the corresponding Toeplitz matrix; a *causal* convolution is a + lower-triangular one. + +Linear vs circular convolution + **Linear** zero-pads the input so the kernel never wraps around + (matches `torch.nn.ConvNd`); **circular** treats the input as periodic + so the kernel wraps at the boundary (useful for PDEs and periodic + signals). See {doc}`ops/README`. + +BHL / BLH + Memory layout. **BHL** is channels-first (`[B, H, *spatial]`, matches + `torch.nn.ConvNd`); **BLH** is channels-last (`[B, *spatial, H]`, common + in transformer code). The FFT is faster on contiguous spatial axes, so + BHL is the fast path; `_w_reshape` wrappers accept BLH and convert. + +Causal + An operator where output position $n$ depends only on inputs at + positions $\le n$, with no leakage from the future. Required for + autoregressive 1D sequence modelling. + +Mixer + An operator with the shared $(q, k, v)$ signature that + {class}`nvsubquadratic.modules.sequence_mixer.QKVSequenceMixer` + dispatches over (Hyena, attention, CKConv, or Mamba), so a network can + swap one for another from the config. + +CKConv + Continuous-Kernel Convolution: a convolution whose kernel is an + {term}`implicit filter` $k_\theta(p)$. A close relative of Hyena's + long-conv path; see {mod}`nvsubquadratic.modules.ckconv_nd`. + +LazyConfig + The library's deferred-instantiation system: a config object that + records *what* to build and *how* without building it yet, so example + recipes describe a whole experiment as a tree of configs. See + {class}`nvsubquadratic.lazy_config.LazyConfig`. +``` + + diff --git a/docs/how_hyenand_works.md b/docs/how_hyenand_works.md new file mode 100644 index 00000000..5bf85a13 --- /dev/null +++ b/docs/how_hyenand_works.md @@ -0,0 +1,249 @@ +# How HyenaND works + +This page builds up the operator at the heart of the library starting from +**attention**, which most readers already know, and shows how HyenaND keeps +what makes attention work while shedding its costs. For reference material, +{doc}`architecture` documents the stack and {doc}`ops/README` documents the +kernels. + +## Start from attention + +Self-attention is the operator everything else is measured against. Two +properties make it work: + +- **Global receptive field.** Every position can look at every other + position, so a token's representation can depend on context arbitrarily + far away. +- **Data-dependence.** The mixing weights are computed *from the input*: + the attention matrix $A(x)$ is built on the fly for each sequence, rather + than being a fixed set of learned weights. + +The price for both is the $N \times N$ attention matrix: compute and memory +grow as $O(N^2)$ in the number of tokens $N$. For a 256×256 image that is +already 65k tokens; for video or a 3D volume it is hopeless. Attention also +has no native notion of 2D/3D geometry. To apply it to an image you flatten +the grid into a 1D sequence and let the model relearn that two pixels are +neighbours. + +**HyenaND's goal is to keep the global receptive field and the +data-dependence, but pay $O(N \log N)$ instead of $O(N^2)$, and to do it +directly on the data's native 2D/3D geometry.** + +```{list-table} +--- +header-rows: 1 +widths: 28 24 24 24 +--- +* - + - Attention + - Mamba + - **HyenaND** (ours) +* - Receptive field + - global + - global (via scan) + - global +* - Cost in tokens $L$ + - $O(L^2)$ + - $O(L)$ + - $O(L \log L)$ +* - Native dimensionality + - any (but geometry-blind) + - 1D only + - **native 1D / 2D / 3D** +* - Data-dependent mixing + - attention matrix $A(x)$ + - selective state + - **gating** +``` + +Mamba is the other popular subquadratic option, but it is inherently 1D: to +process an image it has to pick an ad-hoc raster scan order, and no single +1D ordering respects 2D locality. HyenaND is global and subquadratic while +operating directly on multi-dimensional data. + +## How HyenaND gets the global receptive field cheaply + +A global receptive field means convolving each position with a filter as +large as the whole input. Done naively, a convolution of an $N$-element +signal with an $N$-element kernel costs $O(N^2)$, so we have not saved +anything. Two ideas fix this. + +### 1. An *implicit* filter + +Instead of storing one learnable weight per kernel tap (which would be $N$ +parameters for a global kernel, and a different count for every input +size), HyenaND **generates** the kernel from a small neural network, a +SIREN MLP $f_\theta$, evaluated on the grid coordinates: + +$$ +K(\mathbf{p}) = f_\theta(\mathbf{p}), \qquad +\mathbf{p} \in \text{grid coordinates}. +$$ + +The filter is a continuous *function* of position, not a table of numbers. +This is the difference between storing a line as the equation $y = mx + b$ +versus listing every point on it. Because $f_\theta$ is continuous, the +**same learned kernel can be sampled on a grid of any size or aspect +ratio**: train at 64×64, evaluate at 256×256, no retraining. A learned +Gaussian window $w$ multiplies the filter so its influence can taper with +distance. + +### 2. The convolution theorem (FFT) + +Even with a compact parametrisation, *applying* a global kernel by sliding +it across the input is still $O(N^2)$. The convolution theorem turns the +spatial convolution into an element-wise product in the frequency domain: + +$$ +y = \mathcal{F}^{-1}\!\bigl( \mathcal{F}(x) \odot \mathcal{F}(K) \bigr). +$$ + +The two forward FFTs and the inverse each cost $O(N \log N)$, the +element-wise product is $O(N)$, and, crucially, **the total cost is +independent of kernel size.** A global kernel costs no more than a tiny +one. In $N$ dimensions the FFT runs on the native grid, so a 2D image or +3D volume is convolved on its real geometry with no flattening: + +$$ +\text{cost} = O\!\left(\textstyle\prod_n L_n \;\log \prod_n L_n\right). +$$ + +This frequency-domain step is the **FFT convolution (FFTConv)** that the +{doc}`ops/README` primitives implement, and it is what makes a global-kernel +sequence model subquadratic. The full math primer (linear vs circular +boundaries, precision, the decision tree for picking an op) lives there. + +## How HyenaND stays data-dependent: gating + +Attention is data-dependent because it *builds* the mixing matrix $A(x)$ +from the input. HyenaND never materialises such a matrix. Instead it +interleaves the convolution with **element-wise gating**: multiplying the +signal by a data-derived mask. A convolution is a fixed (Toeplitz) linear +map; multiplying its input and output by input-dependent gates makes the +*effective* operator depend on the data, at the cost of a handful of +element-wise products rather than an $N \times N$ matmul. + +Just as attention forms three projections $q, k, v$ from the input, +HyenaND forms its own projections and threads them through a gate → +long-convolution → gate sandwich. This is the same $(q, k, v)$ mixer +signature that {class}`nvsubquadratic.modules.sequence_mixer.QKVSequenceMixer` +dispatches over, which is why swapping Hyena for attention, CKConv, or +Mamba in a network is a one-line config change. + +## Putting it together: the HyenaND operator + +The diagram below is the operator that gives the stack its name. It maps +1:1 onto the `Short Conv → First Gate → Long Conv → Second Gate` block you +see throughout the network code. + +![The HyenaND operator: kernel-synthesis path on top, data path on the bottom](_static/architecture.png) + +Two paths run and meet at the long convolution: + +- **Kernel synthesis (top), "what filter to use."** Grid coordinates + feed the SIREN MLP $f_\theta$, are masked by the learned Gaussian window + $w$, and are FiLM-conditioned on a control variable $z(\mathbf{x})$ + pooled from the input's register tokens. The result is an + input-dependent, implicitly-parameterised $N$D kernel $K(\mathbf{x})$ that + is global, freely learned, and computed once per input. +- **Data path (bottom), "what to filter."** The input is projected into + $\mathbf{q}, \mathbf{k}, \mathbf{v}$ (with a depthwise short conv for + local context). The **inner gate** $Z = \mathbf{q} \odot + \mathrm{SiLU}(\mathbf{k})$ is convolved with $K(\mathbf{x})$ via the $N$D + FFTConv above, then the **outer gate** $O = H \odot \mathrm{SiLU}(\mathbf + {v})$ conditions the result before a final norm. + +Read top-to-bottom: synthesise an input-dependent global kernel, gate the +signal, convolve it cheaply in the frequency domain, gate again. This +gives attention's two properties, a global receptive field (the long +implicit conv) and data-dependence (the two gates), at $O(N \log N)$ on +native geometry. + +## A worked trace + +The snippet below runs the **actual `Hyena` operator** from the diagram on a +batch of 32×32 images, with no random stand-in kernel. Every argument +maps onto a box in the block above: the `CKConvND` global conv *is* the SIREN +kernel-synthesis + FFTConv path (top), the depthwise `Conv2d` is the short +conv, the two `SiLU`s are the inner and outer gates, and the `LayerNorm`s are +the PixelHyena and output norms. It runs on CPU, so no GPU is needed: + +```python +import torch +from nvsubquadratic.lazy_config import LazyConfig +from nvsubquadratic.modules.ckconv_nd import CKConvND +from nvsubquadratic.modules.hyena_nd import Hyena +from nvsubquadratic.modules.kernels_nd import SIRENKernelND +from nvsubquadratic.utils.qk_norm import L2Norm + +H, X, Y = 64, 32, 32 # hidden channels, 2D grid + +# Top path of the diagram: CKConvND synthesises an implicit kernel with a +# SIREN MLP, then applies it via the ND FFTConv — one O(N log N) step. +global_conv_cfg = LazyConfig(CKConvND)( + data_dim=2, + hidden_dim=H, + kernel_cfg=LazyConfig(SIRENKernelND)( + out_dim=H, + data_dim=2, + mlp_hidden_dim=32, + num_layers=2, + embedding_dim=32, + omega_0=10.0, + L_cache=X, + use_bias=True, + hidden_omega_0=1.0, + ), + mask_cfg=LazyConfig(torch.nn.Identity)(), # Gaussian-window slot (off here) + grid_type="double", + fft_padding="zero", + is_causal=False, +) + +hyena = Hyena( + global_conv_cfg=global_conv_cfg, + short_conv_cfg=LazyConfig(torch.nn.Conv2d)( # depthwise short conv on [Q;K;V] + in_channels=H * 3, + out_channels=H * 3, + kernel_size=3, + padding=1, + groups=H * 3, + ), + gate_nonlinear_cfg=LazyConfig(torch.nn.SiLU)(), # the two gates + pixelhyena_norm_cfg=LazyConfig(torch.nn.LayerNorm)(normalized_shape=H), + qk_norm_cfg=LazyConfig(L2Norm)(), + output_norm_cfg=LazyConfig(torch.nn.LayerNorm)(normalized_shape=H), +).eval() + +B = 2 +x = torch.randn(B, X, Y, H) # channels-last image batch [B, X, Y, C] +y = hyena(x, x, x) # q = k = v = x → self-mixing +print(y.shape) # torch.Size([2, 32, 32, 64]) +``` + +Reading it against the diagram: `hyena(q, k, v)` takes the three projections +(here all equal to `x`, the self-mixing case), runs the depthwise short conv +for local context, forms the inner gate $Z = \mathbf{q} \odot +\mathrm{SiLU}(\mathbf{k})$, convolves $Z$ with the SIREN-synthesised kernel +inside `CKConvND` (the single $O(N \log N)$ FFTConv, the only expensive +step), then applies the outer gate $O = H \odot \mathrm{SiLU}(\mathbf{v})$ and +the output norm. In a full network the three projections come from separate +linear layers, exactly as attention forms its own $q, k, v$. + +For the kernel-level view, one tensor pushed through the bare FFTConv +primitive, see the minimal forward pass in {doc}`getting_started`. + +## Where to go next + +- {doc}`architecture`: the three-library stack (nvSubquadratic / + subquadratic-ops / megatron-core) and the BHL/BLH layout conventions. +- {doc}`ops/README`: the FFT-convolution math primer and the decision + tree for choosing a primitive (linear vs circular, fp32 vs fp16, + chunking, fused CUDA paths). +- [`examples/`](https://github.com/NVIDIA-BioNeMo/nvSubquadratic/tree/main/examples): + native 2D/3D global context applied to images, video, and PDE + rollouts. +- {doc}`glossary`: quick definitions for SIREN, FiLM, implicit filter, + Toeplitz, register tokens, BHL/BLH. + + diff --git a/docs/index.rst b/docs/index.rst index b37cf38e..c0026864 100644 --- a/docs/index.rst +++ b/docs/index.rst @@ -3,14 +3,80 @@ nvSubquadratic Documentation ============================ -``nvsubquadratic`` is a unified PyTorch-native library for subquadratic -alternatives to quadratic attention. It consolidates efforts from across -NVIDIA Research teams (nvResearch, NeMo, BioNeMo) into a single, consistent -API. The current release supports multi-dimensional (1D, 2D, 3D) Hyena -operators backed by optimized CUDA kernels from -:mod:`subquadratic_ops_torch`. Hyena operators provide subquadratic -alternatives to attention, achieving ``O(N log N)`` complexity compared with -``O(N^2)`` for traditional attention. +Attention is global, but it is quadratic and it ignores geometry. +Every token attends to every other token, so compute and memory grow as +``O(N^2)``. A 256×256 image is already 65k tokens, and video or 3D +volumes are out of reach. To apply attention to an image at all you have +to flatten the grid into a 1D sequence and let the model relearn that +neighbouring pixels are neighbours. + +``nvsubquadratic`` is a unified, PyTorch-native library for subquadratic +alternatives to attention. They keep its global receptive field while +running in ``O(N log N)``, directly on native 1D / 2D / 3D +geometry. It consolidates efforts from across NVIDIA Research teams +(nvResearch, NeMo, BioNeMo) into a single, consistent API. The current +release centres on multi-dimensional **HyenaND** operators backed by +optimized CUDA kernels from :mod:`subquadratic_ops_torch`. + +The figure below summarises the trade-off. *(Left)* Attention is natively +multi-dimensional but scales quadratically. Mamba is subquadratic but +inherently 1D, so it needs an ad-hoc 1D scan order to touch +multi-dimensional data, and no single ordering respects 2D locality. +HyenaND is global, natively multi-dimensional, and subquadratic at +the same time. *(Right)* That ``O(N log N)`` complexity is real wall-clock time: +HyenaND scales to million-token sequences while attention +collapses at long context. + +If you are new to the library, start with :doc:`how_hyenand_works`, which builds +the operator up from attention in a few minutes. Then come back for install and the +package tour. + +.. raw:: html + +
+
+
+ + + + + + + + + + + + + + + + + +
AttentionMambaHyenaND (Ours)
+ Attention receptive field + + Mamba scan order 1 + + Mamba scan order 2 + + HyenaND receptive field +
𝒪(L²)𝒪(L)𝒪(L log L)
+
+
+ Forward time vs sequence length +
+
+
+ Figure 1. + (Left) Receptive field and complexity of global operators by + token count L: Attention 𝒪(L²), Mamba 𝒪(L), + HyenaND 𝒪(L log L). + (Right) Forward-pass time vs. sequence length for + flash-attention, the official mamba_chunk_scan_combined + Mamba2 kernel, and nSubQ (HyenaND). +
+
Installation ------------ @@ -30,36 +96,44 @@ To enable the optional fused RMSNorm kernel on Hopper / Blackwell GPUs: Requirements ------------ -- CUDA-compatible NVIDIA GPU (Ampere or newer) +- CUDA-compatible NVIDIA GPU - CUDA Toolkit 12.0 or higher - Python 3.11 or higher Where to go next ---------------- -- **Getting Started** — install, requirements, and a minimal "Hello, - Hyena" forward pass. -- **Architecture** — the three-layer nvSubquadratic / subquadratic-ops / - megatron-core story and the BHL/BLH naming conventions. -- **Package Overview** — bottom-up tour of what's inside - ``nvsubquadratic/`` (ops / modules / networks / parallel / utils). -- **Examples** — per-dataset training recipes under ``examples/``. -- **Benchmarks** — ViT-5-Small throughput tables and FLOP scaling. -- **Reports** — long-form technical reports backed by reproducible - scripts and figures. -- **Ops Overview** — math primer and decision tree for the FFT - convolution primitives. -- **API Reference** — auto-generated reference for the curated public - surface organised by package (ops, modules, networks, parallel, core, - experiments). +- :doc:`How HyenaND Works `: the conceptual on-ramp. + It builds the operator up from attention (global receptive field + + data-dependence) and shows how it gets both for ``O(N log N)`` via + implicit kernels, the FFT, and gating. +- :doc:`Getting Started `: install, requirements, and a + minimal "Hello, Hyena" forward pass. +- :doc:`Architecture `: the three-layer nvSubquadratic / + subquadratic-ops / megatron-core story and the BHL/BLH naming + conventions. +- :doc:`Repository Overview `: bottom-up tour of + what's inside ``nvsubquadratic/`` (ops / modules / networks / parallel / + utils). +- :doc:`Lazy-Config System `: how every run is described by + one config file, with deferred instantiation, ``${...}`` interpolation, and + the base-config + ablation workflow. +- :doc:`Benchmarks `: FLOP scaling, kernel speedups, and a + worked ViT-5-Small ImageNet training optimization case study. +- :doc:`Reports `: long-form technical reports backed by + reproducible scripts and figures. +- :doc:`Glossary `: quick definitions for SIREN, FiLM, implicit + filter, Toeplitz, register tokens, BHL/BLH. +- :doc:`API Reference `: auto-generated reference for + the curated public surface organised by package (ops, modules, networks, + parallel, core, experiments), opening with the FFT-convolution **ops + primer** (math motivation + function decision tree). Contributor docs ---------------- -- `CONVENTIONS.md `_ — +- `CONVENTIONS.md `_: Google-style docstring guide and PR checklist (lives at the repo root). -- `docs-tracker.md `_ — - documentation coverage status per file. Related projects ---------------- @@ -67,7 +141,7 @@ Related projects ``nvsubquadratic`` is the high-level PyTorch interface; the underlying CUDA kernels live in a separate library: -- `subquadratic-ops `_ — +- `subquadratic-ops `_: optimized CUDA kernels (causal conv1d, FFT conv1d/2d, B2B causal conv1d, implicit filters, rearrange) that nvSubquadratic delegates to via :mod:`subquadratic_ops_torch`. Refer to its API reference for kernel-level @@ -81,12 +155,15 @@ CUDA kernels live in a separate library: .. toctree:: :maxdepth: 2 + :hidden: + How HyenaND Works Getting Started Architecture Repository Overview + Lazy-Config System Examples Benchmarks Reports - Ops Overview + Glossary API Reference diff --git a/docs/lazy_config.md b/docs/lazy_config.md new file mode 100644 index 00000000..bd2cdada --- /dev/null +++ b/docs/lazy_config.md @@ -0,0 +1,429 @@ +# The lazy-config system + +Every training run in this repository is described, end to end, by a single +Python config file. The network architecture, the dataset, the optimizer, +the Lightning wrapper, and the schedule are all data, not code wired +together at the call site. The mechanism that makes this work is the +**lazy-config** system in {py:mod}`nvsubquadratic.lazy_config`: a tiny +deferred-instantiation layer (think of it as a few-hundred-line stand-in for +Hydra / detectron2 lazy configs) that lets you *declare* an object (what +class to build and with what arguments) without *building* it yet. + +This page explains why we do it this way, how the machinery works, and walks +through the patterns we actually use in the +[`examples/`](https://github.com/NVIDIA-BioNeMo/nvSubquadratic/tree/main/examples) +configs. + +## Why configs? + +We are moving all model building into config files on purpose. Three +properties motivate it: + +- **One config per experiment makes research backtrackable.** A run is one + file, not a pile of CLI flags and a commit hash you have to reconstruct + months later. Every architectural choice + (kernel size, number of blocks, whether QK-norm is on, which mixer) lives + in one place, under version control, next to the runs it produced. To + reproduce a result you point `run.py` at the same file; to understand what + a run did you read one file top to bottom. + +- **Base configs + overrides make ablations cheap and honest.** You write a + base config once and define each ablation as a small file that *overwrites + parts of it*. A dropout sweep or a mixer swap is its own file, so every run + still has a complete, self-contained config. + There is no "remember to also pass `--dropout 0.1`" footgun: the ablation + is the file. This is exactly how the + [`examples/`](https://github.com/NVIDIA-BioNeMo/nvSubquadratic/tree/main/examples) + configs are organised (see + [the base-config and ablation pattern](#the-base-config-and-ablation-pattern) + below). + +- **The full architecture is serialisable and diffable.** Because a config + is plain data (an OmegaConf tree of `__target__` dicts), it can be dumped + to YAML, logged to W&B, printed as a tree, and *diffed* against another + run. "What changed between run A and run B?" becomes a text diff of two + configs. + +The cost is that the config files look a little unusual the first time you +see one: deeply nested `LazyConfig(...)` calls with `"${...}"` strings +sprinkled through them. The rest of this page makes that syntax legible. + +## The core idea: declare now, build later + +A {py:class}`~nvsubquadratic.lazy_config.LazyConfig` wraps a target class or +callable. *Calling* it with keyword arguments does **not** construct the +object; it produces an OmegaConf `DictConfig` carrying a `__target__` key +plus the arguments: + +```python +from nvsubquadratic.lazy_config import LazyConfig, instantiate +import torch + +cfg = LazyConfig(torch.nn.LayerNorm)(normalized_shape=768, eps=1e-6) +# cfg is a DictConfig: +# {"__target__": "torch.nn.LayerNorm", "normalized_shape": 768, "eps": 1e-6} + +norm = instantiate(cfg) # NOW the LayerNorm is actually built +isinstance(norm, torch.nn.LayerNorm) # True +``` + +Two steps, deliberately separated: + +1. **Declare.** `LazyConfig(target)(**kwargs)` records *what* to build. No + import of heavy framework code is forced at this point; the target can + even be a dotted string like `"torch.nn.LayerNorm"`. +1. **Build.** {py:func}`~nvsubquadratic.lazy_config.instantiate` resolves + `__target__` via `importlib`, processes the arguments, and calls the + target. + +That separation is the point. Between declare and build, the config is +just data you can edit, override, interpolate, serialise, and diff. + +## Nesting + +Configs nest. A `LazyConfig` result can be passed as an argument to another +`LazyConfig`, so an entire module tree is one expression: + +```python +block = LazyConfig(ResidualBlock)( + sequence_mixer_cfg=LazyConfig(QKVSequenceMixer)( + hidden_dim=160, + mixer_cfg=LazyConfig(Hyena)( + global_conv_cfg=LazyConfig(CKConvND)(...), + short_conv_cfg=LazyConfig(torch.nn.Conv2d)(...), + ), + ), + mlp_cfg=LazyConfig(MLP)(dim=160, activation="glu"), +) +``` + +When you `instantiate` the outer config, nested configs are resolved +recursively. There is one subtlety worth knowing +([deferred vs. eager nesting](#deferred-vs-eager-instantiation) below): by +default, nested **`nn.Module`** configs are passed *as config* to their +parent rather than pre-built, so each module constructs its own children and +can inspect or modify their configs first. Non-module callables (e.g. weight +init factories) are instantiated eagerly. + +## Interpolation: `"${...}"` + +Configs are OmegaConf trees, so any value can reference another value by its +dotted path. We use this constantly so that a dimension is written *once* and +flows everywhere: + +```python +config.net = LazyConfig(ClassificationResNet)( + hidden_dim=160, + data_dim=2, + in_proj_cfg=LazyConfig(torch.nn.Linear)( + in_features="${net.in_channels}", + out_features="${net.hidden_dim}", # tracks net.hidden_dim + ), + norm_cfg=LazyConfig(torch.nn.LayerNorm)( + normalized_shape="${net.hidden_dim}", + ), + ..., +) +``` + +Change `hidden_dim` in one place and every `"${net.hidden_dim}"` follows. +Interpolations are resolved at instantiation / override time, not when the +config is declared, which is why they survive being edited and overridden. + +You can also reference across top-level sections. For example, a kernel's +cache length tracks the dataset's canvas size: + +```python +L_cache = "${dataset.canvas_size}" +``` + +### Inline arithmetic + +Two small conveniences let you do math inside configs: + +- **Plain arithmetic strings** are evaluated by `instantiate`. A value like + `"3 * ${net.hidden_dim}"` resolves the interpolation and then evaluates + the arithmetic (`"3 * 160"` → `480`). We use this for the Hyena short + conv, which operates on a 3×-width tensor: + + ```python + short_conv_cfg = LazyConfig(torch.nn.Conv2d)( + in_channels="3 * ${net.hidden_dim}", + out_channels="3 * ${net.hidden_dim}", + groups="3 * ${net.hidden_dim}", + ..., + ) + ``` + +- **The `${eval:...}` resolver** handles arithmetic in CLI overrides and + trainer interpolations, e.g. + `"${eval:'${trainer.samples_per_epoch} // (${train.batch_size} * 2)'}"`. + Only arithmetic on literal numbers is permitted, with no function calls or + attribute access, so configs stay safe to load. (The `${eval:...}` resolver + additionally allows `**`; plain arithmetic strings support `+ - * / // %` + but not power.) + +## `PLACEHOLDER`: a hole to be filled later + +{py:data}`~nvsubquadratic.lazy_config.PLACEHOLDER` is a sentinel marking a +field whose value isn't known yet at declaration time. It plays two roles: + +1. **A required slot to be filled later.** A base config marks a field + `PLACEHOLDER` to say "this *must* be supplied before the object is built." + The hole is filled in one of two ways: + + - *By an experiment file, before running.* The spatial-recall base config + sets `sequence_mixer_cfg=PLACEHOLDER`; each ablation file then asserts the + hole is still empty and plugs in a mixer, a self-documenting contract: + + ```python + block_cfg = LazyConfig(ResidualBlock)( + sequence_mixer_cfg=PLACEHOLDER, ... # filled in by the experiment file + ) + ``` + + ```python + assert config.net.block_cfg.sequence_mixer_cfg == PLACEHOLDER + config.net.block_cfg.sequence_mixer_cfg = get_hyena_mixer_cfg() + ``` + + - *By code, at build time.* The optimizer is declared with + `params=PLACEHOLDER` because the parameters don't exist until the network + is constructed. The Lightning wrapper fills the slot when it builds the + optimizer: `_build_optimizer` in + {py:mod}`experiments.lightning_wrappers.base_lightning_wrapper` resolves + the config to a dict and overwrites `params` with the real parameter + groups. (Note this path constructs the optimizer directly rather than + through `instantiate`.) + +1. **A "don't build me yet" guard.** While `instantiate` walks an argument + tree, any *nested* config that still contains a `PLACEHOLDER` is left as a + config dict rather than constructed, so a half-specified subtree is never + handed to a constructor mid-build. For example, a `block_cfg` whose + `sequence_mixer_cfg` hole hasn't been filled is passed through as config + instead of built. This guard is a check on *nested* configs only: a bare + top-level value like the optimizer's `params=PLACEHOLDER` is not itself + guarded, which is why that slot is filled in by code (role 1) before the + object is built. + +## A full example, end to end + +[`examples/mnist_classification/ccnn_4_160_hyena_rope_qknorm.py`](https://github.com/NVIDIA-BioNeMo/nvSubquadratic/blob/main/examples/mnist_classification/ccnn_4_160_hyena_rope_qknorm.py) +is a complete, self-contained config. Every config file exposes a +`get_config()` that returns an +{py:class}`~experiments.default_cfg.ExperimentConfig`. The skeleton: + +```python +def get_config() -> ExperimentConfig: + config = ExperimentConfig() # typed dataclass of sensible defaults + + # 1. Dataset — a LazyConfig pointing at a LightningDataModule + config.dataset = LazyConfig(MNISTDataModule)( + data_dir=".data/mnist", + batch_size=BATCH_SIZE, + seed=config.seed, + task="classification", + ..., + ) + + # 2. Network — one big nested LazyConfig tree (the architecture) + config.net = LazyConfig(ClassificationResNet)( + in_channels=INPUT_CHANNELS, + hidden_dim=NUM_HIDDEN_CHANNELS, + data_dim=DATA_DIM, + in_proj_cfg=LazyConfig(torch.nn.Linear)( + in_features="${net.in_channels}", out_features="${net.hidden_dim}" + ), + norm_cfg=LazyConfig(torch.nn.LayerNorm)(normalized_shape="${net.hidden_dim}"), + block_cfg=LazyConfig(ResidualBlock)( + sequence_mixer_cfg=LazyConfig(QKVSequenceMixer)( + hidden_dim="${net.hidden_dim}", + mixer_cfg=LazyConfig(Hyena)(...), + ), + mlp_cfg=LazyConfig(MLP)(dim="${net.hidden_dim}", activation="glu", ...), + ..., + ), + ) + + # 3. Lightning wrapper, optimizer (note params=PLACEHOLDER) + config.lightning_wrapper_class = LazyConfig(ClassificationWrapper)() + config.optimizer = LazyConfig(torch.optim.AdamW)( + params=PLACEHOLDER, lr=LEARNING_RATE, weight_decay=WEIGHT_DECAY + ) + + # 4. Plain typed sub-configs for training / schedule / logging + config.train = TrainConfig(batch_size="${dataset.batch_size}", iterations=...) + config.scheduler = SchedulerConfig( + name="cosine", total_iterations="${train.iterations}" + ) + config.wandb = WandbConfig(job_group="mnist_classification", ...) + + return config +``` + +Notice the mix: `dataset`, `net`, `lightning_wrapper_class`, and `optimizer` +are **`LazyConfig`s** (objects built lazily), while `train`, `scheduler`, and +`wandb` are **plain typed dataclasses** from +{py:mod}`experiments.default_cfg` (values read directly). You only set what +differs from the defaults. + +### How `run.py` consumes it + +[`experiments/run.py`](https://github.com/NVIDIA-BioNeMo/nvSubquadratic/blob/main/experiments/run.py) +is the entry point. It: + +```bash +PYTHONPATH=. python experiments/run.py \ + --config examples/mnist_classification/ccnn_4_160_hyena_rope_qknorm.py \ + dataset.batch_size=64 optimizer.lr=3e-4 +``` + +1. Loads the file and calls `get_config()`. + +1. Applies the `key=value` CLI overrides (after checking none of them clobber + an interpolated field, as described below). + +1. Builds the objects exactly when needed: + + ```python + datamodule = instantiate(config.dataset) + network = instantiate(config.net) + model = instantiate(config.lightning_wrapper_class, network=network, cfg=config) + ``` + +The config tree is also serialised and logged to W&B and printed to the +console as a Rich tree, so the exact specification of every run is captured. + +## The base-config and ablation pattern + +It is worth studying the +[`examples/spatial_recall_2d/`](https://github.com/NVIDIA-BioNeMo/nvSubquadratic/tree/main/examples/spatial_recall_2d) +directory to see this pattern in practice. +Instead of copy-pasting a 150-line config per ablation, we factor the shared +structure into helper functions and keep each experiment file tiny. + +[`spatial_recall_2d/base_config.py`](https://github.com/NVIDIA-BioNeMo/nvSubquadratic/blob/main/examples/spatial_recall_2d/base_config.py) +exposes `base_experiment_config(...)` which returns a fully-formed +`ExperimentConfig` with the network, optimizer, scheduler, and callbacks all +wired, but with the sequence mixer and the dataset left as `PLACEHOLDER`. +[`mixer_defaults.py`](https://github.com/NVIDIA-BioNeMo/nvSubquadratic/blob/main/examples/spatial_recall_2d/mixer_defaults.py) +provides `get_hyena_mixer_cfg()`, `get_mamba_mixer_cfg()`, and +`get_attention_mixer_cfg()`, each a `LazyConfig` for one mixer family. + +An individual ablation file is then short but still self-contained: + +```python +# examples/spatial_recall_2d/emnist_regression_color_selection/ccnn_hyena_s.py +def get_config() -> ExperimentConfig: + config = spatial_recall_2d_base_experiment_config( + in_channels=3, + out_channels=1, + hidden_dim=256, # the "S" size + training_iterations=20_000, + wandb_job_group="spatial_recall_2d_emnist_color_selection_s", + ) + + # Fill the mixer hole — swap this line for get_mamba/attention to ablate + assert config.net.block_cfg.sequence_mixer_cfg == PLACEHOLDER + config.net.block_cfg.sequence_mixer_cfg = get_hyena_mixer_cfg() + + # Fill the dataset hole + assert config.dataset == PLACEHOLDER + config.dataset = base_emnist_spatial_recall_2d_dataset_config( + target_size=16, + canvas_size=64, + batch_size=64, + use_colored_frames=True, + num_items=4, + placement="random", + with_mask=False, + normalize_input=True, + ) + return config +``` + +To ablate the mixer, you change one line and the filename (`ccnn_hyena_s.py` +→ `ccnn_mamba_xs.py`). To ablate size, you change `hidden_dim`. Each variant +is its own file, so each run still carries a *complete* config, yet the diff +between any two variants is one or two lines. This is the backtrackable, +ablation-friendly workflow the config system is built for. + +Because helpers return `LazyConfig` trees built from `"${...}"` +interpolations, the swapped-in mixer automatically picks up `hidden_dim`, +`data_dim`, `num_blocks`, and `canvas_size` from the surrounding config, so +you never restate them. + +## Overriding from the command line + +Any field can be overridden with `key=value` positional arguments to +`run.py`. Values are auto-typed (`int` → `float` → `None` → `bool` → tuple → +list → `str`), and dotted paths reach into nested configs: + +```bash +PYTHONPATH=. python experiments/run.py --config \ + train.batch_size=32 \ + optimizer.lr=3e-4 \ + net.hidden_dim=256 +``` + +Two guardrails are worth knowing: + +- **You cannot override an interpolated field.** If a field's current value + is a `"${...}"` string, overriding it directly is rejected + (`verify_no_interpolator_overwrites`). Override the *source* of the + interpolation instead: set `net.hidden_dim=256`, not the dozen places + that read `"${net.hidden_dim}"`. +- **Add genuinely new keys with `+`.** `key=value` requires the key to exist + (typo protection); Hydra-style `+key=value` force-adds it. + +Overrides also feed the deterministic run name, so a sweep over +`optimizer.lr` produces distinctly named, individually reproducible runs. + +## Two details that explain the rest + +### Deferred vs. eager instantiation + +When `instantiate` walks the argument tree it decides, per nested config, +whether to build it now or pass it through as config: + +- **`nn.Module` subclasses are passed through as config** (a `DictConfig`), + not pre-built. The parent module receives the child's config and + constructs it itself. This lets a network inspect or tweak block configs + (injecting `drop_path_rate`, reading `"${net.num_blocks}"`) before building + them. +- **Non-module callables are instantiated eagerly.** For example, a + weight-init factory like `partial_wang_init_fn_with_num_layers(num_layers=...)` + is resolved to a function and handed to the module ready to use. + +Passing `recursive_instantiate=True` overrides this and builds everything +top-down; the default (`False`) is what the module tree relies on. + +### Serialisation + +{py:func}`~nvsubquadratic.lazy_config.save_config` / +{py:func}`~nvsubquadratic.lazy_config.load_config` round-trip a config to +YAML via OmegaConf, and `config_to_dict` (used by `run.py`) flattens the +whole tree (`LazyConfig`s, dataclasses, function references) into a +JSON-serialisable dict for W&B and the console tree. This is what makes a run +fully recoverable from its logged config. + +## Mental model / cheat sheet + +- `LazyConfig(Target)(**kwargs)` → a config dict; nothing is built yet. +- `instantiate(cfg)` → the actual object. +- Nest `LazyConfig`s to describe a whole module tree as one expression. +- `"${a.b.c}"` references another field; resolved at build/override time. +- `"3 * ${net.hidden_dim}"` does inline arithmetic after interpolation. +- `PLACEHOLDER` marks a hole that must be filled and blocks premature builds. +- One file per experiment; a base helper + a one-line swap per ablation. +- Override with `key=value` on the CLI; never override a `"${...}"` field + directly. Change its source. + +For the bigger picture of where configs sit in the stack, see +{doc}`architecture`; for runnable recipes, the +[`examples/`](https://github.com/NVIDIA-BioNeMo/nvSubquadratic/tree/main/examples) +configs. + +``` + +``` diff --git a/docs/ops/MIXED_BC_PLAN.md b/docs/ops/MIXED_BC_PLAN.md new file mode 100644 index 00000000..260285c1 --- /dev/null +++ b/docs/ops/MIXED_BC_PLAN.md @@ -0,0 +1,349 @@ +--- +orphan: true +--- + +# Mixed Boundary-Condition FFT Convolution — Plan & Tracker + +**Status:** In progress (v1 ops + tests) +**Branch:** `dwromero/mixed-bc-fftconv` +**Owner:** dwromero +**Started:** 2026-05-20 + +This is the working plan and tracker for adding **per-axis boundary +condition** support to the FFT-based convolution operators +(`nvsubquadratic/ops/`) and the modules that consume them +(`nvsubquadratic/modules/ckconv_nd.py`, etc.). + +If you pick this up later, read the [survey & decisions](#1-context--decisions) +section, then check [§5 Tracked questions](#5-tracked-questions--revisit-items) +for what still needs to be done. + +______________________________________________________________________ + +## 1. Context & decisions + +### Motivation + +Several Well PDE datasets have boundaries that are **periodic on some axes +and non-periodic (WALL or OPEN) on others**, e.g. + +| Dataset | x | y | z | +| ------------------------------ | -------- | -------- | ---- | +| `rayleigh_benard` | periodic | wall | — | +| `viscoelastic_instability` | periodic | wall | — | +| `turbulent_radiative_layer_2D` | periodic | open | — | +| `turbulent_radiative_layer_3D` | periodic | periodic | open | +| `rayleigh_taylor_instability` | periodic | periodic | wall | +| `helmholtz_staircase` | open | per-face | — | +| `acoustic_scattering_maze` | per-face | per-face | — | + +Today the FFT-conv operators expose only two **global** modes +(`fft_padding="zero"` or `"circular"`) selected per module. There is no way +to say "periodic on x, zero-padded on y" for one and the same conv. + +### Decisions (locked in) + +1. **API** — extend `fft_padding` to accept either a **single mode + string** (applies to every axis) or a **list of mode strings** (one + per spatial axis): + + ```python + fft_padding: str | Sequence[str] = "zero" + # "zero" -> all axes zero-padded. + # "circular" -> all axes periodic. + # ["circular", "zero"] -> 2D, x periodic + y zero-padded. + # ["zero", "circular", "zero"] -> 3D, etc. + # ("circular", "zero") -> tuple form is equivalent. + ``` + + Internally everything normalises to a tuple `periodic: tuple[bool, ...]` + of length `data_dim`. Three inputs are deliberately rejected with an + error that redirects to the canonical form: + + - **Booleans** (`(True, False)`, `True`): the per-axis intent is not + obvious from the boolean values. + - **Comma-separated strings** (`"circular, zero"`): redundant with the + list form and gives two ways to say the same thing; we keep one + canonical per-axis form. + +1. **WALL vs OPEN** — both treated as **zero-padded linear** at the + conv level. Physical distinctions are handled elsewhere (data + normalisation, loss, etc.). Per-face BC (different BC on opposite faces + of the same axis) is **out of scope** for v1. + +1. **Kernel size per axis** — auto-derived from the per-axis BC, **not** a + new knob: + + | axis BC | grid_lens per axis (in `CKConvND.forward`) | SIREN kernel size on that axis | + | ------------ | -------------------------------------------- | ------------------------------ | + | periodic | `(s+1)//2` (≡ today's `grid_type="single"`) | `≈ s` | + | non-periodic | `s` (≡ today's `grid_type="double"`) | `≈ 2s − 1` | + + When `fft_padding` is a tuple, the legacy `grid_type` argument must be + `None` (or omitted) — raise on conflict, no silent overrides. + +1. **First-PR scope** — **ops + tests only.** Module wiring + (`CKConvND`/`CKConvMultiheadND`), Well config updates, fp16, multihead, + and the `subq_ops` CUDA path are explicitly deferred (see §4). + +1. **`subq_ops` CUDA kernel** — left at zero-only. Any + `fft_backend="subq_ops"` + mixed BC will raise in the future + `CKConvND` wiring PR. + +______________________________________________________________________ + +## 2. Algorithm — per-axis recipe + +The mixed N-D FFT convolution applies, **independently per spatial axis**: + +| axis is | FFT length `F_d` | post-IFFT crop range | phase ramp shift on that axis | +| ------------ | ------------------------------ | -------------------------------- | ----------------------------- | +| periodic | `N_d` (no padding) | `0 : N_d` (no crop) | `−(K_d − 1) // 2` | +| non-periodic | `min(N_d + (K_d+1)//2, 2·N_d)` | `K_d//2 : K_d//2 + N_d` (center) | `0` (no shift) | + +The whole conv is still **one** `rfftn` / `irfftn` over all spatial dims; +the per-axis recipe just feeds different per-axis `F_d` values to `s=` and +different per-axis slices to the post-IFFT crop. Phase ramps are the +product of per-axis 1-D ramps (length-1 broadcast on non-periodic axes). + +Edge behaviour required by the tests: + +- All `periodic == False` → bit-identical to existing `fftconv*` linear op. +- All `periodic == True` → bit-identical to existing `circular_fftconv*` op. + +______________________________________________________________________ + +## 3. v1 deliverables (this PR) + +### Code + +- [x] `nvsubquadratic/ops/mixed_fftconv.py` (fp32): + - Self-contained per-axis 1-D phase-ramp LRU cache; N-D ramp built on + demand by broadcasted multiplication so non-periodic axes contribute + nothing (skipped, not just length-1). + - 1D / 2D / 3D BHL variants. + - BHL `_w_reshape` wrappers (BLH inputs). + - Channel-chunked variants. + - Automatic dispatch to the existing linear / circular ops in the + all-False / all-True cases (no perf cost for legacy usage). +- [x] No `nvsubquadratic/ops/__init__.py` exists — callers import from + submodules directly (matches existing convention). + +> **Note:** an earlier draft of this plan called for extracting +> `_PhaseRampCache1D/2D/3D` from `circular_fftconv.py` into a shared +> `_phase_ramp.py`. We **dropped** that refactor for v1 because the existing +> caches hard-code `FFT_shape == input_shape` per axis, which is only true +> for the all-circular case. There is nothing to share without first +> generalising the API. We may revisit this as a unification refactor +> (see §5 Q4). + +### Tests + +- [x] `tests/ops/test_mixed_fftconv.py` — 76 tests, all passing on H100: + - Reference comparison against + `F.pad(x, mode="circular"|"constant")` + `F.conv{1,2,3}d(padding=0)` + for every per-axis combo across 1D / 2D / 3D, including the K==N + edge case and even-K kernels (asymmetric "same" padding). + - Sanity: all-False matches existing `fftconv*`; all-True matches + existing `circular_fftconv*`. + - BHL ↔ BLH wrapper equivalence. + - Chunked vs non-chunked equivalence. + - Backward / gradient equivalence vs the spatial reference (1D, 2D, 3D). + - `use_phase_shift=False` (roll on periodic axes only) matches + `use_phase_shift=True` for every combo. + - Shortcut term equivalence and dtype preservation (fp32, bf16). + - Validation errors: wrong `periodic` length, oversized kernel, + mismatched shortcut dtype. +- [x] Re-ran existing FFT-conv suites — **101 passed**, no regressions: + - `tests/ops/test_fftconv.py` + - `tests/ops/test_circular_fftconv.py` + - `tests/ops/test_fftconv_chunked.py` + +### Docs + +- [x] Updated `docs/ops/README.md` "File map" with the new + `mixed_fftconv.py` row. + +______________________________________________________________________ + +## 4. Deferred — follow-up work + +Each item below is intentionally **not** part of this PR. They are +listed so we don't lose track. + +### 4.1 Module wiring — `CKConvND` ✅ DONE (2026-05-20) + +- [x] `CKConvND` (`nvsubquadratic/modules/ckconv_nd.py`): + - Accepts `fft_padding: str | Sequence[str]` in two forms: a single + mode string (`"zero"` / `"circular"`) that applies to every axis, or + a list of mode strings (e.g. `["circular", "zero"]`) — one per axis. + - Resolves to a normalised per-axis `periodic` tuple via + `_resolve_periodic` (length checked against `data_dim`). + - When `fft_padding` is a per-axis list, **requires** `grid_type=None` + and raises `ValueError` otherwise. When it's a single mode string, + `grid_type` is required as before. + - Boolean inputs (e.g. `(True, False)`) and comma-separated strings + (`"circular, zero"`) are explicitly rejected with errors that + redirect to the list form. + - Per-axis `grid_lens` and per-axis `L_cache` halving auto-derived in + tuple mode (halve only on periodic axes); helper + `_grid_is_single_per_axis(grid_type, periodic)` is the single source + of truth used by both `__init__` (L_cache) and `forward`/`flop_count`. + - Dispatch: tuple mode routes through `MIXED_FFT_FUNCTIONS[_CHUNKED]` + (wrapped by `_wrap_mixed_op` to bind `periodic`). All-False / all-True + tuples internally fall back to the legacy linear / circular ops + bit-identically (verified by tests). String mode keeps the legacy + `FFT_FUNCTIONS` tables unchanged. + - Validation: + - `is_causal=True` + any periodic axis → `ValueError`. + - `fft_backend="subq_ops"` + tuple `fft_padding` → `ValueError`. + - `use_fp16_fft=True` + tuple `fft_padding` → `NotImplementedError` + (planned for v2; see §4.2). + - `use_chunked_fftconv` allowed with tuple `fft_padding` for **all** + per-axis combos (including all-True — new capability vs the legacy + string-mode where circular + chunked was an error). + - `flop_count` uses per-axis padded sizes: `s` on periodic axes, + `min(s + (k+1)//2, 2*s)` on non-periodic axes. +- [x] `mixed_fftconv.py` op: `K <= N` assertion relaxed on non-periodic + axes to `K <= 2*N` to match the "double-grid" SIREN kernel size + (`2N - 1`) that `CKConvND` produces on non-periodic axes. +- [x] `mixed_fftconv.py`: added `*_w_reshape_chunked` BLH wrappers so the + module dispatch table is symmetric with the legacy chunked path. +- [x] Module-level tests in `tests/modules/test_ckconv_nd_mixed_bc.py`: + resolver / helper unit tests, validation errors, per-axis kernel + shape, tuple-vs-string bit-identical equivalence, mixed-mode + forward correctness (matches the underlying op called directly), + BHL/BLH layout equivalence, chunked-vs-non-chunked, and FLOP + accounting (mixed sits between all-zero and all-circular). + +### 4.2 fp16 variant (deferred) + +### 4.2 fp16 variant + +- [ ] `nvsubquadratic/ops/mixed_fftconv_fp16.py`: + - Per-axis cuFFT-fp16 constraints: pad-up to power-of-2 on linear axes; + require input dim power-of-2 on periodic axes (fallback to fp32 with a + warning otherwise). + - Reuse centering / DC-correction logic from `circular_fftconv_fp16` + on the periodic axes only. +- [ ] `tests/ops/test_mixed_fftconv_fp16.py`. +- [ ] `CKConvND` fp16 dispatch update. + +### 4.3 2D multi-head variant + +- [ ] `fftconv2d_multihead_mixed_bhl` (and `_bhi`) in + `nvsubquadratic/ops/fftconv_multihead.py`. +- [ ] `CKConvMultiheadND` wiring + tests (same shape of changes as + `CKConvND` above; 2D only). + +### 4.4 Well experiment configs + +- [ ] Add Hyena variants with mixed BC for the datasets that need them. + Start with the ones we actually run: + - `rayleigh_benard` — periodic on x → `(True, False)` + - `viscoelastic_instability` — periodic on x → `(True, False)` + - `turbulent_radiative_layer_2D` — periodic on x → `(True, False)` + - `turbulent_radiative_layer_3D` — periodic on x,y → `(True, True, False)` + - `rayleigh_taylor_instability` — periodic on x,y → `(True, True, False)` +- [ ] Decide whether to fix `examples/well/v1/supernova_explosion_64/cfg_hyena.py` + which uses `FFT_PADDING="circular"` but the dataset is all-OPEN. v2 is + already corrected. +- [ ] Longer-term: read Well HDF5 `boundary_conditions` from the + datamodule and auto-derive `periodic_axes` instead of hard-coding it + per config. + +### 4.5 Per-face BCs + +- [ ] Investigate whether `acoustic_scattering_maze` and + `helmholtz_staircase` benefit from a *per-face* BC treatment + (different BC on opposite faces of the same axis). +- [ ] If yes: design a follow-up that goes beyond per-axis circular/linear. + v1 maps any non-periodic axis to symmetric zero-pad. + +### 4.6 Custom CUDA path (`subq_ops`) + +- [ ] If/when we want mixed-BC to use the + `subquadratic_ops_torch.fft_conv2d` fast path, the upstream kernel + must grow per-axis BC support. v1 leaves the kernel zero-only and + raises in `CKConvND`. + +### 4.7 SIREN kernel generator anisotropy + +- [ ] Re-verify that the SIREN kernel module + (`nvsubquadratic/modules/kernels_nd.py`) actually handles + anisotropic `grid_lens` (e.g. `(64, 128)`) cleanly end-to-end — + including positional embedding, masks, monitors, and FLOP accounting. + Used today for the all-isotropic case; we need it for the per-axis + grid in the mixed path. + +______________________________________________________________________ + +## 5. Tracked questions & "revisit" items + +These are not bugs we plan to fix in this PR, just things we noticed +along the way that may want attention later. + +### Q1. Silent mutation of user-provided `L_cache` in `CKConvND.__init__` + +`ckconv_nd.py` (~L266–289) does `copy.deepcopy(kernel_cfg)` and then +silently halves `L_cache` when `grid_type=="single"` so the SIREN +positional grid spans `[-1, 1]` over the actual kernel size. The +behaviour is *intentional* (per the in-code comment) but the **silent +mutation of a user-provided config** is mildly surprising. Cleaner +pattern would be to compute the effective `L_cache` at construction +time without round-tripping through the config object. Not a v1 bug — +revisit in a separate refactor. + +For the mixed path, the same adjustment must become **per-axis** (halve +only on periodic axes); design that in the module-wiring PR (§4.1). + +### Q2. `subq_ops` 2D linear kernel — make it BC-aware? + +Out of scope for v1, but worth a conversation with the kernel authors +before we commit to a divergent fast path that only supports zero-pad. + +### Q3. Auto-wiring from Well HDF5 boundary_conditions + +`experiments/datamodules/pde/well.py` can return per-sample BC metadata +from HDF5, but no model code consumes it. Long-term it would be cleaner +to derive the `periodic_axes` tuple from the dataset rather than +hard-coding in each config. + +### Q4. Unify the phase-ramp cache across `circular_fftconv` and `mixed_fftconv` + +After v1, the codebase will have two parallel phase-ramp caches: + +- The original `_PhaseRampCache1D/2D/3D` in `circular_fftconv.py`, which + hard-codes `FFT_shape == input_shape` per axis (only valid for the + all-circular case). +- A new general N-D cache local to `mixed_fftconv.py`, which also handles + per-axis padded `F_d` and zero shifts on linear axes. + +These can be unified into a single shared helper once we are happy with +the mixed op's API. Doing so is a pure refactor (no behaviour change for +either op) and should be its own PR. + +______________________________________________________________________ + +## 6. Changelog + +- **2026-05-20** — Plan written, feature branch created. +- **2026-05-20** — v1 ops landed: `mixed_fftconv.py` (fp32, 1D/2D/3D, + BHL + BLH wrappers + channel-chunked), 76-test suite passing, + no regressions in existing FFT-conv tests. +- **2026-05-20** — Tolerance tightening + analytical-truth tests + (impulse response, DC response) added; 102 op tests passing. +- **2026-05-20** — `CKConvND` integration landed: per-axis + `fft_padding: Sequence[bool]` API, auto-derived per-axis grid + + `L_cache` halving, unified dispatch via `MIXED_FFT_FUNCTIONS*`, + module-level test suite (`tests/modules/test_ckconv_nd_mixed_bc.py`). + Full regression sweep on `tests/ops + tests/modules` → 802 passed. +- **2026-05-21** — Public API revised on PR review: `fft_padding` now + accepts mode-name strings only. Two forms: a single mode string + (`"zero"` / `"circular"`) that applies to every axis, or a list of mode + strings (e.g. `["circular", "zero"]`) — one per axis. Comma-separated + strings and bool tuples are both rejected with redirecting errors. + Rationale: `(True, False)` did not convey which axis was periodic, and + having both a comma-string form and a list form was two ways to say + the same thing. The list form reads identically in Python and OmegaConf + / YAML overrides. `rayleigh_benard` config updated to the list form. diff --git a/docs/ops/README.md b/docs/ops/README.md index 63fc706a..b7bbc934 100644 --- a/docs/ops/README.md +++ b/docs/ops/README.md @@ -1,6 +1,6 @@ -# `nvsubquadratic.ops` — FFT convolution primitives +# `nvsubquadratic.ops`: FFT convolution primitives -This folder contains the **lowest-level building blocks** of the library: FFT-based convolution operators that turn an `O(N · K)` spatial convolution into an `O(N log N)` frequency-domain product. They are the workhorses behind every subquadratic mixer in the library (Hyena, CKConv), and are kept here as plain functions — no `nn.Module` state, no learned parameters — so that higher-level modules can compose them freely. +This folder contains the lowest-level building blocks of the library: FFT-based convolution operators that turn an `O(N · K)` spatial convolution into an `O(N log N)` frequency-domain product. Every subquadratic mixer in the library (Hyena, CKConv, multi-head variants) is built on them. They are kept here as plain functions, with no `nn.Module` state and no learned parameters, so that higher-level modules can compose them freely. If you are reading the paper alongside this codebase, this is the file to start with. @@ -14,7 +14,7 @@ $$ y[n] = \sum_{m} x[n - m] \cdot k[m] $$ -costs `O(N · K)` per channel. When `K` is small (e.g. a 3×3 image kernel) that is fine. When `K` is **comparable to `N`** — the regime Hyena-style models live in, where each layer's effective receptive field can span the whole input — the spatial cost grows quadratically with sequence length. +costs `O(N · K)` per channel. When `K` is small (e.g. a 3×3 image kernel) that is fine. The trouble starts when `K` is **comparable to `N`**, the regime Hyena-style models live in, where each layer's effective receptive field can span the whole input. There the spatial cost grows quadratically with sequence length. The **convolution theorem** lets us replace the spatial convolution with an element-wise product in the frequency domain: @@ -28,23 +28,23 @@ Two flavours show up throughout the folder: | Flavour | What it computes | When to use | | ---------------------------------- | ---------------------------------------------------------------------------------------- | ----------------------------------------------------------------------------------------------------------------------------------------- | -| **Linear** (`fftconv*`) | Standard convolution, zero-padded so no wrap-around occurs, then cropped to "same" size. | Default choice — matches `torch.nn.ConvNd` semantics. | +| **Linear** (`fftconv*`) | Standard convolution, zero-padded so no wrap-around occurs, then cropped to "same" size. | Default choice; matches `torch.nn.ConvNd` semantics. | | **Circular** (`circular_fftconv*`) | Periodic convolution where the kernel wraps around the input boundary. | When you want global mixing under periodic boundary conditions, or when input and kernel are the same size (no padding needed → cheaper). | ______________________________________________________________________ ## File map -| File | Precision | Conv type | Channel mixing | When you'd reach for it | -| --------------------------------------------------------------------------- | --------- | ------------- | -------------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | -| [fftconv.py](../../nvsubquadratic/ops/fftconv.py) | fp32 | linear | depthwise | The default. 1D/2D/3D, causal & non-causal. | -| [circular_fftconv.py](../../nvsubquadratic/ops/circular_fftconv.py) | fp32 | circular | depthwise | Periodic boundaries (e.g. PDEs, ARC grids), or when `K = N` so padding is wasteful. | -| [mixed_fftconv.py](../../nvsubquadratic/ops/mixed_fftconv.py) | fp32 | per-axis BC | depthwise | **Mixed boundaries** — periodic on some spatial axes, zero-padded on others (e.g. Well's `rayleigh_benard`, `viscoelastic_instability`, `turbulent_radiative_layer`). Routes to the existing linear/circular ops in the all-False/all-True cases. | -| [fftconv_chunked.py](../../nvsubquadratic/ops/fftconv_chunked.py) | fp32 | linear | depthwise | Memory-constrained training; processes channels in chunks. Has a global flag so models can opt in transparently. | -| [fftconv_custom.py](../../nvsubquadratic/ops/fftconv_custom.py) | fp32 | linear | depthwise | Wraps optional fused CUDA kernels (`subquadratic_ops_torch.fft_conv2d` for 2D non-causal, `fft_causal_conv1d` for 1D causal) behind the same API as `fftconv.py`. | -| [causal_conv1d_custom.py](../../nvsubquadratic/ops/causal_conv1d_custom.py) | fp32 | direct causal | depthwise | Non-FFT 1D causal kernels (`causal_conv1d` short conv, `b2b_causal_conv1d` fused proj-gate-mixer-gate). Use for kernels short enough that FFT overhead dominates, or as a fused-Hyena building block. | +| File | Precision | Conv type | Channel mixing | When you'd reach for it | +| -------------------------------------------------- | --------- | ------------- | -------------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------ | +| {ref}`fftconv.py ` | fp32 | linear | depthwise | The default. 1D/2D/3D, causal & non-causal. | +| {ref}`circular_fftconv.py ` | fp32 | circular | depthwise | Periodic boundaries (e.g. PDEs, ARC grids), or when `K = N` so padding is wasteful. | +| {ref}`mixed_fftconv.py ` | fp32 | per-axis BC | depthwise | **Mixed boundaries**: periodic on some spatial axes, zero-padded on others (e.g. Well's `rayleigh_benard`, `viscoelastic_instability`, `turbulent_radiative_layer`). Routes to the existing linear/circular ops in the all-False/all-True cases. | +| {ref}`fftconv_chunked.py ` | fp32 | linear | depthwise | Memory-constrained training; processes channels in chunks. Has a global flag so models can opt in transparently. | +| {ref}`fftconv_custom.py ` | fp32 | linear | depthwise | Wraps optional fused CUDA kernels (`subquadratic_ops_torch.fft_conv2d` for 2D non-causal, `fft_causal_conv1d` for 1D causal) behind the same API as `fftconv.py`. | +| {ref}`causal_conv1d_custom.py ` | fp32 | direct causal | depthwise | Non-FFT 1D causal kernels (`causal_conv1d` short conv, `b2b_causal_conv1d` fused proj-gate-mixer-gate). Use for kernels short enough that FFT overhead dominates, or as a fused-Hyena building block. | -`mixed_boundary_conditions.md` describes the per-axis boundary-condition support (periodic on some spatial axes, zero-padded on others) used by the Well PDE datasets. +The [FP16 circular FFT convolution report](https://github.com/NVIDIA-BioNeMo/nvSubquadratic/blob/main/reports/fp16_fft_convolution/REPORT.md) contains the full derivation of the numerically stable fp16 circular conv (dual mean-centering + inclusion-exclusion geometric correction). Read it if you are touching the fp16 path or want to understand the math behind those `T1, T2, T3, T4` terms in the code. ______________________________________________________________________ @@ -53,21 +53,21 @@ ______________________________________________________________________ Every function name encodes its contract: ``` -[causal_] fftconv {1d|2d|3d} _ fp32 _ {bhl|blh} [_w_reshape] [_chunked] +[causal_] fftconv {1d|2d|3d} _ {fp32|fp16} _ {bhl|blh} [_w_reshape] [_chunked] ``` | Part | Meaning | | ------------------ | ---------------------------------------------------------------------------------------------------------------------------------------------- | | `causal_` | Output at position `n` only sees inputs at positions `≤ n`. 1D only. | | `1d` / `2d` / `3d` | Spatial rank. | -| `fp32` | Internal compute precision. The output dtype always matches `x.dtype` regardless. | +| `fp32` / `fp16` | Internal compute precision. The output dtype always matches `x.dtype` regardless. | | `bhl` / `blh` | Memory layout. `bhl` = channels-first (`[B, H, *spatial]`). `blh` = channels-last (`[B, *spatial, H]`). | | `_w_reshape` | Wrapper that accepts BLH input, internally reshapes to BHL (faster), and reshapes back. The recommended entry point for channels-last callers. | | `_chunked` | Processes channels in groups to reduce peak GPU memory. | So `causal_fftconv1d_fp32_bhl_w_reshape` is: causal 1D FFT conv, fp32 internal, accepts channels-last input, internally uses the channels-first kernel. -The CUDA-accelerated wrappers in `fftconv_custom.py` drop the `_fp32_` token because the underlying kernel manages its own precision internally — so the same name in `fftconv_custom` is `causal_fftconv1d_bhl_w_reshape`. The direct-conv wrappers in `causal_conv1d_custom.py` (`causal_conv1d`, `b2b_causal_conv1d`) do not follow this scheme because they are thin pass-throughs to the upstream API; see their docstrings for shapes. +The CUDA-accelerated wrappers in `fftconv_custom.py` drop the `_fp32_` / `_fp16_` token because the underlying kernel manages its own precision internally, so the same name in `fftconv_custom` is `causal_fftconv1d_bhl_w_reshape`. The direct-conv wrappers in `causal_conv1d_custom.py` (`causal_conv1d`, `b2b_causal_conv1d`) do not follow this scheme because they are thin pass-throughs to the upstream API; see their docstrings for shapes. ______________________________________________________________________ @@ -78,7 +78,7 @@ Everything in this folder follows two layouts. Pick whichever matches your surro - **BHL** (channels-first): `x: [B, H, *spatial]`, `kernel: [1|B, H, *K_dims]`. Standard for `torch.nn.ConvNd`-style modules. Faster under the hood because FFT runs on contiguous spatial axes without a transpose. - **BLH** (channels-last): `x: [B, *spatial, H]`, `kernel: [1|B, *K_dims, H]`. Common in transformer-style code. Use the `_w_reshape` variants. -The kernel's leading dim is either `1` (shared kernel across the batch — the standard depthwise case) or `B` (per-sample kernel, e.g. FiLM-conditioned Hyena where each sample gets its own kernel). +The kernel's leading dim is either `1` (shared kernel across the batch, the standard depthwise case) or `B` (per-sample kernel, e.g. FiLM-conditioned Hyena where each sample gets its own kernel). ### The shortcut term @@ -88,7 +88,7 @@ $$ y \leftarrow y + \mathrm{shortcut} \odot x $$ -i.e. a per-channel residual scale. This is *not* a generic skip connection — it fuses a specific algebraic shortcut that shows up repeatedly in Hyena-style gating, saving a separate kernel launch. Pass `None` if you don't need it. +i.e. a per-channel residual scale. This is *not* a generic skip connection; it fuses a specific algebraic shortcut that shows up repeatedly in Hyena-style gating, saving a separate kernel launch. Pass `None` if you don't need it. ______________________________________________________________________ @@ -109,9 +109,16 @@ ______________________________________________________________________ - Channels-first (`[B, H, …]`) → use `_bhl` directly. - Channels-last (`[B, …, H]`) → use `_bhl_w_reshape`. Benchmarks show this is faster than a true `_blh` op because the FFT runs on contiguous spatial axes. +1. **What's my precision budget?** + + - fp32 is the default, always correct. + - For aggressive memory/throughput savings on **power-of-2 spatial dims**, use the `*_fp16_*` variant. The fp16 ops use `norm="ortho"` and (for circular) dual mean-centering to stay within fp16 dynamic range; see the [FP16 circular FFT convolution report](https://github.com/NVIDIA-BioNeMo/nvSubquadratic/blob/main/reports/fp16_fft_convolution/REPORT.md). + - If your spatial dims aren't powers of two, stay in fp32 (cuFFT requires power-of-2 for fp16 transforms). + 1. **Am I OOMing?** - - Try `fftconv_chunked` — splits the channel dim into groups to cap peak memory. Default chunk size 128 gives ~26% memory savings for ~11% overhead. + - Try `fftconv_chunked`, which splits the channel dim into groups to cap peak memory. Default chunk size 128 gives ~26% memory savings for ~11% overhead. + - Or combine: `fftconv_fp16.py` already provides `_chunked` variants that stack both savings. 1. **Is there a fused CUDA kernel for my shape?** @@ -123,17 +130,17 @@ ______________________________________________________________________ ## Numerical notes -- All operators **accept any input dtype** but cast to fp32 before the FFT. The output is returned in the **original dtype of `x`** — no need for a manual cast on the caller side. -- The fp32 ops are correct for any input range. +- All operators **accept any input dtype** but cast to the internal compute precision (fp32 or fp16) before the FFT. The output is returned in the **original dtype of `x`**, so no manual cast is needed on the caller side. +- The fp32 ops are correct for any input range. The fp16 ops impose two constraints: spatial dims must be powers of two (cuFFT), and the dynamic range is handled by mean-centering both `x` and `k` (see derivation doc). - The non-causal linear ops match a standard `torch.nn.ConvNd(padding='same')` up to floating-point rounding. The circular ops match `torch.nn.functional.conv*d` after a circular pad. Both are exercised in `tests/`. ______________________________________________________________________ ## Related modules -- **[`nvsubquadratic/modules/kernels_nd.py`](../../nvsubquadratic/modules/kernels_nd.py)** — learned kernel parametrisations that produce the kernels these ops consume. -- **[`nvsubquadratic/modules/hyena_nd.py`](../../nvsubquadratic/modules/hyena_nd.py)** — the Hyena operator, the main consumer of these ops. -- **[`nvsubquadratic/modules/ckconv_nd.py`](../../nvsubquadratic/modules/ckconv_nd.py)** — the CKConv operator that composes these primitives. +- **[`nvsubquadratic/modules/kernels_nd.py`](../../nvsubquadratic/modules/kernels_nd.py)**: learned kernel parametrisations that produce the kernels these ops consume. +- **[`nvsubquadratic/modules/hyena_nd.py`](../../nvsubquadratic/modules/hyena_nd.py)**: the Hyena operator, the main consumer of these ops. +- **[`nvsubquadratic/modules/ckconv_nd.py`](../../nvsubquadratic/modules/ckconv_nd.py)** / **[`ckconv_multihead_nd.py`](../../nvsubquadratic/modules/ckconv_multihead_nd.py)**: CKConv variants that compose these primitives. ```{toctree} --- diff --git a/docs/ops/mixed_boundary_conditions.md b/docs/ops/mixed_boundary_conditions.md index 93653a6e..62d3bcdd 100644 --- a/docs/ops/mixed_boundary_conditions.md +++ b/docs/ops/mixed_boundary_conditions.md @@ -28,8 +28,8 @@ mixed path closes that gap. Wall and open boundaries are both treated as **zero-padded linear** at the convolution level; physical distinctions (if any) are handled elsewhere -(data normalisation, loss). Per-face boundary conditions — a different BC -on opposite faces of the same axis — are not supported (see +(data normalisation, loss). Per-face boundary conditions, meaning a different +BC on opposite faces of the same axis, are not supported (see [Limitations](#limitations)). ______________________________________________________________________ @@ -59,7 +59,7 @@ deliberately rejected with an error that redirects to the canonical form: ### Kernel size per axis -The kernel grid size is auto-derived from the per-axis boundary condition — +The kernel grid size is auto-derived from the per-axis boundary condition; it is **not** a separate knob. When `fft_padding` is a list, the legacy `grid_type` argument must be `None` (a conflict raises rather than silently overriding): @@ -119,6 +119,6 @@ ______________________________________________________________________ - **Custom CUDA path** (`fft_backend="subq_ops"`) supports zero-padding only; combining it with a per-axis `fft_padding` raises. Use `fft_backend="torch_fft"` (the default) for mixed boundaries. -- **Auto-wiring** — the periodic axes are specified per config. The Well +- **Auto-wiring**: the periodic axes are specified per config. The Well datamodule can read `boundary_conditions` from the HDF5 metadata, but model code does not yet consume it to derive `periodic` automatically. diff --git a/docs/reports.md b/docs/reports.md index 35e3a421..eda0b21b 100644 --- a/docs/reports.md +++ b/docs/reports.md @@ -11,12 +11,13 @@ for the regeneration conventions (snake_case topic names, scripts take ## Current topics -| Topic | Summary | -| -------------------------------------------------------------------------------------------------------------------------------------------------------- | ----------------------------------------------------------------------------------------------------------------------------------------------------------------- | -| [`ckconv_block_diagonal_kernel/`](https://github.com/NVIDIA-BioNeMo/nvSubquadratic/blob/main/reports/ckconv_block_diagonal_kernel/REPORT.md) | Block-diagonal multi-ω₀ SIREN kernel + block-aligned Gaussian mask for ViT-5 hybrid Hyena. Resolution scaling rule (`ω₀ ← m·ω₀`) verified across 1×/2×/4× grids. | -| [`siren_omega0_dimensional_scaling/`](https://github.com/NVIDIA-BioNeMo/nvSubquadratic/blob/main/reports/siren_omega0_dimensional_scaling/REPORT.md) | SIREN ω₀ dimensional scaling rule and supporting figures. | -| [`spatial_recall/`](https://github.com/NVIDIA-BioNeMo/nvSubquadratic/blob/main/reports/spatial_recall/REPORT.md) | Qualitative target-vs-prediction snapshots for the 1D/2D/3D EMNIST spatial-recall task suite (simple copy, mask selection, color selection, color conditioning). | -| [`vit5_imagenet_dataloader_profiling/`](https://github.com/NVIDIA-BioNeMo/nvSubquadratic/blob/main/reports/vit5_imagenet_dataloader_profiling/REPORT.md) | Feb-2026 investigation that diagnosed the CPU-decode bottleneck on ViT-5-Small ImageNet and motivated the move to the DALI-fused dataloader. | +| Topic | Summary | +| -------------------------------------------------------------------------------------------------------------------------------------------------------- | ----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | +| [`ckconv_block_diagonal_kernel/`](https://github.com/NVIDIA-BioNeMo/nvSubquadratic/blob/main/reports/ckconv_block_diagonal_kernel/REPORT.md) | Block-diagonal multi-ω₀ SIREN kernel + block-aligned Gaussian mask for ViT-5 hybrid Hyena. Resolution scaling rule (`ω₀ ← m·ω₀`) verified across 1×/2×/4× grids. | +| [`siren_omega0_dimensional_scaling/`](https://github.com/NVIDIA-BioNeMo/nvSubquadratic/blob/main/reports/siren_omega0_dimensional_scaling/REPORT.md) | SIREN ω₀ dimensional scaling rule and supporting figures. | +| [`spatial_recall/`](https://github.com/NVIDIA-BioNeMo/nvSubquadratic/blob/main/reports/spatial_recall/REPORT.md) | Qualitative target-vs-prediction snapshots for the 1D/2D/3D EMNIST spatial-recall task suite (simple copy, mask selection, color selection, color conditioning). | +| [`vit5_imagenet_dataloader_profiling/`](https://github.com/NVIDIA-BioNeMo/nvSubquadratic/blob/main/reports/vit5_imagenet_dataloader_profiling/REPORT.md) | Feb-2026 investigation that diagnosed the CPU-decode bottleneck on ViT-5-Small ImageNet and motivated the move to the DALI-fused dataloader. | +| [`fp16_fft_convolution/`](https://github.com/NVIDIA-BioNeMo/nvSubquadratic/blob/main/reports/fp16_fft_convolution/REPORT.md) | FP16 circular FFT convolution: dual-mean-centering derivation, accuracy + throughput vs the FP32 reference, and why the (stable, exact) FP16 path was kept opt-in rather than shipped as the default. | ## Adding a new report diff --git a/docs/repository_overview.md b/docs/repository_overview.md index fc3316ea..98d1610c 100644 --- a/docs/repository_overview.md +++ b/docs/repository_overview.md @@ -1,185 +1,53 @@ # Repository overview A map of what lives where at the repo root. The library code -(`nvsubquadratic/`) and the training driver (`experiments/`) sit -alongside the per-task configs (`examples/`), the perf measurement -tree (`benchmarks/`), and the supporting infrastructure (`scripts/`, -`tests/`, `reports/`, `docs/`). +(`nvsubquadratic/`) and the training driver (`experiments/`) sit alongside +the per-task configs (`examples/`), the perf-measurement tree +(`benchmarks/`), and the supporting infrastructure (`scripts/`, `tests/`, +`reports/`, `docs/`). ## Layout ```text nvSubquadratic/ -├── nvsubquadratic/ library code — see "library tree" below -├── experiments/ training framework (PyTorch Lightning) -│ ├── run.py CLI entry point -│ ├── trainer.py construct_trainer (checkpoints, precision, W&B) -│ ├── default_cfg.py typed ExperimentConfig dataclasses -│ ├── lightning_wrappers/ task-specific wrappers (classification, diffusion, regression, …) -│ ├── datamodules/ LightningDataModule subclasses (ImageNet, MNIST, WELL, …) -│ ├── callbacks/ FiLM monitor, image-grid viz, EMA, walltime checkpointer, … -│ └── utils/ cli + checkpointing helpers -├── examples/ LazyConfig recipes that feed experiments.run -│ ├── imagenet_classification/ -│ ├── vit5_imagenet/ ViT-5 baseline suite (v1–v5) -│ ├── spatial_recall_{1,2,3}d/ and spatial_recall_v2/ -│ ├── well/ The Well PDE benchmark suite -│ └── overview_tracker.md active experimental roadmap -├── benchmarks/ performance measurement (the canonical home) -│ ├── README.md ViT-5-Small headline throughput tables -│ ├── compare_flops.py FLOP comparison across ViT-5 variants -│ ├── benchmark_imagenet_diffusion_gpu.py -│ ├── benchmark_patch_size_2d.py -│ ├── ops/ op-level benchmarks (fftconv2d / MLP / subq-ops) -│ ├── vit5_imagenet/ ViT-5 throughput, profile, verify, validate -│ └── well/ WELL dataloader / training-step / VRMSE -├── reports/ frozen-in-time investigations with regen scripts -│ ├── ckconv_block_diagonal_kernel/ -│ ├── siren_omega0_dimensional_scaling/ -│ ├── spatial_recall/ -│ └── vit5_imagenet_dataloader_profiling/ -├── scripts/ utilities (data prep, sanity, SLURM, viz) -│ ├── slurm/ SLURM submit scripts (portable wrapper + per-experiment) -│ ├── data/ data prep (ImageNet folder extraction, normalization stats, …) -│ ├── visualization/ kernel viewers + throughput plot -│ └── check_gpu_availability.py, license_check.py, … -├── tests/ correctness tests -│ ├── conftest.py shared fixtures -│ ├── ops/, modules/, networks/, parallel/ per-package test trees -│ └── test_*.py top-level integration tests -├── docs/ Sphinx documentation site (this site) -│ ├── conf.py, index.rst site config + landing -│ ├── getting_started.md, architecture.md, repository_overview.md -│ ├── api_reference/ curated API per area -│ ├── ops/ FFT-ops math primer + FP16 derivation -│ ├── examples/index.md, benchmarks.md, reports.md -│ └── _templates/, _static/ autosummary templates + custom CSS -├── docs-tracker.md documentation coverage status per file -├── CONVENTIONS.md Google-style docstring guide and PR checklist -├── README.md top-level install / overview -├── pyproject.toml project metadata, dependencies, ruff config -├── Dockerfile production container -├── nvsubquadratic.def Apptainer/Singularity recipe -└── setup_conda_env.sh local conda bootstrap +├── nvsubquadratic/ library code — ops, modules, networks, parallel, utils +├── experiments/ PyTorch Lightning training driver (run.py, wrappers, datamodules, callbacks) +├── examples/ per-task LazyConfig training recipes fed to experiments.run +├── benchmarks/ performance measurement (throughput, FLOP / scaling, op-level) +├── reports/ frozen-in-time technical investigations with regen scripts +├── scripts/ utilities (data prep, evaluation, SLURM, visualization) +├── tests/ correctness tests mirroring the library layout +├── docs/ this Sphinx documentation site +├── CONVENTIONS.md docstring style guide and PR checklist +├── README.md top-level install / overview +├── pyproject.toml project metadata, dependencies, ruff config +├── Dockerfile production container +├── nvsubquadratic.def Apptainer / Singularity recipe +└── setup_conda_env.sh local conda bootstrap ``` -## Library tree (`nvsubquadratic/`) +## The library (`nvsubquadratic/`) -The library itself is organised bottom-up: function-only convolution -primitives in `ops/`, then `nn.Module`-shaped mixers and blocks in -`modules/`, then full architectures in `networks/`. +The library is organised bottom-up: function-only convolution primitives, +then `nn.Module`-shaped building blocks, then full architectures. -```text -nvsubquadratic/ -├── lazy_config.py deferred-instantiation system (LazyConfig, instantiate) -├── ops/ function-only convolution primitives -│ ├── fftconv.py fp32 reference FFT conv (1D / 2D / 3D, linear) -│ ├── fftconv_fp16.py half-precision linear-conv variants -│ ├── circular_fftconv.py fp32 periodic-boundary FFT conv -│ ├── circular_fftconv_fp16.py half-precision periodic variants -│ ├── fftconv_multihead.py multi-head + low-rank factorisations -│ ├── fftconv_chunked.py peak-FFT-memory chunking helpers -│ ├── fftconv_custom.py subquadratic_ops_torch CUDA wrappers -│ ├── causal_conv1d_custom.py direct (non-FFT) CUDA causal conv1d wrappers -│ └── mixed_fftconv.py per-axis mixed boundary-condition FFT conv -├── modules/ nn.Module building blocks -│ ├── hyena_nd.py Hyena ND mixer (two-gate sandwich, CP) -│ ├── mamba_nd.py Mamba SSM (ND, selective, raster scan) -│ ├── attention.py multi-head attention (RoPE, ND) -│ ├── vit5_attention.py ViT-5 register-aware attention -│ ├── vit5_hyena_adapter.py Hyena drop-in for ViT-5 -│ ├── sequence_mixer.py operator-agnostic QKV dispatch -│ ├── condition_mixer.py cross-attention conditioning mixer -│ ├── kernels_nd.py SIREN / RFF kernels (multi-ω₀, block-diag) -│ ├── ckconv_nd.py CKConv ND (implicit k_θ(p)) -│ ├── ckconv_multihead_nd.py multi-head CKConv (low-rank) -│ ├── distributed_depthwise_conv_nd.py CP-aware depthwise convs -│ ├── causal_conv1d.py left-only-padded Conv1d wrapper -│ ├── subq_ops_causal_conv1d.py nn.Conv1d-compatible CUDA depthwise -│ ├── residual_block.py pre-norm + mixer + MLP (+ AdaLN-Zero) -│ ├── vit5_residual_block.py ViT-5 residual block (LayerScale, registers) -│ ├── patchify.py strided patch embedding / unpatchify -│ ├── position_encoding.py axis-factorised learned PE -│ ├── masks_nd.py exponential / Gaussian / block-aligned masks -│ ├── mlp.py GELU / SwiGLU / GLU MLP -│ ├── film.py FiLM kernel generator + register pooling -│ ├── grn.py GlobalResponseNorm (ConvNeXt V2) -│ ├── layer_scale.py LayerScale γ·F(x) -│ ├── drop_path.py stochastic depth -│ ├── rms_norm.py RMSNorm + PerHeadRMSNorm -│ ├── rms_norm_channel_first.py channel-first RMSNorm -│ └── schedulers.py ResumableSequentialLR -├── networks/ end-to-end architectures -│ ├── general_purpose_resnet.py ResidualNetwork (LazyConfig stack) -│ ├── classification_resnet.py GAP-readout classification head -│ ├── vit5_classification.py ViT-5 hybrid Hyena/attention backbone -│ ├── jit.py JiT diffusion backbone -│ ├── jit_utils.py JiT helpers (RoPE, RMSNorm, sin-cos PE) -│ ├── huggingface_diffusers.py HF DiT / UVit adapters -│ └── baselines/ -│ ├── unet_convnext.py Well UNet-ConvNeXt baseline -│ └── unet_convnext_v2.py …with fixed finest-skip -├── parallel/ context-parallel primitives -│ ├── a2a_comms.py AllToAllSingle (autograd-aware) -│ └── utils.py init_parallel_state + zigzag split/gather -├── utils/ cross-area utilities -│ ├── init.py weight-init factories -│ ├── qk_norm.py QK normalization (apply + L2Norm) -│ ├── rope.py rotary position embedding (1D / 2D / 3D) -│ └── quack_utils.py QuACK capability probe -└── testing/ - └── utils.py compute_relative_error -``` - -## What each top-level directory does - -**`nvsubquadratic/`** — The library. Function-only ops, `nn.Module` -building blocks, full networks, context-parallel primitives, and a -deferred-instantiation system (`LazyConfig`) that every example -config relies on. See {doc}`api_reference/index` for the curated API. - -**`experiments/`** — The training driver. Lightning wrappers, -datamodules, callbacks, the `construct_trainer` helper, and the -`run.py` CLI entry point. Consumes a network + datamodule + wrapper -via a `LazyConfig` tree from `examples/` and runs it through Lightning. - -**`examples/`** — Per-task training recipes. Each subdirectory is a -config tree (LazyConfig dataclasses) that fully describes one -experiment. Running it is -`python -m experiments.run --config examples/.../.py`. The -live roadmap is at -[`examples/overview_tracker.md`](https://github.com/NVIDIA-BioNeMo/nvSubquadratic/blob/main/examples/overview_tracker.md). - -**`benchmarks/`** — The single home for performance measurement. -Op-level microbenchmarks (`benchmarks/ops/`), end-to-end model -throughput (`vit5_imagenet/`, `well/`), and FLOP / scaling -comparisons. Headline numbers are pulled into the {doc}`benchmarks` -docs page. - -**`reports/`** — Frozen-in-time technical investigations. One -`REPORT.md` per topic alongside the scripts and figures it cites. -Indexed at {doc}`reports`. - -**`scripts/`** — Utility / glue scripts. SLURM submit drivers -(`scripts/slurm/`), data prep (`scripts/data/`), kernel viewers (`scripts/visualization/`), -and standalone sanity scripts. No benchmarks live here — those moved -to `benchmarks/`. - -**`tests/`** — Correctness tests, mirroring the library's per-package -structure: `tests/ops/`, `tests/modules/`, `tests/networks/`, -`tests/parallel/`. +- **`ops/`**: function-only FFT convolution primitives (linear / circular / + mixed boundary, fp32 / fp16, chunked, and fused-CUDA wrappers). +- **`modules/`**: `nn.Module` building blocks: mixers (Hyena, Mamba, + attention, CKConv), learned kernels, residual blocks, norms, and MLPs. +- **`networks/`**: end-to-end architectures (ResNet / CCNN, ViT-5, the JiT + diffusion backbone, and UNet-ConvNeXt baselines). +- **`parallel/`**: context-parallel primitives (`init_parallel_state`, + AllToAll, zigzag split / gather). +- **`utils/`**, **`metrics/`**, **`testing/`**: weight init, RoPE, QK-norm, + and the QuACK probe; FID; relative-error helpers. -**`docs/`** — This Sphinx documentation site. Narrative pages -(Getting Started, Architecture, this Repository Overview, Examples, -Benchmarks, Reports, Ops Overview) plus the curated -{doc}`api_reference/index`. +See {doc}`api_reference/index` for the curated, per-symbol API. ## Where to go next -- {doc}`architecture` — the three-layer +- {doc}`architecture`: the three-layer nvSubquadratic / subquadratic-ops / megatron-core story. -- {doc}`api_reference/index` — the curated API for each +- {doc}`api_reference/index`: the curated API for each `nvsubquadratic/` area. -- {doc}`examples/index` — per-dataset training recipes that compose - these networks. -- {doc}`ops/README` — math primer for the FFT convolution primitives. +- {doc}`ops/README`: math primer for the FFT convolution primitives.