Fix hlo_diff_test to compile the real NNX path#4308
Merged
Conversation
ecnal-cienet
force-pushed
the
fix/hlo-diff-nnx-references
branch
from
d04acaf to
efbac61
Compare
Codecov Report✅ All modified and coverable lines are covered by tests. 📢 Thoughts on this report? Let us know! |
ecnal-cienet
force-pushed
the
fix/hlo-diff-nnx-references
branch
from
7d8b737 to
bff39a0
Compare
ecnal-cienet
requested review from
A9isha,
NicoGrande,
NuojCheng,
RissyRan,
SurbhiJainUSC,
abhinavclemson,
aireenmei,
bvandermoon,
darisoy,
dipannita08,
gagika,
gobbleturk,
hengtaoguo,
igorts-git,
jiangjy1982,
khatwanimohit,
richjames0,
shralex,
suexu1025,
vipannalla and
xibinliu
as code owners
Contributor
|
🤖 Hi @ecnal-cienet, I've received your request, and I'm working on it now! You can track my progress in the logs for more details. |
Contributor
|
🤖 I'm sorry @ecnal-cienet, but I was unable to process your request. Please see the logs for more details. |
xibinliu
reviewed
Jun 30, 2026
xibinliu
left a comment
xibinliu
left a comment
Collaborator
There was a problem hiding this comment.
Can we also change the NNX flag default values to "True", so even without including the base.yaml, the tests can use the True flags.
hlo_diff_test loaded each model yml directly as the top-level config, which skips base.yml and leaves pure_nnx at its pydantic schema default (False), so the test kept silently compiling the Linen path even after the NNX defaults were flipped to True in base.yml. Its reference HLOs are therefore still Linen. Compile via base.yml + model_name (the normal training path) so the config inherits base.yml's logical_axis_rules and exercises the NNX path. The three reference HLOs are regenerated for NNX via the Update HLO References workflow. Also flip the enable_nnx/pure_nnx/pure_nnx_decoder pydantic schema defaults to True so a config built without base.yml no longer silently falls back to the Linen path.
ecnal-cienet
force-pushed
the
fix/hlo-diff-nnx-references
branch
from
bff39a0 to
36e333c
Compare
xibinliu
approved these changes
Jul 1, 2026
NuojCheng
approved these changes
Jul 1, 2026
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
3 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
hlo_diff_testloaded each model yml (e.g.models/llama3-8b.yml) directly as the top-level config. Those model ymls declare nobase_configparent, sobase.ymlwas never merged in. Two things broke as a result:pure_nnxfell back to its pydantic schema default (False) instead ofbase.yml'sTrue, so the test silently kept compiling the Linen path even after the NNX defaults were flipped. Its reference HLOs were therefore stale Linen graphs.logical_axis_ruleslives only inbase.yml, so it was left empty.Compile via
base.yml+model_name(the normal training path) instead, so the config inheritsbase.yml'slogical_axis_rulesand exercises the real NNX path. The three reference HLOs (deepseek3,llama3_8b,qwen3_1.7b) are regenerated for NNX via the Update HLO References workflow.Tests
tests/utils/reference_hlo_*.txtin the GitHub Actions TPU (v6e-4) environment.tests/integration/hlo_diff_test.pypasses against the regenerated references in CI.Checklist
Before submitting this PR, please make sure (put X in square brackets):
gemini-reviewlabel.