tocmo0nlord/axolotl
1
0
Fork 0
Code Issues Pull Requests Actions 2 Packages Projects Releases Wiki Activity

enable tensor parallel

Browse Source
This commit is contained in:
bursteratom
2024-12-11 10:38:14 -05:00
parent d009ead101
commit 42389c1f78
2 changed files with 2 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
src/axolotl/utils/config/models/input/v0_4_1/__init__.py
View File

@@ -393,7 +393,7 @@ class ModelInputConfig(BaseModel):
default=None, json_schema_extra={"description": "transformers processor class"} default=None, json_schema_extra={"description": "transformers processor class"}
) )
trust_remote_code: Optional[bool] = None trust_remote_code: Optional[bool] = None
tensor_parallel: Optional[Union[Literal["auto"], bool]] = "auto"
model_kwargs: Optional[Dict[str, Any]] = None model_kwargs: Optional[Dict[str, Any]] = None
@field_validator("trust_remote_code") @field_validator("trust_remote_code")

1
src/axolotl/utils/models.py
View File

@@ -621,6 +621,7 @@ class ModelLoader:
self.model_kwargs["device_map"] = device_map self.model_kwargs["device_map"] = device_map
self.model_kwargs["torch_dtype"] = self.cfg.torch_dtype self.model_kwargs["torch_dtype"] = self.cfg.torch_dtype
self.model_kwargs["tp_plan"] = self.cfg.tensor_parallel
cur_device = get_device_type() cur_device = get_device_type()
if "mps" in str(cur_device): if "mps" in str(cur_device):