[gemini]: fix loading checkpoints of models with tied parameters

[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 #6160

📝 What does this PR do?

A checkpoint saved by GeminiPlugin cannot be loaded back by GeminiPlugin when the model has tied parameters (e.g. tie_word_embeddings=True: GPT-2, Qwen2.5 base models, ...):

RuntimeError: Error(s) in loading state_dict for GeminiDDP:
	Unexpected key(s) in state_dict: "lm_head.weight".

Cause. The save and load paths of GeminiDDP enumerate parameter names differently:

• save (_save_to_state_dict, state_dict_shard) iterates self.name2param, which is built from every module's own parameters, so a tied parameter is written under all of its names. This matches torch.nn.Module.state_dict.
• the strict check in _load_from_state_dict builds local_state from self.named_parameters(), which deduplicates and keeps only the first name, so the other name is reported as unexpected.

Fix. Treat every name in name2param as expected during the strict check. The save format is unchanged.

Test. The existing Gemini checkpoint tests only use models without tied weights (LLaMA, BERT for sequence classification). Added exam_tied_weights_state_dict: a save/load round trip with a tied-embedding GPT-2, sharded and unsharded, asserting that both names are present in the state dict and that the loaded weights are equal.

Results (4x H200, torch 2.5.1+cu124)

	main	this PR
tied GPT-2 round trip, shard=False	Unexpected key(s) in state_dict: "lm_head.weight"	loads, weights equal
tied GPT-2 round trip, shard=True	same error	loads, weights equal
exam_tied_weights_state_dict (new)	-	passes
exam_state_dict, exam_state_dict_with_origin with use_async=False	-	pass, no regression

Notes on what was not run: the use_async=True parametrizations need tensornvme, which does not build on my machine. The tp_size=2 cases of exam_state_dict additionally need #6260 fixed on a machine without apex, so they were verified together with that fix (submitted separately).

💥 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.