enable gemma3 conversion from maxtext to huggingface

Author: YixuanWang-99

MaxText/tests/hf_checkpoint_conversion_check.py

@@ -13,17 +13,18 @@
 See the License for the specific language governing permissions and
 limitations under the License.
 """
-
-from typing import Sequence
+import os
 import torch
+import torch.nn.functional as F
+import argparse
 from transformers import AutoTokenizer, AutoModelForCausalLM
-import os
-from absl import app  # Removed flags
+from tabulate import tabulate
 
 from MaxText.utils.ckpt_conversion.utils.hf_utils import (
-    check_predicted_tokens_match,
+    # check_predicted_tokens_match,
     check_arrays_match,
 )
+from MaxText import max_logging
 # Read Hugging Face token from environment variable
 hf_token = os.environ.get("HF_AUTH_TOKEN")
 
@@ -40,6 +41,7 @@
     huggingface_hub
     transformers
     accelerate
+    tabulate
 """
 
 
@@ -70,46 +72,180 @@ def get_logits(inputs, model, golden_model):
   return logits, golden_logits
 
 
-def main(argv: Sequence[str]) -> None:
-  # Parse arguments from argv
-  # Default values
-  parsed_args = {"golden_model_id": "google/gemma-2-2b-it", "hf_checkpoint_path": os.path.expanduser("~/.hf_output/")}
-  for arg in argv[1:]:
-    if "=" in arg:
-      key, value = arg.split("=", 1)
-      if key in parsed_args:
-        parsed_args[key] = value
-      else:
-        print(f"Warning: Unknown argument '{key}' found in argv. Ignoring.")
-
-  golden_model = AutoModelForCausalLM.from_pretrained(parsed_args["golden_model_id"], torch_dtype=torch.float32)
-
-  tokenizer = AutoTokenizer.from_pretrained(parsed_args["hf_checkpoint_path"])
-  model = AutoModelForCausalLM.from_pretrained(parsed_args["hf_checkpoint_path"], torch_dtype=torch.float32)
-
-  # TODO: (@yixuannwang) use 3 prompts to verify
-  input_text = "I love to"
-  inputs = tokenizer(input_text, return_tensors="pt")
-  # --- Generate Output ---
-  with torch.no_grad():
-    outputs = model.generate(**inputs, max_new_tokens=8)
-  # --- Decode and Print ---
-  print(tokenizer.decode(outputs[0], skip_special_tokens=True))
-
-  # Check weights match
-  print("########### check weights match ############### ")
-  check_weights_match(model, golden_model)
+def get_top_k_tokens_scores(logits_tensor, tokenizer_instance, k=10, description=""):
+  """Get the top-k tokens and their scores from a given logits tensor."""
+  max_logging.log(f"\n--- {description} top {k} tokens ---")
+  collected_tokens = []
+  tokens = []
+  # Ensure logits_tensor is on CPU for operations like topk and item()
+  logits_tensor = logits_tensor.cpu()
+  topk_results = torch.topk(logits_tensor[0, -1], k=k)
+  for i in range(k):
+    tok_id = topk_results.indices[i].item()
+    score = topk_results.values[i].item()
+    tok = tokenizer_instance.decode(tok_id)
+    collected_tokens.append({"id": int(tok_id), "token": tok.strip(), "score": float(score)})
+    tokens.append({"id": int(tok_id), "token": tok.strip(), "score": float(score)})
+
+  # Prepare data for tabulate: a list of lists
+  table_data = [[d["id"], d["token"], d["score"]] for d in collected_tokens]
+  max_logging.log(tabulate(table_data, headers=["Token ID", "Token", "Score"], tablefmt="orgtbl"))
+  return tokens
+
+
+def compare_top_tokens(converted_tokens, golden_tokens):
+  """
+  Compares two lists of top tokens and calculates similarity metrics.
+
+  Args:
+      converted_tokens: top tokens from the converted model.
+      golden_tokens:  top tokens from the golden model.
+  """
+  # Extract the sets of token IDs for comparison
+  converted_ids = {token["id"] for token in converted_tokens}
+  golden_ids = {token["id"] for token in golden_tokens}
+
+  # --- Metric 1: Overlap Count & Jaccard Similarity ---
+  intersection = converted_ids.intersection(golden_ids)
+  union = converted_ids.union(golden_ids)
+
+  overlap_count = len(intersection)
+  jaccard_similarity = overlap_count / len(union) if union else 0.0
+
+  # --- Metric 2: Rank Agreement ---
+  rank_matches = 0
+  min_len = min(len(converted_tokens), len(golden_tokens))
+  for i in range(min_len):
+    if converted_tokens[i]["id"] == golden_tokens[i]["id"]:
+      rank_matches += 1
+
+  rank_agreement = (rank_matches / min_len) * 100 if min_len > 0 else 0.0
+
+  metrics = {
+      "overlap_count": f"{overlap_count}/{min_len}",
+      "jaccard_similarity": jaccard_similarity,
+      "rank_agreement_percentage": rank_agreement,
+  }
+
+  max_logging.log("\n--- Similarity Metrics of Top Tokens ---")
+  table = [[key, value] for key, value in metrics.items()]
+  max_logging.log(tabulate(table, headers=["Metric", "Value"], tablefmt="orgtbl"))
+
+
+def check_kl_divergence(model_logits, golden_logits, atol=0.02):
+  """
+  Calculates KL divergence D_KL(P_golden || Q_model) over a batch of sequences.
+
+  Args:
+      model_logits: Logits from the converted model (Batch, SeqLen, VocabSize).
+      golden_logits: Logits from the golden model (Batch, SeqLen, VocabSize).
+      token_size: The number of vocabulary entries to consider for the comparison.
+                  (Effectively vocab_size_to_compare).
+  """
+  # 1. Select the relevant vocabulary slice from the logits.
+  token_size = min(model_logits.shape[2], golden_logits.shape[2])
+  model_logits_sliced = model_logits[..., :token_size]
+  golden_logits_sliced = golden_logits[..., :token_size]
+
+  # 2. Reshape
+  b, s, v = model_logits_sliced.shape
+  model_logits_reshaped = model_logits_sliced.view(b * s, v)
+  golden_logits_reshaped = golden_logits_sliced.view(b * s, v)
+
+  # 3. Get the probability distributions.
+  golden_probabilities = F.softmax(golden_logits_reshaped, dim=-1)
+  model_log_probabilities = F.log_softmax(model_logits_reshaped, dim=-1)
+
+  # 4. Calculate avg KL divergence for all token distributions.
+  # use 'batchmean'; the sum of the KL divergences for each token in the batch
+  # and then divides by the number of tokens (b * s)
+  kl_div_value = F.kl_div(
+      input=model_log_probabilities,
+      target=golden_probabilities,
+      reduction="batchmean",  # Use 'batchmean' for the average KL per token.
+      log_target=False,
+  )
+
+  max_logging.log(f"\nAverage KL divergence per token (D_KL(P_golden || Q_model)): {kl_div_value.item():.6f}")
+
+  # To find the max KL divergence for any single token in the set
+  # use reduction='none'.
+  kl_divs_per_token = F.kl_div(
+      input=model_log_probabilities, target=golden_probabilities, reduction="none", log_target=False
+  ).sum(
+      dim=-1
+  )  # Sum over the vocab dim to get a single KL value per token
+
+  max_kl_div = kl_divs_per_token.max()
+  max_logging.log(f"\nMax KL divergence for a single token in the set: {max_kl_div.item():.6f}")
+
+  assert max_kl_div < atol, f"KL divergence values {max_kl_div.item():.6f} exceed the threshold {atol}"
+
+
+def run_prompts(args: argparse.Namespace) -> None:
+  """
+  Args:
+      - golden_model_id (str): HF model ID for the golden model.
+      - hf_checkpoint_path (str): Path to the converted HF checkpoint.
+      - max_kl_div (float): Maximum allowed KL divergence.
+  """
+  golden_model = AutoModelForCausalLM.from_pretrained(args.golden_model_id, torch_dtype=torch.bfloat16)
+  golden_tokenizer = AutoTokenizer.from_pretrained(args.golden_model_id)
+
+  tokenizer = AutoTokenizer.from_pretrained(args.hf_checkpoint_path)
+  model, _ = AutoModelForCausalLM.from_pretrained(
+      args.hf_checkpoint_path, trust_remote_code=True, torch_dtype=torch.bfloat16, output_loading_info=True
+  )
+
+  # max_logging.log(loading_info)
+
+  prompts = ["I love to", "Today is a", "What is the"]
+  for input_text in prompts:
+    max_logging.log(f"\n--- Prompt: {input_text} ---")
+    inputs = tokenizer(input_text, return_tensors="pt")
+    # --- Generate Output ---
+    with torch.no_grad():
+      outputs = model.generate(**inputs, max_new_tokens=15, do_sample=False)
+    # --- Decode and Print ---
+    max_logging.log(f"Output: {tokenizer.decode(outputs[0], skip_special_tokens=True)}")
+
+    # --- Compare tokens ---
+    model_logits, golden_model_logits = get_logits(inputs, model, golden_model)
+    tokens = get_top_k_tokens_scores(model_logits, tokenizer, k=10, description="converted model")
+    golden_tokens = get_top_k_tokens_scores(golden_model_logits, golden_tokenizer, k=10, description="golden model")
+    compare_top_tokens(converted_tokens=tokens, golden_tokens=golden_tokens)
+
+    check_kl_divergence(model_logits, golden_model_logits, atol=args.max_kl_div)
+
+  """
+  if the model's structure is exactly the same as the golden model (layers, vocab_size, etc.), 
+  you can check more weights details using the following steps:
 
-  # Run forward pass to get logits
-  logits, golden_logits = get_logits(inputs, model, golden_model)
+  check_weights_match(model, golden_model)
 
   # Check logits from the first 5 tokens match
-  print("########### check logits match ############### ")
-  check_arrays_match(logits[0, :5, :], golden_logits[0, :5, :], atol=0.2)
+  check_arrays_match(model_logits[0, :5, :], golden_model_logits[0, :5, :], atol=0.2)
 
-  print("########### check predicted token match ############### ")
-  check_predicted_tokens_match(logits, golden_logits)
+  check_predicted_tokens_match(model_logits, golden_model_logits)
+  """
 
 
 if __name__ == "__main__":
-  app.run(main)
+  parser = argparse.ArgumentParser(description="Verify HuggingFace checkpoints converted from MaxText.")
+  parser.add_argument(
+      "--golden_model_id",
+      type=str,
+      default="google/gemma-2-2b-it",
+      help="The HuggingFace model ID for the golden/reference model.",
+  )
+  parser.add_argument(
+      "--hf_checkpoint_path",
+      type=str,
+      default=os.path.expanduser("~/.hf_output/"),
+      help="Path to the converted HuggingFace checkpoint directory.",
+  )
+  parser.add_argument("--max_kl_div", type=float, default=0.02, help="Maximum allowed KL divergence between model logits.")
+
+  parsed_args = parser.parse_args()
+
+  run_prompts(parsed_args)

MaxText/utils/ckpt_conversion/examples/convert_gemma2_to_hf.sh

@@ -12,7 +12,7 @@ MAXTEXT_CHECKPOINT_DIR="gs://maxtext-model-checkpoints/gemma2-2b-it/2025-02-20-1
 LOCAL_HF_CHECKPOINT_DIR="/tmp/hf_gemma2-2b_output" # HF requires a local dir
 GOLDEN_MODEL_ID="google/gemma-2-2b-it"
 
-CONVERT_MODULE="MaxText.ckpt_conversion.to_huggingface"
+CONVERT_MODULE="MaxText.utils.ckpt_conversion.to_huggingface"
 CONVERT_ARGS=(
     "MaxText/configs/base.yml"
     "model_name=gemma2-2b"
@@ -29,7 +29,7 @@ CONVERT_ARGS=(
     "base_output_directory=${HF_CHECKPOINT_GCS_PATH}"
 )
 
-VERIFY_MODULE="MaxText.tests.huggingface_ckpt_conversion_check"
+VERIFY_MODULE="MaxText.tests.hf_ckpt_conversion_check"
 
 VERIFY_ARGS=(
     "golden_model_id=${GOLDEN_MODEL_ID}"

MaxText/utils/ckpt_conversion/examples/convert_gemma3_to_hf.sh

@@ -0,0 +1,83 @@
+#!/bin/bash
+
+# Exit immediately if a command exits with a non-zero status.
+set -e
+
+export HF_AUTH_TOKEN=""
+
+DATE=$(date +%Y-%m-%d)
+# Define variables for paths and arguments
+HF_CHECKPOINT_GCS_PATH="gs://maxtext-model-checkpoints/HuggingFace/gemma3-4b/${DATE}" # (optional)GCS path for HF model
+MAXTEXT_CHECKPOINT_DIR="gs://maxtext-model-checkpoints/gemma3-4b/2025-03-18-19-03/unscanned/checkpoints/0/items"
+LOCAL_HF_CHECKPOINT_DIR="/tmp/hf_gemma3-4b_output" # HF requires a local dir
+GOLDEN_MODEL_ID="google/gemma-3-4b-it"
+
+CONVERT_MODULE="MaxText.utils.ckpt_conversion.to_huggingface"
+CONVERT_ARGS=(
+    "MaxText/configs/base.yml",
+    "model_name=gemma3-4b",
+    "tokenizer_path=assets/tokenizer.gemma3",
+    "load_parameters_path=${MAXTEXT_CHECKPOINT_DIR}",
+    "per_device_batch_size=1",
+    "run_name=ht_test",
+    "max_prefill_predict_length=8",
+    "max_target_length=16",
+    "steps=1",
+    "async_checkpointing=false",
+    "prompt='I love to'",
+    "scan_layers=false",
+    "attention='dot_product'",
+    "base_output_directory=${HF_CHECKPOINT_GCS_PATH}"
+)
+
+VERIFY_MODULE="MaxText.tests.hf_ckpt_conversion_check"
+
+VERIFY_ARGS=(
+    "--golden_model_id=${GOLDEN_MODEL_ID}"
+    "--hf_checkpoint_path=${LOCAL_HF_CHECKPOINT_DIR}" # Updated to local path
+)
+
+
+# --- Step 1: Run the Hugging Face Conversion ---
+echo "Starting Hugging Face model conversion for gemma2-2b..."
+cd "$MAXTEXT_PROJECT_DIR"
+
+# Construct the command
+CONVERT_CMD=("python3" -m "$CONVERT_MODULE")
+for arg in "${CONVERT_ARGS[@]}"; do
+    CONVERT_CMD+=("$arg")
+done
+
+# Execute the command
+"${CONVERT_CMD[@]}"
+
+echo "Hugging Face model conversion finished."
+
+# --- Step 2: Run the Verification Script ---
+echo "Starting verification for the converted gemma2-2b model..."
+
+# Create local directory for checkpoints and download from GCS
+echo "Creating local directory for HF checkpoints: ${LOCAL_HF_CHECKPOINT_DIR}"
+mkdir -p "${LOCAL_HF_CHECKPOINT_DIR}"
+echo "Downloading HF checkpoints from ${HF_CHECKPOINT_GCS_PATH} to ${LOCAL_HF_CHECKPOINT_DIR}..."
+gsutil -m cp -r "${HF_CHECKPOINT_GCS_PATH}/*" "${LOCAL_HF_CHECKPOINT_DIR}/"
+echo "Download complete."
+
+# Construct the command
+VERIFY_CMD=("python3" -m "$VERIFY_MODULE")
+if [ ${#VERIFY_ARGS[@]} -ne 0 ]; then
+    for arg in "${VERIFY_ARGS[@]}"; do
+        VERIFY_CMD+=("$arg")
+    done
+fi
+
+# Execute the command
+"${VERIFY_CMD[@]}"
+
+# Optional: Clean up the local checkpoint directory
+echo "Cleaning up local HF checkpoint directory: ${LOCAL_HF_CHECKPOINT_DIR}"
+rm -rf "${LOCAL_HF_CHECKPOINT_DIR}"
+echo "Cleanup complete."
+
+echo "Verification script finished. Please check the above generated text"
+echo "All steps completed."

MaxText/utils/ckpt_conversion/to_huggingface.py

@@ -34,7 +34,7 @@
 )
 from MaxText.utils.ckpt_conversion.utils.shape_mapping import SHAPE_MAPPING
 from MaxText.utils.ckpt_conversion.utils.hf_model_configs import HF_MODEL_CONFIGS
-from MaxText.utils.ckpt_conversion.utils.utils import (process_leaf_param, save_model_files, TOKENIZER_HF_IDS)
+from MaxText.utils.ckpt_conversion.utils.utils import (process_leaf_param, save_model_files, HF_IDS)
 
 """Convert MaxText unscanned ckpt into HF format"""
 
@@ -83,9 +83,9 @@ def main(argv: Sequence[str]) -> None:
   hf_config_obj = HF_MODEL_CONFIGS[model_key]
 
   # 2. Load Tokenizer
-  if model_key not in TOKENIZER_HF_IDS:
+  if model_key not in HF_IDS:
     raise ValueError(f"HF Tokenizer ID not found for model key: {model_key}")
-  hf_tokenizer_id = TOKENIZER_HF_IDS[model_key]
+  hf_tokenizer_id = HF_IDS[model_key]
   tokenizer = AutoTokenizer.from_pretrained(hf_tokenizer_id, token=hf_token)
 
   # 3. Get parameter mappings