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

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):

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
@@ -23,7 +23,7 @@
 
 ASYNC_MODES = [False, True] if HAS_TENSORNVME else [False]
 try:
-    import apex.amp
+    pass
 
     APEX_AVAILABLE = True
 except ImportError:
@@ -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

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: