Add a complete DDP program, including the process of collect and world_model/llm training.

Author: xiongjyu

lzero/entry/utils.py

@@ -528,9 +528,11 @@ def calculate_update_per_collect(
             collected_transitions_tensor
         ).item()
         updates = int(total_collected_transitions * cfg.policy.replay_ratio)
+        print(f"total_collected_transitions={total_collected_transitions}\tupdates={updates}")
     else:
         # In a single-process setup.
         updates = int(collected_transitions_num * cfg.policy.replay_ratio)
+        print(f"collected_transitions_num={collected_transitions_num}\tupdates={updates}")
 
     return max(1, updates) # Ensure at least one update.
 

zoo/jericho/priorzero/priorzero_config.py

@@ -135,7 +135,7 @@ class PriorZeroLLMConfig:
     rft_kl_coef: float = 0.01
     kl_estimator: str = "k3"
 
-    train_llm_after_wm_warm_step: int = int(1e3)
+    train_llm_after_wm_warm_step: int = int(1e2)
     value_norm_cfg: Optional[EasyDict] = field(default_factory=lambda: EasyDict({
         'enable_stability_optimizer': True,
         'value_norm_init_momentum': 0.9,        # Fast adaptation in early training
@@ -153,6 +153,7 @@ def get_priorzero_config(
     exp_name: str = None,
     use_cot: bool = False,
     model_key: Optional[str] = None,
+    multi_gpu: bool = False
 ) -> Tuple[EasyDict, EasyDict]:
     """
     Generate complete PriorZero configuration with automatic model configuration.
@@ -218,7 +219,7 @@ def get_priorzero_config(
     )
     policy_config = dict(
         type='priorzero',
-        multi_gpu=False,  
+        multi_gpu=multi_gpu,  
         use_wandb=False,
         learn=dict(
                 learner=dict(

zoo/jericho/priorzero/priorzero_datafactory.py

@@ -215,7 +215,7 @@ def build_llm_samples(self,
                 )
         return samples
 
-    def make_llm_train_samples(self, priorzero_batch) -> List[Dict[str, Any]]:
+    def make_llm_train_samples(self, priorzero_batch, ddp: bool = False) -> List[Dict[str, Any]]:
         """
         Convert PriorZero batch to LLM training samples.
 
@@ -235,14 +235,17 @@ def make_llm_train_samples(self, priorzero_batch) -> List[Dict[str, Any]]:
         samples = self.build_llm_samples(
             raw_obs_list, history_obs_list, action_logprob_list, target_value, cot_prefix_list
         )
-        per_rank = len(samples) // self.world_size
-        start = self.rank * per_rank
-        end = (self.rank + 1) * per_rank if self.rank != self.world_size - 1 else len(samples)
-        print(f"[Rank {self.rank}] process {start}: {end} samples, total {len(samples)} samples.")
-        real_samples = samples[start:end]
+        if ddp:
+            print(f"[Rank {self.rank}] process {len(samples)} samples collected by Rank {self.rank}")
+            real_samples = samples
+        else:
+            per_rank = len(samples) // self.world_size
+            start = self.rank * per_rank
+            end = (self.rank + 1) * per_rank if self.rank != self.world_size - 1 else len(samples)
+            print(f"[Rank {self.rank}] process {start}: {end} samples. Total {len(samples)} samples collected by Rank 0.")
+            real_samples = samples[start:end]
 
         prompts_only = [s["prompt"] for s in real_samples]
-
         if self.use_cot:
             targets_only = [s["prefix_cot"] + " " + s["target"] + self.tokenizer.eos_token for s in real_samples]
             if self.args.reward_func.format_reward: