Skip to content

[PyTorch] Refactor GroupedLinear eager forward and backward - #50

Draft
pggPL wants to merge 1 commit into
mainfrom
grouped_linear_refactor
Draft

pggPL wants to merge 1 commit into
mainfrom
grouped_linear_refactor

Conversation

@pggPL

@pggPL pggPL commented Sep 10, 2026

Copy link
Copy Markdown
Owner

Description

Extract GroupedLinear forward/backward computation and saved-state handling from the autograd wrapper. Both split-quantize and fused grouped-tensor execution expose shared helpers that the compile integration in #32 can reuse.

This is the first of two stacked PRs; it contains the eager refactor.

Type of change

  • Code refactoring

Changes

  • Introduce explicit forward/backward argument containers and thin autograd wrappers for both execution paths. Share backward state cleanup and MCore main-grad handling; omit unused argument fields.
  • Move computation and context setup into shared functions, including fused grouped GEMM.
  • Preserve eager saved-tensor layout, packed gradients, grouped parameters, save-original-input, delayed wgrad and main-grad accumulation.

Validation

RTX Ada workstation, current source:

  • test_grouped_linear.py: 718 passed, 558 skipped.
  • 12 comparisons against main across BF16/FP8, save-original-input and gradient settings: bit-exact outputs/gradients and identical saved-tensor shapes, strides, offsets and storage-alias groups.
  • 72 BF16 fused comparisons against main across grouped weight/bias flags, save-original-input, gradient settings and normal/main-grad/delayed wgrad: bit-exact values and saved metadata. Grouped GEMM was replaced with matmuls for this temporary probe.
  • Changed-file pylint, Black and diff checks pass.

Native fused GEMM requires Hopper/Blackwell and was not exercised on RTX Ada. Full Python lint reports the existing import-order diagnostic in unchanged attention/fused_mla_q_uproj.py.

Checklist

  • I have read and followed the contributing guidelines
  • Existing eager tests pass on the available GPU
  • Changes introduce no lint diagnostics
  • Compile integration is isolated in the follow-up PR

@pggPL
pggPL force-pushed the grouped_linear_refactor branch from 50ca413 to d007d1c Compare September 10, 2026 11:24
@pggPL pggPL changed the title [PyTorch] Refactor GroupedLinear forward and backward argument handling [PyTorch] Refactor GroupedLinear eager forward and backward Sep 10, 2026
Extract shared computation and backward state for both split-quantize and grouped-tensor execution. Preserve eager saved-tensor layout and gradient allocation so compile adapters can reuse these functions in a follow-up change.

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
@pggPL
pggPL force-pushed the grouped_linear_refactor branch from d007d1c to a3c005d Compare September 11, 2026 11:34
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant