Transformers Community Reports Gemma Tensor Parallelism Bug; Mixing Tensor and DTensor May Interrupt Inference

The Transformers community reported a tensor-parallel compatibility issue affecting Gemma multimodal models on October 12 and added a proposed direction for a fix the following day. The issue involves shared embedding weights and Gemma 4’s per-layer input embedding path: some operations receive both regular Tensors and distributed DTensors, causing an exception. The reported environment was Transformers 5.18.0.dev0 and PyTorch 2.14.0; this does not establish that all stable releases are affected. Issue report

The two failures occur at different points. For Gemma 3 and 4, the output layer lm_head may perform matrix operations using regular hidden states and shared DTensor weights. When per-layer embeddings are enabled in Gemma 4, regular Tensors may be mixed with DTensor-filled embeddings in torch.where. The attached reduced-model script creates a two-rank CPU simulation process group, applies tensor parallelism, and rebinds the weights; the error can then be triggered with text tokens alone, lowering the hardware barrier to investigating the issue. Reproduction script and error stack trace

This kind of defect matters for multi-GPU deployment. Tensor parallelism shards each layer’s weights across multiple accelerators, allowing models that exceed the memory of a single device to run. Transformers uses DeviceMesh and DTensor to represent sharding, replication, and reduction layouts. In practice, engineers should check the input and output layouts at module boundaries, and verify that the types of shared parameters still meet the requirements of the operations after rebinding. This is a suggested investigation based on the official architecture documentation. Official tensor parallelism documentation

The issue reporter said on October 13 that existing PR #48938 could address the current model errors, but that the approach was not yet fully general and still needed adjustment. This article has not independently verified the PR’s full contents or merge status, so the author’s report of a partial result should not be treated as a released fix. The public case also does not yet provide a complete validation matrix covering full checkpoints, multimodal generation, or different hardware. Author’s follow-up

Teams preparing to scale Gemma inference can pin their package versions and validate shared-embedding and per-layer-embedding configurations on their actual checkpoints, then test both the forward pass and generation workflow. Follow-up points include whether the fix covers different tensor layouts and whether regression tests can verify model loading, weight sharing, and multi-GPU execution together.