chore: Add unit test for full entropy calculation #1405
Conversation
Signed-off-by: Dheeraj Peri <[email protected]>
📝 WalkthroughWalkthroughA new Ray-based test actor Changes
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~20 minutes Pre-merge checks and finishing touches✅ Passed checks (4 passed)
✨ Finishing touches
🧪 Generate unit tests (beta)
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. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (4)
tests/unit/distributed/test_model_utils.py (4)
613-616: Avoid negative-index gather for padded targets.Using -1 in gather relies on Python negative indexing and can surprise future readers; multiply-by-mask zeroes it later but still indexes the last vocab entry. Prefer clamping/branching to avoid out-of-range indexing and keep gradients strictly zero for pads.
Apply:
- target_mask = target >= 0 # Valid targets (assuming -1 or similar for padding) - log_probs = torch.gather(log_softmax, -1, target.unsqueeze(-1)).squeeze(-1) - log_probs = log_probs * target_mask.float() + # Treat negative targets as padding; avoid negative-index gather + target_valid = target >= 0 + safe_target = torch.where(target_valid, target, torch.zeros_like(target)) + log_probs = torch.gather(log_softmax, -1, safe_target.unsqueeze(-1)).squeeze(-1) + log_probs = log_probs.masked_fill(~target_valid, 0.0)
1135-1144: Uniform‑case comment and minor cleanup.For uniform distributions, sum p log p = -log(V). Also, expand_as is redundant since shapes already match.
- # For uniform distribution over V items: H = -sum(1/V * log(1/V)) = log(V) + # For uniform over V items: sum p log p = log(1/V) = -log(V) @@ - expected_uniform_entropy.expand_as(entropy_uniform), + expected_uniform_entropy,
1211-1237: Enforce numerical thresholds in the harness (not just inside the actor).Actor asserts already gate failures, but adding harness-level asserts keeps parity with the gather-logprob tests and surfaces diffs clearly if tolerances change.
results = ray.get(futures) - for i, result in enumerate(results): - if "forward_max_diff" in result: - print(f"Worker {i} forward max diff: {result['forward_max_diff']:.2e}") - if "grad_max_diff" in result and "entropy_max_diff" in result: - print( - f"Worker {i} gradient max diff: {result['grad_max_diff']:.2e}, " - f"entropy max diff: {result['entropy_max_diff']:.2e}" - ) + for i, result in enumerate(results): + print(f"Worker {i} forward max diff: {result['forward_max_diff']:.2e}") + print( + f"Worker {i} gradient max diff: {result['grad_max_diff']:.2e}, " + f"entropy max diff: {result['entropy_max_diff']:.2e}" + ) + assert result["forward_max_diff"] < 1e-4, ( + f"Worker {i} forward diff too large: {result['forward_max_diff']}" + ) + assert result["grad_max_diff"] < 1e-4, ( + f"Worker {i} grad diff too large: {result['grad_max_diff']}" + ) + assert result["entropy_max_diff"] < 1e-4, ( + f"Worker {i} entropy diff too large: {result['entropy_max_diff']}" + )
1176-1184: Optional: add a “single‑chunk” path to the param grid.Include a chunk_size larger than seq_len to exercise the single-chunk path in ChunkedDistributedEntropy.
@pytest.mark.parametrize( "tp_size, chunk_size", [ (1, 4), (2, 4), (1, 1), (2, 1), + (1, 16), # chunk_size > seq_len (single-chunk path) + (2, 16), ], )
📜 Review details
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (1)
tests/unit/distributed/test_model_utils.py(3 hunks)
🧰 Additional context used
📓 Path-based instructions (1)
**/*.py
📄 CodeRabbit inference engine (CODING_GUIDELINES.md)
**/*.py: Follow the Google Python Style Guide for all Python code
Target Python 3.12+ for all Python code in NeMo-RL
Indent Python code with 4 spaces; do not use tabs
Python filenames should be snake_case (e.g., some_file.py)
Class names should be PascalCase
Function and method names should be snake_case
Local variable names should be snake_case; if starting with a number, prefix with k (e.g., k_99th_percentile)
Global variables should be UPPER_SNAKE_CASE and prefixed with G_ (e.g., G_MY_GLOBAL)
Constants should be UPPER_SNAKE_CASE
Avoid shadowing variables declared in an outer scope
Initialize all externally visible members of a class in the constructor
For public interfaces used outside a file, prefer docstrings over comments
Use comments mainly for code within a function or interfaces local to a file
Commented-out code must include a nearby comment explaining usage and why it is commented out; otherwise remove before merging
Use Google-style docstrings for classes and functions (Sphinx-parseable)
Avoid using reflection when functionality can be easily achieved without it
Limit except clauses to the smallest specific set of exceptions possible
For duck-typing via try/except, keep the try body minimal and use else for main logic
Add the NVIDIA copyright header (with current year) at the top of all Python files, excluding tests/ and test-only scripts
Files:
tests/unit/distributed/test_model_utils.py
🧬 Code graph analysis (1)
tests/unit/distributed/test_model_utils.py (4)
nemo_rl/distributed/model_utils.py (6)
ChunkedDistributedEntropy(985-1056)backward(101-140)backward(210-255)backward(317-381)backward(746-775)backward(1026-1056)nemo_rl/distributed/virtual_cluster.py (2)
PY_EXECUTABLES(42-58)RayVirtualCluster(177-435)nemo_rl/distributed/named_sharding.py (3)
NamedSharding(19-222)layout(99-101)names(84-86)nemo_rl/distributed/worker_groups.py (3)
RayWorkerBuilder(130-300)RayWorkerGroup(303-1004)run_all_workers_single_data(728-772)
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (2)
- GitHub Check: Post submodule check comment / Comment on PR
- GitHub Check: Post automodel integration comment / Comment on PR
| def _torch_baseline_entropy(self, full_logits): | ||
| """Single-GPU PyTorch baseline implementation for entropy computation.""" | ||
| # Compute log softmax and softmax using standard PyTorch | ||
| log_probs = torch.nn.functional.log_softmax(full_logits, dim=-1) | ||
| probs = torch.exp(log_probs) | ||
|
|
||
| # Compute entropy: H = -sum(p * log(p)) = -sum(p * log_p) | ||
| entropy = (probs * log_probs).sum(dim=-1) # [B, S] | ||
|
|
||
| return entropy |
There was a problem hiding this comment.
Entropy docstring/sign confusion.
The implementation returns sum_v p_v log p_v (non-positive), not the conventional positive Shannon entropy -sum p log p. Adjust the docstring/comments to match to avoid confusion.
- def _torch_baseline_entropy(self, full_logits):
- """Single-GPU PyTorch baseline implementation for entropy computation."""
+ def _torch_baseline_entropy(self, full_logits):
+ """Single-GPU baseline for sum_v p_v log p_v (non-positive 'entropy' used by ChunkedDistributedEntropy)."""
@@
- # Compute entropy: H = -sum(p * log(p)) = -sum(p * log_p)
- entropy = (probs * log_probs).sum(dim=-1) # [B, S]
+ # Compute H_all = sum_v p_v log p_v (<= 0)
+ entropy = (probs * log_probs).sum(dim=-1) # [B, S]🤖 Prompt for AI Agents
In tests/unit/distributed/test_model_utils.py around lines 971 to 980, the
docstring and inline comment claim the function computes "entropy" but the code
returns sum_v p_v * log p_v (a non-positive value), i.e. the negative of
conventional Shannon entropy; update the function docstring and the inline
comment to state explicitly that the function returns the negative Shannon
entropy (sum_v p_v * log p_v, ≤ 0) or "negative entropy" (or change the sign in
the computation if you prefer positive Shannon entropy), and make the comment
above the entropy computation consistent with that description.
1 participant
What does this PR do ?
I was trying to implement full entropy by understanding
ChunkedDistributedEntropyvia a test case and later I've observed this PR #1200 :) So this PR is just a test case. Feel free to close it if it isn't needed. cc: @parthchadhaIssues
List issues that this PR closes (syntax):
Usage
# Add a code snippet demonstrating how to use thisBefore your PR is "Ready for review"
Pre checks:
Additional Information
Summary by CodeRabbit