Transformers Proposes Tensor Parallelism Fix to Prevent Gradient Hooks from Slowing Compiled Inference

Hugging Face Transformers has a new proposed performance fix for its tensor-parallel implementation. On October 26, a project contributor submitted PR #49975, which changes when gradient synchronization hooks are registered for ReplicatedWithGradAllReduce. The change is intended to prevent inference that does not require gradients from carrying module-level backward hooks that interfere with torch.compile. As of review, the page showed the proposal in the merge queue; it had not yet been merged.

The original issue was reported on October 23, with two small reproductions using randomly initialized Qwen3 models running in two CPU processes. After enabling tp_plan="auto", Dynamo recorded 11 graph breaks; removing the backward hooks reduced that number to zero. The reporter also said that, when Qwen3-8B was distributed across two Neuron cores on AWS Trainium, decoding took 140 milliseconds per token with the hooks enabled and 37.6 milliseconds after removing them. These are measurements reported for a specific environment and should not be generalized to other hardware or models. Issue and reproduction scripts

The hook has a training purpose. The official documentation explains that the relevant strategy replicates parameters and performs an all-reduce on gradients. It is intended for normalization operations located between different partitioned layers, where each process can see only some of the attention heads. The problem is that module backward hooks that have already been registered can still cause graph breaks in the compiler, even when using torch.no_grad(). Disabling gradient computation alone may therefore not eliminate the added cost.

The proposal checks torch.is_grad_enabled() during the forward pass and installs gradient-accumulation post-hooks on the parameters actually loaded only when gradients are needed. Deferring registration also avoids a problem where placeholder parameters are replaced during weight loading, which can cause their hooks to be lost. The author also says this addresses a potential case in which the old hooks ran before parameter gradients had been produced, causing synchronization to be skipped. Implementation details

For deployment teams, the change points to a performance bottleneck worth investigating: when tensor parallelism and compilation are both enabled, examine the causes of graph breaks instead of comparing only the number of devices. Teams should track the merge and release status, then validate decoding latency and training gradient synchronization on their actual models. The public proposal includes a test script, but there are not yet comprehensive results to support claims of a generally applicable speedup.