[ckpt] feat: add NCCL M2N Reshard layout primitives

[ss16118]

What does this PR do?

Adds transport-independent tensor-layout primitives for the staged NCCL M2N Reshard checkpoint path.

veRL engine producers export rank-local parameters as (name, tensor, ShardSpec), while M2N requires explicit source and destination mesh dimensions, placements, local shapes, and communicator rank ranges. This PR introduces:

• LocalWeightDesc for normalized rank-local exports;
• ReshardLayout for one side of an M2N transfer;
• local_weight_desc_from_shard_api() for validated ShardSpec conversion; and
• build_reshard_layouts() for constructing matching source and destination layouts.

The reusable API uses replica/shard terminology rather than assigning DP or TP roles to its mesh axes.

This is the foundational PR in a staged integration. It intentionally contains no NCCL calls, checkpoint-backend registration, configuration, or rollout changes; those will be added in follow-up PRs.

This does not duplicate an existing PR. Searches for NCCL M2N and reshard layout found no matching implementation. The closest related work—#4876, #7263, #7108, and #7227—implements different consumers, transports, or wire formats rather than this layout-conversion layer.

Related to #400 and #1063.

Checklist Before Starting

• Searched for similar PRs and documented the results above.
• Formatted the title as [{modules}] {type}: {description}.

Test

pytest -q tests/checkpoint_engine/test_reshard_layout_on_cpu.py
# 12 passed

pre-commit run --all-files --show-diff-on-failure --color=always
# All 13 hooks passed

The CPU tests cover sharded and replicated layouts, one- and two-dimensional source meshes, topology mismatches, invalid local tensor sizes, unsupported placements, and deferred Hugging Face conversion metadata.

No GPU test is needed for this PR because it does not initialize NCCL or invoke M2N.

API and Usage Example

This adds an internal checkpoint-engine API and does not change user-facing configuration.

weight = local_weight_desc_from_shard_api(
    exported,
    destination_shard_dim=1,
)

source_layout, destination_layout = build_reshard_layouts(
    weight,
    source_replica_size=2,
    source_shard_size=16,
    destination_replica_size=8,
    destination_shard_size=4,
)

Design & Code Changes

• Convert the current engine ShardSpec contract into a transport-independent descriptor.
• Support replicated tensors and one sharded axis over a one- or two-dimensional source mesh.
• Validate source replica/shard sizes and rank-local tensor element counts.
• Fail closed on deferred conversion, explicit placement metadata, multiple sharded axes, and non-innermost source sharding.
• Represent replicated tensors as (world_size, 1) so all weights share the same two-axis representation.
• Place source and destination ranges contiguously in the dedicated M2N communicator.

Acknowledgements

This layout layer is adapted from the original NCCL M2N Reshard prototype and RFC by Alexandre Chidiac (@achidiac). Thanks to Ke Wen (@kwen2501) and Pouya Kousha (@pkousha) for design and code-review feedback.

AI assistance disclosure

OpenAI Codex was used for code exploration, implementation assistance, test orchestration, and preparation of this PR.

The human submitter must review every changed line and understand and be able to defend the change before requesting upstream review.

Checklist Before Submitting

• Read the Contribute Guide.
• Applied the complete pre-commit suite; all 13 hooks passed.
• Add / update documentation. Not needed for this internal helper. Its API and scope are documented with docstrings and CPU tests.
• Added unit tests discoverable by the CPU CI workflow.
• Request upstream CI in the veRL community channel when ready.
• Recipe submodule update. Not applicable.