Skip to content

add fp64 gradcheck for geglu activation and gate+up paths - #31

Open
tejasexpress wants to merge 1 commit into
Forge-researchlab:mainfrom
tejasexpress:fix/geglu-fp64-gradcheck
Open

add fp64 gradcheck for geglu activation and gate+up paths #31
tejasexpress wants to merge 1 commit into
Forge-researchlab:mainfrom
tejasexpress:fix/geglu-fp64-gradcheck

Conversation

@tejasexpress

Copy link
Copy Markdown

Base branch: please target dev, not main. GitHub defaults the base to
main; change it in the dropdown above if it is not already dev.

dev doesn't exist on the repo right now, only main, so this targets main.

What this changes

fixes #11 .. adds torch.autograd.gradcheck at fp64 for both geglu paths := ForgeGEGLUFunction (activation-only, separate gate/up) and ForgePackedGEGLUFunction (gate+up fused). runs through the Triton kernel, similar approach as RoPE's fp64 gradcheck. small shapes (2,3,8) and both approximate="tanh" and "none".

Numbers

n/a, no math changed

Correctness

  • uv run ruff check .
  • Relevant pytest suites pass
  • Numerics checked against a reference, not just "it runs"

Anything a reviewer should know

using preserve_inputs=True in both new tests. With the default False,
gradcheck fails with "backward is not reentrant" becaus geglu_backward writes
the gradient into the same storage it saved for backward (dgate = gate
when preserve_inputs=False), so calling .backward() twice on the same
graph reads back corrupted values on the second call. that's the
default geglu() uses, so retain_graph=True + two backward calls would
silently give wrong gradients. didn't do anything abt it here, out of scope for
#11 , can open a separate issue if wanted.

tested on Colab Tesla T4

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.

GeGLU: fp64 gradcheck is deferred

1 participant