Lightning-AI / Lightning-AI/litgpt

Gemma 2B weights seem to have changed

Open
#1,665 2 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

bug help wanted
Dominant language
Python
Stars
13.7k
Forks
1.5k
Avg merge
15h 37m
Merged PRs (30d)
1

Description

### Bug description

It seems that they updated the Gemma v1 2B weights. Something to look into:

```
⚡ main ~/litgpt litgpt chat checkpoints/google/gemma-2b
{'access_token': None,
'checkpoint_dir': PosixPath('checkpoints/google/gemma-2b'),
'compile': False,
'max_new_tokens': 50,
'multiline': False,
'precision': None,
'quantize': None,
'temperature': 0.8,
'top_k': 50,
'top_p': 1.0}
Traceback (most recent call last):
File "/home/zeus/miniconda3/envs/cloudspace/bin/litgpt", line 8, in
sys.exit(main())
File "/teamspace/studios/this_studio/litgpt/litgpt/__main__.py", line 71, in main
CLI(parser_data)
File "/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/jsonargparse/_cli.py", line 119, in CLI
return _run_component(component, init.get(subcommand))
File "/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/jsonargparse/_cli.py", line 204, in _run_component
return component(**cfg)
File "/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/torch/utils/_contextlib.py", line 115, in decorate_context
return func(*args, **kwargs)
File "/teamspace/studios/this_studio/litgpt/litgpt/chat/base.py", line 260, in main
load_checkpoint(fabric, model, checkpoint_path)
File "/teamspace/studios/this_studio/litgpt/litgpt/utils.py", line 362, in load_checkpoint
model.load_state_dict(state_dict, strict=strict)
File "/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/torch/nn/modules/module.py", line 2153, in load_state_dict
raise RuntimeError('Error(s) in loading state_dict for {}:\n\t{}'.format(
RuntimeError: Error(s) in loading state_dict for GPT:
Missing key(s) in state_dict: "transformer.h.0.norm_2.weight", "transformer.h.1.norm_2.weight", "transformer.h.2.norm_2.weight", "transformer.h.3.norm_2.weight", "transformer.h.4.norm_2.weight", "transformer.h.5.norm_2.weight", "transformer.h.6.norm_2.weight", "transformer.h.7.norm_2.weight", "transformer.h.8.norm_2.weight", "transformer.h.9.norm_2.weight", "transformer.h.10.norm_2.weight", "transformer.h.11.norm_2.weight", "transformer.h.12.norm_2.weight", "transformer.h.13.norm_2.weight", "transformer.h.14.norm_2.weight", "transformer.h.15.norm_2.weight", "transformer.h.16.norm_2.weight", "transformer.h.17.norm_2.weight".
Unexpected key(s) in state_dict: "transformer.h.0.post_attention_norm.weight", "transformer.h.1.post_attention_norm.weight", "transformer.h.2.post_attention_norm.weight", "transformer.h.3.post_attention_norm.weight", "transformer.h.4.post_attention_norm.weight", "transformer.h.5.post_attention_norm.weight", "transformer.h.6.post_attention_norm.weight", "transformer.h.7.post_attention_norm.weight", "transformer.h.8.post_attention_norm.weight", "transformer.h.9.post_attention_norm.weight", "transformer.h.10.post_attention_norm.weight", "transformer.h.11.post_attention_norm.weight", "transformer.h.12.post_attention_norm.weight", "transformer.h.13.post_attention_norm.weight", "transformer.h.14.post_attention_norm.weight", "transformer.h.15.post_attention_norm.weight", "transformer.h.16.post_attention_norm.weight", "transformer.h.17.post_attention_norm.weight".
```

We can either fix or remove these. Because there's Gemma 2, not sure why someone would care about Gemma 1. What do you think @Andrei-Aksionov ?

### What operating system are you using?

Unknown

### LitGPT Version

```

```

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Reproduce loading checkpoints/google/gemma-2b with the litgpt chat entry point in litgpt/chat/base.py. Trace load_checkpoint in litgpt/utils.py and the GPT state-dict names, then determine whether Gemma v1 compatibility should be restored or the checkpoint support removed; done means the intended behavior is confirmed by a loading or regression test.

Written by the indexing model from the issue text.

Assessment

Tech stack
python, pytorch
Domain
ai, machine-learning
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Stale
Clarity
Needs clarification
Newbie friendliness
28/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.