[cuda backend] reduce memory consumption on gemma4_31b by running embedding in int8 (#20351)

Decode the tied GGUF token_embd/lm_head weight to a gatherable int8
IntxUnpackedToInt8Tensor for the embedding instead of dequantizing the
whole ~1.4B-element weight to bf16 to reduce the memory consumption
during exporation
Halves the embedding's host + GPU-constant footprint (2 -> 1 B/elem) and
reduce about 1.5 gb vram memory.

The decode is lossless for Q6_K; for Q4_K the int8 embedding matches the
precision of the tied lm_head and the other q4_k linears (same bf16
zero-point reconstruction), so it adds no error beyond the model's
existing 4-bit level. Also create test case to verify it.

lm_head keeps its packed format (Q6_K -> CudaDp4aPlanarInt6Tensor from
the same int8 decode; Q4_K -> CudaCoalescedInt4Tensor via
_resolve_tied_lm_head). MLX and non-Q4_K/Q6_K paths unchanged.
 
 Memory consumption: 28.3 GB with context length == 128k + turboquant

Author: Gasoonjia

examples/models/gemma4_31b/gguf_loader.py

@@ -17,9 +17,12 @@
   linear and embedding. ``embed_tokens`` and ``lm_head`` stay tied -- they share
   the one quantized tensor.
 * **CUDA**: Q4_K -> ``Int4Tensor``, Q6_K -> ``CudaDp4aPlanarInt6Tensor`` (a genuine
-  6-bit packed weight, lossless, symmetric); ``lm_head`` keeps the quantized
-  tensor but the token embedding is dequantized to bf16 (the packed tensors can't
-  gather), so they are untied.
+  6-bit packed weight, lossless, symmetric). ``embed_tokens`` and ``lm_head`` are
+  untied: ``lm_head`` keeps a packed (int6/int4) matmul weight, while the token
+  embedding becomes a gatherable ``IntxUnpackedToInt8Tensor`` (int8) -- the truly
+  packed int4/int6 tensors can't gather. For the Q6_K tied weight the decode is
+  done once and shared between the two, avoiding a whole-tensor bf16 dequant and
+  a second decode (see ``_untie_embed_lm_head``).
 
 Usage:
     model, config = load_gguf_model("model.gguf", backend="cuda")
@@ -116,6 +119,55 @@ def _resolve_tied_lm_head(model, lm_head_weight, packers):
         )
 
 
+def _untie_embed_lm_head(model, gtensor, weight, backend):
+    """Untie the GGUF token-embed / lm_head weight, returning ``(embed, lm_head)``.
+
+    GGUF ties ``embed_tokens`` and ``lm_head`` to one quantized weight. The
+    returned ``lm_head`` is packed into ``model.lm_head`` after the streaming loop
+    (``_resolve_tied_lm_head``), or is ``None`` when this function already
+    assigned it.
+
+    * **MLX**: keep both tied on the raw ``ExportableGGUFTensor``.
+    * **CUDA** (Q6_K or Q4_K): untie so ``lm_head`` keeps a packed low-bit matmul
+      weight while the token embedding becomes a gatherable int8
+      ``IntxUnpackedToInt8Tensor`` -- the truly packed int4/int6 tensors can't
+      gather. Instead of dequantizing the whole ~1.4 B-element weight to bf16
+      (2 B/elem), decode it once to int8 (1 B/elem; the decode is lossless so the
+      result is numerically identical), halving the embedding's host + GPU-constant
+      footprint. The token embedding (Q4_K for the Gemma checkpoint) is the single
+      biggest weight, so this is the dominant saving vs the bf16 path. ``lm_head``:
+        - Q6_K -> ``CudaDp4aPlanarInt6Tensor`` from the *same* int8 decode and
+          assigned here (``pack_linear_for_cuda`` would mis-route an int8 tensor to
+          the int8 path), so the post-loop resolve is a no-op.
+        - Q4_K -> kept as the native ``Int4Tensor`` and returned, so
+          ``_resolve_tied_lm_head`` packs it to ``CudaCoalescedInt4Tensor`` (same
+          as a regular Q4_K linear).
+    * **CUDA, other types**: fall back to the bf16 embedding.
+    """
+    if backend == "mlx":
+        return weight, gtensor
+
+    if gtensor.ggml_type in ("q6_k", "q4_k"):
+        intx = gtensor.to_intx_unpacked_to_int8_tensor()
+        if gtensor.ggml_type == "q6_k":
+            import torch.nn as nn
+            from executorch.backends.cuda.dp4a_planar_int6_tensor import (
+                CudaDp4aPlanarInt6Tensor,
+            )
+
+            model.lm_head.weight = nn.Parameter(
+                CudaDp4aPlanarInt6Tensor._from_intx_int8(intx), requires_grad=False
+            )
+            return intx, None
+        # Q4_K: ``weight`` is the native Int4Tensor; let _resolve_tied_lm_head
+        # pack it to CudaCoalescedInt4Tensor. Only the embedding switches to int8.
+        return intx, weight
+
+    from executorch.examples.models.gemma4_31b.quant import dequantize_weight
+
+    return dequantize_weight(weight, torch.bfloat16), weight
+
+
 def load_gguf_model(
     gguf_path: str,
     max_seq_len: int = 4096,
@@ -140,7 +192,7 @@ def load_gguf_model(
         Gemma4_31BConfig,
         materialize_runtime_buffers,
     )
-    from executorch.examples.models.gemma4_31b.quant import dequantize_weight, pack_one
+    from executorch.examples.models.gemma4_31b.quant import pack_one
     from executorch.extension.llm.export.gguf import ExportableGGUFTensor, iter_gguf
 
     if backend == "cuda":
@@ -161,7 +213,7 @@ def load_gguf_model(
     with torch.device("meta"):
         model = Gemma4_31B(config)
 
-    lm_head_weight = None  # weight reused for a tied lm_head
+    lm_head_weight = None  # tied weight resolved into lm_head after the loop
     n_processed = 0
 
     print(f"Streaming GGUF from {gguf_path}...")
@@ -173,11 +225,9 @@ def load_gguf_model(
         if isinstance(value, ExportableGGUFTensor):
             weight = _convert_weight(model, model_key, value, backend)
             if model_key == "embed_tokens.weight":
-                # Tied lm_head reuses the embedding weight: MLX wants the raw
-                # ExportableGGUFTensor (linear pattern), CUDA the quant tensor.
-                lm_head_weight = value if backend == "mlx" else weight
-                if backend == "cuda":
-                    weight = dequantize_weight(weight, torch.bfloat16)
+                weight, lm_head_weight = _untie_embed_lm_head(
+                    model, value, weight, backend
+                )
             value = weight
         elif value.dtype == torch.float32:
             value = value.to(torch.bfloat16)

examples/models/gemma4_31b/tests/test_cuda_pipeline.py

@@ -246,13 +246,14 @@ def _load(self, tmp):
 
     def test_load_converts_weights(self):
         """GGUF -> CUDA: Q4_K -> CudaCoalescedInt4Tensor, Q6_K -> CudaDp4aPlanarInt6Tensor,
-        embedding bf16."""
+        embedding int8 (gatherable)."""
         from executorch.backends.cuda.coalesced_int4_tensor import (
             CudaCoalescedInt4Tensor,
         )
         from executorch.backends.cuda.dp4a_planar_int6_tensor import (
             CudaDp4aPlanarInt6Tensor,
         )
+        from torchao.quantization import IntxUnpackedToInt8Tensor
 
         with tempfile.TemporaryDirectory() as tmp:
             model, _ = self._load(tmp)
@@ -263,11 +264,49 @@ def test_load_converts_weights(self):
         self.assertIsInstance(
             model.layers[0].mlp.down_proj.weight.data, CudaDp4aPlanarInt6Tensor
         )
-        # Tied lm_head is repacked to int6 by pack_cuda (it keeps quantization,
-        # unlike the token embedding which is dequantized for the gather).
+        # Tied lm_head keeps a packed int6 matmul weight.
         self.assertIsInstance(model.lm_head.weight.data, CudaDp4aPlanarInt6Tensor)
-        # Token embedding is dequantized to bf16 (Int4/packed-int6 can't gather).
-        self.assertEqual(model.embed_tokens.weight.dtype, torch.bfloat16)
+        # Token embedding is decoded to a gatherable int8 tensor (not bf16): the
+        # Q6_K decode is lossless and shared with lm_head. Keeping it int8 (vs
+        # bf16) avoids a ~5.6 GB fp32 dequant transient and ~1.4 GB resident at
+        # export time.
+        self.assertIsInstance(model.embed_tokens.weight.data, IntxUnpackedToInt8Tensor)
+
+    def test_int8_embedding_matches_bf16(self):
+        """Guard the bf16 -> int8 token-embedding switch.
+
+        The embedding is now loaded as a gatherable int8 ``IntxUnpackedToInt8Tensor``
+        instead of being dequantized to bf16. Its gathered rows must match the bf16
+        dequant of the *source* GGUF token embedding -- i.e. exactly what the old
+        ``dequantize_weight(..., bf16)`` path returned. The GGUF decode is lossless,
+        so they agree to bf16 precision.
+        """
+        from executorch.examples.models.gemma4_31b.gguf_loader import gguf_to_model_key
+        from executorch.extension.llm.export.gguf import ExportableGGUFTensor, iter_gguf
+        from torchao.quantization import IntxUnpackedToInt8Tensor
+
+        with tempfile.TemporaryDirectory() as tmp:
+            path = os.path.join(tmp, "tiny.gguf")
+            build_gguf_checkpoint(path)
+            # Reference = bf16 dequant of the source GGUF token embedding (the
+            # tensor the previous bf16 embedding path materialized).
+            ref_bf16 = None
+            for name, val in iter_gguf(path):
+                if gguf_to_model_key(name) == "embed_tokens.weight":
+                    self.assertIsInstance(val, ExportableGGUFTensor)
+                    ref_bf16 = val.dequantize(torch.bfloat16)
+                    break
+            self.assertIsNotNone(ref_bf16, "token_embd.weight not found in GGUF")
+            model, _ = load_gguf_model(path, backend="cuda", config=GGUF_CONFIG)
+
+        self.assertIsInstance(model.embed_tokens.weight.data, IntxUnpackedToInt8Tensor)
+
+        ids = torch.tensor([0, 1, 7, GGUF_CONFIG.vocab_size - 1])
+        out = model.embed_tokens(ids)  # int8 gather + dequant
+        ref = ref_bf16[ids]
+        self.assertEqual(out.shape, ref.shape)
+        rel_err = (out.float() - ref.float()).abs().mean() / ref.float().abs().mean()
+        self.assertLess(rel_err.item(), 0.02)
 
     def test_generate(self):
         """GGUF -> CUDA -> eager generate produces valid tokens (inference.py)."""