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

ddp shampoo

Browse Source
This commit is contained in:
Wing Lian
2024-09-18 10:50:46 -07:00
parent 69a29382e1
commit beaee36191
1 changed files with 2 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
src/axolotl/core/trainer_builder.py
View File

@@ -539,7 +539,7 @@ class AxolotlTrainer(SchedulerMixin, Trainer):
distributed_config = None distributed_config = None
if self.args.world_size > 1: if self.args.world_size > 1:
if self.args.fsdp_config: if self.args.fsdp and self.args.fsdp_config:
distributed_config = FSDPShampooConfig( distributed_config = FSDPShampooConfig(
param_to_metadata=compile_fsdp_parameter_metadata( param_to_metadata=compile_fsdp_parameter_metadata(
self.model_wrapped self.model_wrapped
@@ -547,7 +547,7 @@ class AxolotlTrainer(SchedulerMixin, Trainer):
) )
else: else:
distributed_config = DDPShampooConfig( distributed_config = DDPShampooConfig(
communication_dtype=CommunicationDType.BFLOAT16, communication_dtype=CommunicationDType.BF16,
num_trainers_per_group=self.args.world_size, num_trainers_per_group=self.args.world_size,
communicate_params=False, communicate_params=False,
) )