Spaces:
Sleeping
Sleeping
Update app.py
Browse files
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))
|