(model)
| 31 | |
| 32 | |
| 33 | def dispatch_quantized_model(model): |
| 34 | num_devices = torch.cuda.device_count() |
| 35 | device_map = {"model.embed_tokens": 0, "model.norm": num_devices - 1, "lm_head": 0} |
| 36 | num_layers = len(get_layers(model)) |
| 37 | layers_per_device = math.ceil(num_layers / num_devices) |
| 38 | for layer_id in range(num_layers): |
| 39 | device_id = layer_id // layers_per_device |
| 40 | device_map[f"model.layers.{layer_id}"] = device_id |
| 41 | model = dispatch_model(model, device_map) |
| 42 | # for some reason dispatch doesn't put this modules on needed device |
| 43 | model.model.embed_tokens = model.model.embed_tokens.to("cuda:0") |
| 44 | model.lm_head = model.lm_head.to("cuda:0") |
| 45 | return model |
| 46 | |
| 47 | |
| 48 | def get_model( |
no test coverage detected