fix(pu): fix unizero_multitask ddp barrier bug

Author: puyuan1996

lzero/entry/train_unizero_multitask_segment_ddp.py

@@ -622,12 +622,15 @@ def train_unizero_multitask_segment_ddp(
             print('=' * 20)
             print(f'Starting collection for Rank {rank} task_id: {cfg.policy.task_id}...')
             print(f'Rank {rank}: cfg.policy.task_id={cfg.policy.task_id} ')
+            logging.info(f'Rank {rank}: Starting data collection for task {cfg.policy.task_id} at train_iter {learner.train_iter}')
 
             # Reset initial data before each collection, crucial for multi-task settings.
             collector._policy.reset(reset_init_data=True, task_id=cfg.policy.task_id)
             # Collect data.
             new_data = collector.collect(train_iter=learner.train_iter, policy_kwargs=collect_kwargs)
 
+            logging.info(f'Rank {rank}: Finished data collection for task {cfg.policy.task_id}, collected {len(new_data[0]) if new_data else 0} segments')
+
             # Update the replay buffer.
             replay_buffer.push_game_segments(new_data)
             replay_buffer.remove_oldest_data_to_fit()
@@ -648,6 +651,19 @@ def train_unizero_multitask_segment_ddp(
             # Log after data collection.
             logging.info(f'Rank {rank}: Completed data collection for task {cfg.policy.task_id}')
 
+        # ========== CRITICAL FIX: Synchronize all ranks after data collection ==========
+        # Wait for all ranks to complete their data collection before proceeding.
+        # This prevents fast-collecting ranks from reaching barriers/all_gather calls
+        # while slow-collecting ranks are still in the collection loop.
+        try:
+            logging.info(f'Rank {rank}: Waiting at post-collection barrier...')
+            dist.barrier()
+            logging.info(f'Rank {rank}: All ranks completed data collection, proceeding...')
+        except Exception as e:
+            logging.error(f'Rank {rank}: Post-collection barrier failed, error: {e}')
+            raise e
+        # ===============================================================================
+
         # Check if there is enough data for training.
         not_enough_data = any(
             replay_buffer.get_num_of_transitions() < cfgs[0].policy.total_batch_size / world_size
@@ -662,7 +678,9 @@ def train_unizero_multitask_segment_ddp(
             # Calculate task weights.
             try:
                 # Gather task rewards.
+                logging.info(f'Rank {rank}: Entering evaluation synchronization barrier at train_iter {learner.train_iter}')
                 dist.barrier()
+                logging.info(f'Rank {rank}: Passed evaluation barrier, gathering task returns')
                 all_task_returns = [None for _ in range(world_size)]
                 dist.all_gather_object(all_task_returns, task_returns)
                 # Merge task rewards.

lzero/model/unizero_world_models/lpips.py

@@ -13,7 +13,6 @@
 from torchvision import models
 from tqdm import tqdm
 
-
 class LPIPS(nn.Module):
     # Learned perceptual metric
     def __init__(self, use_dropout: bool = True):

zoo/atari/config/atari_unizero_multitask_segment_ddp_config.py

@@ -384,6 +384,13 @@ def create_env_manager() -> EasyDict:
     reanalyze_partition = 0.75
 
     # ==================== Training Loop ====================
+    # Set NCCL timeout to prevent watchdog hang due to unbalanced data collection speeds
+    # Different games (e.g., Pong vs Seaquest) have vastly different episode lengths,
+    # which can cause some ranks to finish collection much faster than others.
+    # Default timeout is 30 minutes; we increase it to 60 minutes for safety.
+    os.environ.setdefault('NCCL_TIMEOUT', '3600')  # 60 minutes in seconds
+    os.environ.setdefault('NCCL_BLOCKING_WAIT', '1')  # Enable blocking wait for better error messages
+
     for seed in [0]:
         configs = generate_configs(
             env_id_list, action_space_size, collector_env_num, n_episode, evaluator_env_num,

zoo/atari/config/atari_unizero_segment_config.py

@@ -68,7 +68,8 @@ def main(env_id, seed):
                 world_model_cfg=dict(
                     latent_recon_loss_weight=0.1, # TODO
                     perceptual_loss_weight=0.1,
-                    use_new_cache_manager=False, # TODO
+                    # use_new_cache_manager=False, # TODO
+                    use_new_cache_manager=True, # ==============TODO==============
 
                     norm_type=norm_type,
                     final_norm_option_in_obs_head='LayerNorm',
@@ -260,7 +261,7 @@ def main(env_id, seed):
 
     # ============ use muzero_segment_collector instead of muzero_collector =============
     from lzero.entry import train_unizero_segment
-    main_config.exp_name = f'data_unizero_st_1226/{env_id[3:-3]}/{env_id[3:-3]}_uz_head-clip-p_target005_allhead4_targetentropy-alpha-500k-098-005-min005_mse-loss2_rec01_poli-clip10_pol-smo-005_pol-loss-tmp-1.5_brf{buffer_reanalyze_freq}-rbs{reanalyze_batch_size}-rp{reanalyze_partition}_nlayer{num_layers}_numsegments-{num_segments}_gsl{game_segment_length}_rr{replay_ratio}_Htrain{num_unroll_steps}-Hinfer{infer_context_length}_bs{batch_size}_seed{seed}'
+    main_config.exp_name = f'data_unizero_st_1226_2/{env_id[3:-3]}/{env_id[3:-3]}_uz_newkv_head-clip-p_target005_allhead4_targetentropy-alpha-500k-098-005-min005_mse-loss2_rec01_poli-clip10_pol-smo-005_pol-loss-tmp-1.5_brf{buffer_reanalyze_freq}-rbs{reanalyze_batch_size}-rp{reanalyze_partition}_nlayer{num_layers}_numsegments-{num_segments}_gsl{game_segment_length}_rr{replay_ratio}_Htrain{num_unroll_steps}-Hinfer{infer_context_length}_bs{batch_size}_seed{seed}'
     # main_config.exp_name = f'data_unizero/{env_id[3:-3]}/{env_id[3:-3]}_uz_brf{buffer_reanalyze_freq}-rbs{reanalyze_batch_size}-rp{reanalyze_partition}_nlayer{num_layers}_numsegments-{num_segments}_gsl{game_segment_length}_rr{replay_ratio}_Htrain{num_unroll_steps}-Hinfer{infer_context_length}_bs{batch_size}_seed{seed}'
 
     train_unizero_segment([main_config, create_config], seed=seed, model_path=main_config.policy.model_path, max_env_step=max_env_step)
@@ -274,9 +275,10 @@ def main(env_id, seed):
     args = parser.parse_args()
 
     # Test environments from atari8 base set
-    # args.env = 'ALE/Pong-v5'               # Memory-planning environment with sparse rewards
+    args.env = 'ALE/Pong-v5'               # Memory-planning environment with sparse rewards
     # args.env = 'ALE/Qbert-v5'               # Memory-planning environment with sparse rewards
-    args.env = 'ALE/MsPacman-v5'               # Memory-planning environment with sparse rewards
+    
+    # args.env = 'ALE/MsPacman-v5'               # Memory-planning environment with sparse rewards
 
     main(args.env, args.seed)