LoufAn commited on
Commit
3453c0f
·
1 Parent(s): db7fd06

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +1 -0
app.py CHANGED
@@ -9,6 +9,7 @@ def split_model(model_path):
9
  from transformers import AutoConfig
10
  device_map = {}
11
  world_size = torch.cuda.device_count()
 
12
  config = AutoConfig.from_pretrained(model_path, trust_remote_code=True)
13
  num_layers = config.llm_config.num_hidden_layers
14
  num_layers_per_gpu = math.ceil(num_layers / (world_size - 0.5))
 
9
  from transformers import AutoConfig
10
  device_map = {}
11
  world_size = torch.cuda.device_count()
12
+ print(f"world_size:{world_size}")
13
  config = AutoConfig.from_pretrained(model_path, trust_remote_code=True)
14
  num_layers = config.llm_config.num_hidden_layers
15
  num_layers_per_gpu = math.ceil(num_layers / (world_size - 0.5))