Repository navigation
[Bug]: _gradient_checkpointing Error in Inference time #321
Description
Bug Description
NextDiTCrossAttn.__init__() crashes with TypeError when using diffusers >= 0.31 (afaik) due to a breaking API change in _set_gradient_checkpointing(). The custom LuminaNextDiT2DModel in nextdit_traj.py implements the old signature (self, module, value=False), but newer diffusers calls it with (enable=True, gradient_checkpointing_func=...).
The gradient checkpointing flag is hardcoded to True in NextDiTCrossAttnConfig (line 28 of nextdit_crossattn_traj.py), but the if path (nextdit_traj.py:328) only activates during training (self.training and self.gradient_checkpointing), so it has no effect during inference.
I would suggest adding False as the default flag.
Steps to Reproduce
Please try to provide a minimal example to reproduce the bug. Error messages and stack traces are also helpful.
# Using the Docker internutopia environment (v1.2) with no code modifications
conda activate internutopia
python scripts/eval/eval.py --config scripts/eval/configs/h1_internvla_n1_async_cfg.py
Traceback (most recent call last):
File "/root/InternNav/scripts/eval/eval.py", line 49, in <module>
main()
...
File
"/root/InternNav/internnav/model/basemodel/internvla_n1/nextdit_crossattn_traj.py",
line 76, in __init__
self.model.enable_gradient_checkpointing()
File ".../diffusers/models/modeling_utils.py", line 328, in
enable_gradient_checkpointing
self._set_gradient_checkpointing(enable=True,
gradient_checkpointing_func=gradient_checkpointing_func)
TypeError: LuminaNextDiT2DModel._set_gradient_checkpointing() got an unexpected
keyword argument 'enable'
For more information on this, check: https://www.markdownguide.org/extended-syntax/#fenced-code-blocks
-->
Expected Behavior
InternVLA-N1 model should instantiate and run evaluation with the default confg.
Screenshots/Videos
Environment
- OS: Native Ubuntu 22.04 (w/ Docker image internnav:v1.2)
- GPU/CPU: NVIDIA RTX A2000 12GB
- GPU-driver version: 570.211.01 CUDA 12.8
Release version or Commit ID
7a5c624 (latest)
Additional Context
No response
Activity
Tai-Wang commentedon Mar 31, 2026
Thanks for your report, and it is caused by the new version of diffusers indeed. Have you ever tried to fix this bug by adjusting the default setting? How much is the memory used affected and is it feasible for your hardware (I notice your GPU has only 12GB of memory)? If the effect is minor, we can revise the default setting for better compatibility with the requirement.
Felipe-Tommaselli commentedon Mar 31, 2026
I tested an L40S with 48GB of RAM and am currently experimenting with an aggressive quantized version on the A200 with 12GB of RAM. For both, I applied the fix by setting it to
Falseby default. I didn't observe any RAM overhead from this change. Do you recommend any specific isolated tests I should perform? Thank you!Tai-Wang commentedon Apr 7, 2026
Sorry for the late reply. Could you please help test the difference in peak memory used by setting the argument to True/False? If the result is acceptable, we may consider changing the default setting. Or, at least we can add the guidance in the documentation and comments to that piece of code.