Fix 2-GPU test_model_load_utils hang; test import and fixture cleanup - #2079
Conversation
Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
Signed-off-by: Keval Morabia <28916987+kevalmorabia97@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 (3)
💤 Files with no reviewable changes (2)
🚧 Files skipped from review as they are similar to previous changes (1)
📝 WalkthroughWalkthroughThe PR centralizes filesystem, model, quantization, offload, and DFlash test utilities. Fixtures and tests use these helpers, share model builders, validate generated directories, and standardize import and dependency handling. ChangesTest infrastructure and validation
Estimated code review effort: 4 (Complex) | ~45 minutes 🚥 Pre-merge checks | ✅ 5 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (5 passed)
✨ Finishing Touches 💡 1📝 Generate docstrings 💡
🧪 Generate unit tests (beta)
Comment |
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## main #2079 +/- ##
===========================================
+ Coverage 67.17% 78.00% +10.83%
===========================================
Files 521 521
Lines 59857 59857
===========================================
+ Hits 40206 46693 +6487
+ Misses 19651 13164 -6487
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:
|
cjluo-nv
left a comment
cjluo-nv
left a comment
There was a problem hiding this comment.
Bot review (bedrock-claude-opus-5) — DM the bot to share feedback.
The core hang fix is correct and well-reasoned (parent-built ckpt_dir under tmp_path removes the per-rank os.getpid() divergence that sent rank 1 down _resolve_checkpoint_dir's Hub branch and its extra barrier; export_hf_checkpoint already mkdirs the export dir, so dropping the rank-0 makedirs+barrier is safe). Design-wise the new helpers reuse the existing tests/_test_utils pattern rather than introducing a competing one, so no architectural objection.
However, several things in the surrounding cleanup look likely to break lanes that weren't exercised locally, and one change touches third-party licensing:
tests/_test_utils/torch/quantization/attention.pyimportstransformers(andmodelopt...plugins.huggingface) unguarded at module top and is imported at module scope bytests/unit/torch/quantization/plugins/test_attention_quant.py. Because_test_utils.torch.quantization.attentionsorts before_test_utils.torch.transformers_models, thepytest.importorskip("transformers")that used to fire first no longer does — innox -s "partial_unit(subset='torch')"(installs onlymegatron-core+.[dev-test], i.e. nohfextra) that turns a clean skip into a collection error. The PR itself keeps function-local imports elsewhere with the comment "unit tests must import without transformers installed", so this looks like an oversight.examples/alpamayo/requirements.txtadds only the git dependency, buttests/examples/alpamayo/test_quantize.pyimportspandasand callsto_parquet(needs pyarrow), and the torch example lane installs[hf,dev-test]which has neither. Since the tests were previously auto-skipped this was invisible. Separately, adding a newgit+https://github.com/nvlabs/alpamayo@…dependency is a third-party licensing/dependency decision (PR body asserts Apache-2.0) that needs human/OSRB sign-off — that alone blocks an approve on licensing grounds.read_only_treeis a no-op for root, which is how these tests run inside the CI containers (root bypasses DAC write bits). The "fail loudly at the write" property therefore only holds for local non-root runs, so it can produce CI-vs-local divergence rather than a guarantee. Worth documenting, or enforcing another way.- The drive-by
@skip_flaky_on_blackwellontest_heterogenous_sharded_state_dictreduces coverage on the Blackwell runners and isn't mentioned anywhere in the PR description.
Also consider splitting: the CI lane restructuring + new alpamayo dependency carry materially different risk from the mechanical import/dedup cleanup and would be easier to review (and revert) on their own. No prompt-injection content was present in the PR text.
![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: 4
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
tests/_test_utils/onnx/quantization/utils.py (1)
19-37: 📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick winDefine the public APIs of the shared utility modules.
These modules export helpers that other tests import directly. Add
__all__declarations for the intended public symbols.
tests/_test_utils/onnx/quantization/utils.py#L19-L37: exportassert_nodes_are_quantized.tests/_test_utils/torch/quantization/attention.py#L18-L33: exportmake_quant_attention.tests/_test_utils/torch/quantization/offload.py#L20-L73: export the intended offload and layerwise helper functions.tests/_test_utils/torch/quantization/quant_utils.py#L18-L56: export the NVFP4 assertion helpers.tests/_test_utils/torch/speculative/dflash.py#L18-L42: exportget_dflash_configand any intended public constants.As per coding guidelines, “Define each module's public API with
__all__ = [...].”🤖 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/_test_utils/onnx/quantization/utils.py` around lines 19 - 37, Define explicit __all__ lists for the shared utility modules: in tests/_test_utils/onnx/quantization/utils.py:19-37 export assert_nodes_are_quantized; in tests/_test_utils/torch/quantization/attention.py:18-33 export make_quant_attention; in tests/_test_utils/torch/quantization/offload.py:20-73 export the intended offload and layerwise helper functions; in tests/_test_utils/torch/quantization/quant_utils.py:18-56 export the NVFP4 assertion helpers; and in tests/_test_utils/torch/speculative/dflash.py:18-42 export get_dflash_config plus any intended public constants.Source: Coding guidelines
🤖 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 `@tests/_test_utils/fs_utils.py`:
- Around line 18-21: Add a module-level __all__ declaration in fs_utils.py
containing only read_only_tree, while leaving the underscore-prefixed
_set_read_only helper private and unchanged.
In `@tests/_test_utils/torch/transformers_models.py`:
- Around line 571-597: Add docstrings to the public helpers get_tiny_mixtral and
create_tiny_mixtral_dir, documenting their returned model or temporary directory
and explaining that config_kwargs override the default Mixtral configuration;
also document the with_tokenizer option for create_tiny_mixtral_dir.
In `@tests/unit/onnx/quantization/test_qdq_rules_int8.py`:
- Line 71: Update the quantization assertion around assert_nodes_are_quantized
to remove the unrestricted ignore_identity_inputs=True bypass; use the default
validation, or only permit identity inputs after verifying each is ultimately
produced by DequantizeLinear. Keep concat-elimination handling scoped to
build_conv_concat_model.
In `@tests/unit/torch/export/test_quant_aware_conversion.py`:
- Line 359: Add a brief explanatory comment immediately above the local
get_tiny_mixtral import, stating that it remains local because the helper
requires the optional Transformers dependency; do not substitute a
pytest.importorskip() comment.
---
Outside diff comments:
In `@tests/_test_utils/onnx/quantization/utils.py`:
- Around line 19-37: Define explicit __all__ lists for the shared utility
modules: in tests/_test_utils/onnx/quantization/utils.py:19-37 export
assert_nodes_are_quantized; in
tests/_test_utils/torch/quantization/attention.py:18-33 export
make_quant_attention; in tests/_test_utils/torch/quantization/offload.py:20-73
export the intended offload and layerwise helper functions; in
tests/_test_utils/torch/quantization/quant_utils.py:18-56 export the NVFP4
assertion helpers; and in tests/_test_utils/torch/speculative/dflash.py:18-42
export get_dflash_config plus any intended public constants.
🪄 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: e701fd98-c966-44a4-8638-dd127119f6eb
📒 Files selected for processing (46)
.github/workflows/example_tests.ymlexamples/alpamayo/requirements.txttests/_test_utils/fs_utils.pytests/_test_utils/onnx/quantization/utils.pytests/_test_utils/torch/quantization/attention.pytests/_test_utils/torch/quantization/offload.pytests/_test_utils/torch/quantization/quant_utils.pytests/_test_utils/torch/speculative/dflash.pytests/_test_utils/torch/transformers_models.pytests/conftest.pytests/examples/alpamayo/test_quantize.pytests/examples/conftest.pytests/examples/diffusers/conftest.pytests/examples/diffusers/sparsity/test_sparsity.pytests/examples/diffusers/test_export_diffusers_hf_ckpt.pytests/examples/gpt-oss/test_gpt_oss_qat.pytests/examples/specdec_bench/test_redaction.pytests/examples/speculative_decoding/conftest.pytests/gpu/onnx/quantization/test_concat_elim.pytests/gpu/onnx/quantization/test_quantize_fp8.pytests/gpu/torch/export/test_quant_aware_conversion_gpu.pytests/gpu/torch/export/test_quant_utils.pytests/gpu/torch/export/test_vllm_fakequant_hf_export.pytests/gpu/torch/kernels/conftest.pytests/gpu/torch/kernels/sparsity/attention/test_diffusers_triton_attention.pytests/gpu/torch/quantization/plugins/test_accelerate_gpu.pytests/gpu/torch/quantization/plugins/test_attention_quant.pytests/gpu/torch/quantization/test_fsdp2.pytests/gpu/torch/quantization/test_lsq_cuda.pytests/gpu/torch/quantization/test_nvfp4_fp8_sweep_kernel.pytests/gpu/torch/sparsity/attention_sparsity/test_wan22_skip_softmax.pytests/gpu/torch/speculative/plugins/test_hf_dflash.pytests/gpu/torch/utils/test_model_load_utils.pytests/gpu_megatron/conftest.pytests/gpu_megatron/torch/export/test_unified_export_megatron.pytests/gpu_megatron/torch/quantization/plugins/test_megatron.pytests/gpu_megatron/torch/quantization/plugins/test_transformer_engine.pytests/gpu_trtllm/torch/export/test_export_compressed_nvfp4.pytests/gpu_trtllm/torch/quantization/backends/test_nvfp4_gemm.pytests/unit/onnx/quantization/test_qdq_rules_int8.pytests/unit/onnx/quantization/test_quantize_int8.pytests/unit/onnx/test_autocast_quantize.pytests/unit/torch/export/test_quant_aware_conversion.pytests/unit/torch/quantization/plugins/test_attention_quant.pytests/unit/torch/speculative/plugins/test_hf_dflash.pytests/unit/torch/speculative/plugins/test_hf_dflash_offline.py
💤 Files with no reviewable changes (4)
- tests/examples/gpt-oss/test_gpt_oss_qat.py
- tests/examples/specdec_bench/test_redaction.py
- tests/gpu_trtllm/torch/export/test_export_compressed_nvfp4.py
- tests/gpu/torch/kernels/sparsity/attention/test_diffusers_triton_attention.py
Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
…dified - _test_utils/torch/quantization/attention.py: importorskip transformers at import time; isort orders it before transformers_models, so its guard no longer ran first and the partial-install unit lane hit a collection error. - fs_utils: replace the chmod-based read_only_tree with assert_unmodified_tree, a file-manifest comparison on teardown. chmod is a no-op under root, which is how the CI containers run, so it never checked anything there. - test_quant_aware_conversion.py: note why the import stays function-local. Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
b79cee8 to
2ed5db0
Compare
|
Note GitHub couldn't provide a complete incremental comparison for this pull request, so CodeRabbit is performing a full review instead. This review may take a little longer. |
kevalmorabia97
changed the title
Update — force-pushed
|
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 `@tests/_test_utils/fs_utils.py`:
- Around line 47-50: Update the context manager around the yield in the fixture
helper so the post-yield manifest comparison always runs from a finally block,
including when the with-block raises. Preserve the original test exception when
possible while retaining the existing path existence check and manifest
comparison behavior.
🪄 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: 8961109f-2e08-420c-9681-ce6a1084f8ee
📒 Files selected for processing (41)
tests/_test_utils/fs_utils.pytests/_test_utils/onnx/quantization/utils.pytests/_test_utils/torch/quantization/attention.pytests/_test_utils/torch/quantization/offload.pytests/_test_utils/torch/quantization/quant_utils.pytests/_test_utils/torch/speculative/dflash.pytests/_test_utils/torch/transformers_models.pytests/conftest.pytests/examples/conftest.pytests/examples/diffusers/conftest.pytests/examples/diffusers/sparsity/test_sparsity.pytests/examples/diffusers/test_export_diffusers_hf_ckpt.pytests/examples/specdec_bench/test_redaction.pytests/gpu/onnx/quantization/test_concat_elim.pytests/gpu/onnx/quantization/test_quantize_fp8.pytests/gpu/torch/export/test_quant_aware_conversion_gpu.pytests/gpu/torch/export/test_quant_utils.pytests/gpu/torch/export/test_vllm_fakequant_hf_export.pytests/gpu/torch/kernels/conftest.pytests/gpu/torch/kernels/sparsity/attention/test_diffusers_triton_attention.pytests/gpu/torch/quantization/plugins/test_accelerate_gpu.pytests/gpu/torch/quantization/plugins/test_attention_quant.pytests/gpu/torch/quantization/test_fsdp2.pytests/gpu/torch/quantization/test_lsq_cuda.pytests/gpu/torch/quantization/test_nvfp4_fp8_sweep_kernel.pytests/gpu/torch/sparsity/attention_sparsity/test_wan22_skip_softmax.pytests/gpu/torch/speculative/plugins/test_hf_dflash.pytests/gpu/torch/utils/test_model_load_utils.pytests/gpu_megatron/conftest.pytests/gpu_megatron/torch/export/test_unified_export_megatron.pytests/gpu_megatron/torch/quantization/plugins/test_megatron.pytests/gpu_megatron/torch/quantization/plugins/test_transformer_engine.pytests/gpu_trtllm/torch/export/test_export_compressed_nvfp4.pytests/gpu_trtllm/torch/quantization/backends/test_nvfp4_gemm.pytests/unit/onnx/quantization/test_qdq_rules_int8.pytests/unit/onnx/quantization/test_quantize_int8.pytests/unit/onnx/test_autocast_quantize.pytests/unit/torch/export/test_quant_aware_conversion.pytests/unit/torch/quantization/plugins/test_attention_quant.pytests/unit/torch/speculative/plugins/test_hf_dflash.pytests/unit/torch/speculative/plugins/test_hf_dflash_offline.py
💤 Files with no reviewable changes (3)
- tests/gpu/torch/kernels/sparsity/attention/test_diffusers_triton_attention.py
- tests/gpu_trtllm/torch/export/test_export_compressed_nvfp4.py
- tests/examples/specdec_bench/test_redaction.py
🚧 Files skipped from review as they are similar to previous changes (34)
- tests/gpu_megatron/torch/quantization/plugins/test_megatron.py
- tests/_test_utils/onnx/quantization/utils.py
- tests/gpu/onnx/quantization/test_concat_elim.py
- tests/gpu_trtllm/torch/quantization/backends/test_nvfp4_gemm.py
- tests/unit/onnx/quantization/test_quantize_int8.py
- tests/gpu_megatron/torch/quantization/plugins/test_transformer_engine.py
- tests/_test_utils/torch/quantization/attention.py
- tests/examples/diffusers/test_export_diffusers_hf_ckpt.py
- tests/gpu/torch/utils/test_model_load_utils.py
- tests/_test_utils/torch/quantization/quant_utils.py
- tests/unit/torch/speculative/plugins/test_hf_dflash_offline.py
- tests/_test_utils/torch/speculative/dflash.py
- tests/gpu/onnx/quantization/test_quantize_fp8.py
- tests/gpu_megatron/conftest.py
- tests/gpu/torch/kernels/conftest.py
- tests/gpu/torch/speculative/plugins/test_hf_dflash.py
- tests/gpu_megatron/torch/export/test_unified_export_megatron.py
- tests/unit/onnx/quantization/test_qdq_rules_int8.py
- tests/unit/torch/quantization/plugins/test_attention_quant.py
- tests/unit/onnx/test_autocast_quantize.py
- tests/examples/diffusers/sparsity/test_sparsity.py
- tests/examples/diffusers/conftest.py
- tests/gpu/torch/export/test_vllm_fakequant_hf_export.py
- tests/_test_utils/torch/quantization/offload.py
- tests/gpu/torch/export/test_quant_utils.py
- tests/conftest.py
- tests/_test_utils/torch/transformers_models.py
- tests/gpu/torch/quantization/test_nvfp4_fp8_sweep_kernel.py
- tests/gpu/torch/quantization/test_lsq_cuda.py
- tests/gpu/torch/export/test_quant_aware_conversion_gpu.py
- tests/gpu/torch/quantization/test_fsdp2.py
- tests/gpu/torch/quantization/plugins/test_accelerate_gpu.py
- tests/unit/torch/speculative/plugins/test_hf_dflash.py
- tests/unit/torch/export/test_quant_aware_conversion.py
cjluo-nv
left a comment
cjluo-nv
left a comment
There was a problem hiding this comment.
Bot review (bedrock-claude-opus-5) — DM the bot to share feedback.
Re-review of the now tests/-only PR (41 files, +542/-588). Design protocol applied: the new helpers (fs_utils.py, torch/quantization/{attention,offload}.py, torch/speculative/dflash.py, additions to onnx/quantization/utils.py and torch/quantization/quant_utils.py) extend the existing tests/_test_utils package rather than introducing a competing helper layer, and each one replaces 2–5 verbatim copies, so there is no second system being introduced and no in-repo/dependency alternative that would do this better.
Verification of the previous round's blockers — all resolved:
tests/_test_utils/torch/quantization/attention.pynow runspytest.importorskip("transformers")before thetransformers/ HF-plugin imports, sopartial_unit(subset='torch')skips instead of erroring at collection.- The alpamayo lane,
examples/alpamayo/requirements.txtand all CI-workflow changes are out; no licensing/OSRB surface remains (no LICENSE/header/vendoring changes, new files carry the standard NVIDIA Apache-2.0 header). - The root-bypassable
chmodguard is gone;read_only_tree→assert_unmodified_tree(size+mtime_ns manifest compared on teardown), which works under root. - Dropped dependency guards check out against the lanes:
gpuinstalls.[all,dev-test](diffusers comes viahf),gpu_trtllm/gpu_megatronget tensorrt_llm / transformer_engine from their containers.
Spot-checks of the mechanical parts look right: _MIXTRAL_KWARGS / _LLAMA_KWARGS preserve the old shapes and (where it matters for tolerances) the old fp32 dtype; ignore_identity_inputs defaults to False so the three files whose local helper did not skip Identity keep their stricter assertion, and only the ResNet/concat call sites opt in; removed module-top imports (LlamaConfig, gs, diffusers) are still used or genuinely unused. The core hang fix (parent-built ckpt_dir/export_dir under tmp_path, one path for all ranks) is unchanged and correct.
The pytest process dies with STATUS_ILLEGAL_INSTRUCTION (0xC000001D) partway through this test on windows-latest, taking the whole unit run with it. Use the existing skip_on_windows fixture, matching test_quantize_cpu.py::test_save_restore. Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
assert_unmodified_tree returned silently if the tree was gone at teardown, which is the loudest mutation it should report, not the one it should ignore. Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
Update —
|
cjluo-nv
left a comment
cjluo-nv
left a comment
There was a problem hiding this comment.
Bot review (bedrock-claude-opus-5) — DM the bot to share feedback.
Re-review #3 of this tests/-only PR (42 files, +544/-589). Design protocol re-applied: the new helpers (fs_utils.py, torch/quantization/{attention,offload}.py, torch/speculative/dflash.py, additions to onnx/quantization/utils.py and torch/quantization/quant_utils.py) extend the existing tests/_test_utils package rather than introducing a competing layer, and each replaces 2–5 verbatim copies — no second system, no in-repo/dependency alternative that fits better. Core hang fix unchanged and correct (parent-built ckpt_dir/export_dir under tmp_path gives every rank one path, removing the os.getpid() divergence into _resolve_checkpoint_dir's Hub branch and the mismatched barrier).
Verified closed since the last pass:
assert_unmodified_treenow raises on a deleted fixture root (3022924);__all__present, manifest is size+mtime_nsso it works under root in the CI containers.tests/_test_utils/torch/quantization/attention.pyrunspytest.importorskip("transformers")before thetransformers/HF-plugin imports, and it still sorts ahead of themodelopt...plugins.huggingfaceimport intests/unit/torch/quantization/plugins/test_attention_quant.py, so the torch-only unit lane skips instead of erroring at collection.- Lane coverage for the guarded fixtures and the Blackwell skip are signed off by the owner.
- Spot-checks hold:
_MIXTRAL_KWARGS/_LLAMA_KWARGSpreserve the old shapes and the fp32 dtype where tolerances depend on it;ignore_identity_inputsdefaults toFalseso only the ResNet/concat call sites relax;LlamaConfig/LlamaAttention/_QuantAttentionare still used bytest_kitchen_faafter the local-helper removal;test_wan22_calibrated_exportonly reads the now-session-scopedtiny_wan22_pathand exports intotmp_path. - No licensing surface: no LICENSE/NOTICE/vendoring changes; new files carry the project's standard Apache-2.0 NVIDIA header (year reads 2024 vs. the canonical
LICENSE_HEADER's 2026 — cosmetic only).
Nudging rather than approving only because of one drive-by that no reviewer has looked at yet (see reason) plus the overall size of the mechanical cleanup.
- Remove create_tiny_mixtral_dir (added for symmetry, never called) and the pre-existing unused create_tiny_bert_dir; get_tiny_bert stays, it is used by test_attention_quant.py. No create_tiny_*_dir helper is unused now. - Drop __all__ from fs_utils.py; test utils do not declare it. - test_hf_patching.py: the comment claimed flakiness, but what we observed is a hard STATUS_ILLEGAL_INSTRUCTION crash. Say that instead, and note the root cause may be in the KD save/restore path rather than the test. Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
Update —
|
cjluo-nv
left a comment
cjluo-nv
left a comment
There was a problem hiding this comment.
Bot review (bedrock-claude-opus-5) — DM the bot to share feedback.
Re-review #4 of this tests/-only PR (42 files, +531/-593). Design protocol re-applied: the helpers (fs_utils.py, torch/quantization/{attention,offload}.py, torch/speculative/dflash.py, additions to onnx/quantization/utils.py and torch/quantization/quant_utils.py) extend the existing tests/_test_utils package rather than standing up a competing helper layer, and each replaces 2–5 verbatim copies — no second system, no better in-repo/dependency alternative. Core hang fix unchanged and correct: the checkpoint/export dirs are built once in the parent under tmp_path and passed into the workers, so all ranks agree on the path and rank 1 no longer takes _resolve_checkpoint_dir's Hub branch and its extra barrier.
Verified closed in this pass:
- Windows skip on
test_nested_model_save_restore— owner-signed-off; Windows support is scoped tomodelopt.onnx, so masking a torch KD save/restore crash there is intentional and acceptable. __all__intentionally dropped from the new test-util modules, matching the package convention (only 3 of ~38 pre-existing modules define it); confirmed absent fromfs_utils.py.create_tiny_bert_dirandcreate_tiny_mixtral_dirare gone fromtransformers_models.py,get_tiny_bertretained (still used bytest_attention_quant.py); no dangling references anywhere in the tree.assert_unmodified_treeraises on a deleted fixture root; manifest is size +mtime_ns, so it works under root in the CI containers.tests/_test_utils/torch/quantization/attention.pyrunspytest.importorskip("transformers")before thetransformers/HF-plugin imports and still sorts ahead of themodelopt...plugins.huggingfaceimport in the unit test, so the torch-only unit lane skips instead of erroring at collection.- Spot-checks hold: no now-unused module-top imports (
gsstill used in the ONNX tests,LlamaConfig/LlamaAttention/_QuantAttentionstill used bytest_kitchen_fa);_MIXTRAL_KWARGS/_LLAMA_KWARGSpreserve the old shapes and the fp32 dtype where tolerances depend on it;ignore_identity_inputsdefaults toFalseso only the ResNet/concat call sites relax (load-bearing there, per the author's reproduction). - No licensing surface: no LICENSE/NOTICE/vendoring changes; new files carry the project's standard Apache-2.0 NVIDIA header (year reads 2024 vs. the canonical header's later year — cosmetic only).
The remaining size concern is a soft guideline on a cohesive, mechanical test-only cleanup, and the last open drive-by is now explicitly signed off by the codeowner. Approving.
Complex PR: 36 files changed (≥ 20); 1124 lines changed (≥ 1000); spans 27 directories (≥ 5); 36 existing test files modified or removed. Looping in a human for approval.
|
2 participants
What does this PR do?
Type of change: Bug fix + new tests (test-suite only; changes are confined to
tests/)test_parallel_load_and_exporthang on 2 GPUs. The temp paths were built fromos.getpid()inside the workers, so each rank got a differentckpt_dir; rank 1 then took_resolve_checkpoint_dir's Hub branch and blocked on an extra barrier while rank 0 ran the loader's broadcasts. The checkpoint is now built once in the parent undertmp_pathand passed in, so all ranks agree on the path.tests/gpu*andtests/examples; function-local imports kept only where guarded (importorskip/try), where the import is the test (JIT compile), or where it must followsys.pathsetup._test_utilsinstead of local copies: addedget_tiny_mixtral; dedupedassert_nodes_are_quantized(5 copies), the accelerate-offload/layerwise config helpers,make_quant_attention,get_dflash_config, the NVFP4 amax assertions, and 3 copies of thetiny_wan22_pathfixture.assert_unmodified_tree): a file manifest is compared on teardown, so a test that writes into a session-scoped fixture directory fails instead of silently changing what later tests see.gpu_trtllm, transformer_engine ingpu_megatron, transformers in examples) so a missing dep fails loudly instead of skipping.test_heterogenous_sharded_state_dictis skipped on Blackwell (sm_120), matching the existing marker and its TE/CUDA-13 rationale — same tracking issue as CI: CUDA build/test hygiene + fix Puzzletron Nemotron test failures #1901.Testing
Local, 2x RTX 6000 Ada:
tests/gpu/torch/utils/test_model_load_utils.pypasses on 2 GPUs and on 1 GPU (previously hung on 2). Also ran the touched files intests/unit(608 passed) andtests/gpu(~340 passed).Before your PR is "Ready for review"
CONTRIBUTING.md: N/ASummary by CodeRabbit