transformers - 💡(How to fix) Fix [Gemma4] `Gemma4VisionPatchEmbedder._position_embeddings` materializes a ~19 GiB one-hot tensor that's mathematically a 2-row embedding lookup

def _position_embeddings(self, pixel_position_ids, padding_positions):
    """Prepare patch positions map for matmul with positon embedding table."""
    clamped_positions = pixel_position_ids.clamp(min=0)
    one_hot = F.one_hot(clamped_positions, num_classes=self.position_embedding_size)
    one_hot = one_hot.permute(0, 2, 1, 3).to(self.position_embedding_table)
    position_embeddings = one_hot @ self.position_embedding_table
    position_embeddings = position_embeddings.sum(dim=1)
    position_embeddings = torch.where(padding_positions.unsqueeze(-1), 0.0, position_embeddings)
    return position_embeddings

---

[1] alloc peak 36.00 GiB 
      19.57 GiB  (54.4%)  x10    vision tower (vision_model)
      14.09 GiB  (39.1%)        ZeRO-3 
       ...

---

one_hot @ line 581:   15.38 GiB (single int64 tensor, [40, 2520, 2, 10240])
one_hot.to bf16 @ 582: 3.85 GiB (bf16 cast of the same shape)

---

_position_embeddings @ modeling_gemma4.py:581
forward @ modeling_gemma4.py:596

---

def _position_embeddings(self, pixel_position_ids, padding_positions):
    """Prepare patch positions map via direct embedding lookup.

    Mathematically equivalent to the original one_hot @ table + sum,
    but avoids the [batch, num_patches, 2, position_embedding_size]
    one-hot temporary (~15 GiB int64 + ~3.85 GiB bf16 cast at bs=4,
    n_img=10, position_embedding_size=10240).
    """
    # Original .clamp(min=0) only bounds the lower side; one_hot() did
    # an implicit upper-bound check. F.embedding would read OOB on
    # overflow, so clamp both sides.
    clamped = pixel_position_ids.clamp(0, self.position_embedding_size - 1)
    x_emb = F.embedding(clamped[..., 0], self.position_embedding_table[0])
    y_emb = F.embedding(clamped[..., 1], self.position_embedding_table[1])
    position_embeddings = x_emb + y_emb
    # The original's trailing .sum(dim=1) promotes bf16 -> fp32 under
    # autocast (reduction-dtype rule); plain `+` doesn't, so cast
    # explicitly only inside autocast to keep the downstream
    # `hidden_states + pos` dtype contract.
    if torch.is_autocast_enabled() and position_embeddings.dtype != torch.float32:
        position_embeddings = position_embeddings.to(torch.float32)
    position_embeddings = torch.where(padding_positions.unsqueeze(-1), 0.0, position_embeddings)
    return position_embeddings