Add Qwen 3.5 0.8B Puzzletron configs - #2116
Conversation
Signed-off-by: Johannes Rausch <jrausch@nvidia.com>
Cover the complete model metadata and axis domains so incomplete recipe changes fail in CPU CI. Signed-off-by: Johannes Rausch <jrausch@nvidia.com>
Keep the default aligned with the tracked 0.8B FFN campaign while making the broader 9B-derived search an explicit, unvalidated opt-in. Signed-off-by: Johannes Rausch <jrausch@nvidia.com>
|
Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually. Contributors can view more details about this message here. |
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
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)
📝 WalkthroughWalkthroughAdds a pinned Qwen3.5 0.8B model configuration, pruning and search spaces, revision-aware Hugging Face conversion, transactional checkpoint publication, and unit tests for configuration and conversion behavior. ChangesQwen3.5 0.8B configuration and conversion
Estimated code review effort: 4 (Complex) | ~45 minutes Suggested reviewers: Sequence Diagram(s)sequenceDiagram
participant convert_stage
participant HuggingFace
participant TemporaryCheckpoint
participant PublishedCheckpoint
convert_stage->>HuggingFace: Resolve source with configured revision
convert_stage->>TemporaryCheckpoint: Convert and write metadata
TemporaryCheckpoint->>convert_stage: Validate checkpoint layout and metadata
convert_stage->>PublishedCheckpoint: Atomically publish validated checkpoint
🚥 Pre-merge checks | ✅ 6✅ Passed checks (6 passed)
✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Comment |
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## feature/puzzletron_v2 #2116 +/- ##
=========================================================
+ Coverage 54.76% 55.26% +0.50%
=========================================================
Files 704 704
Lines 90869 90960 +91
=========================================================
+ Hits 49761 50269 +508
+ Misses 41108 40691 -417
Flags with carried forward coverage won't be shown. Click here to find out more. ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Warning
CodeRabbit couldn't request changes on this pull request because it doesn't have sufficient GitHub permissions.
Please grant CodeRabbit Pull requests: Read and write permission and re-run the review.
Actionable comments posted: 2
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@examples/puzzletron/configs/families/qwen3_5/qwen3p5_0p8b/advanced.yaml`:
- Around line 45-50: Disable the unvalidated gdn_key_head_dim candidate by
setting its enabled flag to false or removing 96 from values, while preserving
the teacher value. Update the matching expected domain assertion in
test_qwen3p5_0p8b_example.py so it reflects the disabled or removed candidate.
In `@examples/puzzletron/configs/families/qwen3_5/qwen3p5_0p8b/model.yaml`:
- Around line 3-10: Update the Hugging Face source-resolution flow for
input_hf_model_path to read model_info.hf_revision and pass it as the revision
argument to snapshot_download alongside the repository ID, ensuring the resolved
teacher checkpoint uses the configured pinned revision.
🪄 Autofix
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 6e6e202c-5d2d-4f31-8a7e-d0ce33d32e2e
📒 Files selected for processing (4)
examples/puzzletron/configs/families/qwen3_5/qwen3p5_0p8b/advanced.yamlexamples/puzzletron/configs/families/qwen3_5/qwen3p5_0p8b/model.yamltests/unit/torch/puzzletron/test_portable_configs.pytests/unit/torch/puzzletron/test_qwen3p5_0p8b_example.py
Prevent the blocked GDN target from entering advanced searches and pin Hub downloads to the configured model revision. Signed-off-by: Johannes Rausch <jrausch@nvidia.com>
There was a problem hiding this comment.
Warning
CodeRabbit couldn't request changes on this pull request because it doesn't have sufficient GitHub permissions.
Please grant CodeRabbit Pull requests: Read and write permission and re-run the review.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@modelopt/torch/puzzletron/stages/convert.py`:
- Line 281: Update _is_complete_checkpoint and the surrounding convert_stage
resume logic to include model_cfg.get("revision") in the teacher checkpoint
identity, storing and comparing it before marking the stage skipped. Ensure
checkpoints from different source revisions are not reused, and add a regression
test covering the revision mismatch.
🪄 Autofix
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 806e8377-150e-4e6b-9cca-9504f635fa0a
📒 Files selected for processing (5)
examples/puzzletron/configs/families/qwen3_5/qwen3p5_0p8b/advanced.yamlexamples/puzzletron/configs/families/qwen3_5/qwen3p5_0p8b/model.yamlmodelopt/torch/puzzletron/stages/convert.pytests/unit/torch/puzzletron/test_convert_anymodel.pytests/unit/torch/puzzletron/test_qwen3p5_0p8b_example.py
🚧 Files skipped from review as they are similar to previous changes (3)
- examples/puzzletron/configs/families/qwen3_5/qwen3p5_0p8b/advanced.yaml
- examples/puzzletron/configs/families/qwen3_5/qwen3p5_0p8b/model.yaml
- tests/unit/torch/puzzletron/test_qwen3p5_0p8b_example.py
Signed-off-by: Johannes Rausch <jrausch@nvidia.com>
|
@coderabbitai review |
✅ Action performedReview finished.
|
There was a problem hiding this comment.
Warning
CodeRabbit couldn't request changes on this pull request because it doesn't have sufficient GitHub permissions.
Please grant CodeRabbit Pull requests: Read and write permission and re-run the review.
Actionable comments posted: 1
🧹 Nitpick comments (3)
tests/unit/torch/puzzletron/test_convert_anymodel.py (2)
52-57: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueConsider using the real
disthelpers in single-process tests.
modelopt/torch/utils/distributed.pyalready degrades to single-process behavior:size()returns 1,broadcastreturns the object unchanged,barrierreturns early, andrank()returns 0. The fakes foris_master,broadcast, andbarriertherefore add no isolation, and they hide call-signature drift. For example,lambda value, src: valuebreaks if production later callsdist.broadcast(obj)positionally withoutsrc.Keep the
_distributed_if_neededand_register_automodel_config_aliasespatches, and drop the threedistpatches so the tests exercise the real single-rank path.♻️ Proposed simplification
def _patch_single_rank_convert(monkeypatch): monkeypatch.setattr(convert_stage_module, "_register_automodel_config_aliases", lambda: None) monkeypatch.setattr(convert_stage_module, "_distributed_if_needed", nullcontext) - monkeypatch.setattr(convert_stage_module.dist, "is_master", lambda: True) - monkeypatch.setattr(convert_stage_module.dist, "broadcast", lambda value, src: value) - monkeypatch.setattr(convert_stage_module.dist, "barrier", lambda: None)Based on learnings from the coding guidelines: "For backend or runtime behavior, include an end-to-end test using the real implementation rather than replacing it with a fake."
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/unit/torch/puzzletron/test_convert_anymodel.py` around lines 52 - 57, Update _patch_single_rank_convert to retain only the _register_automodel_config_aliases and _distributed_if_needed patches. Remove the dist.is_master, dist.broadcast, and dist.barrier monkeypatches so single-process tests use the real distributed helpers and their production call signatures.Source: Coding guidelines
409-436: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winConsider covering the ambiguous transaction state.
This test covers both crash windows of
_publish_conversion_transaction. It does not cover the third branch, where the backup directory, the checkpoint directory, and the transaction directory all exist. That branch raisesRuntimeError("ambiguous conversion transaction state...")and protects against deleting the wrong checkpoint. Add a short case so a future change cannot silently turn the guard into a delete.💚 Proposed additional coverage
convert_stage_module._recover_conversion_transaction(teacher_dir) assert (teacher_dir / "published-shard.bin").read_text() == "published" assert not backup_dir.exists() + + backup_dir.mkdir() + transaction_dir.mkdir() + with pytest.raises(RuntimeError, match="ambiguous"): + convert_stage_module._recover_conversion_transaction(teacher_dir) + assert (teacher_dir / "published-shard.bin").read_text() == "published"🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/unit/torch/puzzletron/test_convert_anymodel.py` around lines 409 - 436, Add a focused case to test_conversion_transaction_recovers_interrupted_swap that creates the teacher, backup, checkpoint, and transaction directories simultaneously, then assert _recover_conversion_transaction raises RuntimeError with the ambiguous conversion transaction message and preserves all directories and their contents.modelopt/torch/puzzletron/stages/convert.py (1)
192-217: 🗄️ Data Integrity & Integration | 🔵 Trivial | ⚡ Quick winConsider guarding the transaction with a file lock.
_prepare_conversion_transactionand_publish_conversion_transactionhave no mutual exclusion across processes. Within one job,dist.is_master()limits writers to one rank. If two jobs share the sameteacher_dir, the two runs can interleave: both create the transaction directory, and the secondpath.replace(backup_dir)can hit the ambiguous state or delete the other run's backup. The repository already providesFileLockinmodelopt/torch/utils/distributed.py, which usesO_EXCLfor this purpose.Wrap prepare/publish in a lock on a sibling lock file if concurrent jobs on one experiment directory are supported.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@modelopt/torch/puzzletron/stages/convert.py` around lines 192 - 217, Guard the conversion transaction lifecycle in _prepare_conversion_transaction and _publish_conversion_transaction with the repository’s existing FileLock, using a lock file adjacent to the target path so concurrent jobs sharing the same teacher directory cannot interleave. Hold the lock across recovery, transaction-directory creation, publishing, rollback, and backup cleanup, while preserving the current behavior once the lock is acquired.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@modelopt/torch/puzzletron/stages/convert.py`:
- Around line 427-471: Wrap the master-only conversion and publishing block
around _prepare_conversion_transaction through _publish_conversion_transaction
in coordinated error handling that captures any exception on the master,
broadcasts the failure to all ranks, and raises it on every rank before reaching
dist.barrier(). Preserve the barrier only for successful conversions and ensure
non-master ranks do not wait when the master fails.
---
Nitpick comments:
In `@modelopt/torch/puzzletron/stages/convert.py`:
- Around line 192-217: Guard the conversion transaction lifecycle in
_prepare_conversion_transaction and _publish_conversion_transaction with the
repository’s existing FileLock, using a lock file adjacent to the target path so
concurrent jobs sharing the same teacher directory cannot interleave. Hold the
lock across recovery, transaction-directory creation, publishing, rollback, and
backup cleanup, while preserving the current behavior once the lock is acquired.
In `@tests/unit/torch/puzzletron/test_convert_anymodel.py`:
- Around line 52-57: Update _patch_single_rank_convert to retain only the
_register_automodel_config_aliases and _distributed_if_needed patches. Remove
the dist.is_master, dist.broadcast, and dist.barrier monkeypatches so
single-process tests use the real distributed helpers and their production call
signatures.
- Around line 409-436: Add a focused case to
test_conversion_transaction_recovers_interrupted_swap that creates the teacher,
backup, checkpoint, and transaction directories simultaneously, then assert
_recover_conversion_transaction raises RuntimeError with the ambiguous
conversion transaction message and preserves all directories and their contents.
🪄 Autofix
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 72169ca5-cb5e-469e-9270-ebb4248f50e4
📒 Files selected for processing (2)
modelopt/torch/puzzletron/stages/convert.pytests/unit/torch/puzzletron/test_convert_anymodel.py
Signed-off-by: Johannes Rausch <jrausch@nvidia.com>
|
@coderabbitai review |
✅ Action performedReview finished.
|
Signed-off-by: Johannes Rausch <jrausch@nvidia.com>
|
@coderabbitai review |
✅ Action performedReview finished.
|
2 participants
What does this PR do?
Type of change: New example and tests.
Qwen 3.5 0.8B does not have a version-pinned Puzzletron model example. It also needs a clear separation between a conservative default and broader targets that have not yet been validated end to end.
This PR pins the public checkpoint identity and geometry, then provides two search choices:
model.yamlis the default and searches only FFN intermediate sizes[3072, 2048].advanced.yamlis an explicit opt-in. It inherits the pinned model config and exposes six pruning axes. The axis structure is adapted from the existing Qwen 3.5 9B config, while the concrete targets are derived from the pinned 0.8B geometry. These targets were not selected from completed 0.8B search results and have not been fully runtime-validated.Hugging Face conversion now passes the configured model revision to checkpoint resolution, so a repository ID resolves to the pinned teacher snapshot instead of the repository's current revision. Converted teachers record that source and revision locally and are reused only when the configured identity matches. Replacements are built and validated in a clean sibling directory before publication, preventing stale shards and preserving the previous teacher if conversion or the directory swap is interrupted. In distributed runs, rank-zero failures are broadcast before later barriers so peers exit promptly instead of waiting for a timeout. A pinned legacy teacher without this record is converted once to establish the identity; unpinned legacy checkpoints retain the existing resume behavior.
The outcome is a narrow default for normal example reuse plus a clearly labeled advanced surface for follow-up validation, without presenting derived targets as measured results. The proposed advanced
gdn_key_head_dimchange from128to96remains documented but disabled with no executable candidate values until physical runtime-equivalence evidence is available.Usage
Select the default model config in a run:
- /families/qwen3_5/qwen3p5_0p8b/model@_global_Replace that line with the advanced choice when evaluating the broader validated search surface:
- /families/qwen3_5/qwen3p5_0p8b/advanced@_global_Testing
The dedicated Puzzletron v2 CPU suite covers the pinned public model identity, exact default and advanced domains, portable configuration discovery, Hydra inheritance and revision resolution, blocked-axis behavior, offline propagation of optional Hugging Face revisions through conversion, and rejection of a converted teacher from a different revision. No GPU smoke has been run, so the advanced configuration remains experimental.
Summary by CodeRabbit
New Features
Bug Fixes
Tests