From eb662557a7304485534779876f49793df0ec6a6a Mon Sep 17 00:00:00 2001 From: Ed Sealing <97403758+drikster80@users.noreply.github.com> Date: Fri, 11 Jul 2025 13:59:59 -0700 Subject: [PATCH] Register Plugins in Ray Workers (#2901) [skip ci] * Access plugins in ray cluster * Add comment * chore: lint --------- Co-authored-by: Ed Sealing Co-authored-by: Wing Lian --- src/axolotl/cli/train.py | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/src/axolotl/cli/train.py b/src/axolotl/cli/train.py index 28d4d543a..d0cf8455b 100644 --- a/src/axolotl/cli/train.py +++ b/src/axolotl/cli/train.py @@ -109,6 +109,13 @@ def ray_train_func(kwargs: dict): # initialize accelerator before model instantiation 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 do_train(**kwargs)