Skip to content

Support current gradient checkpointing APIs - #367

Open
zjn20030811 wants to merge 1 commit into
InternRobotics:mainfrom
zjn20030811:fix-gradient-checkpointing-api
Open

Support current gradient checkpointing APIs#367
zjn20030811 wants to merge 1 commit into
InternRobotics:mainfrom
zjn20030811:fix-gradient-checkpointing-api

Conversation

@zjn20030811

Copy link
Copy Markdown

Summary

  • Support both per-module and model-level gradient-checkpointing hook signatures used by diffusers releases.
  • Propagate the checkpoint function supplied by newer releases to checkpointable submodules.
  • Add a small CPU regression test that enables and disables checkpointing through NextDiTCrossAttn.

Problem

LuminaNextDiT2DModel implemented the older
_set_gradient_checkpointing(module, value=False) hook. Newer diffusers
calls the hook with enable= and gradient_checkpointing_func=, so the
default InternVLA-N1 configuration failed during model construction whenever
gradient checkpointing was enabled.

The adapter accepts both conventions while preserving the existing behavior
on older releases.

Validation

  • PYTHONPATH=F:\\contributor\\.deps-wan-audit;F:\\contributor\\audit-internnav-next2 python -m pytest tests/unit_test -q (5 passed, 1 skipped)
  • The focused regression passes with diffusers 0.35.2 and 0.32.1.
  • python -m py_compile internnav/model/basemodel/internvla_n1/nextdit_traj.py tests/unit_test/test_gradient_checkpointing_compat.py
  • git diff --check

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