Fix training crash in addmm_act by falling back to unfused ops when grad is enabled - #613
Conversation
…rad is enabled The fused _addmm_activation path detaches weights and casts to bf16, so it only works under no_grad. Since the SAM 3.1 update, vitdet's Mlp.forward() calls addmm_act unconditionally, which makes any training/fine-tuning run (e.g. the documented Roboflow configs) fail with 'ValueError: Expected grad to be disabled.' (facebookresearch#610). Fall back to the equivalent standard linear+activation ops when gradients are enabled; inference keeps the fused fast path unchanged.
|
Hi @Rudr4Khunt1! Thank you for your pull request and welcome to our community. Action RequiredIn order to merge any pull request (code, docs, etc.), we require contributors to sign our Contributor License Agreement, and we don't seem to have one on file for you. ProcessIn order for us to review and merge your suggested changes, please sign at https://code.facebook.com/cla. If you are contributing on behalf of someone else (eg your employer), the individual CLA may not be sufficient and your employer may need to sign the corporate CLA. Once the CLA is signed, our tooling will perform checks and validations. Afterwards, the pull request will be tagged with If you have received this in error or have any questions, please contact us at cla@meta.com. Thanks! |
|
Thank you for signing our Contributor License Agreement. We can now accept your code for this (and any) Meta Open Source project. Thanks! |
meta-cla
Bot
added
the
CLA Signed
1 participant
Problem
Since the SAM 3.1 update,
vitdet.py::Mlp.forward()unconditionally callsperflib/fused.py::addmm_act(), which is inference-only (it detaches weights and casts to bf16) and explicitly raises when gradients are enabled. As a result, any training or fine-tuning run on currentmain— including the documented Roboflow configs (roboflow_v100_full_ft_100_images.yaml) — crashes in the first vision-backbone forward pass with:This is the crash reported in #610 (closed with a "pin to the pre-3.1 commit" workaround) and related to the question in #509. The workaround forfeits the 3.1 improvements; this PR makes training work on current
main.Fix
When
torch.is_grad_enabled(), fall back to the mathematically equivalent standard ops (linear+relu/gelu), which support autograd. Inference is untouched — the fused fast path still runs whenever grads are disabled.Validation
With this patch, a full 20-epoch fine-tune of the Roboflow config (custom 200-image single-class COCO dataset, 1× L4,
--use-cluster 0 --num-gpus 1) runs to completion: losses converge normally (composite 214 → 61; bbox L1 halved; GIoU −43%), and the resulting checkpoint improves held-out detection F1 from 0.394 (zero-shot) to 0.520 — confirming gradients flow correctly through the patched path. Inference before/after the patch is unchanged.Happy to rework if you'd prefer gating at the module level (e.g. branching in
vitdet.Mlp.forward) instead of insideaddmm_act.