From beaee36191a2c67408639d7f161beca0f2e059d2 Mon Sep 17 00:00:00 2001 From: Wing Lian Date: Wed, 18 Sep 2024 10:50:46 -0700 Subject: [PATCH] ddp shampoo --- src/axolotl/core/trainer_builder.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/axolotl/core/trainer_builder.py b/src/axolotl/core/trainer_builder.py index 02b68c7d7..52ac22b7d 100755 --- a/src/axolotl/core/trainer_builder.py +++ b/src/axolotl/core/trainer_builder.py @@ -539,7 +539,7 @@ class AxolotlTrainer(SchedulerMixin, Trainer): distributed_config = None if self.args.world_size > 1: - if self.args.fsdp_config: + if self.args.fsdp and self.args.fsdp_config: distributed_config = FSDPShampooConfig( param_to_metadata=compile_fsdp_parameter_metadata( self.model_wrapped @@ -547,7 +547,7 @@ class AxolotlTrainer(SchedulerMixin, Trainer): ) else: distributed_config = DDPShampooConfig( - communication_dtype=CommunicationDType.BFLOAT16, + communication_dtype=CommunicationDType.BF16, num_trainers_per_group=self.args.world_size, communicate_params=False, )