New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

[nnx] improve Optimizer metadata propagation #4180

Merged

copybara-service merged 1 commit into main from nnx-fix-optimizer

Sep 9, 2024

Collaborator

cgarciae commented Sep 8, 2024 •

edited

Loading

Changes

Creates two new OptState subtypes OptArray and OptVariable.
Improves Optimizer to properly propagate sharding metadata from model to opt_state by wrapping VariableState leaves in OptVariable.
Small cleanups to nnx.with_partitioning and Optimizer.


          [nnx] improve Optimizer metadata propagation

4c47287

cgarciae force-pushed the nnx-fix-optimizer branch from ffe04bd to 4c47287 Compare

September 8, 2024 12:53

cgarciae added the pull ready label

IvyZX approved these changes

View reviewed changes

copybara-service bot merged commit 671130b into main

18 checks passed

copybara-service bot deleted the nnx-fix-optimizer branch

September 9, 2024 19:10

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels