deepspeedai / deepspeedai/DeepSpeedExamples

Same model Llama 7B, why does zero3 initialize different parameter sizes?

Open
#742 1 comment 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

Dominant language
Python
Stars
6.8k
Forks
1.1k
Avg merge
2d 16h
Merged PRs (30d)
1

Description

When I use Zero3, in initializing the network, if my Llama is rewritten by inheritance as follows:

class FlashLlamaModel(LlamaModel):
   def __init__(self, config: LlamaConfig):
       super().__init__(config) 

class FlashLlamaForCausalLM(LlamaForCausalLM):
   def __init__(self, config):
       super().__init__(config)
       self.model = FlashLlamaModel(config) 

actor_model = create_hf_model(
           model_class= FlashLlamaForCausalLM,  
           model_name_or_path=actor_model_name_or_path,
           tokenizer=self.tokenizer,
           ds_config=ds_config, 
           disable_dropout=self.args.disable_actor_dropout,
           debug=debug)

Log output: initialize about 14B parameters.
'[2023-09-27 01:48:00,149] [INFO] [partition_parameters.py:454:__exit__] finished initializing model with 13.63B parameters'

But when not overriding the parent self.model, the log output goes back to 7B.

class FlashLlamaModel(LlamaModel):
    def __init__(self, config: LlamaConfig):
        super().__init__(config) 

class FlashLlamaForCausalLM(LlamaForCausalLM):
    def __init__(self, config):
        super().__init__(config) 

actor_model = create_hf_model(
            model_class= FlashLlamaForCausalLM,  
            model_name_or_path=actor_model_name_or_path,
            tokenizer=self.tokenizer,
            ds_config=ds_config, 
            disable_dropout=self.args.disable_actor_dropout,
            debug=debug)

Log output: initialize about 7B parameters.
[2023-09-27 01:52:23,475] [INFO] [partition_parameters.py:454:__exit__] finished initializing model with 6.93B parameters

Contributor guide

No contributing guide indexed for this repository

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

Start at create_hf_model and compare the two FlashLlamaForCausalLM constructors, focusing on the explicit self.model assignment and the Zero3 initialization logs. Trace how each model instance is registered and verify what explains the change from about 7B to 14B parameters; done means the cause is documented or a focused fix is identified.

Written by the indexing model from the issue text.

Assessment

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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.