Expand description
Serial-over-chunks SSD with a custom recompute backward. Serial SSD with a custom, memory-efficient backward.
The forward + Mamba2BackendExt impl live in serial_recalculated; the
registered autodiff [backward] node and the recompute-based gradient math
in [combined_backward] together mirror the official ssd_combined.py,
saving ~⅓ of the training memory versus storing every intermediate.
Modules§
- backward
- The registered custom
Backwardnode (autodiff op). - combined_
backward - Recompute-based gradient math (the memory-efficient backward).
- serial_
recalculated 🔒 - Serial SSD with a custom, memory-efficient backward (Mamba-2)
Traits§
- Mamba2
Autodiff Backend Ext - Marker for an autodiff-capable backend that also implements the SSD extension trait (so the custom memory-efficient backward is available).
- Mamba2
Backend Ext - Extends the backend and wraps it for
burn.