(checkpoint_path: Union[str, os.PathLike], num_gpus: int = 2,
device_map: Optional[Dict[str, int]] = None, **kwargs)
| 43 | |
| 44 | |
| 45 | def load_model_on_gpus(checkpoint_path: Union[str, os.PathLike], num_gpus: int = 2, |
| 46 | device_map: Optional[Dict[str, int]] = None, **kwargs) -> Module: |
| 47 | if num_gpus < 2 and device_map is None: |
| 48 | model = AutoModel.from_pretrained(checkpoint_path, trust_remote_code=True, **kwargs).half().cuda() |
| 49 | else: |
| 50 | from accelerate import dispatch_model |
| 51 | |
| 52 | model = AutoModel.from_pretrained(checkpoint_path, trust_remote_code=True, **kwargs).half() |
| 53 | |
| 54 | if device_map is None: |
| 55 | device_map = auto_configure_device_map(num_gpus) |
| 56 | |
| 57 | model = dispatch_model(model, device_map=device_map) |
| 58 | |
| 59 | return model |
no test coverage detected