fix(minimax-h3): materialize TP column shards

[Leonccaa]

Summary

• materialize MiniMax-H3 dim-0 tensor-parallel shards with clone
• add a regression test for shard values, contiguity, storage ownership, and storage size

Why

torch.chunk(..., dim=0) returns a contiguous view. Calling .contiguous() on that view does not allocate new storage, so a rank-local CPU shard can keep the complete checkpoint tensor backing storage alive.

The local shard now owns only its expected storage. This does not change shard values, shapes, dtypes, non-TP loading, or inference numerics.

Upstream alignment

Rebased onto current upstream main at f8aee98b5462cca8d7288888146ebd95592bf266 after MiniMax-H3 AdaLN caching landed in #1413. The shard fix remains independent and applies cleanly.

Validation

• repository-pinned pre-commit run --all-files
• Ruff 0.11 lint and format checks
• Python bytecode compilation
• PyTorch 2.11 CPU storage regression check against _select_tensor_parallel_shard

V100 TP4 integration validation

Validated PR head 16e3a5fce94ad14b66c12c463c0fc2f67d69959d on four Tesla V100 PCIe 32 GB GPUs with a full MiniMax-H3 workload (864x480, 124 frames, 20 evaluations). The downstream V100 runtime uses the exact loader expression introduced by this PR:

torch.chunk(tensor, self.tp_size, dim=0)[self.tp_rank].clone(
    memory_format=torch.contiguous_format
)

Results were consistent across all four ranks:

• source tensor storage: 40,225,668,128 bytes (37.463 GiB) per rank before local materialization
• retained tensors: 10,543,386,144 bytes (9.819 GiB) per rank, or 26.21% of source storage; the excess over exactly 25% comes from intentionally replicated tensors
• CUDA allocated after load: 9.840 GiB per rank
• TP4 NVML peak: 14,078 MiB per GPU
• 20/20 model evaluations and 1000/1000 finite checks; latent finite; media decode passed
• TP4 process exit 0, OOMKilled=false; corrected and uncorrected volatile ECC remained 0

This confirms that each TP rank owns its local shard instead of retaining the complete 37.46 GiB source checkpoint storage.

Scope note: the end-to-end run used a downstream branch that also contains separate V100 compatibility work. This validates the identical shard-materialization path on V100; it does not claim that this PR alone provides all changes required for complete V100 inference.

[Leonccaa]

V100 TP4 integration results for the current head (16e3a5fce94ad14b66c12c463c0fc2f67d69959d) are now included in the PR body, including per-rank retained storage, CUDA/NVML peaks, finite checks, process status, and ECC results.

@gushiqiao, could you please review this TP shard storage-ownership fix when convenient? It is adjacent to the MiniMax-H3 TP checkpoint-loading path covered by #1353.

The current GitHub Actions suite is action_required with zero check runs, so it still needs maintainer approval before CI can start.

@codex review

[chatgpt-codex-connector]

Codex Review: Didn't find any major issues. Swish!

Reviewed commit: 16e3a5fce9

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".