make sure everything stays in the same dtype when using dpo + FSDP (#1559)

src/axolotl/core/trainer_builder.py

@@ -54,6 +54,7 @@ from axolotl.utils.collators import (
     MambaDataCollator,
     V2BatchSamplerDataCollatorForSeq2Seq,
 )
+from axolotl.utils.models import ensure_dtype
 from axolotl.utils.samplers import MultipackBatchSampler, get_dataset_lengths
 from axolotl.utils.schedulers import (
     get_cosine_schedule_with_min_lr,
@@ -1569,6 +1570,9 @@ class HFRLTrainerBuilder(TrainerBuilderBase):
             callbacks=self.get_callbacks(),
             **dpo_trainer_kwargs,
         )
+        if self.cfg.fsdp:
+            ensure_dtype(dpo_trainer.model, dtype=self.cfg.torch_dtype)
+
         dpo_trainer = self.hook_post_create_trainer(dpo_trainer)
         for callback in self.get_post_trainer_create_callbacks(dpo_trainer):
             dpo_trainer.add_callback(callback)

src/axolotl/utils/models.py

@@ -993,3 +993,13 @@ def load_lora(model, cfg, inference=False, config_only=False):
         setup_quantized_peft_meta_for_training(model)
 
     return model, lora_config
+
+
+def ensure_dtype(model, dtype=torch.bfloat16):
+    for name, module in model.named_modules():
+        try:
+            if module.weight.dtype != dtype:
+                print(f"Converting module {name}: {module.weight.dtype} -> {dtype}")
+                module.to(dtype)
+        except AttributeError:
+            pass