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..1869a4efc0a4 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 @@ -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 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: