Skip to content

Commit 244f0ca

Browse files
committed
feat(peft): add fused LoRA SwiGLU/ReLU² MLP with recompute backward
Fuses gate+up+down+activation into a single autograd Function that saves only (x, gate_out, up_out) and recomputes the activation and down-projection input in backward, roughly halving MLP activation memory at equal speed during LoRA SFT. SwiGLU forward/backward use elementwise Triton kernels (with in-place backward buffer reuse) and a pure-torch fallback when Triton is unavailable; matmuls stay on cuBLAS. Covers SiLU-SwiGLU (gate/up/down) and non-gated ReLU² (e.g. Nemotron-H dense) MLPs. install_fused_lora_mlp() swaps each LoRA-applied MLP's forward and falls back to the per-linear path at runtime under DTensor (TP/EP), DoRA, or active dropout, keeping it correct under sharding. Already wired from lora.py; opt out via NEMO_AUTOMODEL_DISABLE_FUSED_LORA_MLP=1. Activation recompute follows Megatron-Core's SwiGLUFunction; the fused LoRA-MLP and in-place buffer reuse follow Unsloth's LoRA_MLP (both Apache-2.0). Signed-off-by: Alexandros Koumparoulis <akoumparouli@nvidia.com>
1 parent ae64756 commit 244f0ca

2 files changed

Lines changed: 658 additions & 0 deletions

File tree

0 commit comments

Comments
 (0)