save memory by deleting params. lint. #277

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

Merged

entrpn merged 1 commit into main from deallocate_params_tree

Oct 24, 2025

Collaborator

entrpn commented Oct 23, 2025

Deletes params pytree after creating TrainState to reduce HBM usage. Lints the code.

The only change is in this line: https://github.com/AI-Hypercomputer/maxdiffusion/compare/deallocate_params_tree?expand=1#diff-69cda939e98b489aca1cc8aa543ecc537d9f4bc6c58fa767cbe01cd3636aacf3R311


          save memory by deleting params. lint.

95b0c98

entrpn requested a review from susanbao

October 23, 2025 22:56

github-actions bot commented Oct 23, 2025

e2e testgrid: https://8bcf50593faf4ea38060e236169827e5-dot-us-central1.composer.googleusercontent.com/dags/maxdiffusion_tpu_e2e/grid

susanbao approved these changes

View reviewed changes

entrpn merged commit c9229c3 into main

3 of 4 checks passed

entrpn deleted the deallocate_params_tree branch

October 24, 2025 01:30

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

Labels

None yet