Back Home

GitHub Repo

Transformers proposes speech model compile fix, removing legacy gradient checkpointing compatibility logic

On November 10, the Transformers community proposed a patch to remove gradient-setting logic that blocks full-graph compilation in ten speech models, including Wav2Vec2. The proposal has not been merged and changes reentrant gradient checkpointing behavior; users should verify that the first convolution still receives gradients before adopting it.

Marcus Qwertyus · Public domain · Image source
zh-Hant

On November 9, the Transformers community reported a compilation blocker in speech model training and submitted a patch on November 10. When the Wav2Vec2 feature encoder is unfrozen and in training mode, every forward pass sets requires_grad=True on an intermediate tensor—even when gradient checkpointing is disabled. This causes torch.compile(fullgraph=True) to reject compilation. Issue report

The reporter reproduced the error with a small randomly initialized model by compiling only the feature encoder and then running a backward pass. The environment used Transformers 5.23.0.dev0 and PyTorch 2.15.0. This reproduction points to the convolutional feature extraction path, which engineers can inspect independently when debugging. Reproduction script

The logic was originally intended to support reentrant gradient checkpointing. Checkpointing saves memory by recomputing intermediate results during backpropagation, but the two implementations have different constraints: PyTorch officially recommends use_reentrant=False; the reentrant version requires at least one input and output to require gradients, or that checkpointed region will not receive gradients. PyTorch documentation

Following the direction discussed by repository contributors, the new proposal simply removes the compatibility logic and an unused flag, and applies the change to ten models, including HuBERT, WavLM, SpeechT5, and Data2Vec-Audio. The tradeoff is explicit: when using use_reentrant=True, the first convolution will no longer receive gradients. The patch therefore affects both compilation compatibility and migration of training settings. Patch proposal

The author reported that the reproduction script passes. Across ten model test suites, 1,955 tests passed and 2,438 were skipped; two other SpeechT5 integration tests also fail on the original main branch. As of this check, the proposal is still under review, and CI is awaiting maintainer approval. Next, track the merged release, actual GPU training performance, and gradient integrity with existing checkpoint settings; there is currently no evidence quantifying throughput gains. Test and review status

Sources

  1. Wav2Vec2-family feature encoder breaks torch.compile(fullgraph=True) — Issue #50527
  2. Removing legacy requires_grad workaround — Pull Request #50544
  3. torch.utils.checkpoint — PyTorch 2.15 documentation