update

Author: Yunnglin

src/twinkle/model/megatron/multi_lora_megatron.py

@@ -170,7 +170,7 @@ def set_lr_scheduler(self, scheduler_cls: Union[LRScheduler, Type[LRScheduler],
         self._check_adapter_valid(kwargs.get('adapter_name'))
         return super().set_lr_scheduler(scheduler_cls, **kwargs)
 
-    @remote_function(dispatch='all', sync=True)
+    @remote_function(dispatch='all', collect='first', sync=True)
     def save(self, name, output_dir: Optional[str] = None, interval=1, **kwargs):
         self._check_adapter_valid(kwargs.get('adapter_name'))
         optimizer_config = self.optimizer_group[kwargs.get('adapter_name')]

src/twinkle/model/transformers/multi_lora_transformers.py

@@ -209,7 +209,7 @@ def get_state_dict(self, **kwargs):
         self._check_adapter_valid(kwargs.get('adapter_name'))
         return self.multi_adapter.get_state_dict(kwargs.get('adapter_name'))
 
-    @remote_function()
+    @remote_function(collect='first')
     def save(self, name, output_dir: Optional[str] = None, interval=1, **kwargs):
         self._check_adapter_valid(kwargs.get('adapter_name'))
         with self.multi_adapter.save_context(kwargs.get('adapter_name')):

src/twinkle/model/transformers/transformers.py

@@ -781,7 +781,7 @@ def apply_patch(self, patch_cls: Union[Patch, Type[Patch], str], **kwargs):
     def __del__(self):
         HubOperation.wait_for()
 
-    @remote_function()
+    @remote_function(collect='first')
     def save(self, name: Optional[str] = None, output_dir: Optional[str] = None, interval: int = 1, **kwargs):
         """Save model.