[tensor]: fix saving checkpoints of padded parameters with tp_size=1 (keep detach/clone of padded tensors unpaddable)

[LiRunGuo]

📌 Checklist before creating the PR

• I have created an issue for this PR for traceability
• The title follows the standard format: [doc/gemini/tensor/...]: A concise description
• I have added relevant tags if possible for us to better distinguish different PRs
• I have installed pre-commit: pip install pre-commit && pre-commit install (ran the pinned hooks manually instead: black 24.10.0 --line-length=120, isort 5.13.2 --profile black, autoflake 2.3.1, all clean)

🚨 Issue number

fixed #6253

📝 What does this PR do?

booster.save_model with HybridParallelPlugin(tp_size=1) crashes for any model whose vocab size is not a multiple of make_vocab_size_divisible_by (64), e.g. GPT-2 (50257), sharded or not, with or without pipeline parallelism:

AttributeError: 'Tensor' object has no attribute '_unpad_detach'

Cause. With tp_size=1 the embedding / lm_head weights are padded tensors (PaddingEmbedding / PaddingLMHead). detach() and clone() of a padded tensor copy the padding metadata to the new tensor but do not hijack the new tensor's own detach / clone. The copy therefore looks like a padded tensor to is_padded_tensor, but to_unpadded_tensor fails on it because _hijack_back_detach_and_clone looks for _unpad_detach. HybridParallelCheckpointIO._model_sharder always unpads param.detach(), so every save hits this.

With tp_size > 1 the parameter is a distributed tensor and takes a different path, which is why the existing checkpoint tests (LLaMA, vocab 32000, a multiple of 64) never hit it.

Fix. Hijack detach / clone of the copies too, so they are complete padded tensors.

Tests.

• test_padded_tensor.py: detach(), clone() and detach().clone() of a padded tensor can be unpadded and give the original values.
• test_hybrid_parallel_plugin_checkpoint_io.py: save/load round trip of a GPT-2 with vocab 1000 for tp=1, pp=1 and tp=1, pp=2, sharded and unsharded, also checking that the checkpoint holds the unpadded vocab size.

Results (4x H200, torch 2.5.1+cu124)

	main	this PR
new unit test	AttributeError: '_unpad_detach'	passes
save_model, GPT-2 vocab 1000, tp=1 pp=1 / tp=1 pp=2, sharded and unsharded	crash	saves, reloads, shapes unpadded
tp=2 (distributed + padded tensor)	works	still works
test_padded_tensor (4 GPUs), exam_state_dict in the hybrid checkpoint test with use_async=False	-	pass

The use_async=True parametrizations were not run: tensornvme does not build on my machine.

💥 Checklist before requesting a review

• I have linked my PR to an issue (instruction)
• My issue clearly describes the problem/feature/proposal, with diagrams/charts/table/code if possible
• I have performed a self-review of my code
• I have added thorough tests.
• I have added docstrings for all the functions/methods I implemented

⭐️ Do you enjoy contributing to Colossal-AI?

• 🌝 Yes, I do.
• 🌚 No, I don't.