Add Decoding Verification for the Checkpoint Validation Pipeline - #4728
Open
olufiyin19 wants to merge 7 commits into
Open
Add Decoding Verification for the Checkpoint Validation Pipeline#4728olufiyin19 wants to merge 7 commits into
olufiyin19 wants to merge 7 commits into
Conversation
…dation for the Checkpoint Validation Agent
olufiyin19
requested review from
A9isha,
NuojCheng,
RissyRan,
SurbhiJainUSC,
abhinavclemson,
aireenmei,
bvandermoon,
darisoy,
dipannita08,
gagika,
gobbleturk,
hengtaoguo,
huytransformer,
igorts-git,
jiangjy1982,
khatwanimohit,
richjames0,
shralex,
vipannalla and
xibinliu
as code owners
August 4, 2026 18:03
|
Caution The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased. |
olufiyin19
changed the title
Codecov Report✅ All modified and coverable lines are covered by tests. 📢 Thoughts on this report? Let us know! |
4 tasks
Contributor
|
🤖 Hi @entrpn, I've received your request, and I'm working on it now! You can track my progress in the logs for more details. |
![github-actions[bot]](https://avatars.githubusercontent.com/in/15368?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
This Pull Request successfully introduces a comprehensive, automated autoregressive decoding and generative validation pipeline via decode_validator.py, completing the fourth stage of the post-checkpoint conversion validation suite. The design of standardizing reporting to GCS in JSON format and streaming logs is a valuable addition to enable automated fail-fast diagnostics under Airflow orchestration.
🔍 General Feedback
- Excellent Architectural Consistency: The scripts align nicely with the transition to pure NNX modules while maintaining backwards compatibility where appropriate.
- Robust Exception Handling Pattern: While the standard execution flow is clean, the error handling around initial file reading and subprocess timeouts should be reinforced to ensure that JSON reports are guaranteed to write and upload even under catastrophic failures.
- Process Cleanup Discipline: It is crucial to clean up global monkeypatches in Python
finallyblocks when scripts are executed in-process, to prevent polluting the shared runtime process space and introducing flaky test behaviors.
Comment on lines
+69
to
+72
| # applying a monkeypatch to maxtext's model_creation_utils because it has a bug where | ||
| # it cannot resolve SequenceKey (list indices) to string keys in Linen checkpoints. | ||
|
|
||
| source = inspect.getsource(model_creation_utils._fix_restore_args_for_shape_mismatch) # pylint: disable=protected-access |
Contributor
There was a problem hiding this comment.
🟠 Save a reference to the original `_fix_restore_args_for_shape_mismatch` function before overriding/monkeypatching it. This is necessary so that we can restore the original implementation in the `finally` block and prevent global namespace/monkeypatch pollution of the Python process space.
Suggested change
| # applying a monkeypatch to maxtext's model_creation_utils because it has a bug where | |
| # it cannot resolve SequenceKey (list indices) to string keys in Linen checkpoints. | |
| source = inspect.getsource(model_creation_utils._fix_restore_args_for_shape_mismatch) # pylint: disable=protected-access | |
| # applying a monkeypatch to maxtext's model_creation_utils because it has a bug where | |
| # it cannot resolve SequenceKey (list indices) to string keys in Linen checkpoints. | |
| _original_fix_restore = model_creation_utils._fix_restore_args_for_shape_mismatch | |
| source = inspect.getsource(_original_fix_restore) # pylint: disable=protected-access |
Comment on lines
+301
to
+307
| if _orig_array_delete is not None: | ||
| jax.Array.delete = _orig_array_delete | ||
| transformers.AutoTokenizer.from_pretrained = _orig_from_pretrained | ||
| sys.stdout = old_stdout | ||
| sys.stderr = old_stderr | ||
| os.chdir(old_cwd) | ||
|
|
Contributor
There was a problem hiding this comment.
🔴 **Critical Bug Fix:** Restore the original implementations of `ocp.Checkpointer.restore` and `model_creation_utils._fix_restore_args_for_shape_mismatch` in the `finally` block. Because this validation script is run in-process using `runpy.run_path()`, failing to clean up global monkeypatches pollutes the process namespace, leading to unexpected behavior and hard-to-debug failures in any subsequent tasks or test executions running in the same context.
Suggested change
| if _orig_array_delete is not None: | |
| jax.Array.delete = _orig_array_delete | |
| transformers.AutoTokenizer.from_pretrained = _orig_from_pretrained | |
| sys.stdout = old_stdout | |
| sys.stderr = old_stderr | |
| os.chdir(old_cwd) | |
| finally: | |
| ocp.Checkpointer.restore = _original_restore | |
| model_creation_utils._fix_restore_args_for_shape_mismatch = _original_fix_restore | |
| if _orig_array_delete is not None: | |
| jax.Array.delete = _orig_array_delete | |
| transformers.AutoTokenizer.from_pretrained = _orig_from_pretrained | |
| sys.stdout = old_stdout | |
| sys.stderr = old_stderr | |
| os.chdir(old_cwd) |
Comment on lines
+137
to
+142
| target_lookup = r" def _lookup_stored_meta\(path\):[\s\S]*?(?=\n\s*mismatched_paths_sharded = \[\])" | ||
| patched_source = re.sub(target_lookup, new_lookup, source) | ||
|
|
||
| env = dict(model_creation_utils.__dict__) | ||
| exec(patched_source, env) # pylint: disable=exec-used | ||
| model_creation_utils._fix_restore_args_for_shape_mismatch = env[ # pylint: disable=protected-access |
Contributor
There was a problem hiding this comment.
🟡 **Fragile Implementation:** Modifying `_fix_restore_args_for_shape_mismatch` via a regular expression find-and-replace (`re.sub`) and then calling `exec` is highly brittle. This is vulnerable to code updates in `model_creation_utils.py` (such as formatting changes, changing variable names, or modifying the target function's structure), which would cause the regex to silently fail to match and not apply the required fix. Consider refactoring `model_creation_utils.py` directly to support these key mappings or subclassing/wrapping the checkpointer rather than performing string manipulation on source code.
Comment on lines
+70
to
+85
| # run subprocess (from the top level repo directory) | ||
| result = subprocess.run(command, text=True, capture_output=True, check=False, cwd=repo_root) | ||
| if result.stdout: | ||
| logger.info("=== Subprocess Stdout ===") | ||
| logger.info(result.stdout) | ||
|
|
||
| # generate report | ||
| report = { | ||
| "run_name": run_name, | ||
| "model": internal_model_name, | ||
| "status": "SUCCESS" if result.returncode == 0 else "FAILED", | ||
| "success": result.returncode == 0, # if returncode is 0, command worked | ||
| "stdout": result.stdout, # store standard output (contains generated text like "Input ... -> ...") | ||
| "stderr": (result.stderr if result.returncode != 0 else "Success"), # store error message if there's a failure | ||
| "checkpoint_used": checkpoint_path, | ||
| } |
Contributor
There was a problem hiding this comment.
🟠 **Buffering Bug:** The PR description states that this PR implements real-time subprocess streaming of stdout so that sequence output, token latency, and throughput metrics are visible in real time. However, using `subprocess.run(capture_output=True)` buffers all output in memory and blocks until the subprocess completely finishes before printing anything. If the decode script hangs, or is slow to initialize, no output will be streamed. Using `subprocess.Popen` with line-by-line reading of standard output resolves this and ensures real-time visibility.
Suggested change
| # run subprocess (from the top level repo directory) | |
| result = subprocess.run(command, text=True, capture_output=True, check=False, cwd=repo_root) | |
| if result.stdout: | |
| logger.info("=== Subprocess Stdout ===") | |
| logger.info(result.stdout) | |
| # generate report | |
| report = { | |
| "run_name": run_name, | |
| "model": internal_model_name, | |
| "status": "SUCCESS" if result.returncode == 0 else "FAILED", | |
| "success": result.returncode == 0, # if returncode is 0, command worked | |
| "stdout": result.stdout, # store standard output (contains generated text like "Input ... -> ...") | |
| "stderr": (result.stderr if result.returncode != 0 else "Success"), # store error message if there's a failure | |
| "checkpoint_used": checkpoint_path, | |
| } | |
| # run subprocess with real-time streaming (from the top level repo directory) | |
| logger.info("=== Subprocess Stdout ===") | |
| try: | |
| with subprocess.Popen( | |
| command, | |
| stdout=subprocess.PIPE, | |
| stderr=subprocess.PIPE, | |
| text=True, | |
| bufsize=1, | |
| cwd=repo_root, | |
| ) as proc: | |
| stdout_lines = [] | |
| for line in proc.stdout: | |
| logger.info(line.rstrip()) | |
| stdout_lines.append(line) | |
| stderr_str = proc.stderr.read() | |
| proc.wait(timeout=1800) # 30 minutes timeout | |
| returncode = proc.returncode | |
| stdout_str = "".join(stdout_lines) | |
| except subprocess.TimeoutExpired as e: | |
| proc.kill() | |
| stdout_str, stderr_str = proc.communicate() | |
| returncode = -1 | |
| logger.error("Subprocess decode.py timed out after 30 minutes!") | |
| except Exception as e: | |
| returncode = -1 | |
| stdout_str = "" | |
| stderr_str = str(e) | |
| # generate report | |
| report = { | |
| "run_name": run_name, | |
| "model": internal_model_name, | |
| "status": "SUCCESS" if returncode == 0 else "FAILED", | |
| "success": returncode == 0, # if returncode is 0, command worked | |
| "stdout": stdout_str, # store standard output (contains generated text like "Input ... -> ...") | |
| "stderr": (stderr_str if returncode != 0 else "Success"), # store error message if there's a failure | |
| "checkpoint_used": checkpoint_path, | |
| } |
Comment on lines
+40
to
+57
| """Computes cosine similarity between two layer activation arrays.""" | ||
| try: | ||
| a = np.asarray(arr1, dtype=np.float64).flatten() | ||
| b = np.asarray(arr2, dtype=np.float64).flatten() | ||
| min_len = min(a.size, b.size) | ||
| if min_len == 0: | ||
| return 0.0 | ||
| a = a[:min_len] | ||
| b = b[:min_len] | ||
| norm_a = np.linalg.norm(a) | ||
| norm_b = np.linalg.norm(b) | ||
| if norm_a == 0 or norm_b == 0: | ||
| return 0.0 | ||
| return float(np.dot(a, b) / (norm_a * norm_b)) | ||
| except Exception: # pylint: disable=broad-exception-caught | ||
| return 0.0 | ||
|
|
||
|
|
Contributor
There was a problem hiding this comment.
🟠 **Mathematical Error:** Computing cosine similarity by flattening different-sized tensors and then slicing them to a common minimum size `min_len` is mathematically incorrect and misleading. This compares unrelated spatial or channel dimensions when the tensor structures differ (e.g. comparing model outputs with mismatched sequence lengths or hidden dimensions). Slicing the flattened vectors hides serious architectural or layer layout bugs while returning a meaningless similarity score. The function should check for identical shapes and return `0.0` or raise a clear warning if they diverge.
Suggested change
| """Computes cosine similarity between two layer activation arrays.""" | |
| try: | |
| a = np.asarray(arr1, dtype=np.float64).flatten() | |
| b = np.asarray(arr2, dtype=np.float64).flatten() | |
| min_len = min(a.size, b.size) | |
| if min_len == 0: | |
| return 0.0 | |
| a = a[:min_len] | |
| b = b[:min_len] | |
| norm_a = np.linalg.norm(a) | |
| norm_b = np.linalg.norm(b) | |
| if norm_a == 0 or norm_b == 0: | |
| return 0.0 | |
| return float(np.dot(a, b) / (norm_a * norm_b)) | |
| except Exception: # pylint: disable=broad-exception-caught | |
| return 0.0 | |
| def compute_cosine_similarity(arr1: Any, arr2: Any) -> float: | |
| """Computes cosine similarity between two layer activation arrays.""" | |
| try: | |
| a = np.asarray(arr1, dtype=np.float64) | |
| b = np.asarray(arr2, dtype=np.float64) | |
| if a.shape != b.shape: | |
| logger.warning("Shape mismatch in cosine similarity: %s vs %s", a.shape, b.shape) | |
| return 0.0 | |
| a = a.flatten() | |
| b = b.flatten() | |
| if a.size == 0: | |
| return 0.0 | |
| norm_a = np.linalg.norm(a) | |
| norm_b = np.linalg.norm(b) | |
| if norm_a == 0 or norm_b == 0: | |
| return 0.0 | |
| return float(np.dot(a, b) / (norm_a * norm_b)) | |
| except Exception: # pylint: disable=broad-exception-caught | |
| return 0.0 |
Comment on lines
+95
to
+101
| if hf_arr is not None and mt_arr is not None: | ||
| cos_sim = compute_cosine_similarity(hf_arr, mt_arr) | ||
| row["cosine_similarity"] = cos_sim | ||
| if cos_sim < 0.98 and first_divergence_layer is None: | ||
| first_divergence_layer = idx | ||
| else: | ||
| row["cosine_similarity"] = None |
Contributor
There was a problem hiding this comment.
🟠 **Silent NaN Failure:** In Python, any comparison with a NaN value (e.g. `cos_sim < 0.98` where `cos_sim` is `nan`) always evaluates to `False`. If activations contain NaNs, the cosine similarity returns `nan`, which bypasses the divergence check completely and fails to set `first_divergence_layer`. We should check for `math.isnan` (or `np.isnan`) and register NaNs as a divergence layer.
Suggested change
| if hf_arr is not None and mt_arr is not None: | |
| cos_sim = compute_cosine_similarity(hf_arr, mt_arr) | |
| row["cosine_similarity"] = cos_sim | |
| if cos_sim < 0.98 and first_divergence_layer is None: | |
| first_divergence_layer = idx | |
| else: | |
| row["cosine_similarity"] = None | |
| if hf_arr is not None and mt_arr is not None: | |
| cos_sim = compute_cosine_similarity(hf_arr, mt_arr) | |
| row["cosine_similarity"] = cos_sim | |
| import math | |
| is_diverged = (cos_sim < 0.98 or math.isnan(cos_sim)) if cos_sim is not None else False | |
| if is_diverged and first_divergence_layer is None: | |
| first_divergence_layer = idx | |
| else: | |
| row["cosine_similarity"] = None |
Comment on lines
+80
to
+91
| ideal_shapes = load_shapes(args.ideal_shapes_path) | ||
| actual_shapes = load_shapes(args.actual_shapes_path) | ||
|
|
||
| _has_mismatch, _mismatched_layers = check_mismatches(ideal_shapes, actual_shapes) | ||
|
|
||
| report = { | ||
| "task": "checkpoint_shape_validation", | ||
| "timestamp": time.time(), | ||
| "status": "FAILURE" if _has_mismatch else "SUCCESS", | ||
| "mismatches_found": _has_mismatch, | ||
| "mismatched_layers": _mismatched_layers, | ||
| } |
Contributor
There was a problem hiding this comment.
🟠 **Failing Silently without Report:** If the shape files at `--ideal_shapes_path` or `--actual_shapes_path` are missing or fail to load, `load_shapes` raises an unhandled `FileNotFoundError`, causing the script to crash immediately. This prevents the schema-compliant JSON validation report (`shape_validation_report_*.json`) from being created and uploaded to GCS. Because Airflow fail-fast DAGs and Overwatch Agents expect a valid JSON report to diagnose and react to failures, we should wrap the shape file loading in a try-except block, log the failure, and write/upload a failure report to GCS before exiting.
Suggested change
| ideal_shapes = load_shapes(args.ideal_shapes_path) | |
| actual_shapes = load_shapes(args.actual_shapes_path) | |
| _has_mismatch, _mismatched_layers = check_mismatches(ideal_shapes, actual_shapes) | |
| report = { | |
| "task": "checkpoint_shape_validation", | |
| "timestamp": time.time(), | |
| "status": "FAILURE" if _has_mismatch else "SUCCESS", | |
| "mismatches_found": _has_mismatch, | |
| "mismatched_layers": _mismatched_layers, | |
| } | |
| try: | |
| ideal_shapes = load_shapes(args.ideal_shapes_path) | |
| actual_shapes = load_shapes(args.actual_shapes_path) | |
| _has_mismatch, _mismatched_layers = check_mismatches(ideal_shapes, actual_shapes) | |
| status = "FAILURE" if _has_mismatch else "SUCCESS" | |
| except FileNotFoundError as e: | |
| logger.error("ERROR: Shape file not found: %s", e) | |
| _has_mismatch = True | |
| _mismatched_layers = [] | |
| status = "FAILURE" | |
| report = { | |
| "task": "checkpoint_shape_validation", | |
| "timestamp": time.time(), | |
| "status": status, | |
| "mismatches_found": _has_mismatch, | |
| "mismatched_layers": _mismatched_layers, | |
| } |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
This PR introduces
decode_validator.py, adding automated autoregressive decoding and generative inference validation to the automated post checkpoint conversion validation pipeline. This is the fourth phase in the stacked checkpoint validation series (building on PR #4462, PR #4726, and #4727).While static shape checks and single-step forward-pass logit verification confirm that converted checkpoints (e.g. Gemma, Llama, Deepseek, Mistral, Qwen) align with expected network geometries and logit distributions, they do not validate multi-step autoregressive generation. Converted checkpoints can pass single-step forward checks yet still fail during text generation due to KV-cache allocation errors, RoPE/positional embedding misconfigurations, attention mask incompatibilities, or tokenizer decode incompatibilities.
This PR standardizes generative validation by utilizing MaxText's autoregressive decoding pipeline (
src/maxtext/inference/decode.py), streaming subprocess stdout/stderr directly into application logs for remote observability, and publishing schema-compliant JSON reports (report_{run_name}.json) to Google Cloud Storage for Airflow fail-fast orchestration.Key Implementation Details
validate_checkpoint)src/maxtext/inference/decode.pythat executes cleanly from the repository root directory (repo_root).tokenizer_pathandscan_layers), raising immediate descriptive exceptions if required checkpoint metadata is omitted from Airflow or CLI overrides.=== Subprocess Stdout ===) directly intoabsl.logging/ MaxText logger so that generated prompt-response sequences (Input ... -> ...), token generation latency, and throughput metrics are visible in real time within Airflow and CI task logs.SUCCESSvsFAILED), and error diagnostics into a structured JSON report (report_{run_name}.json)../reports/report_{run_name}.json) and uploads them to Google Cloud Storage (report_gcs_dir), enabling Airflow fail-fast pipelines and the automated Fixer Agent Sidecar to inspect decoding fidelity and trigger self-healing loops upon failure.Tests
Verified autoregressive decoding execution, real-time subprocess logging, and JSON report generation across 4 distinct models. The script correctly enforces required overrides, parses custom GCS checkpoint/tokenizer paths, and uploads structured reports to GCS upon completion.
1. Qwen3-14b
Airflow Trigger JSON:
{ "checkpoint_gcs_path": "gs://maxtext-model-checkpoints/qwen3-14b/scanned/2026-04-12/0/items", "hf_config_url": "https://huggingface.co/Qwen/Qwen3-14B/raw/main/config.json", "hf_model_path": "Qwen/Qwen3-14B", "hf_ref_code_url": "https://raw.githubusercontent.com/huggingface/transformers/main/src/transformers/models/qwen3/modeling_qwen3.py", "hf_token": "***", "maxtext_branch": "feature/checkpoint-validation-clean", "maxtext_commit_hash": null, "maxtext_model_name": "qwen3-14b", "maxtext_overrides": { "attention": "dot_product", "debug_tensors": true, "max_prefill_predict_length": 16, "max_target_length": 128, "per_device_batch_size": 4, "prompt": "I love to ", "scan_layers": true, "tokenizer_path": "Qwen/Qwen3-14B", "tokenizer_type": "huggingface" }, "report_gcs_dir": "gs://maxtext-validation-agent-reports/", "run_name": "master-qwen3-14b-full-run" }Decoding Validation Task Execution Log:
https://paste.googleplex.com/5045773933871104
2. Mistral-7b
Airflow Trigger JSON:
{ "checkpoint_gcs_path": "gs://maxtext-model-checkpoints/mistral-7b/2025-01-23-19-04/unscanned/checkpoints/0/items", "hf_config_url": "https://huggingface.co/mistralai/Mistral-7B-v0.1/raw/main/config.json", "hf_model_path": "mistralai/Mistral-7B-v0.1", "hf_ref_code_url": "https://raw.githubusercontent.com/huggingface/transformers/main/src/transformers/models/mistral/modeling_mistral.py", "hf_token": "***", "maxtext_branch": "feature/checkpoint-validation-clean", "maxtext_commit_hash": null, "maxtext_model_name": "mistral-7b", "maxtext_overrides": { "attention": "dot_product", "debug_tensors": true, "dtype": "bfloat16", "max_prefill_predict_length": 16, "max_target_length": 128, "per_device_batch_size": 4, "prompt": "I love to ", "scan_layers": false, "tokenizer_path": "mistralai/Mistral-7B-v0.1", "tokenizer_type": "huggingface" }, "report_gcs_dir": "gs://maxtext-validation-agent-reports/", "run_name": "master-mistral-7b-full-run" }Decoding Validation Task Execution Log:
https://paste.googleplex.com/4642963094372352
3. Qwen3-8b
Airflow Trigger JSON:
{ "checkpoint_gcs_path": "gs://maxtext-model-checkpoints/qwen3-8b/unscanned/0/items", "hf_config_url": "https://huggingface.co/Qwen/Qwen3-8B/raw/main/config.json", "hf_model_path": "Qwen/Qwen3-8B", "hf_ref_code_url": "https://raw.githubusercontent.com/huggingface/transformers/main/src/transformers/models/qwen3/modeling_qwen3.py", "hf_token": "***", "maxtext_branch": "feature/checkpoint-validation-clean", "maxtext_commit_hash": null, "maxtext_model_name": "qwen3-8b", "maxtext_overrides": { "attention": "dot_product", "debug_tensors": true, "max_prefill_predict_length": 16, "max_target_length": 128, "per_device_batch_size": 4, "prompt": "I love to ", "scan_layers": false, "tokenizer_path": "Qwen/Qwen3-8B", "tokenizer_type": "huggingface" }, "report_gcs_dir": "gs://maxtext-validation-agent-reports/", "run_name": "master-qwen3-8b-full-run" }Decoding Validation Task Execution Log:
https://paste.googleplex.com/6236433022058496
4. Gemma3-4b
Airflow Trigger JSON:
{ "checkpoint_gcs_path": "gs://maxtext-model-checkpoints/gemma3-4b/2025-03-18-19-03/unscanned/checkpoints/0/items", "hf_config_url": "https://huggingface.co/google/gemma-3-4b-it/raw/main/config.json", "hf_model_path": "google/gemma-3-4b-it", "hf_ref_code_url": "https://raw.githubusercontent.com/huggingface/transformers/main/src/transformers/models/gemma3/modeling_gemma3.py", "hf_token": "***", "maxtext_branch": "feature/checkpoint-validation-clean", "maxtext_commit_hash": null, "maxtext_model_name": "gemma3-4b", "maxtext_overrides": { "attention": "dot_product", "debug_tensors": true, "max_prefill_predict_length": 16, "max_target_length": 128, "per_device_batch_size": 4, "prompt": "I love to ", "scan_layers": false, "tokenizer_path": "google/gemma-3-4b-it", "tokenizer_type": "huggingface" }, "report_gcs_dir": "gs://maxtext-validation-agent-reports/", "run_name": "master-gemma3-4b-full-run" }Decoding Validation Task Execution Log:
https://paste.googleplex.com/4656868940185600
Stack
4th PR in stack. Depends on the preceding upstream PR: #4727
Checklist
Before submitting this PR, please make sure (put X in square brackets):
gemini-reviewlabel.