fix(train): check trainable layers separately during checkpointing

[MrCapricornLiu]

PR type

• Bug Fix
• New Feature
• Document Updates
• More Models or Datasets Support

PR information

Dynamic gradient checkpointing shares one requires_grad cache across every wrapped layer. If the first layer is frozen, later trainable layers also see False. With reentrant checkpointing and inputs that do not require gradients, those layers receive no parameter gradients. This affects partially frozen towers, such as a vision tower with a frozen prefix. The cache also becomes stale when a layer is unfrozen later.

Check the current layer's trainable parameters when its input needs to be made differentiable. This removes the shared cache while preserving the existing checkpointing behavior.

Experiment results

python -m unittest discover -s tests -p test_dynamic_checkpointing.py -v

Both tests pass on CPU and a single CUDA device. They use the public dynamic checkpointing helper with a small ten-layer tower and compare outputs and every parameter gradient against an eager copy. Coverage includes a frozen prefix, an all-trainable control, non-reentrant checkpointing, and unfreezing a layer after the first forward. On the original code, 12 parameter-gradient checks fail across the two tests.

pre-commit run --all-files passed. Tested with PyTorch 2.12.0 and Transformers 5.3.0; no pretrained vision model or distributed training run was used.

[MrCapricornLiu]

The lint job failed while importing YAPF, with EOFError: Ran out of input when reading its generated grammar cache. No source-formatting failure was reported, and the other hooks passed. The same cache initialization issue occurred locally; after warming the cache, the full pre-commit run passed. Could the failed lint job be rerun? My account cannot rerun workflows in this repository.