Fix test assertions for 2-gpu

[kevalmorabia97]

• Assertions should account for layers split across PP ranks

Summary by CodeRabbit

• Tests
  • Improved GPU test handling by capping usage to at most 2 GPUs.
  • Enhanced validation logic for distributed training scenarios with dynamic assertion handling.
  • Refined pruning score assertion checks to properly handle multi-GPU distributed validation paths.

_{✏️ Tip: You can customize this high-level summary in your review settings.}

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 74.64%. Comparing base (83ac3b1) to head (1b9d5cd).
⚠️ Report is 1 commits behind head on feature/compress.

Additional details and impacted files

@@                 Coverage Diff                  @@
##           feature/compress     #772      +/-   ##
====================================================
+ Coverage             74.62%   74.64%   +0.01%     
====================================================
  Files                   192      192              
  Lines                 18989    18989              
====================================================
+ Hits                  14171    14174       +3     
+ Misses                 4818     4815       -3     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[coderabbitai]

Important

Review skipped

Auto incremental reviews are disabled on this repository.

Please check the settings in the CodeRabbit UI or the .coderabbit.yaml file in this repository. To trigger a single review, invoke the @coderabbitai review command.

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

📝 Walkthrough

Walkthrough

This change modifies a GPU-accelerated torch compression test to cap GPU usage to at most 2 and restructures pruning score assertions. The validation logic is refactored to handle distributed scenarios by splitting assertion paths based on the distributed size and rank of processes.

Changes

Cohort / File(s)	Summary
Distributed Compression Test Updates tests/gpu/torch/_compress/test_compress.py	Caps GPU usage to 2 in tests. Restructures pruning score assertion validation to be size and rank-aware, introducing conditional paths for layer-specific validation. Moves assertion call to execute unconditionally post-compression while maintaining size-scaled assertion counts.

Estimated code review effort

🎯 2 (Simple) | ⏱️ ~12 minutes

🚥 Pre-merge checks | ✅ 3

✅ Passed checks (3 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title 'Fix test assertions for 2-gpu' clearly and concisely describes the main change: updating test assertions to handle 2-GPU scenarios where layers are distributed across pipeline-parallel ranks.
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands and usage tips.}