Fix DeepSpeed SP loss aggregation and LocalLayerRepository kwargs
Author: sshivampeta
Pull request: yes

<!-- ci-dashboard-badge:start -->
[![CI](https://transformers-ci.lor-e.huggingface.cool/badge/pr?pr=47073)](https://transformers-ci.lor-e.huggingface.cool/d/pytest-observability-by-pr/pytest-observability-branch?var-pr=47073)
<!-- ci-dashboard-badge:end -->

## What does this PR do?
Fixes weighted loss aggregation in DeepSpeed sequence parallelism and unblocks make check-repo by aligning LocalLayerRepository calls with the installed kernels API.

DeepSpeed SP loss (#47068): deepspeed_sp_compute_loss used a Python if good_tokens_per_rank[rank] > 0 loop over gathered tensors. That forces a CUDA host sync on every step and lets NaN losses from zero-token ranks (e.g. SFT ranks with only masked prompt tokens) pollute the aggregated loss. This PR replaces that loop with pure tensor ops via a small _aggregate_weighted_sp_loss helper using torch.where, preserving the same weighted average while staying on-device and ignoring zero-token ranks safely.

Kernels build fix: kernel_config.py and hub_kernels.py passed an unsupported package_name kwarg to LocalLayerRepository. The installed kernels package (0.15.x / 0.16.x) only accepts repo_path and layer_name, which caused ty to fail during make check-repo and would raise TypeError at runtime for local kernel loading. Removed the invalid argument to match the upstream API.

Fixes #47068

Motivation
Reported in #47068: DeepSpeed sequence-parallel training can produce incorrect or NaN aggregated loss when some SP ranks have zero valid tokens. The root cause is the Python conditional over CUDA tensors in the loss aggregation path.
The LocalLayerRepository change is included because without it the repo does not pass type checking locally, blocking validation of this fix.
Changes
•Add _aggregate_weighted_sp_loss() in src/transformers/integrations/deepspeed.py
•Use it from deepspeed_sp_compute_loss() instead of the Python sum/conditional loop
•Remove invalid package_name from LocalLayerRepository calls in kernel_config.py and hub_kernels.py
•Add unit tests in tests/integrations/test_deepspeed_sp_loss.py

Test plan
[x] pytest tests/integrations/test_deepspeed_sp_loss.py -v — 5/5 passed
[x] make check-repo — 24/24 passed

AI assistance
AI tools were used to draft and iterate on this change. The submitter reviewed all changed lines and ran the tests above locally.
Code Agent Policy
[x] I confirm that this is not a pure code agent PR.
Before submitting
[ ] This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
[x] Did you read the contributor guideline and the Pull Request checks?
[ ] Was this discussed/approved via a Github issue or the forum? Please add a link to it if that's the case.
[ ] Did you make sure to update the documentation with your changes according to the guidelines?
[x] Did you write any new necessary tests?


## Who can review?
@SunMarc @3outeille @vasqu


- AMD ROCm: @ivarflakstad
- Intel XPU: @IlyasMoutawwakil
- Ascend NPU: @ivarflakstad 
- Documentation: @stevhliu

Comments:
Thanks for working on this so quickly.

I originally reported this issue in #47068 after identifying the repeated 8-byte D2H memcpy and host-side `pthread_cond_wait` through Nsight Systems profiling. The GPU-side aggregation direction in this PR matches the approach I had in mind.

If you think my issue analysis and proposed direction warrant it, I’d be grateful if you would consider adding me as a co-author. If not, no worries — I appreciate you implementing the fix.

Co-authored-by: Myeongsu Koo <rnaudtndi@naver.com>

Thanks for the detailed profiling in #47068 — the Nsight trace identifying the repeated 8-byte D2H memcpy and the host-side pthread_cond_wait was what made the root cause clear. Happy to add you as co-author; I'll update the commit with the Co-authored-by trailer.

> Thanks for the detailed profiling in #47068 — the Nsight trace identifying the repeated 8-byte D2H memcpy and the host-side pthread_cond_wait was what made the root cause clear. Happy to add you as co-author; I'll update the commit with the Co-authored-by trailer.

Thanks a lot! I really appreciate the co-author credit. Glad the profiling was helpful!

@SunMarc Could you review when you have a chance? This addresses the D2H sync in `deepspeed_sp_compute_loss` reported in #47068 (cc'd to you by @Rocketknight1). Happy to adjust based on feedback.

> Thanks, just some nits. Does this fixes your issue @hdimmfh ?

Yes, this resolves the issue I reported. Thanks for reviewing!

The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/transformers/pr_47073). All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

Addressed both comments @SunMarc :
•Inlined the weighted aggregation into deepspeed_sp_compute_loss (removed _aggregate_weighted_sp_loss).
•Removed tests/integrations/test_deepspeed_sp_loss.py per your note that existing tests should cover this.
@hdimmfh — does this match what you saw in your Nsight profiling?

Thanks for the review, @SunMarc! CI is green — ready to merge when you have a moment.

<!-- ci-dashboard-recap:start -->

---

### CI recap

**Dashboard:** [View test results in Grafana](https://transformers-ci.lor-e.huggingface.cool/d/pytest-observability-by-pr/pytest-observability-branch?var-pr=47073)
**Latest run:** [28883268585:2](https://github.com/huggingface/transformers/actions/runs/28883268585)
**Result:** `success` | **Jobs:** 15 | **Tests:** 172,522 | **Failures:** 0 | **Duration:** 18h 10m

<!-- ci-dashboard-recap:end -->

Yes, this matches what I observed in my Nsight profiling. The changes still seem appropriate for addressing the issue. Thanks for addressing the review comments!
@sshivampeta 

This looks like a flaky CI issue (dist.barrier() → Gloo Connection reset by peer) rather than a regression from this PR. Could we rerun or re-enqueue it @SunMarc  ?

@SunMarc All checks are green and the branch is mergeable — ready to merge whenever you have a moment. Thanks!