[https://nvbugs/5154414][fix] Balanced layer to PP rank assignment

[amukkara]

Description

For some model and PP size combinations, num_hidden_layers % pp_size != 0. This PR creates a balanced assignment of layers to PP ranks in such cases, with few ranks assigned just one extra layer.

For example, Deepseek-V3 with 61 layers, pp size = 8, 61 % 8 = 5:
First 5 ranks get 8 layers each, last 3 ranks get 7 layers each.

Before this change, first 7 ranks get 7 layers each, and last rank get 12 layers causing OOM on last rank.

[amukkara]

/bot run

[tensorrt-cicd]

PR_Github #3247 [ run ] triggered by Bot

[tensorrt-cicd]

PR_Github #3247 [ run ] completed with state SUCCESS
/LLM/main/L0_MergeRequest_PR pipeline #2255 completed with status: 'FAILURE'

[amukkara]

/bot run

[tensorrt-cicd]

PR_Github #3305 [ run ] triggered by Bot

[tensorrt-cicd]

PR_Github #3305 [ run ] completed with state SUCCESS
/LLM/main/L0_MergeRequest_PR pipeline #2303 completed with status: 'FAILURE'

[amukkara]

/bot run

[tensorrt-cicd]

PR_Github #3339 [ run ] triggered by Bot

[tensorrt-cicd]

PR_Github #3339 [ run ] completed with state FAILURE
/LLM/main/L0_MergeRequest_PR pipeline #2330 completed with status: 'FAILURE'

[amukkara]

/bot run

[amukkara]

/bot run

[tensorrt-cicd]

PR_Github #3648 [ run ] triggered by Bot

[chang-l]

Can you use tensor_split (ref) here for simplicity? I think it should do the same thing.

[tensorrt-cicd]

PR_Github #3648 [ run ] completed with state SUCCESS
/LLM/main/L0_MergeRequest_PR pipeline #2581 completed with status: 'FAILURE'

[amukkara]

Can you use tensor_split (ref) here for simplicity? I think it should do the same thing.

@chang-l done in 1c9388eecbba058020795df4901d5338c56d1389

[amukkara]

/bot run

[tensorrt-cicd]

PR_Github #3928 [ run ] triggered by Bot

[tensorrt-cicd]

PR_Github #3928 [ run ] completed with state SUCCESS
/LLM/main/L0_MergeRequest_PR pipeline #2787 completed with status: 'FAILURE'

[amukkara]

/bot run

[tensorrt-cicd]

PR_Github #3952 [ run ] triggered by Bot

[tensorrt-cicd]

PR_Github #3952 [ run ] completed with state SUCCESS
/LLM/main/L0_MergeRequest_PR pipeline #2802 completed with status: 'FAILURE'

[amukkara]

/bot run

[tensorrt-cicd]

PR_Github #3974 [ run ] triggered by Bot

[tensorrt-cicd]

PR_Github #3974 [ run ] completed with state SUCCESS
/LLM/main/L0_MergeRequest_PR pipeline #2819 completed with status: 'FAILURE'

[amukkara]

Change included in #4399 and #4034