[BugFix] Detect empty distillation sequences in multidimensional batches

[aswanth-07]

Description

Fixes DistillationLoss's empty-sequence validation for multidimensional TensorDict batches. The previous mask.any() check collapsed both the token and remaining batch dimensions after the outer iteration, so an empty sequence was accepted whenever a sibling sequence selected at least one token.

The validation now reduces only the token dimension and checks every resulting per-sequence flag. The same token-dimension value is reused by the existing KL and sequence-length reductions. A transformers-free regression covers a [2, 2] batch where exactly one sequence is empty.

Closes #4081

Motivation and Context

An accepted empty sequence silently received a zero loss because normalization clamps the zero token count. Grouped LLM batches such as [prompts, generations] could therefore hide malformed masks instead of raising the documented validation error.

• I have raised an issue to propose this change: [BUG] DistillationLoss misses empty sequences in multidimensional batches #4081

Types of changes

• Bug fix (non-breaking change which fixes an issue)

Validation

• python -m pytest test/llm/test_llm_objectives.py -q -k "distillation and not integration": 16 passed
• python -m pytest test/llm/test_llm_objectives.py -q -m "not slow and not integration": 69 passed, 1 skipped, 2 deselected
• Repository pre-commit hooks on both changed files: all passed (ufmt, flake8, pydocstyle, pyupgrade, codespell, autoflake, docstring checks, and file checks)
• git diff --check: passed

Checklist

• I have read the CONTRIBUTING guide and CLAUDE.md.
• I have updated the tests accordingly.
• My change requires a documentation update.
• I have updated the documentation accordingly.

AI assistance disclosure

I used OpenAI Codex to help inspect the implementation, construct and run the pristine reproduction, search issues/pull requests/commits and the introducing PR for duplicates, implement the focused fix and regression, run validation, and draft this pull request. I reviewed the diff and reproduced the behavior before submission.

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/rl/4082

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

✅ No Failures

As of commit 0fe08b2 with merge base d7659c7 ():
💚 Looks good so far! There are no failures yet. 💚

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[vmoens]

Reviewed rebased head 0fe08b2dd1 plus the current PR discussion (there are no human review threads yet). The fix is correctly scoped: each as_list=True mask retains the remaining batch dimensions plus tokens, so reducing dim=tensordict.ndim - 1 handles only the token axis and the outer .any() detects an empty sequence anywhere in the residual batch. Reusing the same axis for KL summation and normalization keeps those paths aligned. The new multidimensional regression passes locally (1 passed), and git diff --check is clean. No blocking finding from this review.

[vmoens]

LGTM