add trainable_only to gather_cpu_state_dict

pytorch · Dec 13, 2024 · b5a8e63 · b5a8e63
1 parent e168f07
commit b5a8e63
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 0 deletions.
diff --git a/recipes/lora_dpo_distributed.py b/recipes/lora_dpo_distributed.py
@@ -550,6 +550,7 @@ def save_checkpoint(
             self._model,
             self._is_rank_zero,
             device=self._device,
+            trainable_only=self._save_adapter_weights_only,
         )
         if intermediate_checkpoint:
             opt_state_dict = training.get_full_optimizer_state_dict(

diff --git a/recipes/qat_lora_finetune_distributed.py b/recipes/qat_lora_finetune_distributed.py
@@ -699,6 +699,7 @@ def save_checkpoint(
             self._model,
             self._is_rank_zero,
             device=self._device,
+            trainable_only=self._save_adapter_weights_only,
         )
         if self._is_rank_zero:
             log.info(