Returns the kwargs needed to apply `accelerate` in `AutoModel.from_pretrained`.
(
device_map_option: Optional[str] = "auto",
max_memory_per_gpu: Optional[Union[int, str]] = None,
max_cpu_memory: Optional[Union[int, str]] = None,
offload_folder: Optional[str] = "./offload",
)
| 19 | |
| 20 | |
| 21 | def _get_accelerate_args( |
| 22 | device_map_option: Optional[str] = "auto", |
| 23 | max_memory_per_gpu: Optional[Union[int, str]] = None, |
| 24 | max_cpu_memory: Optional[Union[int, str]] = None, |
| 25 | offload_folder: Optional[str] = "./offload", |
| 26 | ) -> dict: |
| 27 | """Returns the kwargs needed to apply `accelerate` in `AutoModel.from_pretrained`.""" |
| 28 | max_memory = {} |
| 29 | if max_memory_per_gpu is not None: |
| 30 | max_memory_per_gpu_map = { |
| 31 | device_idx: max_memory_per_gpu |
| 32 | for device_idx in range(torch.cuda.device_count()) |
| 33 | } |
| 34 | max_memory.update(max_memory_per_gpu_map) |
| 35 | if max_cpu_memory is not None: |
| 36 | max_memory["cpu"] = max_cpu_memory |
| 37 | |
| 38 | args = {} |
| 39 | if max_memory: |
| 40 | args["max_memory"] = max_memory |
| 41 | args["device_map"] = device_map_option |
| 42 | args["offload_folder"] = offload_folder |
| 43 | return args |
| 44 | |
| 45 | |
| 46 | def _get_dtype( |