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

Register Plugins in Ray Workers (#2901) [skip ci]

Browse Source 
* Access plugins in ray cluster

* Add comment

* chore: lint

---------

Co-authored-by: Ed Sealing <ed.sealing@patapsco.ai>
Co-authored-by: Wing Lian <wing@axolotl.ai>
This commit is contained in:
Ed Sealing
2025-07-11 13:59:59 -07:00
committed by GitHub
parent 03b2a113fe
commit eb662557a7
1 changed files with 7 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

7
src/axolotl/cli/train.py
View File

@@ -109,6 +109,13 @@ def ray_train_func(kwargs: dict):
# initialize accelerator before model instantiation # initialize accelerator before model instantiation
Accelerator(gradient_accumulation_steps=cfg.gradient_accumulation_steps) Accelerator(gradient_accumulation_steps=cfg.gradient_accumulation_steps)
# Register plugins in Ray workers
if cfg.get("plugins"):
from axolotl.cli.config import plugin_set_cfg, prepare_plugins
prepare_plugins(cfg)
plugin_set_cfg(cfg)
kwargs["cfg"] = cfg kwargs["cfg"] = cfg
do_train(**kwargs) do_train(**kwargs)