myownskyW7 commited on
Commit
e782b1f
1 Parent(s): aa00681

optimize memory usage

Browse files
Files changed (1) hide show
  1. modeling_InternLM_XComposer.py +2 -2
modeling_InternLM_XComposer.py CHANGED
@@ -81,8 +81,8 @@ conversation
81
  # speed up init llm
82
  with torch.device('meta'):
83
  self.internlm_model = InternLMForCausalLM._from_config(config)
84
- self.internlm_model.to_empty(device=config.device).to(
85
- torch.float16)
86
  for n, m in self.internlm_model.named_modules():
87
  if 'lora' in n:
88
  m.float()
 
81
  # speed up init llm
82
  with torch.device('meta'):
83
  self.internlm_model = InternLMForCausalLM._from_config(config)
84
+ self.internlm_model.to_empty(device='cpu').to(torch.float16)
85
+ self.internlm_model.to(config.device)
86
  for n, m in self.internlm_model.named_modules():
87
  if 'lora' in n:
88
  m.float()