From 5a622f0a385cfe9e2156b77b07260e9348415a6e Mon Sep 17 00:00:00 2001 From: RunguoLi Date: Mon, 21 Sep 2026 16:55:41 -0500 Subject: [PATCH 1/2] [gemini] accept tied parameter names when loading a state dict Fixes #6160. GeminiDDP builds `name2param` from every module's own parameters, so a tied parameter (e.g. `lm_head.weight` / `transformer.wte.weight` with `tie_word_embeddings=True`) is saved under all of its names, which matches `torch.nn.Module.state_dict`. The strict check in `_load_from_state_dict`, however, builds `local_state` from `named_parameters()`, which deduplicates and only keeps the first name. Loading a checkpoint that Gemini itself just saved therefore fails with Unexpected key(s) in state_dict: "lm_head.weight" for both sharded and unsharded checkpoints. Treat every name in `name2param` as expected. The existing Gemini checkpoint tests only use models without tied weights (LLaMA, BERT for sequence classification), so add a save/load round trip with a tied-embedding GPT-2. --- colossalai/zero/gemini/gemini_ddp.py | 4 ++- .../test_gemini_checkpoint_io.py | 28 ++++++++++++++++++- 2 files changed, 30 insertions(+), 2 deletions(-) diff --git a/colossalai/zero/gemini/gemini_ddp.py b/colossalai/zero/gemini/gemini_ddp.py index 9e89e88272e0..d27b5cdc1c08 100644 --- a/colossalai/zero/gemini/gemini_ddp.py +++ b/colossalai/zero/gemini/gemini_ddp.py @@ -813,7 +813,9 @@ def load_parameter(chunk_slice, data): for key in state_dict.keys(): if key.startswith(prefix) and key != extra_state_key: input_name = key[len(prefix) :] - if input_name not in local_state: + # tied parameters are saved under every name they have (same as torch.nn.Module.state_dict), + # while local_state only holds the first one, so the other names are expected too + if input_name not in local_state and input_name not in self.name2param: unexpected_keys.append(key) def _init_chunks(self, param_order, strict_ddp_mode: bool, cpu_offload: bool, pin_memory: bool): diff --git a/tests/test_checkpoint_io/test_gemini_checkpoint_io.py b/tests/test_checkpoint_io/test_gemini_checkpoint_io.py index ba60de0bd1c1..db1ac52f4fdc 100644 --- a/tests/test_checkpoint_io/test_gemini_checkpoint_io.py +++ b/tests/test_checkpoint_io/test_gemini_checkpoint_io.py @@ -3,7 +3,7 @@ import pytest import torch import torch.distributed as dist -from transformers import LlamaForCausalLM +from transformers import GPT2Config, GPT2LMHeadModel, LlamaForCausalLM from utils import shared_tempdir import colossalai @@ -216,6 +216,31 @@ def exam_lazy_from_pretrained(): check_state_dict_equal(state_dict, orig_state_dict, ignore_dtype=True) +@clear_cache_before_run() +@parameterize("shard", [True, False]) +def exam_tied_weights_state_dict(shard: bool): + # lm_head.weight and transformer.wte.weight are the same parameter, which is saved under both names + config = GPT2Config(n_layer=2, n_head=4, n_embd=64, vocab_size=512, n_positions=64) + + def model_fn(): + model = GPT2LMHeadModel(config) + assert model.lm_head.weight is model.transformer.wte.weight + return model + + booster = Booster(plugin=GeminiPlugin(precision="fp16", initial_scale=(2**14))) + model, *_ = booster.boost(model_fn()) + new_model, *_ = booster.boost(model_fn()) + + with shared_tempdir() as tempdir: + model_ckpt_path = f"{tempdir}/model" + booster.save_model(model, model_ckpt_path, shard=shard) + dist.barrier() + booster.load_model(new_model, model_ckpt_path) + state_dict = model.state_dict(only_rank_0=False) + assert "lm_head.weight" in state_dict and "transformer.wte.weight" in state_dict + check_state_dict_equal(state_dict, new_model.state_dict(only_rank_0=False), ignore_dtype=True) + + def run_dist(rank, world_size, port): colossalai.launch(rank=rank, world_size=world_size, host="localhost", port=port, backend="nccl") exam_state_dict() @@ -224,6 +249,7 @@ def run_dist(rank, world_size, port): exam_lazy_from_pretrained() elif rank == 0: print("Skipping lazy from_pretrained checkpoint coverage: LLAMA_PATH is not configured.") + exam_tied_weights_state_dict() @pytest.mark.dist From 636ffdbbc3579e767220130d51ad4452aa4ab9f4 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Wed, 14 Oct 2026 01:02:56 +0000 Subject: [PATCH 2/2] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- tests/test_checkpoint_io/test_gemini_checkpoint_io.py | 2 +- tests/test_zero/test_gemini/test_optim.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/test_checkpoint_io/test_gemini_checkpoint_io.py b/tests/test_checkpoint_io/test_gemini_checkpoint_io.py index db1ac52f4fdc..1869a4efc0a4 100644 --- a/tests/test_checkpoint_io/test_gemini_checkpoint_io.py +++ b/tests/test_checkpoint_io/test_gemini_checkpoint_io.py @@ -23,7 +23,7 @@ ASYNC_MODES = [False, True] if HAS_TENSORNVME else [False] try: - import apex.amp + pass APEX_AVAILABLE = True except ImportError: diff --git a/tests/test_zero/test_gemini/test_optim.py b/tests/test_zero/test_gemini/test_optim.py index b89c591882d6..a0970e43db9e 100644 --- a/tests/test_zero/test_gemini/test_optim.py +++ b/tests/test_zero/test_gemini/test_optim.py @@ -15,7 +15,7 @@ from tests.kit.model_zoo import model_zoo, run_fwd_bwd try: - import apex.amp + pass APEX_AMP_AVAILABLE = True except ImportError: