teowu commited on
Commit
d705cd9
1 Parent(s): 613f33e

Update modeling_InternLM_XComposer.py

Browse files

model.to_empty() on CPU will help the model to successfully be loaded on 3090/4090.

Files changed (1) hide show
  1. modeling_InternLM_XComposer.py +2 -1
modeling_InternLM_XComposer.py CHANGED
@@ -74,7 +74,8 @@ class InternLMXComposerForCausalLM(PreTrainedModel):
74
  # speed up init llm
75
  with torch.device('meta'):
76
  self.internlm_model = InternLMForCausalLM._from_config(config)
77
- self.internlm_model.to_empty(device=config.device).to(torch.float16)
 
78
  for n, m in self.internlm_model.named_modules():
79
  if 'lora' in n:
80
  m.float()
 
74
  # speed up init llm
75
  with torch.device('meta'):
76
  self.internlm_model = InternLMForCausalLM._from_config(config)
77
+ self.internlm_model.to_empty(device='cpu').to(torch.float16)
78
+ self.internlm_model.to(config.device)
79
  for n, m in self.internlm_model.named_modules():
80
  if 'lora' in n:
81
  m.float()