`modeling_nemotron_h.py` missing `supports_gradient_checkpointing = True` (same as BF16)
#4
by IshiiKosuke - opened
Hi NVIDIA team — this model ships the same modeling_nemotron_h.py as the BF16 variant, and has the same missing flag. Cross-linking to the discussion I just opened on the BF16 card so both repos pick up the patch:
→ https://huggingface.co/nvidia/Nemotron-3-Nano-Omni-30B-A3B-Reasoning-BF16/discussions/12
TL;DR: add supports_gradient_checkpointing = True to NemotronHPreTrainedModel. NemotronHBlock already inherits from GradientCheckpointingLayer, so no other change is needed.
class NemotronHPreTrainedModel(PreTrainedModel):
config: NemotronHConfig
base_model_prefix = "backbone"
+ supports_gradient_checkpointing = True
_no_split_modules = ["NemotronHBlock"]
_skip_keys_device_placement = ["past_key_values"]
_supports_flash_attn = True
The canonical transformers/models/nemotron_h/modeling_nemotron_h.py in the transformers library already has this flag — the trust-remote-code copies in the Omni repos diverged.
Repro + verification details on the BF16 discussion. Thanks!