diff --git a/funasr/bin/train_ds.py b/funasr/bin/train_ds.py index b32881d99..d1f25973e 100644 --- a/funasr/bin/train_ds.py +++ b/funasr/bin/train_ds.py @@ -40,6 +40,37 @@ try: except: deepspeed = None +_DISTRIBUTED_TRAIN_CONF_KEYS = ( + "use_ddp", + "use_fsdp", + "use_deepspeed", + "deepspeed_config", +) + + +def _resolve_distributed_config(kwargs, world_size): + """Resolve distributed settings with top-level values taking precedence.""" + train_conf = dict(kwargs.get("train_conf") or {}) + + def get_setting(name, default): + if name in kwargs: + return kwargs[name] + return train_conf.get(name, default) + + use_fsdp = get_setting("use_fsdp", False) + use_deepspeed = get_setting("use_deepspeed", False) + deepspeed_config = get_setting("deepspeed_config", "") + if use_deepspeed and use_fsdp: + raise ValueError("use_deepspeed and use_fsdp cannot be enabled at the same time") + + trainer_conf = { + key: value + for key, value in train_conf.items() + if key not in _DISTRIBUTED_TRAIN_CONF_KEYS + } + use_ddp = world_size > 1 and not use_deepspeed and not use_fsdp + return use_ddp, use_fsdp, use_deepspeed, deepspeed_config, trainer_conf + @hydra.main(config_name=None, version_base=None) def main_hydra(kwargs: DictConfig): @@ -83,9 +114,13 @@ def main(**kwargs): if local_rank == 0: tables.print() - use_ddp = world_size > 1 - use_fsdp = kwargs.get("use_fsdp", False) - use_deepspeed = kwargs.get("use_deepspeed", False) + ( + use_ddp, + use_fsdp, + use_deepspeed, + deepspeed_config, + trainer_conf, + ) = _resolve_distributed_config(kwargs, world_size) if use_deepspeed: logging.info(f"use_deepspeed: {use_deepspeed}") deepspeed.init_distributed(dist_backend=kwargs.get("backend", "nccl")) @@ -143,10 +178,12 @@ def main(**kwargs): world_size=world_size, use_ddp=use_ddp, use_fsdp=use_fsdp, + use_deepspeed=use_deepspeed, + deepspeed_config=deepspeed_config, device=kwargs["device"], excludes=kwargs.get("excludes", None), output_dir=kwargs.get("output_dir", "./exp"), - **kwargs.get("train_conf"), + **trainer_conf, ) model = trainer.warp_model(model, **kwargs) diff --git a/tests/test_train_ds_distributed_config.py b/tests/test_train_ds_distributed_config.py new file mode 100644 index 000000000..20537d160 --- /dev/null +++ b/tests/test_train_ds_distributed_config.py @@ -0,0 +1,62 @@ +import unittest + +from funasr.bin.train_ds import _resolve_distributed_config + + +class TestDistributedConfig(unittest.TestCase): + def test_nested_train_conf_is_supported(self): + use_ddp, use_fsdp, use_deepspeed, config, trainer_conf = ( + _resolve_distributed_config( + { + "train_conf": { + "use_deepspeed": True, + "deepspeed_config": "nested_ds.json", + "log_interval": 10, + } + }, + world_size=2, + ) + ) + + self.assertFalse(use_ddp) + self.assertFalse(use_fsdp) + self.assertTrue(use_deepspeed) + self.assertEqual(config, "nested_ds.json") + self.assertEqual(trainer_conf, {"log_interval": 10}) + + def test_top_level_values_override_train_conf(self): + use_ddp, use_fsdp, use_deepspeed, config, trainer_conf = ( + _resolve_distributed_config( + { + "use_deepspeed": False, + "deepspeed_config": "top_level_ds.json", + "train_conf": { + "use_deepspeed": True, + "deepspeed_config": "nested_ds.json", + "use_ddp": True, + "log_interval": 10, + }, + }, + world_size=2, + ) + ) + + self.assertTrue(use_ddp) + self.assertFalse(use_fsdp) + self.assertFalse(use_deepspeed) + self.assertEqual(config, "top_level_ds.json") + self.assertEqual(trainer_conf, {"log_interval": 10}) + + def test_deepspeed_and_fsdp_are_mutually_exclusive(self): + with self.assertRaisesRegex(ValueError, "cannot be enabled"): + _resolve_distributed_config( + { + "use_deepspeed": True, + "train_conf": {"use_fsdp": True}, + }, + world_size=2, + ) + + +if __name__ == "__main__": + unittest.main()