[Bug]: Fix sharding unit tests regression

[MrGeva]

System Info

Investigate and fix new regressions:
https://nvbugs/5766986
https://nvbugs/5766982
https://nvbugs/5766974

Who can help?

No response

Information

• The official example scripts
• My own modified scripts

Tasks

• An officially supported task in the examples folder (such as GLUE/SQuAD, ...)
• My own task or dataset (give details below)

Reproduction

na

Expected behavior

na

actual behavior

na

additional notes

na

Before submitting a new issue...

• Make sure you already searched for relevant issues, and checked the documentation and examples for answers to frequently asked questions.

[coderabbitai]

📝 CodeRabbit Plan Mode

Generate an implementation plan and prompts that you can use with your favorite coding agent.

• Create Plan

Examples

• Example 1
• Example 2

---

🔗 Similar Issues

Possible Duplicates

• ISSUE-10086

Related Issues

• [Bug]: hangs in allreduce tests #10086
• [Bug]: Crash and hang when using AUTO allreduce strategy #9149
• [Bug]: perf regression seen in the dashboard #8540
• [TEST_BUG][Bug]: Compilation error in uballocator when ENABLE_MULTI_DEVICE=0 #6804
• it run good on rtx 5060 ti 16 gb #8821

👤 Suggested Assignees

• MrGeva
• poweiw
• kmnnmk212-source

---

🧪 Issue enrichment is currently in open beta.

You can configure auto-planning by selecting labels in the issue_enrichment configuration.

To disable automatic issue enrichment, add the following to your .coderabbit.yaml:

issue_enrichment:
  auto_enrich:
    enabled: false

💬 Have feedback or questions? Drop into our discord or schedule a call!

[MrGeva]

fixed: #10271