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

use skip_move_to_device for all cases (#3015)

Browse Source 
* use skip_move_to_device for all cases

* use experimental option for skip move
This commit is contained in:
Wing Lian
2025-08-06 00:13:12 -04:00
committed by GitHub
parent 42f5e6f9e9
commit 8021c718ce
2 changed files with 11 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
src/axolotl/loaders/model.py
View File

@@ -845,6 +845,9 @@ class ModelLoader:
self.model._tp_size = self.cfg.tensor_parallel_size self.model._tp_size = self.cfg.tensor_parallel_size
self.model._device_mesh = self.model_kwargs["device_mesh"] self.model._device_mesh = self.model_kwargs["device_mesh"]
if self.cfg.experimental_skip_move_to_device is not None:
skip_move_to_device = self.cfg.experimental_skip_move_to_device
return skip_move_to_device return skip_move_to_device
def _set_z3_leaf_modules(self): def _set_z3_leaf_modules(self):

8
src/axolotl/utils/schemas/model.py
View File

@@ -62,6 +62,14 @@ class ModelInputConfig(BaseModel):
json_schema_extra={"description": "Trust remote code for untrusted source"}, json_schema_extra={"description": "Trust remote code for untrusted source"},
) )
experimental_skip_move_to_device: bool | None = Field(
default=None,
json_schema_extra={
"description": "Don't move the model to the device before sharding. "
"This is an experimental feature that may be included in the future as the default."
},
)
@field_validator("trust_remote_code") @field_validator("trust_remote_code")
@classmethod @classmethod
def hint_trust_remote_code(cls, trust_remote_code): def hint_trust_remote_code(cls, trust_remote_code):