Support non orthogonal parallel axes and explicit replication annotation in dump comparator#19679
Merged
fzyzcjy merged 539 commits intosgl-project:mainfrom Mar 2, 2026
Merged
Support non orthogonal parallel axes and explicit replication annotation in dump comparator#19679fzyzcjy merged 539 commits intosgl-project:mainfrom
fzyzcjy merged 539 commits intosgl-project:mainfrom
Conversation
…verrideConfig→MetaOverrideConfig
…h dims+side field, move e2e tests to test_entrypoint.py
…omparator callsite
…ests, extract helpers
…or future extensibility
Move existing token_aligner logic into smart/ subdirectory. Add concat.py for simple BS=1 concat+truncate mode. AlignerPlan gains token_aligner_mode field for mode dispatch. Default mode is concat (no aux tensors needed).
TokenAlignerMode is an implementation detail of the token_aligner package. Outer modules (bundle_comparator, planner, entrypoint) now use Optional[str] instead of the typed alias.
Tests in TestEntrypointAlignment rely on aux tensor filtering which only happens in smart mode. With the new concat default, these tests need to explicitly request smart mode.
THD zigzag tests require thd_global_seq_lens from aux tensor processing, which only happens in smart mode.
11 new tests in TestEntrypointConcatMode covering: - multi-step concat with different data + truncation - concat combined with TP/CP unshard - unequal step counts between baseline and target - token dim resolution (t / s / fallback dim 0) - step ordering preservation - aux tensors not filtered in concat mode - aligner_plan field verification - failure path validation Also adds two helpers: _create_multi_step_rank_dump and _create_multi_step_tp_sharded_dumps for tests needing different tensors per step.
…ides, _run_concat
layer_id=[012] is not valid Python (leading zero in integer literal). Change to layer_id in [0, 1, 2] which eval() can handle correctly.
Add ParallelAxis.DP, ParallelState enum (SHARDED/CONCATED), DimSpec.parallel_state field, and axis:state colon parsing (e.g. "t(dp:concated)") to support DP Attention scenarios.
Concated axes (parallel_state == CONCATED) are excluded from both sharded and replicated sets, so they don't participate in unshard, validation, or coord tracking. This makes DP transparent when tensors are already dp_gathered (concated).
- test_dims: axis:state parsing, DP axis, error cases, dp_attn_mlp_dims - test_planner: TestConcatedAxes (skipped/empty/not-replicated/not-in-info) - test_parallel_info: test_dp_info - entrypoint test_planner: dp_concated_with_tp_sharded, _make_meta dp support
TestEntrypointDpConcated: test_dp2_tp2_concated and test_dp2_tp4_concated verify that DP concated axis is skipped by unsharder while TP is correctly concatenated, producing unified_shape matching full tensor.
# Conflicts: # test/registered/debug_utils/comparator/aligner/unsharder/test_executor.py
# Conflicts: # python/sglang/srt/debug_utils/comparator/bundle_comparator.py # python/sglang/srt/debug_utils/comparator/dims.py # python/sglang/srt/debug_utils/comparator/entrypoint.py # test/registered/debug_utils/comparator/aligner/reorderer/test_planner.py # test/registered/debug_utils/comparator/aligner/unsharder/test_executor.py # test/registered/debug_utils/comparator/test_dims.py # test/registered/debug_utils/comparator/test_entrypoint.py # test/registered/debug_utils/comparator/test_model_validation.py # test/registered/debug_utils/source_patcher/test_code_patcher.py # test/registered/debug_utils/source_patcher/test_dumper_integration.py # test/registered/debug_utils/source_patcher/test_source_editor.py
# Conflicts: # test/registered/debug_utils/comparator/test_entrypoint.py
# Conflicts: # test/registered/debug_utils/comparator/test_dims.py
Contributor
Summary of ChangesHello, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed! This pull request significantly upgrades the dump comparator's ability to handle complex tensor layouts and parallelization strategies. By introducing a new dimension specification system, it enables the tool to accurately align and compare tensors with fused dimensions and explicitly declared replicated axes. The refactored logging and CLI presets further enhance the usability and diagnostic capabilities of the comparator, making it more adaptable to diverse debugging needs. Highlights
🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console. Changelog
Using Gemini Code AssistThe full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips. Invoking Gemini You can request assistance from Gemini at any point by creating a comment using either
Customization To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a Limitations & Feedback Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here. You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension. Footnotes
|
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Code Review
This pull request introduces a significant and well-executed refactoring of the dump comparator utility. The core of the changes is the new dims_spec package, which provides a more structured and robust way to parse and handle tensor dimension specifications, including support for fused dimensions and explicit replication annotations. The logging system has also been improved by replacing the warning_sink with a more expressive log_sink that distinguishes between errors and informational messages. Additionally, the command-line interface is now more user-friendly with the introduction of presets. Overall, these changes greatly enhance the functionality and maintainability of the comparator. I've found one minor issue with a duplicated function.
Comment on lines
+157
to
+175
| def _maybe_load_tokenizer(args: argparse.Namespace) -> Any: | ||
| tokenizer_path: Optional[str] = getattr(args, "tokenizer", None) | ||
|
|
||
| if tokenizer_path is None: | ||
| for directory in [Path(args.baseline_path), Path(args.target_path)]: | ||
| tokenizer_path = read_tokenizer_path(directory) | ||
| if tokenizer_path is not None: | ||
| break | ||
|
|
||
| if tokenizer_path is None: | ||
| return None | ||
|
|
||
| try: | ||
| from transformers import AutoTokenizer | ||
|
|
||
| return AutoTokenizer.from_pretrained(tokenizer_path) | ||
| except Exception: | ||
| return None | ||
|
|
Contributor
There was a problem hiding this comment.
The function _maybe_load_tokenizer is defined twice in this file. This second definition appears to be a copy of the one defined on lines 137-155. This duplication should be removed.
# Conflicts: # python/sglang/srt/debug_utils/comparator/dims.py # test/registered/debug_utils/comparator/test_dims.py
Kangyan-Zhou
pushed a commit
to Kangyan-Zhou/sglang
that referenced
this pull request
Mar 4, 2026
…ion in dump comparator (sgl-project#19679)
magicYang1573
pushed a commit
to magicYang1573/sglang
that referenced
this pull request
Mar 9, 2026
…ion in dump comparator (sgl-project#19679)
Wangzheee
pushed a commit
to Wangzheee/sglang
that referenced
this pull request
Mar 21, 2026
…ion in dump comparator (sgl-project#19679)
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
1 participant
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.
Motivation
Modifications
Accuracy Tests
Benchmarking and Profiling
Checklist
Review Process
/tag-run-ci-label,/rerun-failed-ci,/tag-and-rerun-ci