Pccl base

configs/debug/diloco.toml

@@ -9,7 +9,7 @@ micro_bs = 8
 [optim]
 batch_size = 16
 warmup_steps = 10
-total_steps = 4
+total_steps = 128
 
 [data]
 fake = true

configs/debug/normal.toml

@@ -9,7 +9,7 @@ micro_bs = 8
 [optim]
 batch_size = 16
 warmup_steps = 10
-total_steps = 4
+total_steps = 128
 
 [data]
 fake = true

pyproject.toml

@@ -20,6 +20,7 @@ dependencies = [
     "psutil",
     "torch-shampoo @ git+https://github.com/facebookresearch/optimizers.git@main",
     "liger-kernel-nightly>=0.5.2.dev20250122195349",
+    "pccl @ git+https://github.com/PrimeIntellect-ai/pccl.git@main#subdirectory=python/framework"
 ]
 
 [project.optional-dependencies]

scripts/simulate_multi_node_diloco.sh

@@ -71,7 +71,7 @@ export GLOO_SOCKET_IFNAME=lo
 for i in $(seq 0 $(($N - 1 )))
 do
     > logs/log$i.log
-    WANDB_MODE=$([ $i -eq 0 ] && echo "online" || echo "offline") GLOBAL_UNIQUE_ID=$i GLOBAL_RANK=$i CUDA_VISIBLE_DEVICES=$(get_cuda_devices $NUM_GPU $i) uv run torchrun --nproc_per_node=$NUM_GPU --node-rank 0 --rdzv-endpoint localhost:$((BASE_PORT + $i)) --nnodes=1  $@ --data.data_rank $i --data.data_world_size $N > logs/log$i.log 2>&1 &
+    WANDB_MODE=$([ $i -eq 0 ] && echo "online" || echo "offline") GLOBAL_UNIQUE_ID=$i GLOBAL_RANK=$i CUDA_VISIBLE_DEVICES=$(get_cuda_devices $NUM_GPU $i) OMP_NUM_THREADS=1 PCCL_LOG_LEVEL=DEBUG uv run torchrun --nproc_per_node=$NUM_GPU --node-rank 0 --rdzv-endpoint localhost:$((BASE_PORT + $i)) --nnodes=1  $@ --data.data_rank $i --data.data_world_size $N > logs/log$i.log 2>&1 &
     child_pids+=($!)
 done
 

scripts/skip_data.py

@@ -24,7 +24,7 @@
 
 from zeroband.data import get_dataloader
 
-from zeroband.utils.world_info import get_world_info
+from zeroband.utils.world_info import get_local_world_info
 from zeroband.utils.logger import get_logger
 
 
@@ -79,7 +79,7 @@ def skip_data(config: Config):
 if __name__ == "__main__":
     torch.manual_seed(42)
 
-    world_info = get_world_info()
+    world_info = get_local_world_info()
     logger = get_logger()
 
     config = Config(**parse_argv())