Synchronous Multi-node RL

[rafapi]

Synchronous Multi-node RL

Distributed Training Enhancements

• Robust distributed training support with multiple nodes
• New DistributedManager class to handle distributed operations
• File-based synchronization mechanisms for non-distributed phases
• Handling of multi-node NCCL configurations

DeepSpeed Integration

• New DeepSpeed configuration for multi-node scenarios (ds_multinode.json)
• Fixed DeepSpeed checkpoint handling and model saving
• Enhanced optimizer configurations for distributed training
• Added support for CPU Adam optimizer in distributed settings

Data Processing Improvements

• Implemented data splitting across nodes for better load distribution
• Added sample gathering from multiple ranks
• Improved error handling and retry mechanisms for data loading

Memory Management

• Added GPU resource cleanup functionality
• Improved memory status monitoring
• Added CUDA cache management

Logging and Monitoring

• Enhanced W&B logging for distributed training
• Added rank-aware logging
• Improved error reporting across distributed processes
• Added execution time tracking for different training phases

Synchronization Mechanisms

• Added robust barrier implementations with timeout and retry logic
• Implemented file-based synchronization for non-distributed phases
• Added checks to ensure all nodes complete critical sections
• Enhanced coordination between nodes during data processing and training

Best squeeze of throughput

[rizar]

Good stuff! Please see my comments below.

[rafapi]

Refactoring to mitigate NCCL timeouts and handle multi-node rollout generation

https://wandb.ai/rafa/tapeagents/runs/rl_gsm8k_deepspeed_llama31_70b_new_save_rl_no_kl_lr1e-6_2nodes_take5_finetune?nw=nwuserrafa

https://wandb.ai/rafa/tapeagents/runs/rl_gsm8k_deepspeed_llama31_70b_new_save_rl_no_kl_lr1e-6_2nodes_take5?nw=nwuserrafa

[rafapi]

https://wandb.ai/rafa/tapeagents/runs/rl_gsm8k_deepspeed_llama31_70b_rl_no_kl_lr1e-6_2nodes_no_gather_file_sync_finetune?nw=nwuserrafa

[rafapi]