Merge pull request #32 from Qmeiyi/main

update save_interval

Author: Qmeiyi

src/dataflex/configs/components.yaml

@@ -8,6 +8,7 @@ selectors:
       gradient_type: adam
       proj_dim: 4096
       seed: 123
+      save_interval: 16
       reward_model_backend: local_vllm   # choices: [local_vllm, api]
       reward_backend_params:
         local_vllm:
@@ -37,6 +38,7 @@ selectors:
       gradient_type: adam
       proj_dim: 4096
       seed: 123
+      save_interval: 16
 
   loss:
     name: loss

src/dataflex/train/selector/nice_selector.py

@@ -266,6 +266,7 @@ def __init__(self,
                  reward_backend_params: Optional[Dict[str, Any]] = None,
                  gradient_type: str = "adam",
                  proj_dim: int = 8192,
+                 save_interval: int = 16,
                  seed: int = 42,
                  mc_samples: int = 4,
                  max_new_tokens: int = 512,
@@ -280,6 +281,7 @@ def __init__(self,
         self.eval_dataset = eval_dataset
         self.gradient_type = gradient_type
         self.proj_dim = proj_dim
+        self.save_interval = save_interval
         self.seed = seed
         self.mc_samples = mc_samples
         self.max_new_tokens = max_new_tokens
@@ -679,7 +681,7 @@ def indexed_collator_wrapper(features):
         dataloader = self.accelerator.prepare(dataloader)
 
         # 4) 设置保存间隔
-        save_interval = 64
+        save_interval = self.save_interval
 
         # 5) 断点续传
         max_index = self._get_max_saved_index(save_dir=save_dir)