Skip to content

Add optimizer-owned gradient clipping extension point - #10

Draft
ryan-li-ryanllli-aws wants to merge 1 commit into
mainfrom
ryanllli/compiled-grad-clip-extension
Draft

ryan-li-ryanllli-aws wants to merge 1 commit into
mainfrom
ryanllli/compiled-grad-clip-extension

Conversation

@ryan-li-ryanllli-aws

@ryan-li-ryanllli-aws ryan-li-ryanllli-aws commented Oct 14, 2026 •

Copy link
Copy Markdown
Collaborator

Summary

  • add a configurable GradientClipper interface owned by the optimizer component
  • add an optional Optim.Config.gradient_clipper and build it once for the optimizer lifetime
  • preserve the existing eager clipping path when no override is configured
  • add CPU unit coverage for component construction, shared dispatch, and eager compatibility

Motivation

Gradient clipping operates on optimizer-owned mutable gradient state, but there
was no scoped component seam for a private device backend to provide its own
implementation. This change makes that behavior explicit in optimizer
configuration without replacing the eager utility process-wide.

The private device backend implementation remains outside TorchTitan core.

User story: integrate a private device backend

A private device backend has its own gradient clipping implementation,
including any compile support it needs. Its integration package supplies a
concrete GradientClipper component and opts in through optimizer
configuration:

config.optim.gradient_clipper = PrivateDeviceGradientClipper.Config()

The optimizer constructs the configured component once and uses the same
dispatch from its standard and specialized update paths. When no component is
configured, the existing eager gradient clipping behavior is unchanged.

Testing

  • SKIP=pyrefly-check pre-commit run --all-files
  • focused Pyrefly check for the changed production files
  • Python 3.11 bytecode compilation for the changed Python files

The added unit tests were not executed locally because the available
environment does not contain the repository's current dependency set. They are
included for execution in an environment with that stack.

@ryan-li-ryanllli-aws
ryan-li-ryanllli-aws force-pushed the ryanllli/compiled-grad-clip-extension branch 2 times, most recently from 376ea8e to 2f7c6ae Compare October 14, 2026 18:05
@ryan-li-ryanllli-aws
ryan-li-ryanllli-aws force-pushed the ryanllli/compiled-grad-clip-extension branch from 2f7c6ae to aec3a81 Compare October 14, 2026 18:11
@ryan-li-ryanllli-aws ryan-li-ryanllli-aws changed the title Add optimizer-owned compiled gradient clipping extension point Add optimizer-owned gradient clipping extension point Oct 14, 2026
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