|
| 1 | +"""Tests for HLO Graph Diff Verification. |
| 2 | +
|
| 3 | +Integrates validation checks system that detects unintended compiler graph transformations |
| 4 | +or model graph deviations without breaking isolated PR constraints validations. |
| 5 | +
|
| 6 | +This is part 1 of the automated compiler validation pipeline: |
| 7 | +- Ensures the PR isn't changing compiler performance and model graph generation unintentionally. |
| 8 | +- Checks current runtime dumps against base reference checkpoints (in tests/utils/reference_hlo_*.txt). |
| 9 | +""" |
| 10 | + |
| 11 | +import os |
| 12 | +import unittest |
| 13 | +import shutil |
| 14 | +import jax |
| 15 | +import pytest |
| 16 | +from maxtext.trainers.pre_train import train_compile |
| 17 | +from absl import flags |
| 18 | +import difflib |
| 19 | +import re |
| 20 | + |
| 21 | +# Maximum number of filtered lines from HLO to compare or store in the reference file |
| 22 | +MAX_LINES = 2000 |
| 23 | + |
| 24 | +# Matches operation sections content like: 1234 {"sharding parameters", ...} or 1234 "operation metadata" |
| 25 | +SECTION_CONTENT_REGEX = re.compile(r'^\d+ (?:"[^"]*"|\{[^\}]*\})$') |
| 26 | + |
| 27 | + |
| 28 | +# Filters and normalizes an HLO line to ignore environment-specific details. |
| 29 | +def filter_line(line): |
| 30 | + # Matches operation sections content like: 1234 {"sharding parameters", ...} or 1234 "operation metadata" |
| 31 | + if SECTION_CONTENT_REGEX.match(line): |
| 32 | + return None |
| 33 | + return line |
| 34 | + |
| 35 | + |
| 36 | +@pytest.mark.tpu_backend |
| 37 | +class TestHloDiff: |
| 38 | + """Tests for HLO Graph Diff Verification.""" |
| 39 | + |
| 40 | + def setup_method(self): |
| 41 | + # Disable cache to ensure compilation occurs every time |
| 42 | + jax.config.update("jax_enable_compilation_cache", False) |
| 43 | + |
| 44 | + def check_files_equal_ignoring_sections(self, file_path1, file_path2): |
| 45 | + """Asserts that two text files are identical, ignoring specific sections.""" |
| 46 | + |
| 47 | + def get_filtered_lines(file_path): |
| 48 | + with open(file_path, "r", encoding="utf-8") as f: |
| 49 | + lines = [] |
| 50 | + for line in f: |
| 51 | + line = filter_line(line) |
| 52 | + if line is None: |
| 53 | + continue |
| 54 | + lines.append(line) |
| 55 | + return lines |
| 56 | + |
| 57 | + lines1 = get_filtered_lines(file_path1)[:MAX_LINES] |
| 58 | + lines2 = get_filtered_lines(file_path2)[:MAX_LINES] |
| 59 | + |
| 60 | + if lines1 == lines2: |
| 61 | + return True |
| 62 | + |
| 63 | + print("\n" + "=" * 20 + " HLO Diff " + "=" * 20) |
| 64 | + for line in difflib.unified_diff(lines1, lines2, fromfile=file_path1, tofile=file_path2): |
| 65 | + print(line, end="") |
| 66 | + print("=" * 50 + "\n") |
| 67 | + return False |
| 68 | + |
| 69 | + @pytest.mark.parametrize("test_id, config_file, overrides, reference_suffix", [ |
| 70 | + ( |
| 71 | + "deepseek3", |
| 72 | + "src/maxtext/configs/models/deepseek3-test.yml", |
| 73 | + { |
| 74 | + "compile_topology": "v6e-4", |
| 75 | + "base_num_decoder_layers": 4, |
| 76 | + "per_device_batch_size": 1, |
| 77 | + "max_target_length": 128, |
| 78 | + }, |
| 79 | + "deepseek3", |
| 80 | + ), |
| 81 | + ]) |
| 82 | + def test_hlo_diff(self, test_id, config_file, overrides, reference_suffix): |
| 83 | + """Test HLO diff for parameterized configurations.""" |
| 84 | + local_landing_dir = os.path.join(os.path.dirname(__file__), f"hlo_diff_dump_{test_id}") |
| 85 | + |
| 86 | + # Clean up before run |
| 87 | + if os.path.exists(local_landing_dir): |
| 88 | + shutil.rmtree(local_landing_dir) |
| 89 | + os.makedirs(local_landing_dir, exist_ok=True) |
| 90 | + |
| 91 | + base_dir = os.path.abspath(os.path.join(os.path.dirname(__file__), "../..")) |
| 92 | + config_path = os.path.join(base_dir, config_file) |
| 93 | + |
| 94 | + # Arguments for train_compile |
| 95 | + test_args = [ |
| 96 | + None, |
| 97 | + config_path, |
| 98 | + "steps=1", |
| 99 | + "dataset_type=synthetic", |
| 100 | + "override_model_config=true", |
| 101 | + "compile_topology_num_slices=1", |
| 102 | + ] |
| 103 | + |
| 104 | + for key, value in overrides.items(): |
| 105 | + test_args.append(f"{key}={value}") |
| 106 | + |
| 107 | + test_args.append( |
| 108 | + f'compile_xla_flags="--xla_dump_to={local_landing_dir} ' |
| 109 | + f'--xla_dump_hlo_as_text=True --xla_dump_hlo_module_re=jit_train_step"' |
| 110 | + ) |
| 111 | + |
| 112 | + print(f"Running train_compile for {test_id} with args:", test_args) |
| 113 | + |
| 114 | + # Initialize flags if not already done to avoid ParseCommandLine error |
| 115 | + if not flags.FLAGS.is_parsed(): |
| 116 | + flags.FLAGS(["dummy_program", "--grain_train_files=dummy"], known_only=True) |
| 117 | + |
| 118 | + train_compile.main(tuple(test_args)) |
| 119 | + |
| 120 | + print(f"Files in landing dir for {test_id}:", os.listdir(local_landing_dir)) |
| 121 | + |
| 122 | + # Locate the dumped HLO file |
| 123 | + files = os.listdir(local_landing_dir) |
| 124 | + matches = [f for f in files if f.endswith(".after_optimizations.txt")] |
| 125 | + dumped_hlo = matches[0] if matches else None |
| 126 | + |
| 127 | + assert dumped_hlo, f"Dumped HLO file not found for {test_id}!" |
| 128 | + |
| 129 | + dumped_hlo_path = os.path.join(local_landing_dir, dumped_hlo) |
| 130 | + |
| 131 | + reference_hlo_path = os.path.join( |
| 132 | + os.path.dirname(__file__), f"../utils/reference_hlo_{reference_suffix}.txt" |
| 133 | + ) |
| 134 | + |
| 135 | + if not os.path.exists(reference_hlo_path): |
| 136 | + print(f"Reference file not found. Creating it at {reference_hlo_path}") |
| 137 | + with open(dumped_hlo_path, "r", encoding="utf-8") as f_in, open(reference_hlo_path, "w", encoding="utf-8") as f_out: |
| 138 | + count = 0 |
| 139 | + for line in f_in: |
| 140 | + line = filter_line(line) |
| 141 | + if line is None: |
| 142 | + continue |
| 143 | + f_out.write(line) |
| 144 | + count += 1 |
| 145 | + if count >= MAX_LINES: |
| 146 | + break |
| 147 | + print(f"Reference file created for {test_id}. Please commit it.") |
| 148 | + else: |
| 149 | + assert self.check_files_equal_ignoring_sections(dumped_hlo_path, reference_hlo_path), ( |
| 150 | + f"HLO deviation detected in {test_id}! If this is intended, please run the 'Update HLO Reference' workflow " |
| 151 | + f"from Github Actions against your branch to update the reference HLO." |
| 152 | + ) |
| 153 | + print(f"HLO comparison successful against {reference_hlo_path}") |
| 154 | + |
| 155 | + # Clean up |
| 156 | + shutil.rmtree(local_landing_dir) |
| 157 | + |
| 158 | + |
| 159 | +if __name__ == "__main__": |
| 160 | + unittest.main() |
0 commit comments