fix(collector): reject non-divisible DSV4 head geometry - #1611
Conversation
Signed-off-by: git-jxj <65210887+git-jxj@users.noreply.github.com>
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Enterprise Run ID: 📒 Files selected for processing (2)
Included review availability: Your plan provides up to 12 included reviews per hour; 5 remain after this review. 📜 Recent review details⏰ Context from checks skipped due to timeout. (12)
🧰 Additional context used📓 Path-based instructions (10)Enforce the collector rules from `.claude/rules/collector/layer_permissions.md`, `failure_handling.md`, and `case_authoring.md`.⚙️ CodeRabbit configuration file Files:
Check that tests cover the changed behavior rather than only the happy path.⚙️ CodeRabbit configuration file Files:
Prefer applicable inline comments.⚙️ CodeRabbit configuration file Files:
When a declared model row, quant mode, attention/MLA profile, or artifact configuration cannot be resolved, raise an error; never substitute defaults, another model's geometry, or a close-enough quant mode.📄 CodeRabbit inference engine (.claude/rules/collector/case_authoring.md) Files:
A collector has exactly two legal responses to a queued case: **execute it, or raise**.📄 CodeRabbit inference engine (.claude/rules/collector/layer_permissions.md) Files:
Record and classify every worker failure; include the module error record, backend collection summary, case parameters, exception details, and `(model, dtype)` group label before any worker reset.📄 CodeRabbit inference engine (.claude/rules/collector/failure_handling.md) Files:
Collector unit tests should preserve and verify the base-grid/model-shape expansion, deduplication, capability filtering, declaration validation, and loud failure behavior described by the collector rules.📄 CodeRabbit inference engine (.claude/rules/collector/case_authoring.md) Files:
Check the collector when generated configuration formats or generator parameter names change, because it may parse generated configs or reference parameter names.📄 CodeRabbit inference engine (.claude/rules/generator/cross_module_impact.md) Files:
Before making any change under `collector/**`, read `.claude/rules/collector/layer_permissions.md` and `.claude/rules/collector/failure_handling.md` first.📄 CodeRabbit inference engine (AGENTS.md) Files:
When editing `collector/**`, read and follow `.claude/rules/collector/layer_permissions.md` and `.claude/rules/collector/failure_handling.md`; for case YAML work, also read `case_authoring.md`.📄 CodeRabbit inference engine (.claude/rules/repo-guide.md) Files:
🔇 Additional comments (2)
WalkthroughThe DSV4 collector now validates that native attention heads divide evenly across tensor-parallel ranks. Unit tests cover exact division and invalid geometries across supported head reports. ChangesDSV4 head geometry
Estimated code review effort: 2 (Simple) | ~10 minutes Severity of issue fixed: Medium Merge Risk: ⚪ Minimal · up to This change rejects invalid non-divisible DSV4 tensor-parallel head configurations before they can produce truncated persisted head counts, while preserving valid exact-division behavior. The updated validation and coverage leave no identified merge-blocking risk. 🚥 Pre-merge checks | ✅ 4✅ Passed checks (4 passed)
Native heads line up in flight Comment |
|
Rechecked the failing CI on current head |
1 participant
AIC transition scope:
Overview:
Reject non-divisible DSV4 native-head/TP geometry before deriving the rank-local head count, as requested in #1466. For example, 64 native heads with TP=3 currently produces 21 local heads, which reconstructs as a different native identity (63).
This is a configuration-validation correction. Current standard 64/128-head models with the existing power-of-two TP sweep are unaffected; this PR does not claim those models have produced incorrect data.
Details:
Add the requested divisibility check and retain the existing module-head consistency validation. Extend the existing collector contract test with non-divisible and TP-larger-than-head-count cases, covering native, local, and absent module head attributes. Exact division, including one head per rank, remains valid.
Validation:
DID NOT RAISE RuntimeError) and passes after the fix.python -m pytest --confcutdir=tests/unit/collector tests/unit/collector/sglang/test_collect_dsv4_attn_contract.py -q: 6 passed. This follows the existing test's AST extraction of the actual helper and does not require importing SGLang or loading a model.node:path; no JavaScript files are changed.Remote CI baseline comparison: the Rust/Python engine-step parity job reports the same two golden-energy mismatches as the unchanged direct main baseline f254959. Both runs report 2 failed / 62 passed, with identical actual values for GPT-OSS-20B and Nemotron-NAS. This collector-only change does not modify those TRTLLM parity paths.
The unit CI job also matches the same main baseline: 1 failed / 2750 passed / 12 skipped. Its only failure is
test_power_columns_satisfy_energy_model_input_contract, with identical invalidpower_limitdata in 15 B200/TRTLLM parquet files.Where should the reviewer start?
collector/sglang/collect_dsv4_attn.py::_resolve_local_headsRelated Issues:
Fixes #1466
Summary by CodeRabbit
Bug Fixes
Tests