Add a backwards linear function to be used with the fused mla q up-proj - #3330
Add a backwards linear function to be used with the fused mla q up-proj#3330chaseblock wants to merge 4 commits into
Conversation
Signed-off-by: Chase Block <cblock@nvidia.com>
Requires TE PR: NVIDIA/TransformerEngine#3330 Signed-off-by: Chase Block <cblock@nvidia.com>
for more information, see https://pre-commit.ci
Greptile SummaryThe PR adds a reusable linear-backward helper for the fused MLA Q up-projection and exposes dgrad, wgrad, and bias-gradient results.
Confidence Score: 4/5The PR is not yet safe to merge because BF16 fused MLA training still loses the projection bias gradient. The non-FP8 path delegates bias-gradient computation to the weight-gradient GEMM, but the new wrapper supplies no bias tensor, so the binding does not allocate or return db. Files Needing Attention: transformer_engine/pytorch/module/linear.py Important Files Changed
Reviews (3): Last reviewed commit: "Handle biad gradient in lin bwd wrapper." | Re-trigger Greptile |
Signed-off-by: Chase Block <cblock@nvidia.com>
Signed-off-by: Chase Block <cblock@nvidia.com>
| inputmat=x_saved, | ||
| weight_fp8=w_q, | ||
| saved_weight=w_q, | ||
| bias=None, |
There was a problem hiding this comment.
BF16 bias gradient remains missing
When a BF16 fused MLA backward call sets use_bias=True, this wrapper still passes bias=None; the non-FP8 path delegates bgrad to the weight-gradient GEMM, which does not allocate db without a bias tensor, causing backward_linear to return grad_bias=None and preventing the projection bias from being updated.
Knowledge Base Used: PyTorch Fused Modules (transformer_engine/pytorch/module)
Description
Please include a brief summary of the changes, relevant motivation and context.
Fixes # (issue)
Type of change
Changes
Please list the changes introduced in this PR:
Checklist: